{-# LANGUAGE ScopedTypeVariables, BangPatterns #-}
module ECC.Tester (eccMain, eccPrinter, eccMerger) where

import ECC.Types
import System.Random.MWC
import Control.Monad
--import Control.Monad.Primitive (PrimMonad, PrimState)
import Control.Monad.Trans (MonadIO(liftIO))
import Control.Concurrent
import Data.Word
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed  as U
import System.Random.MWC.Distributions
import Data.Monoid
import Statistics.Resampling.Bootstrap
import System.IO
import Data.Char
import System.Environment
import ECC.Estimate
import GHC.Conc
import Control.Concurrent.ParallelIO
import Numeric
import Data.Time.Clock
import qualified Data.UUID as UUID
import System.Directory
import qualified Data.Csv as CSV
import qualified Data.ByteString.Lazy as LBS
import Statistics.Types (Estimate (..), ConfInt(..), CL, confidenceLevel)
import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)
import Data.List as List
import Data.Ratio
import Text.Read
import Data.Word (Word32)
import Data.Maybe (isJust)

data Options = Options
        { codenames :: [String]
        , ebN0s     :: [EbN0]
        , verbose   :: Int
        , enough    :: Enough
        , logDir    :: FilePath -- directory or file to put the log
        , cmp       :: Bool
        , replay    :: Maybe FilePath -- file to run replays
        , seed      :: Maybe Word32
        } deriving Show

data Enough = BitErrorCount Int
            | MessageCount Int
        deriving Show

data TestRun = TestRun
       { test_encode_time    :: !Double  -- in seconds
       , test_decode_time    :: !Double -- in seconds
       , test_ber            :: !BEs
       }
      deriving Show

instance Monoid TestRun where
    mempty = TestRun 0 0 mempty
    TestRun en1 de1 ber1 `mappend` TestRun en2 de2 ber2
        = TestRun (en1 + en2) (de1 + de2) (ber1 `mappend` ber2)

defaultLogDir :: String
defaultLogDir = "log/"

-- | Give a 'Code' (set of possible Error Correcting Codes) and a printer, run the tests.
eccMain :: Code -> (Options -> [String] -> IO (ECC IO -> EbN0 -> TestRun -> IO Bool)) -> IO ()
eccMain code k = do
        args <- getArgs
        if null args
         then error $ "usage: <name> [-v<n>] [-b<n>] [-m<n>] [-l<file or dir/>] [-c] <EbN0_1> <EbN0_2> ... <Code Name> <Code Name>"
                   ++ "\ncodes: " ++ show code
         else eccTester (parseOptions args $ defaultOptions defaultLogDir) code k

parseOptions :: [String] -> Options -> Options
parseOptions (('-':'v':ns):rest) o
        | all isDigit ns = (parseOptions rest o) { verbose = read ns }
parseOptions (('-':'b':ns):rest) o
        | all isDigit ns = (parseOptions rest o) { enough = BitErrorCount (read ns) }
parseOptions (('-':'m':ns):rest) o
        | all isDigit ns = (parseOptions rest o) { enough = MessageCount (read ns) }
parseOptions (('-':'l':lg):rest) o
        | otherwise = (parseOptions rest o) { logDir = lg }
parseOptions (('-':'r':r):rest) o
        | otherwise = (parseOptions rest o) { replay = Just r }
parseOptions (('-':'s':s):rest) o
        | otherwise = (parseOptions rest o) { seed = Just (read s) }
parseOptions ("-c":rest) o
        | otherwise       = (parseOptions rest o) { cmp = True }
parseOptions (arg:rest) o =
        case reads arg of
          [(ebN0::EbN0,"")] -> opts { ebN0s = ebN0 : ebN0s opts }
          _                 -> opts { codenames = arg : codenames opts }
  where
     opts = parseOptions rest o
parseOptions [] o = o

defaultOptions :: String -> Options
defaultOptions d = Options {
    codenames = [],
    ebN0s = [],
    verbose = 0,
    enough = BitErrorCount 1000,
    logDir = d,
    cmp = False,
    replay = Nothing,
    seed = Nothing
  }

-- | A basic printer for our tests. Currently, we report on powers of two,
-- and accept a value if there are at least 1000 bit errors (say).
-- We also output a log file into a diretory, using a UUID style file.
eccPrinter :: Options -> [String] -> IO (ECC f -> EbN0 -> TestRun -> IO Bool)
eccPrinter opts names = do

   let tab1 = maximum (map length names)

   let rjust n xs = take (n - length xs) (cycle " ") ++ xs

   -- This generator is uses for the bootstrapping, and generating the UUID
   gen :: GenIO <- createSystemRandom
   start <- getCurrentTime

   (w1,w2,w3,w4) <- uniform gen
   let uuid = UUID.fromWords w1 w2 w3 w4

   logFileName <- if last (logDir opts) == '/' then do
       createDirectoryIfMissing True $ logDir opts
       return $ logDir opts ++ UUID.toString uuid
     else
       return $ logDir opts

     -- Don't write to the default log directory if a seed was manually
     -- provided (to avoid skewing the statistics, if the same seed is
     -- manually used for a number of runs).
   unless (logDir opts == defaultLogDir && isJust (seed opts))
     $ writeHeader logFileName

   putStrLn $ "#" ++
              rjust 7    "Time" ++ " " ++
              rjust tab1 "ECC" ++ " " ++
              rjust 5    "EbN0" ++ " " ++
              rjust 8    "Packets" ++ " " ++
              rjust 10   "Encode/s" ++ " " ++
              rjust 10   "Decode/s" ++ " " ++
              rjust 8    "Errors" ++ " " ++
              rjust 8    "BER" ++ " " ++
              ""

   return $ \  ecc ebN0 (TestRun tEn tDe bes) -> do
           est <- if sizeBEs bes <= 2
                  then return Nothing
                  else estimate gen 0.95 (message_length ecc) bes
           let accept = case enough opts of
                          BitErrorCount n -> sumBEs bes >= n
                          MessageCount n  -> sizeBEs bes >= n

           now <- getCurrentTime
           let diff = diffUTCTime now start

           putStrLn $
                    rjust 8 (showFFloat (Just 2) (realToFrac diff) "") ++ " " ++
                    rjust tab1 (name ecc) ++ " " ++
                    rjust 5 (showFFloat (Just 2) ebN0 "") ++ " " ++
                    rjust 8 (show (sizeBEs bes)) ++ " " ++
                    rjust 10 (showFFloat (Just 4) (fromIntegral (sizeBEs bes)/tEn) "") ++ " " ++
                    rjust 10 (showFFloat (Just 4) (fromIntegral (sizeBEs bes)/tDe) "") ++ " " ++
                    rjust 8 (show (sumBEs bes)) ++ " " ++
                    (case est of
                      Just e -> " " ++ rjust 20 (showEstimate e)
                      Nothing -> " 0.00e-0") ++
                    if accept then "." else ","
           hFlush stdout

           -- Log file, with raw data

           unless (logDir opts == defaultLogDir && isJust (seed opts))
             $ LBS.appendFile logFileName $ CSV.encode [ Result
                 ( realToFrac diff :: Double )
                 ( name ecc        :: String )
                 ( ebN0            :: Double )
                 ( sizeBEs bes     :: Int )
                 ( tEn        )
                 ( tDe        )
                 ( sumBEs bes      :: Int )
                 ( estPoint <$> est )
                 ( (confIntLDX . estError) <$> est )
                 ( (confIntUDX . estError) <$> est )
                 ( (confidenceLevel . confIntCL . estError) <$> est )
                 ( message_length ecc )
                 ( codeword_length ecc )
                 ( bes )
                 ]
           return accept

-- eccTester runs testECC at different ebN0 levels.
eccTester :: Options -> Code -> (Options -> [String] -> IO (ECC IO -> EbN0 -> TestRun -> IO Bool)) -> IO ()
eccTester opts (Code _ initializeC finalizeC f) k = do
   print opts
   let debug n msg | n <= verbose opts  = putStrLn msg
                   | otherwise  = return ()

   -- This generator is uses for the generation of bits
   gen :: GenIO <- case seed opts of
                     Just s -> initialize (V.singleton s)
                     _      -> createSystemRandom

   let mg :: (Show a, Show b, Eq b, Monad m) => (a -> m b) -> (a -> m b) -> a -> m b
       mg f1 f2 a = do
           r1 <- f1 a
           r2 <- f2 a
           if r1 == r2
            then return r1
            else fail $ "merge of two codes failed at runtime" ++ show (a,r1,r2)

   let jn :: Monad m => ECC m -> ECC m -> ECC m
       jn ecc1 ecc2
            | message_length ecc1 /= message_length ecc2
             = error "trying to combine two codes with different message lengths"
            | codeword_length ecc1 /= codeword_length ecc2
            = error "trying to combine two codes with different codeword lengths"
            | otherwise
            = ECC { name = name ecc1 ++ "+" ++ name ecc2
                  , encode = mg (encode ecc1) (encode ecc2)
                  , decode = mg (decode ecc1) (decode ecc2)
                  , message_length = message_length ecc1
                  , codeword_length = codeword_length ecc2
                  }

   vars <- initializeC

   eccs <- liftM concat
            $ sequence
            $ map (f vars)
            $ map splitCodename
            $ codenames opts

   let eccs' = if cmp opts then [foldr1 jn eccs] else eccs

   k2 <- k opts (map name eccs)

   -- First, try replay (which have fixed EbN0s)
   replay_log <- case replay opts of
                   Nothing -> pure []
                   Just fileName -> do
                        str <- readFile fileName
                        return [ r :: Log
                               | ln <- lines str
                               , r <- case readMaybe ln of
                                        Nothing -> []
                                        Just x -> [x]
                               ]
   let replayMe [] = return ()
       replayMe (Message m:TxCodeword tx:RxCodeword ebN0 rx:rest) = do
           sequence_
               [ do bes <- rxECC (verbose opts) ecc m ebN0 rx
                    k2 ecc ebN0 bes
               | ecc <- eccs'
               ]
           replayMe rest
       replayMe (_:rest) = replayMe rest

   replayMe replay_log

   -- Then, loop over all given EbN0s.
   sequence_
          [ sequence_
                [ do testECC (verbose opts) gen ebN0 ecc k2
                | ebN0 <- ebN0s opts
                ]
          | ecc <- eccs'
          ]

   finalizeC vars

   stopGlobalPool -- TODO: use a local pool instead?

-- Running a *multi* run of an ECC, giving a compounted TestRun result to the
-- continuation, and stopping when the continuation returns True.
testECC :: Int -> GenIO -> EbN0 -> ECC IO -> (ECC IO -> EbN0 -> TestRun -> IO Bool) -> IO ()
testECC verb gen ebN0 ecc k = do
   let debug n msg | n <= verb  = putStrLn msg
                   | otherwise  = return ()

   cap <- getNumCapabilities

   let ourpar = sequence -- parallel
   let cap = 1           -- for now

   let loop (n:ns) !bes = do
        debug 1 $ "trying " ++ show n ++ " messages"
        let real_par = min (cap * 16) n
        bess <- ourpar   [ do foldM (\ !a !_ -> do bes0 <- runECC verb gen ebN0 ecc
                                                   return $ a <> bes0) mempty [1..(n `div` real_par)]
                         | _ <- [1..real_par]
                         ]
        let bes1 = mconcat (bes:bess)
        let errs = sumBEs $ test_ber bes1
        debug 1 $ "found " ++ show errs ++ " so far"
        okay <- k ecc ebN0 bes1
        if okay then return () else loop ns bes1

   -- We run with the cap twice, then double each time
   loop (cap : iterate (*2) cap) mempty

splitCodename :: String -> [String]
splitCodename = words . map (\ c -> if c == '/' then ' ' else c)

-- Running a *single* run of an ECC, getting a single bit error count
runECC :: Int -> GenIO -> EbN0 -> ECC IO -> IO TestRun
runECC verb gen ebN0 !ecc = do

  let debug n msg | n <= verb  = putStrLn msg
                  | otherwise  = return ()

  debug 3 $ "starting message"
  !mess0  <- liftM U.fromList $ sequence [ uniform gen | _ <- [1..message_length ecc]]
  debug 3 $ "generated message"
  debug 4 $ show (Message mess0)

  start_encoding <- getCurrentTime

  !code0  <- encode ecc mess0
  debug 3 $ "encoded message"
  debug 4 $ show (TxCodeword code0)

  end_encoding <- getCurrentTime

  !rx <- txRx_EbN0 ebN0 (rateOf ecc) gen code0
  debug 3 $ "tx/rx'd message"
  debug 4 $ show (RxCodeword ebN0 rx)

  rxECC verb ecc mess0 (realToFrac $ diffUTCTime end_encoding start_encoding) rx

rxECC :: Int -> ECC IO -> U.Vector Bool -> Double -> U.Vector Double -> IO TestRun
rxECC verb ecc mess0 encoding_time rx = do

  let debug n msg | n <= verb  = putStrLn msg
                  | otherwise  = return ()

  start_decoding <- getCurrentTime

  !(!mess1,parity) <- decode ecc rx

  debug 3 $ "decoded message"
  debug 4 $ "parity: " ++ show parity
  debug 4 $ show (Decoded (name ecc) mess1)

  when (U.length mess0 /= U.length mess1) $ do
    error $ "before and after codeword different lengths" ++ show (U.length mess0,U.length mess1)

  let !bitErrorCount = length [ () | (b,a) <- U.toList $ U.zip mess0 mess1, a /= b ]
  debug 2 $ show bitErrorCount ++ " bit errors in message"

  end_decoding <- getCurrentTime

  return $ TestRun encoding_time
                   (realToFrac $ diffUTCTime end_decoding start_decoding)
         $ eventBEs bitErrorCount

{-

    Sigma2 = max(( 1/10^(EbNoDB(jj)/10) )/2, 1e-9);
    Lc    = 2*sqrt(Ec)/Sigma2;

    for ii=1:NumBlocks(jj)

        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        Bits = 0.5 >= rand(1,k);
        CodedBits = mod(Bits*G,2);

        % Modulate the CodedBits
        s     = sqrt(Ec)*(2*CodedBits - 1);
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

        % AWGN channel
        noise  = sqrt(Sigma2)*randn(size(s));
        r(1:n) = s + noise;
-}

-- Adding randomness
txRx_EbN0 :: EbN0 -> Rate -> GenIO -> U.Vector Bool -> IO (U.Vector Double)
txRx_EbN0 ebnoDB rate gen xs
        | isNaN ebnoDB = return $ U.map soft
                                $ xs
txRx_EbN0 ebnoDB rate gen xs = do
        rs :: U.Vector Double  <- U.fromList <$> sequence [ standard gen | _ <- U.toList xs ]
        return $ U.map (* lc)
               $ U.zipWith (+) (U.map (* sqrt sigma2) rs)
                               (U.map (* sqrt ec)
                                  $ U.map soft
                                  $ xs)
     where
         sigma2 = ((1/10) ** (ebnoDB/10)) / 2
         ec     = fromIntegral (numerator rate) / fromIntegral (denominator rate)
         lc     = 2 * sqrt ec / sigma2

-- eccMerge merges all the input files.
eccMerger :: IO ()
eccMerger = do
    args <- getArgs
    if null args
     then error $ "usage: <name> [-v<n>] [-b<n>] [-m<n>] [-l<file or dir/>] [-c] log/<file> log/<file>"
     else do
       let opts = parseOptions args $ defaultOptions "merge/"
--       print opts
       -- now we try load the files.
       logs <- sequence
         [ do bs <- LBS.readFile nm
              case CSV.decode CSV.HasHeader bs of
                Left msg -> fail msg
                Right t  -> return $ V.toList $ (t :: V.Vector Result)
         | nm <- codenames opts
         ]

       -- first, pick only the final results. We can not merge partual with final results.
       let sumMe :: Result -> Map String (Map Double Result) -> Map String (Map Double Result)
           sumMe r m = flip (Map.insert (res_name r)) m $
                   case Map.lookup (res_name r) m of
                     Nothing -> Map.singleton (res_ebN0 r) r
                     Just m' -> flip (Map.insert (res_ebN0 r)) m' $
                       case Map.lookup (res_ebN0 r) m' of
                         Nothing -> r
                         Just r' | res_size r >= res_size r' -> r
                                 | otherwise                 -> r'

       let logs' = map (foldr sumMe Map.empty) logs
       -- next, merge the different runs
       let mergeRes :: Result -> Result -> Result
           mergeRes r1 r2 = r1
              { timestamp = 0.0
              , res_size = res_size r1 + res_size r2
              , res_tEn  = res_tEn r1 + res_tEn r2
              , res_tDe  = res_tDe r1 + res_tDe r2
              , res_sum  = res_sum r1 + res_sum r2
              , res_ber  = Nothing
              , res_ldx  = Nothing
              , res_udx  = Nothing
              , res_conf = Nothing
              , res_bes  = res_bes r1 <> res_bes r2
              }
       let sumLog = Map.unionsWith (Map.unionWith mergeRes) logs'
--       print sumLog

       -- setup the printer

       -- Options -> [String] -> IO (ECC f -> EbN0 -> TestRun -> IO Bool)
       p <- eccPrinter opts (Map.keys sumLog)

       sequence_ [ do
         sequence_ [ p ecc db testrun
                   | db <- List.sort $ Map.keys m
                   , let r = m Map.! db
                     -- These are both fake, for the purposes of the log printing
                   , let ecc = ECC (res_name r) undefined undefined (res_mesg r) (res_code r)
                   , let testrun = TestRun (res_tEn r) (res_tDe r) (res_bes r)
                   ]
                 | (nm,m) <- Map.toList sumLog
                 ]

data Result = Result
 { timestamp  :: Double
 , res_name   :: String
 , res_ebN0   :: Double
 , res_size   :: Int     -- number of packets tested
 , res_tEn    :: Double
 , res_tDe    :: Double
 , res_sum    :: Int     -- total bit errors
 , res_ber    :: Maybe Double
 , res_ldx    :: Maybe Double
 , res_udx    :: Maybe Double
 , res_conf   :: Maybe Double
 , res_mesg   :: Int        -- length of message
 , res_code   :: Int        -- length of codeword
 , res_bes    :: BEs
 } deriving Show

instance CSV.ToRecord Result where
    toRecord (Result a b c d e f g h i j k l m n) = V.fromList [
        CSV.toField a, CSV.toField b, CSV.toField c, CSV.toField d, CSV.toField e, CSV.toField f,
        CSV.toField g, CSV.toField h, CSV.toField i, CSV.toField j, CSV.toField k, CSV.toField l,
        CSV.toField m, CSV.toField (show n)]

instance CSV.FromRecord Result where
    parseRecord v
        | n == 14    = Result        <$> CSV.unsafeIndex v 0
                                     <*> CSV.unsafeIndex v 1
                                     <*> CSV.unsafeIndex v 2
                                     <*> CSV.unsafeIndex v 3
                                     <*> CSV.unsafeIndex v 4
                                     <*> CSV.unsafeIndex v 5
                                     <*> CSV.unsafeIndex v 6
                                     <*> CSV.unsafeIndex v 7
                                     <*> CSV.unsafeIndex v 8
                                     <*> CSV.unsafeIndex v 9
                                     <*> CSV.unsafeIndex v 10
                                     <*> CSV.unsafeIndex v 11
                                     <*> CSV.unsafeIndex v 12
                                     <*> (read <$> CSV.unsafeIndex v 13)
        | otherwise = fail "Fail to read Result"
          where
            n = V.length v

writeHeader :: String -> IO ()
writeHeader logFileName =
   LBS.writeFile logFileName $ CSV.encode [
               [ "Time"
               , "ECC"
               , "EbN0"
               , "Packets"
               , "Encodes"
               , "Decodes"
               , "Errors"
               , "BER"
               , "LDX"
               , "UDX"
               , "Confidence"
               , "Message"
               , "Codeword"
               , "Internal"
               ]]


module PeekyBlinders.Prelude
  ( module Exports,
    showBin,
  )
where

import Control.Applicative as Exports hiding (WrappedArrow (..))
import Control.Arrow as Exports hiding (first, second)
import Control.Category as Exports
import Control.Concurrent as Exports
import Control.Exception as Exports
import Control.Monad as Exports hiding (fail, forM, forM_, mapM, mapM_, msum, sequence, sequence_)
import Control.Monad.Fail as Exports
import Control.Monad.Fix as Exports hiding (fix)
import Control.Monad.IO.Class as Exports
import Control.Monad.ST as Exports
import Data.Bifunctor as Exports
import Data.Bits as Exports
import Data.Bool as Exports
import Data.ByteString as Exports (ByteString)
import Data.ByteString.Short as Exports (ShortByteString)
import Data.Char as Exports
import Data.Coerce as Exports
import Data.Complex as Exports
import Data.Data as Exports
import Data.Dynamic as Exports
import Data.Either as Exports
import Data.Fixed as Exports
import Data.Foldable as Exports hiding (toList)
import Data.Function as Exports hiding (id, (.))
import Data.Functor as Exports
import Data.Functor.Compose as Exports
import Data.Functor.Contravariant as Exports
import Data.IORef as Exports
import Data.Int as Exports
import Data.Ix as Exports
import Data.List as Exports hiding (all, and, any, concat, concatMap, elem, find, foldl, foldl', foldl1, foldr, foldr1, isSubsequenceOf, mapAccumL, mapAccumR, maximum, maximumBy, minimum, minimumBy, notElem, or, product, sortOn, sum, uncons)
import Data.List.NonEmpty as Exports (NonEmpty (..))
import Data.Maybe as Exports
import Data.Monoid as Exports hiding (Alt, (<>))
import Data.Ord as Exports
import Data.Proxy as Exports
import Data.Ratio as Exports
import Data.STRef as Exports
import Data.Semigroup as Exports hiding (First (..), Last (..))
import Data.String as Exports
import Data.Text as Exports (Text)
import Data.Traversable as Exports
import Data.Tuple as Exports
import Data.Unique as Exports
import Data.Version as Exports
import Data.Void as Exports
import Data.Word as Exports
import Debug.Trace as Exports
import Foreign as Exports hiding (void)
import GHC.Conc as Exports hiding (orElse, threadWaitRead, threadWaitReadSTM, threadWaitWrite, threadWaitWriteSTM, withMVar)
import GHC.Exts as Exports (IsList (..), groupWith, inline, lazy, sortWith)
import GHC.Generics as Exports (Generic)
import GHC.IO.Exception as Exports
import GHC.OverloadedLabels as Exports
import Numeric as Exports
import System.Environment as Exports
import System.Exit as Exports
import System.IO as Exports
import System.IO.Error as Exports
import System.IO.Unsafe as Exports
import System.Mem as Exports
import System.Mem.StableName as Exports
import System.Timeout as Exports
import Text.ParserCombinators.ReadP as Exports (ReadP, readP_to_S, readS_to_P)
import Text.ParserCombinators.ReadPrec as Exports (ReadPrec, readP_to_Prec, readPrec_to_P, readPrec_to_S, readS_to_Prec)
import Text.Printf as Exports (hPrintf, printf)
import Text.Read as Exports (Read (..), readEither, readMaybe)
import Unsafe.Coerce as Exports
import Prelude as Exports hiding (all, and, any, concat, concatMap, elem, fail, foldl, foldl1, foldr, foldr1, id, mapM, mapM_, maximum, minimum, notElem, or, product, sequence, sequence_, sum, (.))

showBin :: (Integral a, Show a) => a -> String
showBin a =
  showIntAtBase 2 intToDigit a ""


-- |
-- Copyright: (C) 2013 Amgen, Inc.
--
-- Debugging facilities, in particular to analyze the internal structure of
-- a 'SEXP' as structured JSON.
--
-- This module is intended to be imported qualified.

{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}

module Language.R.Debug
  ( inspect )
  where

import Control.Memory.Region (V)
import qualified Data.Vector.SEXP as Vector
import qualified Foreign.R as R
import Foreign.R (SEXP, SomeSEXP(..), SEXPTYPE, SEXPInfo)
import Foreign.R.Type (IsVector)
import Foreign.Storable
import Language.R.Globals as H
import Language.R.HExp

import Data.Complex
import System.IO.Unsafe ( unsafePerformIO )

import Data.Aeson as A
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.ByteString.Lazy.Char8 as LBS

instance ToJSON SEXPTYPE where
  toJSON = A.String . T.pack . show

instance ToJSON SEXPInfo where
  toJSON x =
    object
      [ "type"  .= R.infoType x
      , "obj"   .= R.infoObj x
      , "named" .= R.infoNamed x
      , "gp"    .= R.infoGp x
      , "mark"  .= R.infoMark x
      , "debug" .= R.infoDebug x
      , "trace" .= R.infoTrace x
      , "spare" .= R.infoSpare x
      , "gcgen" .= R.infoGcGen x
      , "gccls" .= R.infoGcCls x
      ]

instance ToJSON a => ToJSON (Complex a) where
  toJSON (x :+ y) =
    object ["Re" .= x, "Im" .= y]

instance ToJSON (SEXP s a) where
  toJSON x =
      object
        [ "header" .= info
        , "attributes" .= if R.unsexp x == R.unsexp attr then "loop" else toJSON attr
        , tp .= go x
        ]
    where
      vector :: (IsVector a, ToJSON (Vector.ElemRep V a), Storable (Vector.ElemRep V a))
             => Vector.Vector a (Vector.ElemRep V a) -> V.Vector Value
      vector = V.fromList . map toJSON . Vector.toList -- XXX: do not use lists
      ub = R.unsexp H.unboundValue
      nil = R.unsexp H.nilValue
      miss = R.unsexp H.missingArg
      info = unsafePerformIO $ R.peekInfo x
      attr = unsafePerformIO $ R.getAttributes x
      tp = T.pack . show $ R.infoType info
      go :: SEXP s a -> Value
      go y | R.unsexp y == ub   = A.String "UnboundValue"
           | R.unsexp y == nil  = A.String "NilValue"
           | R.unsexp y == miss = A.String "MissingArg"
      go (hexp -> Nil) = A.String "NilValue"
      go (hexp -> Lang i j) =
          object [ "function" .= i
                 , "parameters" .= j
                 ]
      go h@(hexp -> Symbol i j k) =
          object [ "name" .= i
                 , "value" .= if R.unsexp j == R.unsexp h then "loop" else toJSON j
                 , "internal" .= k
                 ]
      go (hexp -> Special i) = object ["index" .= i]
      go (hexp -> Builtin i) = object ["index" .= i]
      go (hexp -> Char v) = A.String (T.pack (Vector.toString v))
      go (hexp -> Int v) = A.Array (vector v)
      go (hexp -> Real v) = A.Array (vector v)
      go (hexp -> Complex v) = A.Array (vector v)
      go (hexp -> Vector _ v) = A.Array (vector v)
--  String    :: {-# UNPACK #-} !(Vector.Vector (SEXP (R.Vector Word8)))
--            -> HExp (R.Vector (SEXP (R.Vector Word8)))
      go (hexp -> List i j k) =
          object [ "value" .= i
                 , "next"  .= j
                 , "tag"   .= k
                 ]
      go (hexp -> Env _ _ _) = A.String "Environment"
      go (hexp -> Closure f b e) =
         object [ "formals" .= f
                , "body" .= b
                , "environment" .= e
                ]
      go (hexp -> Promise vl ex en) =
          object [ "value" .= vl
                 , "expr" .= ex
                 , "environment" .= en
                 ]
      go (hexp -> DotDotDot v) =
          object [ "promises" .= v]
      go (hexp -> Expr _ v)   = A.Array (vector v)
      go (hexp -> Bytecode) = A.String "Bytecode"
      go (hexp -> ExtPtr _ a b) =
          object [ "ptr" .= A.String "<PTR>"
                 , "second" .= a
                 , "symbol" .= b
                 ]
      go (hexp -> WeakRef k v fn nxt) =
          object [ "key" .= k
                 , "value" .= v
                 , "finalizer" .= fn
                 , "next" .= nxt
                 ]
      go (hexp -> Raw _bs) = A.String "<data>"
      go (hexp -> S4 s) =
          object [ "tagval" .= s ]
      go _ = A.String "Unimplemented."

instance ToJSON (SomeSEXP s) where
  toJSON (R.SomeSEXP s) = toJSON s

inspect :: SEXP s a -> String
inspect = LBS.unpack . A.encode


--------------------------------------------------------------------------------

{-# OPTIONS_GHC -fno-warn-orphans #-}

--------------------------------------------------------------------------------

{-# LANGUAGE AllowAmbiguousTypes       #-}
{-# LANGUAGE BangPatterns              #-}
{-# LANGUAGE CPP                       #-}
{-# LANGUAGE EmptyDataDecls            #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE ForeignFunctionInterface  #-}
{-# LANGUAGE FunctionalDependencies    #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE KindSignatures            #-}
{-# LANGUAGE LambdaCase                #-}
{-# LANGUAGE MagicHash                 #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE Rank2Types                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE StandaloneDeriving        #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeFamilyDependencies    #-}
{-# LANGUAGE TypeInType                #-}
{-# LANGUAGE UnboxedTuples             #-}
{-# LANGUAGE UndecidableInstances      #-}

--------------------------------------------------------------------------------

-- | Internal module to Eigen.
--   Here we define all foreign function calls,
--   and some typeclasses integral to the public and private interfaces
--   of the library.
module Eigen.Internal where --   FIXME: Explicit export list

--------------------------------------------------------------------------------

import           Control.Monad            (when)
import           Data.Binary              (Binary(put,get))
import           Data.Binary.Get          (getByteString, getWord32be)
import           Data.Binary.Put          (putByteString, putWord32be)
import           Data.Bits                (xor)
import           Data.Complex             (Complex((:+)))
import           Data.Kind                (Type)
import           Data.Proxy               (Proxy(Proxy))
import           Foreign.C.String         (CString, peekCString)
import           Foreign.C.Types          (CInt(CInt), CFloat(CFloat), CDouble(CDouble), CChar)
import           Foreign.ForeignPtr       (ForeignPtr, castForeignPtr, withForeignPtr)
import           Foreign.Ptr              (Ptr, castPtr, nullPtr, plusPtr)
import           Foreign.Storable         (Storable(sizeOf, alignment, poke, peek, peekByteOff, peekElemOff, pokeByteOff, pokeElemOff))
import           GHC.TypeLits             (natVal, KnownNat, Nat)
import           System.IO.Unsafe         (unsafeDupablePerformIO)
import qualified Data.Vector.Storable     as VS
import qualified Data.ByteString          as BS
import qualified Data.ByteString.Internal as BSI

--------------------------------------------------------------------------------

-- | Like 'Proxy', but specialised to 'Nat'.
data Row (r :: Nat) = Row
-- | Like 'Proxy', but specialised to 'Nat'.
data Col (c :: Nat) = Col

-- | Used internally. Given a 'KnownNat' constraint, turn the type-level 'Nat' into an 'Int'.
natToInt :: forall n. KnownNat n => Int
{-# INLINE natToInt #-}
natToInt = fromIntegral (natVal @n Proxy)

--------------------------------------------------------------------------------

-- | Cast to and from a C-FFI type
--   'Cast' is a closed typeclass with an associated injective type family.
--   It is closed in the sense that we provide only four types
--   with instances for it; and intend for eigen to only be used
--   with those four types. The injectivity of the type family is
--   then useful for avoiding MPTCs. 'Cast' has two functions; 'toC'
--   and 'fromC', where 'toC' goes from a Haskell type to its associated
--   C type for internal use, with the C FFI, and 'fromC' goes from the
--   associated C type to the Haskell type.
class Cast (a :: Type) where
  type family C a = (result :: Type) | result -> a
  toC   :: a -> C a
  fromC :: C a -> a

instance Cast Int where
  type C Int = CInt
  toC = CInt . fromIntegral
  {-# INLINE toC #-}
  fromC (CInt x) = fromIntegral x
  {-# INLINE fromC #-}

instance Cast Float where
  type C Float = CFloat
  toC = CFloat
  {-# INLINE toC #-}
  fromC (CFloat x) = x
  {-# INLINE fromC #-}

instance Cast Double where
  type C Double = CDouble
  toC = CDouble
  {-# INLINE toC #-}
  fromC (CDouble x) = x
  {-# INLINE fromC #-}

instance Cast a => Cast (Complex a) where
  type C (Complex a) = CComplex (C a)
  toC (a :+ b) = CComplex (toC a) (toC b)
  {-# INLINE toC #-}
  fromC (CComplex a b) = (fromC a) :+ (fromC b)
  {-# INLINE fromC #-}

-- | WARNING! 'toC' is lossy for any Int greater than (maxBound :: Int32)!
instance Cast a => Cast (Int, Int, a) where
  type C (Int, Int, a) = CTriplet a
  {-# INLINE toC #-}
  toC (x, y, z) = CTriplet (toC x) (toC y) (toC z)
  {-# INLINE fromC #-}
  fromC (CTriplet x y z) = (fromC x, fromC y, fromC z)

--------------------------------------------------------------------------------

-- | Complex number for FFI with the same memory layout as std::complex\<T\>
data CComplex a = CComplex !a !a deriving (Show)

instance Storable a => Storable (CComplex a) where
    sizeOf _ = sizeOf (undefined :: a) * 2
    alignment _ = alignment (undefined :: a)
    poke p (CComplex x y) = do
        pokeElemOff (castPtr p) 0 x
        pokeElemOff (castPtr p) 1 y
    peek p = CComplex
        <$> peekElemOff (castPtr p) 0
        <*> peekElemOff (castPtr p) 1

--------------------------------------------------------------------------------

-- | FIXME: Doc
data CTriplet a where
  CTriplet :: Cast a => !CInt -> !CInt -> !(C a) -> CTriplet a

deriving instance (Show a, Show (C a)) => Show (CTriplet a)

instance (Storable a, Elem a) => Storable (CTriplet a) where
    sizeOf _ = sizeOf (undefined :: a) + sizeOf (undefined :: CInt) * 2
    alignment _ = alignment (undefined :: CInt)
    poke p (CTriplet row col val) = do
        pokeElemOff (castPtr p) 0 row
        pokeElemOff (castPtr p) 1 col
        pokeByteOff p (sizeOf (undefined :: CInt) * 2) val
    peek p = CTriplet
        <$> peekElemOff (castPtr p) 0
        <*> peekElemOff (castPtr p) 1
        <*> peekByteOff p (sizeOf (undefined :: CInt) * 2)

--------------------------------------------------------------------------------

-- | `Elem` is a closed typeclass that encompasses the properties
--   eigen expects its values to possess, and simplifies the external
--   API quite a bit.
class (Num a, Cast a, Storable a, Storable (C a), Code (C a)) => Elem a

instance Elem Float
instance Elem Double
instance Elem (Complex Float)
instance Elem (Complex Double)

--------------------------------------------------------------------------------

-- | Encode a C Type as a CInt
--
--   Hack used in FFI wrapper functions when constructing FFI calls
class Code a where; code :: a -> CInt
instance Code CFloat             where; code _ = 0
instance Code CDouble            where; code _ = 1
instance Code (CComplex CFloat)  where; code _ = 2
instance Code (CComplex CDouble) where; code _ = 3

-- | Hack used in constructing FFI calls.
newtype MagicCode = MagicCode CInt deriving Eq

instance Binary MagicCode where
    put (MagicCode _code) = putWord32be $ fromIntegral _code
    get = MagicCode . fromIntegral <$> getWord32be

-- | Hack used in constructing FFI calls.
magicCode :: Code a => a -> MagicCode
magicCode x = MagicCode (code x `xor` 0x45696730)

--------------------------------------------------------------------------------

-- | Machine size of a 'CInt'.
intSize :: Int
intSize = sizeOf (undefined :: CInt)

-- | FIXME: Doc
encodeInt :: CInt -> BS.ByteString
encodeInt x = BSI.unsafeCreate (sizeOf x) $ (`poke` x) . castPtr

-- | FIXME: Doc
decodeInt :: BS.ByteString -> CInt
decodeInt (BSI.PS fp fo fs)
    | fs == sizeOf x = x
    | otherwise = error "decodeInt: wrong buffer size"
    where x = performIO $ withForeignPtr fp $ peek . (`plusPtr` fo)

--------------------------------------------------------------------------------

-- | 'Binary' instance for 'Data.Vector.Storable.Mutable.Vector'
instance Storable a => Binary (VS.Vector a) where
    put vs = put (BS.length bs) >> putByteString bs where
        (fp,fs) = VS.unsafeToForeignPtr0 vs
        es = sizeOf (VS.head vs)
        bs = BSI.fromForeignPtr (castForeignPtr fp) 0 (fs * es)

    get = get >>= getByteString >>= \bs -> let
        (fp,fo,fs) = BSI.toForeignPtr bs
        es = sizeOf (VS.head vs)
        -- `plusForeignPtr` is used qualified here to just remind a reader
        -- that it is defined internally within eigen
        vs = VS.unsafeFromForeignPtr0 (Eigen.Internal.plusForeignPtr fp fo) (fs `div` es)
        in return vs

--------------------------------------------------------------------------------

-- | FIXME: Doc
data CSparseMatrix a
-- | FIXME: Doc
type CSparseMatrixPtr a = Ptr (CSparseMatrix a)

-- | FIXME: Doc
data CSolver a
-- | FIXME: Doc
type CSolverPtr a = Ptr (CSolver a)

-- {-# INLINE unholyPerformIO #-}
-- unholyPerformIO :: IO a -> a
-- unholyPerformIO (IO m) = case m realWorld# of (# _, r #) -> r

-- | FIXME: replace with unholyPerformIO (?)
performIO :: IO a -> a
performIO = unsafeDupablePerformIO

-- | FIXME: Doc
plusForeignPtr :: ForeignPtr a -> Int -> ForeignPtr b
plusForeignPtr fp fo = castForeignPtr fp1 where
    vs :: VS.Vector CChar
    vs = VS.unsafeFromForeignPtr (castForeignPtr fp) fo 0
    (fp1, _) = VS.unsafeToForeignPtr0 vs

foreign import ccall "eigen-proxy.h free" c_freeString :: CString -> IO ()

call :: IO CString -> IO ()
call func = func >>= \c_str -> when (c_str /= nullPtr) $
    peekCString c_str >>= \str -> c_freeString c_str >> fail str

foreign import ccall "eigen-proxy.h free" free :: Ptr a -> IO ()

foreign import ccall "eigen-proxy.h eigen_setNbThreads" c_setNbThreads :: CInt -> IO ()
foreign import ccall "eigen-proxy.h eigen_getNbThreads" c_getNbThreads :: IO CInt

--------------------------------------------------------------------------------

#let api1 name, args = "foreign import ccall \"eigen_%s\" c_%s :: CInt -> %s\n%s :: forall a . Code (C a) => %s\n%s = c_%s (code (undefined :: (C a)))", #name, #name, args, #name, args, #name, #name

#api1 random,        "Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 identity,      "Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 add,           "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 sub,           "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 mul,           "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 diagonal,      "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 transpose,     "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 inverse,       "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 adjoint,       "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 conjugate,     "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 normalize,     "Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 sum,           "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 prod,          "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 mean,          "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 norm,          "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 trace,         "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 squaredNorm,   "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 blueNorm,      "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 hypotNorm,     "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 determinant,   "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 rank,          "CInt -> Ptr CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 image,         "CInt -> Ptr (Ptr (C a)) -> Ptr CInt -> Ptr CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 kernel,        "CInt -> Ptr (Ptr (C a)) -> Ptr CInt -> Ptr CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 solve,         "CInt -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 relativeError, "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"

--------------------------------------------------------------------------------

#let api2 name, args = "foreign import ccall \"eigen_%s\" c_%s :: CInt -> %s\n%s :: forall a . Code (C a) => %s\n%s = c_%s (code (undefined :: (C a)))", #name, #name, args, #name, args, #name, #name

#api2 sparse_new,           "CInt -> CInt -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_clone,         "CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_fromList,      "CInt -> CInt -> Ptr (CTriplet a) -> CInt -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_toList,        "CSparseMatrixPtr a -> Ptr (CTriplet a) -> CInt -> IO CString"
#api2 sparse_free,          "CSparseMatrixPtr a -> IO CString"
#api2 sparse_makeCompressed,"CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_uncompress,    "CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_isCompressed,  "CSparseMatrixPtr a -> Ptr CInt -> IO CString"
#api2 sparse_transpose,     "CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_adjoint,       "CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_pruned,        "CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_prunedRef,     "CSparseMatrixPtr a -> Ptr (C a) -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_scale,         "CSparseMatrixPtr a -> Ptr (C a) -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_nonZeros,      "CSparseMatrixPtr a -> Ptr CInt -> IO CString"
#api2 sparse_innerSize,     "CSparseMatrixPtr a -> Ptr CInt -> IO CString"
#api2 sparse_outerSize,     "CSparseMatrixPtr a -> Ptr CInt -> IO CString"
#api2 sparse_coeff,         "CSparseMatrixPtr a -> CInt -> CInt -> Ptr (C a) -> IO CString"
#api2 sparse_coeffRef,      "CSparseMatrixPtr a -> CInt -> CInt -> Ptr (Ptr (C a)) -> IO CString"
#api2 sparse_cols,          "CSparseMatrixPtr a -> Ptr CInt -> IO CString"
#api2 sparse_rows,          "CSparseMatrixPtr a -> Ptr CInt -> IO CString"
#api2 sparse_norm,          "CSparseMatrixPtr a -> Ptr (C a) -> IO CString"
#api2 sparse_squaredNorm,   "CSparseMatrixPtr a -> Ptr (C a) -> IO CString"
#api2 sparse_blueNorm,      "CSparseMatrixPtr a -> Ptr (C a) -> IO CString"
#api2 sparse_add,           "CSparseMatrixPtr a -> CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_sub,           "CSparseMatrixPtr a -> CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_mul,           "CSparseMatrixPtr a -> CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_block,         "CSparseMatrixPtr a -> CInt -> CInt -> CInt -> CInt -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_fromMatrix,    "Ptr (C a) -> CInt -> CInt -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_toMatrix,      "CSparseMatrixPtr a -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api2 sparse_values,        "CSparseMatrixPtr a -> Ptr CInt -> Ptr (Ptr (C a)) -> IO CString"
#api2 sparse_outerStarts,   "CSparseMatrixPtr a -> Ptr CInt -> Ptr (Ptr CInt) -> IO CString"
#api2 sparse_innerIndices,  "CSparseMatrixPtr a -> Ptr CInt -> Ptr (Ptr CInt) -> IO CString"
#api2 sparse_innerNNZs,     "CSparseMatrixPtr a -> Ptr CInt -> Ptr (Ptr CInt) -> IO CString"
#api2 sparse_setZero,       "CSparseMatrixPtr a -> IO CString"
#api2 sparse_setIdentity,   "CSparseMatrixPtr a -> IO CString"
#api2 sparse_reserve,       "CSparseMatrixPtr a -> CInt -> IO CString"
#api2 sparse_resize,        "CSparseMatrixPtr a -> CInt -> CInt -> IO CString"

#api2 sparse_conservativeResize,    "CSparseMatrixPtr a -> CInt -> CInt -> IO CString"
#api2 sparse_compressInplace,       "CSparseMatrixPtr a -> IO CString"
#api2 sparse_uncompressInplace,     "CSparseMatrixPtr a -> IO CString"

--------------------------------------------------------------------------------

#let api3 name, args = "foreign import ccall \"eigen_%s\" c_%s :: CInt -> CInt -> %s\n%s :: forall s a . (Code s, Code (C a)) => s -> %s\n%s s = c_%s (code (undefined :: (C a))) (code s)", #name, #name, args, #name, args, #name, #name

#api3 sparse_la_newSolver,          "Ptr (CSolverPtr a) -> IO CString"
#api3 sparse_la_freeSolver,         "CSolverPtr a -> IO CString"
#api3 sparse_la_factorize,          "CSolverPtr a -> CSparseMatrixPtr a -> IO CString"
#api3 sparse_la_analyzePattern,     "CSolverPtr a -> CSparseMatrixPtr a -> IO CString"
#api3 sparse_la_compute,            "CSolverPtr a -> CSparseMatrixPtr a -> IO CString"
#api3 sparse_la_tolerance,          "CSolverPtr a -> Ptr CDouble -> IO CString"
#api3 sparse_la_setTolerance,       "CSolverPtr a -> CDouble -> IO CString"
#api3 sparse_la_maxIterations,      "CSolverPtr a -> Ptr CInt -> IO CString"
#api3 sparse_la_setMaxIterations,   "CSolverPtr a -> CInt -> IO CString"
#api3 sparse_la_info,               "CSolverPtr a -> Ptr CInt -> IO CString"
#api3 sparse_la_error,              "CSolverPtr a -> Ptr CDouble -> IO CString"
#api3 sparse_la_iterations,         "CSolverPtr a -> Ptr CInt -> IO CString"
#api3 sparse_la_solve,              "CSolverPtr a -> CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
-- #api3 sparse_la_solveWithGuess,     "CSolverPtr a -> CSparseMatrixPtr a -> CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api3 sparse_la_matrixQ,            "CSolverPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api3 sparse_la_matrixR,            "CSolverPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api3 sparse_la_setPivotThreshold,  "CSolverPtr a -> CDouble -> IO CString"
#api3 sparse_la_rank,               "CSolverPtr a -> Ptr CInt -> IO CString"
#api3 sparse_la_matrixL,            "CSolverPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api3 sparse_la_matrixU,            "CSolverPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api3 sparse_la_setSymmetric,       "CSolverPtr a -> CInt -> IO CString"
#api3 sparse_la_determinant,        "CSolverPtr a -> Ptr (C a) -> IO CString"
#api3 sparse_la_logAbsDeterminant,  "CSolverPtr a -> Ptr (C a) -> IO CString"
#api3 sparse_la_absDeterminant,     "CSolverPtr a -> Ptr (C a) -> IO CString"
#api3 sparse_la_signDeterminant,    "CSolverPtr a -> Ptr (C a) -> IO CString"

--------------------------------------------------------------------------------


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}

-- | Marshaling between Feldspar and C99 types
--
module Feldspar.Compiler.Marshal
  ( SA(..)
  , allocSA
  , Marshal(..)
  )
  where

import System.Plugins.MultiStage
import Feldspar.Core.Types (IntN(..), WordN(..), Tuple(..), (:*), TNil)
import Feldspar.Core.NestedTuples (npair)

import Data.Int (Int32)
import Data.Word (Word32)
import Data.Complex (Complex(..))
import Data.Default
import Control.Applicative

import Foreign.Ptr
import Foreign.Marshal (free, new, newArray, peekArray,mallocBytes)
import Foreign.Storable (Storable(..))
import Foreign.Storable.Tuple ()
import qualified Foreign.Storable.Record as Store

instance Reference IntN        where type Ref IntN        = IntN
instance Reference WordN       where type Ref WordN       = WordN
instance Reference (Complex a) where type Ref (Complex a) = Complex a

instance Marshal IntN        where type Rep IntN        = IntN
instance Marshal WordN       where type Rep WordN       = WordN
instance Marshal (Complex a) where type Rep (Complex a) = Complex a

instance Default (Ptr a) where def = nullPtr

instance (Storable (Rep a), Marshal a) => Marshal [a]
  where
    type Rep [a] = SA (Rep a)
    to xs = do
        let len  = fromIntegral $ length xs
        let size = fromIntegral $ sizeOf (undefined :: Rep a)
        ys <- mapM to xs
        buffer <- newArray ys
        return $ SA buffer len
    from p | elems p == 0 = return []
    from p = go p
      where
        go SA{..} = do
          res <- mapM from =<< peekArray (fromIntegral elems) buf
          free buf
          return res

-- | Buffer descriptor for Feldspar arrays
data SA a = SA { buf   :: Ptr a
               , elems :: Int32
               }
  deriving (Eq, Show)

instance Default (SA a) where
    def = SA nullPtr def

allocSA :: forall a. Storable a => Int -> IO (Ptr (SA a))
allocSA len = do
    let size  = fromIntegral $ sizeOf (undefined :: a)
    let bytes = len * size
    buffer <- mallocBytes bytes
    new $ SA buffer (fromIntegral len)

storeSA :: Storable a => Store.Dictionary (SA a)
storeSA = Store.run $ SA
    <$> Store.element buf
    <*> Store.element elems

instance Storable a => Storable (SA a)
  where
    sizeOf    = Store.sizeOf    storeSA
    alignment = Store.alignment storeSA
    peek      = Store.peek      storeSA
    poke      = Store.poke      storeSA

instance Reference (Ptr a)
  where
    type Ref (Ptr a) = Ptr a
    ref   = return
    deref = return

instance (Storable a) => Reference (SA a)
  where
    type Ref (SA a) = Ptr (SA a)
    ref   = new
    deref = peek

instance (Storable (a,b)) => Reference (a,b)
  where
    type Ref (a,b) = Ptr (a,b)
    ref   = new
    deref = peek

instance (Storable (a,b,c)) => Reference (a,b,c)
  where
    type Ref (a,b,c) = Ptr (a,b,c)
    ref   = new
    deref = peek

instance (Storable (a,b,c,d)) => Reference (a,b,c,d)
  where
    type Ref (a,b,c,d) = Ptr (a,b,c,d)
    ref   = new
    deref = peek

instance (Storable (a,b,c,d,e)) => Reference (a,b,c,d,e)
  where
    type Ref (a,b,c,d,e) = Ptr (a,b,c,d,e)
    ref   = new
    deref = peek

instance (Storable (a,b,c,d,e,f)) => Reference (a,b,c,d,e,f)
  where
    type Ref (a,b,c,d,e,f) = Ptr (a,b,c,d,e,f)
    ref   = new
    deref = peek

instance (Storable (a,b,c,d,e,f,g)) => Reference (a,b,c,d,e,f,g)
  where
    type Ref (a,b,c,d,e,f,g) = Ptr (a,b,c,d,e,f,g)
    ref   = new
    deref = peek

instance ( Marshal a
         , Marshal b
         ) => Marshal (a,b)
  where
    type Rep (a,b) = (Rep a,Rep b)
    to (a,b)   = (,) <$> to a <*> to b
    from (a,b) = (,) <$> from a <*> from b

instance ( Marshal a
         , Marshal b
         , Marshal c
         ) => Marshal (a,b,c)
  where
    type Rep (a,b,c) = (Rep a,Rep b,Rep c)
    to (a,b,c)   = (,,) <$> to a <*> to b <*> to c
    from (a,b,c) = (,,) <$> from a <*> from b <*> from c

instance ( Marshal a
         , Marshal b
         , Marshal c
         , Marshal d
         ) => Marshal (a,b,c,d)
  where
    type Rep (a,b,c,d) = (Rep a,Rep b,Rep c,Rep d)
    to (a,b,c,d) =
      (,,,) <$> to a <*> to b <*> to c <*> to d
    from (a,b,c,d) =
      (,,,) <$> from a <*> from b <*> from c <*> from d

instance ( Marshal a
         , Marshal b
         , Marshal c
         , Marshal d
         , Marshal e
         ) => Marshal (a,b,c,d,e)
  where
    type Rep (a,b,c,d,e) = (Rep a,Rep b,Rep c,Rep d,Rep e)
    to (a,b,c,d,e) =
      (,,,,) <$> to a <*> to b <*> to c <*> to d <*> to e
    from (a,b,c,d,e) =
      (,,,,) <$> from a <*> from b <*> from c <*> from d <*> from e

instance ( Marshal a
         , Marshal b
         , Marshal c
         , Marshal d
         , Marshal e
         , Marshal f
         ) => Marshal (a,b,c,d,e,f)
  where
    type Rep (a,b,c,d,e,f) = (Rep a,Rep b,Rep c,Rep d,Rep e,Rep f)
    to (a,b,c,d,e,f) =
      (,,,,,) <$> to a <*> to b <*> to c <*> to d <*> to e <*> to f
    from (a,b,c,d,e,f) =
      (,,,,,) <$> from a <*> from b <*> from c <*> from d <*> from e <*> from f

instance ( Marshal a
         , Marshal b
         , Marshal c
         , Marshal d
         , Marshal e
         , Marshal f
         , Marshal g
         ) => Marshal (a,b,c,d,e,f,g)
  where
    type Rep (a,b,c,d,e,f,g) = (Rep a,Rep b,Rep c,Rep d,Rep e,Rep f,Rep g)
    to (a,b,c,d,e,f,g) =
      (,,,,,,) <$> to a <*> to b <*> to c <*> to d <*> to e <*> to f <*> to g
    from (a,b,c,d,e,f,g) =
      (,,,,,,) <$> from a <*> from b <*> from c <*> from d <*> from e <*> from f <*> from g

instance (Marshal a, Marshal b) => Marshal (Tuple (a :* b :* TNil)) where
  type Rep (Tuple (a :* b :* TNil)) = (Rep a, Rep b)
  to (a :* b :* TNil) = (,) <$> to a <*> to b
  from (a, b) = npair <$> from a <*> from b


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes #-}

module Hw7 where

-- import Hw1
-- import Hw2
-- import Hw3
-- import Hw4
-- import Hw5
-- import Hw6
-- import Hw7
import Data.Time.Clock
import Data.Complex

import Control.Monad.State
import Control.Monad (forM_)
import System.Directory (doesDirectoryExist, listDirectory, Permissions(..),
                         getModificationTime, getPermissions, getFileSize, getDirectoryContents)
import System.FilePath ((</>), takeExtension)
import Control.Exception (handle, SomeException(..))

import Data.Time (UTCTime(..))
import Data.Time.Calendar(toGregorian)
import Data.Time.Clock(utctDay,UTCTime)
import Data.Time.Clock.POSIX(posixSecondsToUTCTime)

epochToUTC :: Integral a => a -> UTCTime
epochToUTC = posixSecondsToUTCTime . fromIntegral

epoch = epochToUTC 0

match :: String -> String -> Bool
match s p
  | length s < l = False
  | otherwise = (take l s) == p || match (drop 1 s) p
  where l = length p

class Applicative m => Predicate m where
   (>?) :: (Ord a) => m a -> a -> m Bool
   p >? a = (> a) <$> p
   infixl 4 >?

   (<?) :: (Ord a) => m a -> a -> m Bool
   p <? a = (< a) <$> p
   infixl 4 <?

   (==?) :: (Ord a) => m a -> a -> m Bool
   p ==? a = (== a) <$> p
   infixl 4 ==?

   (!=?) :: (Ord a) => m a -> a -> m Bool
   p !=? a = (/= a) <$> p
   infixl 4 !=?

   (~?) :: m String -> String -> m Bool
   p ~? s = (`match` s) <$> p
   infixl 4 ~?

   (&&?) :: m Bool -> m Bool -> m Bool
   p1 &&? p2 = (&&) <$> p1 <*> p2
   infixl 3 &&?

   (||?) :: m Bool -> m Bool -> m Bool
   p1 ||? p2 = (||) <$> p1 <*> p2
   infixl 3 ||?

   notP :: m Bool -> m Bool
   notP = fmap (not)

data FileInfo = FileInfo {
                    filePath :: FilePath,
                    filePermissions :: Maybe Permissions,
                    fileSize :: Maybe Integer,
                    fileLastModified :: Maybe UTCTime
                }

newtype FilePredicate a = FileP { runFileP :: FileInfo -> a } deriving (Functor,Applicative)

instance Predicate FilePredicate

pathP :: FilePredicate String
pathP = FileP filePath

extP :: FilePredicate String
extP = takeExtension <$> pathP

sizeP :: FilePredicate Integer
sizeP = maybe (-1) id <$> FileP fileSize

getPermissionP :: (Permissions -> Bool) -> FilePredicate Bool
getPermissionP f = maybe False f <$> FileP filePermissions

searchP:: FilePredicate Bool
searchP = getPermissionP searchable

readP :: FilePredicate Bool
readP = getPermissionP readable

writeP :: FilePredicate Bool
writeP = getPermissionP writable

execP :: FilePredicate Bool
execP = getPermissionP executable

timeP :: FilePredicate UTCTime
timeP = maybe epoch id <$> FileP fileLastModified

maybeIO :: IO a -> IO (Maybe a)
maybeIO act = handle (\(SomeException _) -> return Nothing) (fmap Just act)

getInfo :: FilePath -> IO FileInfo
getInfo path = do
  perms <- maybeIO (getPermissions path)
  size <- maybeIO (getFileSize path)
  modified <- maybeIO (getModificationTime path)
  return (FileInfo path perms size modified)

data FileAction = Done | Skip | Continue

type FileState s =  StateT s IO FileAction

find :: forall s. (FileInfo -> FileState s) -> FilePath -> FileState s

find getState path = do
  names <- lift $ listDirectory path
  iterate names
  where iterate :: [FilePath] -> FileState s
        iterate (name: names) = do
            let path' = path </> name
            let info = getInfo path'
            state <- getState =<< lift info
            isDirectory <- lift $ doesDirectoryExist path'

            case state of
                Done -> return state
                Continue
                    | isDirectory -> do
                        next <- find getState path'
                        case next of
                            Done -> return next
                            _ -> iterate names
                    | otherwise -> iterate names
                Skip -> iterate names

        iterate _ = return Continue

main :: IO ()

main = do
--    let downloads = "C:\\Users\\tzhao\\Downloads"
   let downloads = "/Users/nikolapeevski/Downloads"

   let yearP = ((\(x,_,_) -> x) . toGregorian . utctDay) <$> timeP
   let recurseP = yearP >? 2018

   let filterP = (extP ==? ".pdf" &&? sizeP >? 2^22) ||? execP

   let iter :: FileInfo -> FileState [(FilePath, Integer, UTCTime)]
       iter = \info -> do
            let eval = \p -> runFileP p info

            s <- get
            let r = (eval pathP, eval sizeP, eval timeP)
            let s' = if eval filterP then r : s else s
            put s'

            let a = if length s' >= 10 then Done
                    else if eval recurseP then Continue else Skip
            return a

   results <- execStateT (find iter downloads) []

   forM_ results $ \(p,n,t) ->
                            handle (\(SomeException _) -> return ())
                          $ putStrLn
                          $ (take 100 p) ++ "\t"
                              ++ show (round $ fromInteger n / 2^10) ++ " KB\t"
                              ++ show (utctDay t)


{-# LANGUAGE OverloadedStrings, ScopedTypeVariables #-}

module Main where

import Control.Monad (forM, forM_)
import Data.List (intercalate)
import Statistics.Sample (mean, stdDev)
import Turtle
import qualified Data.Text as T
import qualified Data.Vector as V

data Platform = PlatformLinux | PlatformMacOS deriving (Eq)

main :: IO ()
main = do
    uname <- strict $ inshell "uname" empty
    platform <-
        if "Linux" `T.isInfixOf` uname then return PlatformLinux
        else if "Darwin" `T.isInfixOf` uname then return PlatformMacOS
        else failWithMsg "Currently only Linux and macOS are supported."

    -- Required for pcregrep on Linux.
    when (platform == PlatformLinux) $ export "LC_ALL" "en_US.UTF-8"

    let logFile :: (IsString a) => a
        logFile = "./tmp/bench.sh.log"

    export "logfile" (fromString logFile)
    testfile logFile >>= flip when (rm logFile)

    echo "Compiling C program... "
    shells "mkdir -p ./tmp" empty
    shells ("gcc -O2 -Wall bench-archive.c -I/usr/local/opt/libarchive/include -L/usr/local/opt/libarchive/lib "
            `T.append` "-larchive -o bench-archive >> $logfile 2>&1") empty

    echo "Compiling Haskell programs... "
    shells "stack build >> $logfile 2>&1" empty

    let c_executable = "./bench-archive"
    hs_plain_executable <- T.strip <$> strict (inshell "stack exec -- which bench-archive-plain 2> /dev/null" empty)
    hs_streamly_executable <- T.strip <$> strict (inshell "stack exec -- which bench-archive-streamly 2> /dev/null" empty)

    let fileCountsAndFileSizes :: [(Int,Int)] =
            [(2500000,1), (5000000,1), (20,250000000), (40,250000000), (80,250000000)]

    echoHrule
    answerIsYes "Create archives?" >>= flip when (do
        echo "Creating archives using C program (unless they already exist)..."
        let tmpDir :: (IsString a) => a
            tmpDir = "./tmp"
        mktree tmpDir

        forM_ fileCountsAndFileSizes $ \(fileCount, fileSize') -> do
            let archivePath = format (""%fp%"/test_"%d%"_"%d%".tar") tmpDir fileCount fileSize'
            (_, existingFileCount) <-
                shellStrict (format ("tar -tf "%s%" 2> /dev/null | wc -l | bc") archivePath) empty
            if T.strip existingFileCount == format (""%d%"") fileCount then
                procs "echo" [format ("    Archive "%s%" already exists.") archivePath] empty
            else do
                procs "echo" ["-n", format ("    Creating archive "%s%"... ") archivePath] empty
                shells (format ("rm -rf "%s%" 2> /dev/null") archivePath) empty
                shells (format (""%s%" write "%s%" "%d%" "%d%"") c_executable archivePath fileCount fileSize') empty
                echo "success.")

    echoHrule
    answerIsYes "Measure read performance?" >>= flip when (do
        let tmpDir :: (IsString a) => a
            tmpDir = "./tmp"
        mktree tmpDir

        let csvFile :: (IsString a) => a
            csvFile = "./tmp/read.csv"
        shells (format ("rm -f "%fp%"") csvFile) empty

        output csvFile . return . unsafeTextToLine . T.pack $
            "fileCount,fileSize,"
            ++ "c_mean,c_std,"
            ++ "hsPlain_mean,hsPlain_std,"
            ++ "hsStreamly_mean,hsStreamly_std,"
            ++ "hsPlain_div_c_mean,hsPlain_div_c_std,"
            ++ "hsStreamly_div_c_mean,hsStreamly_div_c_std,"
            ++ "hsStreamly_div_hsPlain_mean,hsStreamly_div_hsPlain_std"

        forM_ fileCountsAndFileSizes $ \(fileCount, fileSize') -> do
            let archivePath = format (""%fp%"/test_"%d%"_"%d%".tar") tmpDir fileCount fileSize'
            let args  = format ("read "%s%"") archivePath
            let expectedInOutput1 = format ("Total filesize: "%d%"") (fileSize' * fileCount)
            let expectedInOutput2 = format ("File count:     "%d%"") fileCount
            let warmCount = 1
            let timeCount = 10
            procs "echo" [format ("    Reading archive "%s%" with C...") archivePath] empty
            cTimes <- timeCommand Nothing (format (""%s%" "%s%"") c_executable args)
                        [expectedInOutput1, expectedInOutput2] warmCount timeCount
            procs "echo" [format ("    Reading archive "%s%" with Haskell (plain)...") archivePath] empty
            hsPlainTimes <- timeCommand Nothing (format (""%s%" "%s%"") hs_plain_executable args)
                        [expectedInOutput1, expectedInOutput2] warmCount timeCount
            procs "echo" [format ("    Reading archive "%s%" with Haskell (streamly)...") archivePath] empty
            hsStreamlyTimes <- timeCommand Nothing (format (""%s%" "%s%"") hs_streamly_executable args)
                        [expectedInOutput1, expectedInOutput2] warmCount timeCount

            let hsPlainDivCTimes = [a / b | a <- hsPlainTimes, b <- cTimes]
            let hsStreamlyDivCTimes = [a / b | a <- hsStreamlyTimes, b <- cTimes]
            let hsStreamlyDivHsPlainTimes = [a / b | a <- hsStreamlyTimes, b <- hsPlainTimes]

            let mean' :: [NominalDiffTime] -> Double
                mean' xs = mean . V.fromList $ map realToFrac xs

            let std' :: [NominalDiffTime] -> Double
                std' xs = stdDev . V.fromList $ map realToFrac xs

            let line = intercalate "," $ [show fileCount, show fileSize'] ++ map show [
                    mean' cTimes, std' cTimes,
                    mean' hsPlainTimes, std' hsPlainTimes,
                    mean' hsStreamlyTimes, std' hsStreamlyTimes,
                    mean' hsPlainDivCTimes, std' hsPlainDivCTimes,
                    mean' hsStreamlyDivCTimes, std' hsStreamlyDivCTimes,
                    mean' hsStreamlyDivHsPlainTimes, std' hsStreamlyDivHsPlainTimes ]
            append csvFile . return . unsafeTextToLine . T.pack $ line
            procs "echo" [format ("    Appended results to "%s%"") csvFile] empty
            return ())

failWithMsg :: Line -> IO a
failWithMsg msg = echo msg >> exit (ExitFailure 1)

echoHrule :: IO ()
echoHrule = echo "----------------------------------------"

answerIsYes :: String -> IO Bool
answerIsYes question = do
    procs "echo" ["-n", T.pack $ question ++ " (y/N) "] empty
    line <- readline
    case line of
        Just answer | answer == "y" -> return True
        _ -> return False

timeCommand' :: Text -> [Text] -> IO NominalDiffTime
timeCommand' cmdAndArgs expectedInOut = do
    (out, t) <- time . strict $ inshell cmdAndArgs empty
    if all (`T.isInfixOf` out) expectedInOut then
        return t
    else
        error $ "all of " ++ show expectedInOut ++ " were not found in output "
            ++ show out ++ " of command " ++ show cmdAndArgs

timeCommand :: Maybe T.Text -> Text -> [Text] -> Int -> Int -> IO [NominalDiffTime]
timeCommand forceRemovalPath cmdAndArgs expectedInOut warmCount timeCount = do
    let removePath = maybe (return ()) (\fp' -> shells (format ("rm -rf "%s%"") fp') empty) forceRemovalPath
    let echoCount c = procs "echo" ["-n", format (""%d%" ") c] empty
    when (warmCount > 0) $ do
        procs "echo" ["-n", "        Warming up... "] empty
        forM_ [1..warmCount] $ \c ->
            echoCount c >> removePath >> timeCommand' cmdAndArgs expectedInOut
        echo "done."
    procs "echo" ["-n", "        Timing... "] empty
    times <- forM [1..timeCount] $ \c ->
        echoCount c >> removePath >> timeCommand' cmdAndArgs expectedInOut
    removePath
    echo "done."
    return times