{-# LANGUAGE GADTs #-}

module Benchmark where

import System.Environment

import Ypnos.Examples.Stencils
  (runAvgY, runAvg, runAvgY', raiseToList, runLife, runLife', runId)

import Criterion
import Criterion.Monad
import Criterion.Analysis
import Criterion.Environment
import Criterion.Config
import Statistics.Resampling.Bootstrap
import Data.Vector.Unboxed hiding (map, mapM, foldr, foldr1, (++))
import Prelude hiding (sum, length)
import Control.Monad.Trans

l  = [1,2,3,4,5,6,7,8,9,10]
l2 = [True, False, False, True, False, True]

avgGPU = runAvgY l
avgCPU = runAvgY' l
avgA = raiseToList runAvg l

lifeCPU = runLife' l2
lifeGPU = runLife  l2

idGPU = runId l

stenBench :: ((Int,Int) -> b) -> [Benchmark]
stenBench f = [ bench "10x10" $ whnf f (10,10)
              , bench "100x10" $ whnf f (100,10)
              , bench "100x100" $ whnf f (100,100)
              --, bench "1000x1000" $ whnf f (1000,1000)
              ]

--main = defaultMain [ bgroup "Ypnos" (stenBench runAvgY')
--                   , bgroup "Accelerate" (stenBench runAvgA')
--                   ]
type Result = (Double, Double, Double)

runB :: Fun -> Int -> IO Result
runB (Fun f) x = let v = do env <- measureEnvironment
                            l <- runBenchmark env (whnf f (x,x))
                            cfg <- getConfig
                            let conf = fromLJ cfgConfInterval cfg
                            let res  = fromLJ cfgResamples cfg
                            s <- liftIO $ analyseSample conf l res
                            let m = anMean s
                            return (estPoint m,
                                    estLowerBound m,
                                    estUpperBound m)

         in  withConfig defaultConfig v

makeSet :: (Int -> IO Result) -> [Int] -> IO [[String]]
makeSet f range = let tup x = do (y, l, u)  <- f x
                                 return [show x, show y, show l, show u]
                  in  mapM tup range

insert i x y = x ++ i ++ y

makeLine [] = ""
makeLine xs = foldr1 (insert ", ") xs

printCSV :: [[String]] -> String
printCSV lls = foldr (insert "\n") "" (map makeLine lls)

type FunS = String
type ImplS = String

data Fun where
  Fun :: ((Int,Int) -> b) -> Fun

getFun :: FunS -> ImplS -> Fun
getFun "avg" "gpu" = Fun avgGPU
getFun "avg" "cpu" = Fun avgCPU
getFun "life" "gpu" = Fun lifeGPU
getFun "life" "cpu" = Fun lifeCPU
getFun "id" "gpu" = Fun idGPU

main = do [function, impl, begin, step, end, filename] <- getArgs
          let fun = getFun function impl
          let b = read begin :: Int
          let s = read step :: Int
          let e = read end :: Int

          a <- makeSet (runB fun) ([b, s.. e])
          let ca = printCSV a
          writeFile filename ca


module Physics.Involutive
  ( Involutive(..)
  ) where

import Data.Complex hiding (conjugate)

class Involutive a where
  conjugate :: a -> a
  conjugate = id

instance Involutive Double
instance Involutive Float
instance (Involutive a, RealFloat a) => Involutive (Complex a) where
  {-# SPECIALIZE instance Involutive (Complex Float) #-}
  {-# SPECIALIZE instance Involutive (Complex Double) #-}
  conjugate (a :+ b) = conjugate a :+ negate b


module Types (
  LispVal(..),
  LispError(..),
  ThrowsError,
  trapError,
  extractValue,
  Env,
  nullEnv,
  IOThrowsError,
  liftThrows,
  runIOThrows,
  showVal
) where

import           Control.Monad.Error
import           Data.Complex
import           Data.IORef
import           System.IO
import           Text.ParserCombinators.Parsec

data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal
             | Number Integer
             | Float Double
             | Ratio Rational
             | Complex (Complex Double)
             | String String
             | Bool Bool
             | Character Char
             | PrimitiveFunc ([LispVal] -> ThrowsError LispVal)
             | Func {
               params  :: [String],
               vararg  :: Maybe String,
               body    :: [LispVal],
               closure :: Env }
             | IOFunc ([LispVal] -> IOThrowsError LispVal)
             | Port Handle

showVal :: LispVal -> String
showVal (String contents) = "\"" ++ contents ++ "\""
showVal (Atom name)       = name
showVal (Number contents) = show contents
showVal (Bool True)       = "#t"
showVal (Bool False)      = "#f"
showVal (List contents)   = "(" ++ unwordsList contents ++ ")"
showVal (DottedList head tail) = "(" ++ unwordsList head ++ " . " ++ showVal tail ++ ")"
showVal (Character content) = case content of
  ' '  -> "#\\space"
  '\n' -> "#\\newline"
  _    -> "#\\" ++ [content]
showVal (Float x) = show x
showVal (Ratio x) = show x
showVal (Complex x) = show x
showVal (PrimitiveFunc _) = "<primitive>"
showVal Func { params = args, vararg = varargs, body = body, closure = env } =
  "(lambda (" ++ unwords (map show args) ++ (
    case varargs of
      Nothing  -> ""
      Just arg -> " . " ++arg
  ) ++ ") ...)"
showVal (Port _) = "<IO port>"
showVal (IOFunc _) = "<IO primitive>"

unwordsList :: [LispVal] -> String
unwordsList = unwords . map showVal

instance Show LispVal where show = showVal
instance Eq LispVal where
  (==) a b = show a == show b

data LispError = NumArgs Integer [LispVal]
               | TypeMismatch String LispVal
               | Parser ParseError
               | BadSpecialForm String LispVal
               | NotFunction String String
               | UnboundVar String String
               | Default String

showError :: LispError -> String
showError (UnboundVar message varname) = message ++ ": " ++ varname
showError (BadSpecialForm message form) = message ++ ": " ++ show form
showError (NotFunction message func) = message ++ ": " ++ show func
showError (NumArgs expected found) =
  "Expected " ++ show expected
  ++ " args; found values " ++ unwordsList found
showError (TypeMismatch expected found) =
  "Invalid type: expected " ++ expected
  ++ ", found " ++ show found
showError (Parser parseErr) = "Parse error at " ++ show parseErr
showError (Default err) = "LispError: " ++ show err

instance Show LispError where show = showError
instance Error LispError where
  noMsg = Default "An error has occurred"
  strMsg = Default

type ThrowsError = Either LispError

trapError action = catchError action (return . show)

extractValue :: ThrowsError a -> a
extractValue (Right val) = val

type Env = IORef [(String, IORef LispVal)]

nullEnv :: IO Env
nullEnv = newIORef []

type IOThrowsError = ErrorT LispError IO

liftThrows :: ThrowsError a -> IOThrowsError a
liftThrows (Left err)  = throwError err
liftThrows (Right val) = return val

runIOThrows :: IOThrowsError String -> IO String
runIOThrows action = fmap extractValue (runErrorT (trapError action))


{-# LANGUAGE QuasiQuotes #-}

{-# OPTIONS_GHC -fwarn-incomplete-patterns #-}

-- | C code generation of primitive Feldspar expressions

module Feldspar.Primitive.Backend.C where

import Data.Complex

import Data.Constraint (Dict (..))
import Data.Proxy

import Language.C.Quote.C
import qualified Language.C.Syntax as C

import Language.C.Monad
import Language.Embedded.Backend.C

import Language.Syntactic

import Feldspar.Primitive.Representation

-- Note: This module assumes a 32-bit target. For example the C function `abs`
-- is used up to 32-bits, and `labs` is used above that.

viewLitPrim :: ASTF (Primitive :&: PrimTypeRep) a -> Maybe a
viewLitPrim (Sym (Lit a :&: _)) = Just a
viewLitPrim (Sym (Pi :&: _))    = Just pi
viewLitPrim _ = Nothing

instance CompTypeClass PrimType'
  where
    compType _ (_ :: proxy a) = case primTypeRep :: PrimTypeRep a of
      BoolT   -> addInclude "<stdbool.h>" >> return [cty| typename bool     |]
      Int8T   -> addInclude "<stdint.h>"  >> return [cty| typename int8_t   |]
      Int16T  -> addInclude "<stdint.h>"  >> return [cty| typename int16_t  |]
      Int32T  -> addInclude "<stdint.h>"  >> return [cty| typename int32_t  |]
      Int64T  -> addInclude "<stdint.h>"  >> return [cty| typename int64_t  |]
      Word8T  -> addInclude "<stdint.h>"  >> return [cty| typename uint8_t  |]
      Word16T -> addInclude "<stdint.h>"  >> return [cty| typename uint16_t |]
      Word32T -> addInclude "<stdint.h>"  >> return [cty| typename uint32_t |]
      Word64T -> addInclude "<stdint.h>"  >> return [cty| typename uint64_t |]
      FloatT  -> return [cty| float |]
      DoubleT -> return [cty| double |]
      ComplexFloatT  -> addInclude "<tgmath.h>" >> return [cty| float  _Complex |]
      ComplexDoubleT -> addInclude "<tgmath.h>" >> return [cty| double _Complex |]

    compLit _ a = case primTypeOf a of
      BoolT   -> do addInclude "<stdbool.h>"
                    return $ if a then [cexp| true |] else [cexp| false |]
      Int8T   -> return [cexp| $a |]
      Int16T  -> return [cexp| $a |]
      Int32T  -> return [cexp| $a |]
      Int64T  -> return [cexp| $a |]
      Word8T  -> return [cexp| $a |]
      Word16T -> return [cexp| $a |]
      Word32T -> return [cexp| $a |]
      Word64T -> return [cexp| $a |]
      FloatT  -> return [cexp| $a |]
      DoubleT -> return [cexp| $a |]
      ComplexFloatT  -> return $ compComplexLit a
      ComplexDoubleT -> return $ compComplexLit a

-- | Compile a complex literal
compComplexLit :: (Eq a, Num a, ToExp a) => Complex a -> C.Exp
compComplexLit (r :+ 0) = [cexp| $r |]
compComplexLit (0 :+ i) = [cexp| $i * I |]
compComplexLit (r :+ i) = [cexp| $r + $i * I |]

addTagMacro :: MonadC m => m ()
addTagMacro = addGlobal [cedecl|$esc:("#define TAG(tag,exp) (exp)")|]

-- | Compile a unary operator
compUnOp :: MonadC m => C.UnOp -> ASTF PrimDomain a -> m C.Exp
compUnOp op a = do
    a' <- compPrim $ Prim a
    return $ C.UnOp op a' mempty

-- | Compile a binary operator
compBinOp :: MonadC m =>
    C.BinOp -> ASTF PrimDomain a -> ASTF PrimDomain b -> m C.Exp
compBinOp op a b = do
    a' <- compPrim $ Prim a
    b' <- compPrim $ Prim b
    return $ C.BinOp op a' b' mempty

-- | Compile a function call
compFun :: MonadC m => String -> Args (AST PrimDomain) sig -> m C.Exp
compFun fun args = do
    as <- sequence $ listArgs (compPrim . Prim) args
    return [cexp| $id:fun($args:as) |]

-- | Compile a call to 'abs'
compAbs :: MonadC m => PrimTypeRep a -> ASTF PrimDomain a -> m C.Exp
compAbs t a = case viewPrimTypeRep t of
    PrimTypeBool     -> error "compAbs: type BoolT not supported"
    PrimTypeIntWord (IntType _)  -> addInclude "<stdlib.h>" >> compFun "abs" (a :* Nil)
    PrimTypeIntWord (WordType _) -> compPrim $ Prim a
    _ -> addInclude "<tgmath.h>" >> compFun "fabs" (a :* Nil)
      -- Floating and complex types

complexSign_def = [cedecl|
double _Complex feld_complexSign(double _Complex c) {
    double z = cabs(c);
    if (z == 0) {
        return 0;
    } else {
        return (creal(c)/z + I*(cimag(c)/z));
    }
}
|]

complexSignf_def = [cedecl|
float _Complex feld_complexSignf(float _Complex c) {
    float z = cabsf(c);
    if (z == 0) {
        return 0;
    } else {
        return (crealf(c)/z + I*(cimagf(c)/z));
    }
}
|]

-- | Compile a call to 'signum'
compSign :: MonadC m => PrimTypeRep a -> ASTF PrimDomain a -> m C.Exp
compSign t a = case viewPrimTypeRep t of
    PrimTypeBool -> error "compSign: type BoolT not supported"
    PrimTypeIntWord (WordType _) -> do
        addTagMacro
        a' <- compPrim $ Prim a
        return [cexp| TAG("signum", $a' > 0) |]
    PrimTypeIntWord (IntType _) -> do
        addTagMacro
        a' <- compPrim $ Prim a
        return [cexp| TAG("signum", ($a' > 0) - ($a' < 0)) |]
    PrimTypeFloating FloatType -> do
        addTagMacro
        a' <- compPrim $ Prim a
        return [cexp| TAG("signum", (float) (($a' > 0) - ($a' < 0))) |]
    PrimTypeFloating DoubleType -> do
        addTagMacro
        a' <- compPrim $ Prim a
        return [cexp| TAG("signum", (double) (($a' > 0) - ($a' < 0))) |]
    PrimTypeComplex ComplexDoubleType -> do
        addInclude "<tgmath.h>"
        addGlobal complexSign_def
        a' <- compPrim $ Prim a
        return [cexp| feld_complexSign($a') |]
    PrimTypeComplex ComplexFloatType -> do
        addInclude "<complex.h>"
        addGlobal complexSignf_def
        a' <- compPrim $ Prim a
        return [cexp| feld_complexSignf($a') |]
  -- TODO The floating point cases give `sign (-0.0) = 0.0`, which is (slightly)
  -- wrong. They should return -0.0. I don't know whether it's correct for other
  -- strange values.

-- | Compile a type casted primitive expression
compCast :: MonadC m => PrimTypeRep a -> ASTF PrimDomain b -> m C.Exp
compCast t a = compPrim (Prim a) >>= compCastExp t

-- | Applies type cast on a compiled expression
compCastExp :: MonadC m => PrimTypeRep a -> C.Exp -> m C.Exp
compCastExp t a
    | Dict <- witPrimType t = do
        t' <- compType (Proxy :: Proxy PrimType') t
        return [cexp|($ty:t') $a|]

-- | Compile a call to 'round'
--   A cast is added to the resulting expression to allow compatibility
--   between different integral and floating return types.
compRound :: (PrimType' a, Num a, RealFrac b, MonadC m) =>
    PrimTypeRep a -> ASTF PrimDomain b -> m C.Exp
compRound t a = do
    addInclude "<tgmath.h>"
    rounded <- case viewPrimTypeRep t of
        PrimTypeIntWord _  -> compFun "lround" (a :* Nil)
        PrimTypeFloating _ -> compFun "round"  (a :* Nil)
        PrimTypeComplex _  -> compFun "round"  (a :* Nil)
        _ -> error $ "compRound: type " ++ show t ++ " not supported"
    compCastExp t rounded

-- Note: There's no problem with including both `tgmath.h` and `math.h`. As long
-- as the former is included, including the latter (before or after) doesn't
-- make a difference.
--
-- See: <https://gist.github.com/emilaxelsson/51310b3353f96914cd9bdb18b10b3103>

div_def = [cedecl|
int feld_div(int x, int y) {
    int q = x/y;
    int r = x%y;
    if ((r!=0) && ((r<0) != (y<0))) --q;
    return q;
}
|]

ldiv_def = [cedecl|
long int feld_ldiv(long int x, long int y) {
    int q = x/y;
    int r = x%y;
    if ((r!=0) && ((r<0) != (y<0))) --q;
    return q;
}
|]

mod_def = [cedecl|
int feld_mod(int x, int y) {
    int r = x%y;
    if ((r!=0) && ((r<0) != (y<0))) { r += y; }
    return r;
}
|]

lmod_def = [cedecl|
long int feld_lmod(long int x, long int y) {
    int r = x%y;
    if ((r!=0) && ((r<0) != (y<0))) { r += y; }
    return r;
}
|]

-- The above C implementations are taken from
-- <http://www.microhowto.info/howto/round_towards_minus_infinity_when_dividing_integers_in_c_or_c++.html>

compDiv :: MonadC m =>
    PrimTypeRep a -> ASTF PrimDomain a -> ASTF PrimDomain b -> m C.Exp
compDiv t a b = case viewPrimTypeRep t of
    PrimTypeIntWord (WordType _) -> compBinOp C.Div a b
    PrimTypeIntWord (IntType Int64Type) -> do
        addGlobal ldiv_def
        compFun "feld_ldiv" (a :* b :* Nil)
    PrimTypeIntWord _ -> do
        addGlobal div_def
        compFun "feld_div" (a :* b :* Nil)
    _ -> error $ "compDiv: type " ++ show t ++ " not supported"

compMod :: MonadC m =>
    PrimTypeRep a -> ASTF PrimDomain a -> ASTF PrimDomain b -> m C.Exp
compMod t a b = case viewPrimTypeRep t of
    PrimTypeIntWord (WordType _) -> compBinOp C.Mod a b
    PrimTypeIntWord (IntType Int64Type) -> do
        addGlobal lmod_def
        compFun "feld_lmod" (a :* b :* Nil)
    PrimTypeIntWord _ -> do
        addGlobal mod_def
        compFun "feld_mod" (a :* b :* Nil)
    _ -> error $ "compMod: type " ++ show t ++ " not supported"

-- | Compile an expression
compPrim :: MonadC m => Prim a -> m C.Exp
compPrim = simpleMatch (\(s :&: t) -> go t s) . unPrim
  where
    go :: forall m sig . MonadC m
        => PrimTypeRep (DenResult sig)
        -> Primitive sig
        -> Args (AST PrimDomain) sig
        -> m C.Exp
    go _ (FreeVar v) Nil = touchVar v >> return [cexp| $id:v |]
    go t (Lit a) Nil
        | Dict <- witPrimType t
        = compLit (Proxy :: Proxy PrimType') a
    go _ Add  (a :* b :* Nil) = compBinOp C.Add a b
    go _ Sub  (a :* b :* Nil) = compBinOp C.Sub a b
    go _ Mul  (a :* b :* Nil) = compBinOp C.Mul a b
    go _ Neg  (a :* Nil)      = compUnOp C.Negate a
    go t Abs  (a :* Nil)      = compAbs t a
    go t Sign (a :* Nil)      = compSign t a

    go _ Quot (a :* b :* Nil) = compBinOp C.Div a b
    go _ Rem  (a :* b :* Nil) = compBinOp C.Mod a b
    go t Div  (a :* b :* Nil) = compDiv t a b
    go t Mod  (a :* b :* Nil) = compMod t a b
    go _ FDiv (a :* b :* Nil) = compBinOp C.Div a b

    go _ Pi Nil = addGlobal pi_def >> return [cexp| FELD_PI |]
      where pi_def = [cedecl|$esc:("#define FELD_PI 3.141592653589793")|]
              -- This is the value of `pi :: Double`.
              -- Apparently there is no standard C99 definition of pi.
    go _ Exp   args = addInclude "<tgmath.h>" >> compFun "exp" args
    go _ Log   args = addInclude "<tgmath.h>" >> compFun "log" args
    go _ Sqrt  args = addInclude "<tgmath.h>" >> compFun "sqrt" args
    go _ Pow   args = addInclude "<tgmath.h>" >> compFun "pow" args
    go _ Sin   args = addInclude "<tgmath.h>" >> compFun "sin" args
    go _ Cos   args = addInclude "<tgmath.h>" >> compFun "cos" args
    go _ Tan   args = addInclude "<tgmath.h>" >> compFun "tan" args
    go _ Asin  args = addInclude "<tgmath.h>" >> compFun "asin" args
    go _ Acos  args = addInclude "<tgmath.h>" >> compFun "acos" args
    go _ Atan  args = addInclude "<tgmath.h>" >> compFun "atan" args
    go _ Sinh  args = addInclude "<tgmath.h>" >> compFun "sinh" args
    go _ Cosh  args = addInclude "<tgmath.h>" >> compFun "cosh" args
    go _ Tanh  args = addInclude "<tgmath.h>" >> compFun "tanh" args
    go _ Asinh args = addInclude "<tgmath.h>" >> compFun "asinh" args
    go _ Acosh args = addInclude "<tgmath.h>" >> compFun "acosh" args
    go _ Atanh args = addInclude "<tgmath.h>" >> compFun "atanh" args

    go _ Complex (a :* b :* Nil) = do
        addInclude "<tgmath.h>"
        a' <- compPrim $ Prim a
        b' <- compPrim $ Prim b
        return $ case (viewLitPrim a, viewLitPrim b) of
            (Just 0, _) -> [cexp| I*$b'       |]
            (_, Just 0) -> [cexp| $a'         |]
            _           -> [cexp| $a' + I*$b' |]
      -- We assume that constant folding has been performed, so that not both
      -- `a` and `b` are constants
    go _ Polar (m :* p :* Nil)
        | Just 0 <- viewLitPrim m = return [cexp| 0 |]
        | Just 0 <- viewLitPrim p = do
            m' <- compPrim $ Prim m
            return [cexp| $m' |]
        | Just 1 <- viewLitPrim m = do
            p' <- compPrim $ Prim p
            return [cexp| exp(I*$p') |]
        | otherwise = do
            m' <- compPrim $ Prim m
            p' <- compPrim $ Prim p
            return [cexp| $m' * exp(I*$p') |]
      -- We assume that constant folding has been performed, so that not both
      -- `m` and `p` are constants
    go _ Real      args = addInclude "<tgmath.h>" >> compFun "creal" args
    go _ Imag      args = addInclude "<tgmath.h>" >> compFun "cimag" args
    go _ Magnitude args = addInclude "<tgmath.h>" >> compFun "cabs"  args
    go _ Phase     args = addInclude "<tgmath.h>" >> compFun "carg"  args
    go _ Conjugate args = addInclude "<tgmath.h>" >> compFun "conj"  args

    go t I2N   (a :* Nil) = compCast t a
    go t I2B   (a :* Nil) = compCast t a
    go t B2I   (a :* Nil) = compCast t a
    go t Round (a :* Nil) = compRound t a

    go _ Not  (a :* Nil)      = compUnOp C.Lnot a
    go _ And  (a :* b :* Nil) = compBinOp C.Land a b
    go _ Or   (a :* b :* Nil) = compBinOp C.Lor a b
    go _ Eq   (a :* b :* Nil) = compBinOp C.Eq a b
    go _ NEq  (a :* b :* Nil) = compBinOp C.Ne a b
    go _ Lt   (a :* b :* Nil) = compBinOp C.Lt a b
    go _ Gt   (a :* b :* Nil) = compBinOp C.Gt a b
    go _ Le   (a :* b :* Nil) = compBinOp C.Le a b
    go _ Ge   (a :* b :* Nil) = compBinOp C.Ge a b

    go _ BitAnd   (a :* b :* Nil) = compBinOp C.And a b
    go _ BitOr    (a :* b :* Nil) = compBinOp C.Or a b
    go _ BitXor   (a :* b :* Nil) = compBinOp C.Xor a b
    go _ BitCompl (a :* Nil)      = compUnOp C.Not a
    go _ ShiftL   (a :* b :* Nil) = compBinOp C.Lsh a b
    go _ ShiftR   (a :* b :* Nil) = compBinOp C.Rsh a b

    go _ (ArrIx arr) (i :* Nil) = do
        i' <- compPrim $ Prim i
        touchVar arr
        return [cexp| $id:arr[$i'] |]

    go _ Cond (c :* t :* f :* Nil) = do
        c' <- compPrim $ Prim c
        t' <- compPrim $ Prim t
        f' <- compPrim $ Prim f
        return $ C.Cond c' t' f' mempty

    go _ s _ = error $ "compPrim: no handling of symbol " ++ renderSym s
      -- Should not occur, but the completeness checker doesn't know that

instance CompExp Prim where compExp = compPrim


-- | Type classes for random generation of values.
{-# LANGUAGE CPP #-}
#ifndef NO_GENERICS
{-# LANGUAGE DefaultSignatures, FlexibleContexts, TypeOperators #-}
#endif
module Test.QuickCheck.Arbitrary
  (
  -- * Arbitrary and CoArbitrary classes
    Arbitrary(..)
  , CoArbitrary(..)

  -- ** Helper functions for implementing arbitrary
  , arbitrarySizedIntegral        -- :: Integral a => Gen a
  , arbitraryBoundedIntegral      -- :: (Bounded a, Integral a) => Gen a
  , arbitrarySizedBoundedIntegral -- :: (Bounded a, Integral a) => Gen a
  , arbitrarySizedFractional      -- :: Fractional a => Gen a
  , arbitraryBoundedRandom        -- :: (Bounded a, Random a) => Gen a
  , arbitraryBoundedEnum          -- :: (Bounded a, Enum a) => Gen a
  -- ** Helper functions for implementing shrink
#ifndef NO_GENERICS
  , genericShrink      -- :: (Generic a, Typeable a, RecursivelyShrink (Rep a), Subterms (Rep a)) => a -> [a]
  , subterms           -- :: (Generic a, Subterms (Rep a)) => a -> [a]
  , recursivelyShrink  -- :: (Generic a, RecursivelyShrink (Rep a)) => a -> [a]
#endif
  , shrinkNothing            -- :: a -> [a]
  , shrinkList               -- :: (a -> [a]) -> [a] -> [[a]]
  , shrinkIntegral           -- :: Integral a => a -> [a]
  , shrinkRealFrac           -- :: RealFrac a => a -> [a]
  , shrinkRealFracToInteger  -- :: RealFrac a => a -> [a]
  -- ** Helper functions for implementing coarbitrary
  , coarbitraryIntegral      -- :: Integral a => a -> Gen b -> Gen b
  , coarbitraryReal          -- :: Real a => a -> Gen b -> Gen b
  , coarbitraryShow          -- :: Show a => a -> Gen b -> Gen b
  , coarbitraryEnum          -- :: Enum a => a -> Gen b -> Gen b
  , (><)

  -- ** Generators which use arbitrary
  , vector       -- :: Arbitrary a => Int -> Gen [a]
  , orderedList  -- :: (Ord a, Arbitrary a) => Gen [a]
  , infiniteList -- :: Arbitrary a => Gen [a]
  )
 where

--------------------------------------------------------------------------
-- imports

import System.Random(Random)
import Test.QuickCheck.Gen
import Test.QuickCheck.Gen.Unsafe

{-
import Data.Generics
  ( (:*:)(..)
  , (:+:)(..)
  , Unit(..)
  )
-}

import Data.Char
  ( chr
  , ord
  , isLower
  , isUpper
  , toLower
  , isDigit
  , isSpace
  )

#ifndef NO_FIXED
import Data.Fixed
  ( Fixed
  , HasResolution
  )
#endif

import Data.Ratio
  ( Ratio
  , (%)
  , numerator
  , denominator
  )

import Data.Complex
  ( Complex((:+)) )

import Data.List
  ( sort
  , nub
  )

import Control.Monad
  ( liftM
  , liftM2
  , liftM3
  , liftM4
  , liftM5
  )

import Data.Int(Int8, Int16, Int32, Int64)
import Data.Word(Word, Word8, Word16, Word32, Word64)

#ifndef NO_GENERICS
import GHC.Generics
import Data.Typeable
#endif

--------------------------------------------------------------------------
-- ** class Arbitrary

-- | Random generation and shrinking of values.
class Arbitrary a where
  -- | A generator for values of the given type.
  arbitrary :: Gen a
  arbitrary = error "no default generator"

  -- | Produces a (possibly) empty list of all the possible
  -- immediate shrinks of the given value. The default implementation
  -- returns the empty list, so will not try to shrink the value.
  --
  -- Most implementations of 'shrink' should try at least three things:
  --
  -- 1. Shrink a term to any of its immediate subterms.
  --
  -- 2. Recursively apply 'shrink' to all immediate subterms.
  --
  -- 3. Type-specific shrinkings such as replacing a constructor by a
  --    simpler constructor.
  --
  -- For example, suppose we have the following implementation of binary trees:
  --
  -- > data Tree a = Nil | Branch a (Tree a) (Tree a)
  --
  -- We can then define 'shrink' as follows:
  --
  -- > shrink Nil = []
  -- > shrink (Branch x l r) =
  -- >   -- shrink Branch to Nil
  -- >   [Nil] ++
  -- >   -- shrink to subterms
  -- >   [l, r] ++
  -- >   -- recursively shrink subterms
  -- >   [Branch x' l' r' | (x', l', r') <- shrink (x, l, r)]
  --
  -- There are a couple of subtleties here:
  --
  -- * QuickCheck tries the shrinking candidates in the order they
  --   appear in the list, so we put more aggressive shrinking steps
  --   (such as replacing the whole tree by @Nil@) before smaller
  --   ones (such as recursively shrinking the subtrees).
  --
  -- * It is tempting to write the last line as
  --   @[Branch x' l' r' | x' <- shrink x, l' <- shrink l, r' <- shrink r]@
  --   but this is the /wrong thing/! It will force QuickCheck to shrink
  --   @x@, @l@ and @r@ in tandem, and shrinking will stop once /one/ of
  --   the three is fully shrunk.
  --
  -- There is a fair bit of boilerplate in the code above.
  -- We can avoid it with the help of some generic functions;
  -- note that these only work on GHC 7.2 and above.
  -- The function 'genericShrink' tries shrinking a term to all of its
  -- subterms and, failing that, recursively shrinks the subterms.
  -- Using it, we can define 'shrink' as:
  --
  -- > shrink x = shrinkToNil x ++ genericShrink x
  -- >   where
  -- >     shrinkToNil Nil = []
  -- >     shrinkToNil (Branch _ l r) = [Nil]
  --
  -- 'genericShrink' is a combination of 'subterms', which shrinks
  -- a term to any of its subterms, and 'recursivelyShrink', which shrinks
  -- all subterms of a term. These may be useful if you need a bit more
  -- control over shrinking than 'genericShrink' gives you.
  --
  -- A final gotcha: we cannot define 'shrink' as simply @'shrink' x = Nil:'genericShrink' x@
  -- as this shrinks @Nil@ to @Nil@, and shrinking will go into an
  -- infinite loop.
  --
  -- If all this leaves you bewildered, you might try @'shrink' = 'genericShrink'@ to begin with,
  -- after deriving @Generic@ and @Typeable@ for your type. However, if your data type has any
  -- special invariants, you will need to check that 'genericShrink' can't break those invariants.
  shrink :: a -> [a]
  shrink _ = []

#ifndef NO_GENERICS
-- | Shrink a term to any of its immediate subterms,
-- and also recursively shrink all subterms.
genericShrink :: (Generic a, Typeable a, RecursivelyShrink (Rep a), Subterms (Rep a)) => a -> [a]
genericShrink x = subterms x ++ recursivelyShrink x

-- | Recursively shrink all immediate subterms.
recursivelyShrink :: (Generic a, RecursivelyShrink (Rep a)) => a -> [a]
recursivelyShrink = map to . grecursivelyShrink . from

class RecursivelyShrink f where
  grecursivelyShrink :: f a -> [f a]

instance (RecursivelyShrink f, RecursivelyShrink g) => RecursivelyShrink (f :*: g) where
  grecursivelyShrink (x :*: y) =
    [x' :*: y | x' <- grecursivelyShrink x] ++
    [x :*: y' | y' <- grecursivelyShrink y]

instance (RecursivelyShrink f, RecursivelyShrink g) => RecursivelyShrink (f :+: g) where
  grecursivelyShrink (L1 x) = map L1 (grecursivelyShrink x)
  grecursivelyShrink (R1 x) = map R1 (grecursivelyShrink x)

instance RecursivelyShrink f => RecursivelyShrink (M1 i c f) where
  grecursivelyShrink (M1 x) = map M1 (grecursivelyShrink x)

instance Arbitrary a => RecursivelyShrink (K1 i a) where
  grecursivelyShrink (K1 x) = map K1 (shrink x)

instance RecursivelyShrink U1 where
  grecursivelyShrink U1 = []

-- | All immediate subterms of a term.
subterms :: (Generic a, Typeable a, Subterms (Rep a)) => a -> [a]
subterms = gsubterms . from

class Subterms f where
  gsubterms :: Typeable b => f a -> [b]

instance (Subterms f, Subterms g) => Subterms (f :*: g) where
  gsubterms (x :*: y) =
    gsubterms x ++ gsubterms y

instance (Subterms f, Subterms g) => Subterms (f :+: g) where
  gsubterms (L1 x) = gsubterms x
  gsubterms (R1 x) = gsubterms x

instance Subterms f => Subterms (M1 i c f) where
  gsubterms (M1 x) = gsubterms x

instance Typeable a => Subterms (K1 i a) where
  gsubterms (K1 x) =
    case cast x of
      Nothing -> []
      Just y -> [y]

instance Subterms U1 where
  gsubterms U1 = []

#endif

-- instances

instance (CoArbitrary a, Arbitrary b) => Arbitrary (a -> b) where
  arbitrary = promote (`coarbitrary` arbitrary)

instance Arbitrary () where
  arbitrary = return ()

instance Arbitrary Bool where
  arbitrary = choose (False,True)
  shrink True = [False]
  shrink False = []

instance Arbitrary Ordering where
  arbitrary = elements [LT, EQ, GT]
  shrink GT = [EQ, LT]
  shrink LT = [EQ]
  shrink EQ = []

instance Arbitrary a => Arbitrary (Maybe a) where
  arbitrary = frequency [(1, return Nothing), (3, liftM Just arbitrary)]

  shrink (Just x) = Nothing : [ Just x' | x' <- shrink x ]
  shrink _        = []

instance (Arbitrary a, Arbitrary b) => Arbitrary (Either a b) where
  arbitrary = oneof [liftM Left arbitrary, liftM Right arbitrary]

  shrink (Left x)  = [ Left  x' | x' <- shrink x ]
  shrink (Right y) = [ Right y' | y' <- shrink y ]

instance Arbitrary a => Arbitrary [a] where
  arbitrary = sized $ \n ->
    do k <- choose (0,n)
       sequence [ arbitrary | _ <- [1..k] ]

  shrink xs = shrinkList shrink xs

-- | Shrink a list of values given a shrinking function for individual values.
shrinkList :: (a -> [a]) -> [a] -> [[a]]
shrinkList shr xs = concat [ removes k n xs | k <- takeWhile (>0) (iterate (`div`2) n) ]
                 ++ shrinkOne xs
 where
  n = length xs

  shrinkOne []     = []
  shrinkOne (x:xs) = [ x':xs | x'  <- shr x ]
                  ++ [ x:xs' | xs' <- shrinkOne xs ]

  removes k n xs
    | k > n     = []
    | null xs2  = [[]]
    | otherwise = xs2 : map (xs1 ++) (removes k (n-k) xs2)
   where
    xs1 = take k xs
    xs2 = drop k xs

{-
  -- "standard" definition for lists:
  shrink []     = []
  shrink (x:xs) = [ xs ]
               ++ [ x:xs' | xs' <- shrink xs ]
               ++ [ x':xs | x'  <- shrink x ]
-}

instance (Integral a, Arbitrary a) => Arbitrary (Ratio a) where
  arbitrary = arbitrarySizedFractional
  shrink    = shrinkRealFracToInteger

instance (RealFloat a, Arbitrary a) => Arbitrary (Complex a) where
  arbitrary = liftM2 (:+) arbitrary arbitrary
  shrink (x :+ y) = [ x' :+ y | x' <- shrink x ] ++
                    [ x :+ y' | y' <- shrink y ]

#ifndef NO_FIXED
instance HasResolution a => Arbitrary (Fixed a) where
  arbitrary = arbitrarySizedFractional
  shrink    = shrinkRealFrac
#endif

instance (Arbitrary a, Arbitrary b)
      => Arbitrary (a,b)
 where
  arbitrary = liftM2 (,) arbitrary arbitrary

  shrink (x, y) =
       [ (x', y) | x' <- shrink x ]
    ++ [ (x, y') | y' <- shrink y ]

instance (Arbitrary a, Arbitrary b, Arbitrary c)
      => Arbitrary (a,b,c)
 where
  arbitrary = liftM3 (,,) arbitrary arbitrary arbitrary

  shrink (x, y, z) =
    [ (x', y', z')
    | (x', (y', z')) <- shrink (x, (y, z)) ]

instance (Arbitrary a, Arbitrary b, Arbitrary c, Arbitrary d)
      => Arbitrary (a,b,c,d)
 where
  arbitrary = liftM4 (,,,) arbitrary arbitrary arbitrary arbitrary

  shrink (w, x, y, z) =
    [ (w', x', y', z')
    | (w', (x', (y', z'))) <- shrink (w, (x, (y, z))) ]

instance (Arbitrary a, Arbitrary b, Arbitrary c, Arbitrary d, Arbitrary e)
      => Arbitrary (a,b,c,d,e)
 where
  arbitrary = liftM5 (,,,,) arbitrary arbitrary arbitrary arbitrary arbitrary

  shrink (v, w, x, y, z) =
    [ (v', w', x', y', z')
    | (v', (w', (x', (y', z')))) <- shrink (v, (w, (x, (y, z)))) ]

-- typical instance for primitive (numerical) types

instance Arbitrary Integer where
  arbitrary = arbitrarySizedIntegral
  shrink    = shrinkIntegral

instance Arbitrary Int where
  arbitrary = arbitrarySizedIntegral
  shrink    = shrinkIntegral

instance Arbitrary Int8 where
  arbitrary = arbitrarySizedBoundedIntegral
  shrink    = shrinkIntegral

instance Arbitrary Int16 where
  arbitrary = arbitrarySizedBoundedIntegral
  shrink    = shrinkIntegral

instance Arbitrary Int32 where
  arbitrary = arbitrarySizedBoundedIntegral
  shrink    = shrinkIntegral

instance Arbitrary Int64 where
  arbitrary = arbitrarySizedBoundedIntegral
  shrink    = shrinkIntegral

instance Arbitrary Word where
  arbitrary = arbitrarySizedBoundedIntegral
  shrink    = shrinkIntegral

instance Arbitrary Word8 where
  arbitrary = arbitrarySizedBoundedIntegral
  shrink    = shrinkIntegral

instance Arbitrary Word16 where
  arbitrary = arbitrarySizedBoundedIntegral
  shrink    = shrinkIntegral

instance Arbitrary Word32 where
  arbitrary = arbitrarySizedBoundedIntegral
  shrink    = shrinkIntegral

instance Arbitrary Word64 where
  arbitrary = arbitrarySizedBoundedIntegral
  shrink    = shrinkIntegral

instance Arbitrary Char where
  arbitrary = chr `fmap` oneof [choose (0,127), choose (0,255)]
  shrink c  = filter (<. c) $ nub
            $ ['a','b','c']
           ++ [ toLower c | isUpper c ]
           ++ ['A','B','C']
           ++ ['1','2','3']
           ++ [' ','\n']
   where
    a <. b  = stamp a < stamp b
    stamp a = ( (not (isLower a)
              , not (isUpper a)
              , not (isDigit a))
              , (not (a==' ')
              , not (isSpace a)
              , a)
              )

instance Arbitrary Float where
  arbitrary = arbitrarySizedFractional
  shrink    = shrinkRealFrac

instance Arbitrary Double where
  arbitrary = arbitrarySizedFractional
  shrink    = shrinkRealFrac

-- ** Helper functions for implementing arbitrary

-- | Generates an integral number. The number can be positive or negative
-- and its maximum absolute value depends on the size parameter.
arbitrarySizedIntegral :: Integral a => Gen a
arbitrarySizedIntegral =
  sized $ \n ->
  inBounds fromInteger (choose (-toInteger n, toInteger n))

inBounds :: Integral a => (Integer -> a) -> Gen Integer -> Gen a
inBounds fi g = fmap fi (g `suchThat` (\x -> toInteger (fi x) == x))

-- | Generates a fractional number. The number can be positive or negative
-- and its maximum absolute value depends on the size parameter.
arbitrarySizedFractional :: Fractional a => Gen a
arbitrarySizedFractional =
  sized $ \n ->
    let n' = toInteger n in
      do a <- choose ((-n') * precision, n' * precision)
         b <- choose (1, precision)
         return (fromRational (a % b))
 where
  precision = 9999999999999 :: Integer

-- Useful for getting at minBound and maxBound without having to
-- fiddle around with asTypeOf.
withBounds :: Bounded a => (a -> a -> Gen a) -> Gen a
withBounds k = k minBound maxBound

-- | Generates an integral number. The number is chosen uniformly from
-- the entire range of the type. You may want to use
-- 'arbitrarySizedBoundedIntegral' instead.
arbitraryBoundedIntegral :: (Bounded a, Integral a) => Gen a
arbitraryBoundedIntegral =
  withBounds $ \mn mx ->
  do n <- choose (toInteger mn, toInteger mx)
     return (fromInteger n)

-- | Generates an element of a bounded type. The element is
-- chosen from the entire range of the type.
arbitraryBoundedRandom :: (Bounded a, Random a) => Gen a
arbitraryBoundedRandom = choose (minBound,maxBound)

-- | Generates an element of a bounded enumeration.
arbitraryBoundedEnum :: (Bounded a, Enum a) => Gen a
arbitraryBoundedEnum =
  withBounds $ \mn mx ->
  do n <- choose (fromEnum mn, fromEnum mx)
     return (toEnum n)

-- | Generates an integral number from a bounded domain. The number is
-- chosen from the entire range of the type, but small numbers are
-- generated more often than big numbers. Inspired by demands from
-- Phil Wadler.
arbitrarySizedBoundedIntegral :: (Bounded a, Integral a) => Gen a
arbitrarySizedBoundedIntegral =
  withBounds $ \mn mx ->
  sized $ \s ->
    do let bits n | n `quot` 2 == 0 = 0
                  | otherwise = 1 + bits (n `quot` 2)
           k  = 2^(s*(bits mn `max` bits mx `max` 40) `div` 100)
       n <- choose (toInteger mn `max` (-k), toInteger mx `min` k)
       return (fromInteger n)

-- ** Helper functions for implementing shrink

-- | Returns no shrinking alternatives.
shrinkNothing :: a -> [a]
shrinkNothing _ = []

-- | Shrink an integral number.
shrinkIntegral :: Integral a => a -> [a]
shrinkIntegral x =
  nub $
  [ -x
  | x < 0, -x > x
  ] ++
  [ x'
  | x' <- takeWhile (<< x) (0:[ x - i | i <- tail (iterate (`quot` 2) x) ])
  ]
 where
   -- a << b is "morally" abs a < abs b, but taking care of overflow.
   a << b = case (a >= 0, b >= 0) of
            (True,  True)  -> a < b
            (False, False) -> a > b
            (True,  False) -> a + b < 0
            (False, True)  -> a + b > 0

-- | Shrink a fraction, but only shrink to integral values.
shrinkRealFracToInteger :: RealFrac a => a -> [a]
shrinkRealFracToInteger x =
  nub $
  [ -x
  | x < 0
  ] ++
  map fromInteger (shrinkIntegral (truncate x))

-- | Shrink a fraction.
shrinkRealFrac :: RealFrac a => a -> [a]
shrinkRealFrac x =
  nub $
  shrinkRealFracToInteger x ++
  [ x - x'
  | x' <- take 20 (iterate (/ 2) x)
  , (x - x') << x ]
 where
  a << b = abs a < abs b

--------------------------------------------------------------------------
-- ** CoArbitrary

-- | Used for random generation of functions.
class CoArbitrary a where
  -- | Used to generate a function of type @a -> b@.
  -- The first argument is a value, the second a generator.
  -- You should use 'variant' to perturb the random generator;
  -- the goal is that different values for the first argument will
  -- lead to different calls to 'variant'. An example will help:
  --
  -- @
  -- instance CoArbitrary a => CoArbitrary [a] where
  --   coarbitrary []     = 'variant' 0
  --   coarbitrary (x:xs) = 'variant' 1 . coarbitrary (x,xs)
  -- @

  coarbitrary :: a -> Gen b -> Gen b

{-# DEPRECATED (><) "Use ordinary function composition instead" #-}
-- | Combine two generator perturbing functions, for example the
-- results of calls to 'variant' or 'coarbitrary'.
(><) :: (Gen a -> Gen a) -> (Gen a -> Gen a) -> (Gen a -> Gen a)
(><) = (.)

instance (Arbitrary a, CoArbitrary b) => CoArbitrary (a -> b) where
  coarbitrary f gen =
    do xs <- arbitrary
       coarbitrary (map f xs) gen

instance CoArbitrary () where
  coarbitrary _ = id

instance CoArbitrary Bool where
  coarbitrary False = variant 0
  coarbitrary True  = variant 1

instance CoArbitrary Ordering where
  coarbitrary GT = variant 0
  coarbitrary EQ = variant 1
  coarbitrary LT = variant 2

instance CoArbitrary a => CoArbitrary (Maybe a) where
  coarbitrary Nothing  = variant 0
  coarbitrary (Just x) = variant 1 . coarbitrary x

instance (CoArbitrary a, CoArbitrary b) => CoArbitrary (Either a b) where
  coarbitrary (Left x)  = variant 0 . coarbitrary x
  coarbitrary (Right y) = variant 1 . coarbitrary y

instance CoArbitrary a => CoArbitrary [a] where
  coarbitrary []     = variant 0
  coarbitrary (x:xs) = variant 1 . coarbitrary (x,xs)

instance (Integral a, CoArbitrary a) => CoArbitrary (Ratio a) where
  coarbitrary r = coarbitrary (numerator r,denominator r)

#ifndef NO_FIXED
instance HasResolution a => CoArbitrary (Fixed a) where
  coarbitrary = coarbitraryReal
#endif

instance (RealFloat a, CoArbitrary a) => CoArbitrary (Complex a) where
  coarbitrary (x :+ y) = coarbitrary x . coarbitrary y

instance (CoArbitrary a, CoArbitrary b)
      => CoArbitrary (a,b)
 where
  coarbitrary (x,y) = coarbitrary x
                    . coarbitrary y

instance (CoArbitrary a, CoArbitrary b, CoArbitrary c)
      => CoArbitrary (a,b,c)
 where
  coarbitrary (x,y,z) = coarbitrary x
                      . coarbitrary y
                      . coarbitrary z

instance (CoArbitrary a, CoArbitrary b, CoArbitrary c, CoArbitrary d)
      => CoArbitrary (a,b,c,d)
 where
  coarbitrary (x,y,z,v) = coarbitrary x
                        . coarbitrary y
                        . coarbitrary z
                        . coarbitrary v

instance (CoArbitrary a, CoArbitrary b, CoArbitrary c, CoArbitrary d, CoArbitrary e)
      => CoArbitrary (a,b,c,d,e)
 where
  coarbitrary (x,y,z,v,w) = coarbitrary x
                          . coarbitrary y
                          . coarbitrary z
                          . coarbitrary v
                          . coarbitrary w

-- typical instance for primitive (numerical) types

instance CoArbitrary Integer where
  coarbitrary = coarbitraryIntegral

instance CoArbitrary Int where
  coarbitrary = coarbitraryIntegral

instance CoArbitrary Int8 where
  coarbitrary = coarbitraryIntegral

instance CoArbitrary Int16 where
  coarbitrary = coarbitraryIntegral

instance CoArbitrary Int32 where
  coarbitrary = coarbitraryIntegral

instance CoArbitrary Int64 where
  coarbitrary = coarbitraryIntegral

instance CoArbitrary Word where
  coarbitrary = coarbitraryIntegral

instance CoArbitrary Word8 where
  coarbitrary = coarbitraryIntegral

instance CoArbitrary Word16 where
  coarbitrary = coarbitraryIntegral

instance CoArbitrary Word32 where
  coarbitrary = coarbitraryIntegral

instance CoArbitrary Word64 where
  coarbitrary = coarbitraryIntegral

instance CoArbitrary Char where
  coarbitrary = coarbitrary . ord

instance CoArbitrary Float where
  coarbitrary = coarbitraryReal

instance CoArbitrary Double where
  coarbitrary = coarbitraryReal

-- ** Helpers for implementing coarbitrary

-- | A 'coarbitrary' implementation for integral numbers.
coarbitraryIntegral :: Integral a => a -> Gen b -> Gen b
coarbitraryIntegral = variant

-- | A 'coarbitrary' implementation for real numbers.
coarbitraryReal :: Real a => a -> Gen b -> Gen b
coarbitraryReal x = coarbitrary (toRational x)

-- | 'coarbitrary' helper for lazy people :-).
coarbitraryShow :: Show a => a -> Gen b -> Gen b
coarbitraryShow x = coarbitrary (show x)

-- | A 'coarbitrary' implementation for enums.
coarbitraryEnum :: Enum a => a -> Gen b -> Gen b
coarbitraryEnum = variant . fromEnum

--------------------------------------------------------------------------
-- ** arbitrary generators

-- these are here and not in Gen because of the Arbitrary class constraint

-- | Generates a list of a given length.
vector :: Arbitrary a => Int -> Gen [a]
vector k = vectorOf k arbitrary

-- | Generates an ordered list of a given length.
orderedList :: (Ord a, Arbitrary a) => Gen [a]
orderedList = sort `fmap` arbitrary

-- | Generate an infinite list.
infiniteList :: Arbitrary a => Gen [a]
infiniteList = infiniteListOf arbitrary

--------------------------------------------------------------------------
-- the end.


-- |
-- Module: Optics.Each.Core
-- Description: An 'IxTraversal' for each element of a (potentially monomorphic) container.
--
-- This module defines the 'Each' class, which provides an 'IxTraversal' that
-- extracts 'each' element of a (potentially monomorphic) container.
--
-- Note that orphan instances for this class are defined in the @Optics.Each@
-- module from @optics-extra@, so if you are not simply depending on @optics@
-- you may wish to import that module instead.
--
{-# LANGUAGE UndecidableInstances #-}
module Optics.Each.Core
  (
  -- * Each
    Each(..)
  ) where

import Data.Array
import Data.Complex
import Data.Functor.Identity
import Data.IntMap as IntMap
import Data.List.NonEmpty
import Data.Map as Map
import Data.Sequence as Seq
import Data.Tree as Tree

import Optics.IxTraversal

-- | Extract 'each' element of a (potentially monomorphic) container.
--
-- >>> over each (*10) (1,2,3)
-- (10,20,30)
--
-- >>> iover each (\i a -> a*10 + succ i) (1,2,3)
-- (11,22,33)
--
class Each i s t a b | s -> i a, t -> i b, s b -> t, t a -> s where
  each :: IxTraversal i s t a b

  default each
    :: (TraversableWithIndex i g, s ~ g a, t ~ g b)
    => IxTraversal i s t a b
  each = itraversed
  {-# INLINE[1] each #-}

-- | @'each' :: 'IxTraversal' 'Int' (a, a) (b, b) a b@
instance
  (a ~ a1,
   b ~ b1
  ) => Each Int (a, a1)
                (b, b1) a b where
  each = itraversalVL $ \f (a0, a1) ->
    (,) <$> f 0 a0 <*> f 1 a1
  {-# INLINE[1] each #-}

-- | @'each' :: 'IxTraversal' 'Int' (a, a, a) (b, b, b) a b@
instance
  (a ~ a1, a ~ a2,
   b ~ b1, b ~ b2
  ) => Each Int (a, a1, a2)
                (b, b1, b2) a b where
  each = itraversalVL $ \f (a0, a1, a2) ->
    (,,) <$> f 0 a0 <*> f 1 a1 <*> f 2 a2
  {-# INLINE[1] each #-}

-- | @'each' :: 'IxTraversal' 'Int' (a, a, a, a) (b, b, b, b) a b@
instance
  (a ~ a1, a ~ a2, a ~ a3,
   b ~ b1, b ~ b2, b ~ b3
  ) => Each Int (a, a1, a2, a3)
                (b, b1, b2, b3) a b where
  each = itraversalVL $ \f (a0, a1, a2, a3) ->
    (,,,) <$> f 0 a0 <*> f 1 a1 <*> f 2 a2 <*> f 3 a3
  {-# INLINE[1] each #-}

-- | @'each' :: 'IxTraversal' 'Int' (a, a, a, a, a) (b, b, b, b, b) a b@
instance
  (a ~ a1, a ~ a2, a ~ a3, a ~ a4,
   b ~ b1, b ~ b2, b ~ b3, b ~ b4
  ) => Each Int (a, a1, a2, a3, a4)
                (b, b1, b2, b3, b4) a b where
  each = itraversalVL $ \f (a0, a1, a2, a3, a4) ->
    (,,,,) <$> f 0 a0 <*> f 1 a1 <*> f 2 a2 <*> f 3 a3 <*> f 4 a4
  {-# INLINE[1] each #-}

-- | @'each' :: 'IxTraversal' 'Int' (a, a, a, a, a, a) (b, b, b, b, b, b) a b@
instance
  (a ~ a1, a ~ a2, a ~ a3, a ~ a4, a ~ a5,
   b ~ b1, b ~ b2, b ~ b3, b ~ b4, b ~ b5
  ) => Each Int (a, a1, a2, a3, a4, a5)
                (b, b1, b2, b3, b4, b5) a b where
  each = itraversalVL $ \f (a0, a1, a2, a3, a4, a5) ->
    (,,,,,) <$> f 0 a0 <*> f 1 a1 <*> f 2 a2 <*> f 3 a3 <*> f 4 a4
            <*> f 5 a5
  {-# INLINE[1] each #-}

-- | @'each' :: 'IxTraversal' 'Int' (a, a, a, a, a, a, a) (b, b, b, b, b, b, b)
-- a b@
instance
  (a ~ a1, a ~ a2, a ~ a3, a ~ a4, a ~ a5, a ~ a6,
   b ~ b1, b ~ b2, b ~ b3, b ~ b4, b ~ b5, b ~ b6
  ) => Each Int (a, a1, a2, a3, a4, a5, a6)
                (b, b1, b2, b3, b4, b5, b6) a b where
  each = itraversalVL $ \f (a0, a1, a2, a3, a4, a5, a6) ->
    (,,,,,,) <$> f 0 a0 <*> f 1 a1 <*> f 2 a2 <*> f 3 a3 <*> f 4 a4
             <*> f 5 a5 <*> f 6 a6
  {-# INLINE[1] each #-}

-- | @'each' :: 'IxTraversal' 'Int' (a, a, a, a, a, a, a, a) (b, b, b, b, b, b,
-- b, b) a b@
instance
  (a ~ a1, a ~ a2, a ~ a3, a ~ a4, a ~ a5, a ~ a6, a ~ a7,
   b ~ b1, b ~ b2, b ~ b3, b ~ b4, b ~ b5, b ~ b6, b ~ b7
  ) => Each Int (a, a1, a2, a3, a4, a5, a6, a7)
                (b, b1, b2, b3, b4, b5, b6, b7) a b where
  each = itraversalVL $ \f (a0, a1, a2, a3, a4, a5, a6, a7) ->
    (,,,,,,,) <$> f 0 a0 <*> f 1 a1 <*> f 2 a2 <*> f 3 a3 <*> f 4 a4
              <*> f 5 a5 <*> f 6 a6 <*> f 7 a7
  {-# INLINE[1] each #-}

-- | @'each' :: 'IxTraversal' 'Int' (a, a, a, a, a, a, a, a, a) (b, b, b, b, b,
-- b, b, b, b) a b@
instance
  (a ~ a1, a ~ a2, a ~ a3, a ~ a4, a ~ a5, a ~ a6, a ~ a7, a ~ a8,
   b ~ b1, b ~ b2, b ~ b3, b ~ b4, b ~ b5, b ~ b6, b ~ b7, b ~ b8
  ) => Each Int (a, a1, a2, a3, a4, a5, a6, a7, a8)
                (b, b1, b2, b3, b4, b5, b6, b7, b8) a b where
  each = itraversalVL $ \f (a0, a1, a2, a3, a4, a5, a6, a7, a8) ->
    (,,,,,,,,) <$> f 0 a0 <*> f 1 a1 <*> f 2 a2 <*> f 3 a3 <*> f 4 a4
               <*> f 5 a5 <*> f 6 a6 <*> f 7 a7 <*> f 8 a8
  {-# INLINE[1] each #-}

-- | @'each' :: 'IxTraversal' 'Int' (a, a, a, a, a, a, a, a, a, a) (b, b, b, b,
-- b, b, b, b, b, b) a b@
instance
  (a ~ a1, a ~ a2, a ~ a3, a ~ a4, a ~ a5, a ~ a6, a ~ a7, a ~ a8, a ~ a9,
   b ~ b1, b ~ b2, b ~ b3, b ~ b4, b ~ b5, b ~ b6, b ~ b7, b ~ b8, b ~ b9
  ) => Each Int (a, a1, a2, a3, a4, a5, a6, a7, a8, a9)
                (b, b1, b2, b3, b4, b5, b6, b7, b8, b9) a b where
  each = itraversalVL $ \f (a0, a1, a2, a3, a4, a5, a6, a7, a8, a9) ->
    (,,,,,,,,,) <$> f 0 a0 <*> f 1 a1 <*> f 2 a2 <*> f 3 a3 <*> f 4 a4
                <*> f 5 a5 <*> f 6 a6 <*> f 7 a7 <*> f 8 a8 <*> f 9 a9
  {-# INLINE[1] each #-}

-- | @'each' :: 'IxTraversal' ('Either' () ()) ('Either' a a) ('Either' b b) a
-- b@
instance
  (a ~ a', b ~ b'
  ) => Each (Either () ()) (Either a a') (Either b b') a b where
  each = itraversalVL $ \f -> \case
    Left  a -> Left  <$> f (Left ())  a
    Right a -> Right <$> f (Right ()) a
  {-# INLINE[1] each #-}

-- | @'each' :: ('RealFloat' a, 'RealFloat' b) => 'IxTraversal' (Either () ())
-- ('Complex' a) ('Complex' b) a b@
instance Each (Either () ()) (Complex a) (Complex b) a b where
  each = itraversalVL $ \f (a :+ b) -> (:+) <$> f (Left ()) a <*> f (Right ()) b
  {-# INLINE[1] each #-}

-- | @'each' :: 'IxTraversal' k ('Map' k a) ('Map' k b) a b@
instance k ~ k' => Each k (Map k a) (Map k' b) a b where
  -- traverseWithKey has best performance for all flavours for some reason.
  each = itraversalVL Map.traverseWithKey
  {-# INLINE[1] each #-}

-- | @'each' :: 'IxTraversal' 'Int' ('IntMap' a) ('IntMap' b) a b@
instance Each Int (IntMap a) (IntMap b) a b where
  -- traverseWithKey has best performance for all flavours for some reason.
  each = itraversalVL IntMap.traverseWithKey
  {-# INLINE[1] each #-}

-- | @'each' :: 'IxTraversal' 'Int' [a] [b] a b@
instance Each Int [a] [b] a b

-- | @'each' :: 'IxTraversal' 'Int' (NonEmpty a) (NonEmpty b) a b@
instance Each Int (NonEmpty a) (NonEmpty b) a b

-- | @'each' :: 'IxTraversal' () ('Identity' a) ('Identity' b) a b@
instance Each () (Identity a) (Identity b) a b

-- | @'each' :: 'IxTraversal' () ('Maybe' a) ('Maybe' b) a b@
instance Each () (Maybe a) (Maybe b) a b

-- | @'each' :: 'IxTraversal' 'Int' ('Seq' a) ('Seq' b) a b@
instance Each Int (Seq a) (Seq b) a b

-- | @'each' :: 'IxTraversal' [Int] ('Tree' a) ('Tree' b) a b@
instance Each [Int] (Tree a) (Tree b) a b

-- | @'each' :: 'Ix' i => 'IxTraversal' i ('Array' i a) ('Array' i b) a b@
instance (Ix i, i ~ j) => Each i (Array i a) (Array j b) a b

-- $setup
-- >>> import Optics.Core


-- |
-- Module      : Data.Express.Name
-- Copyright   : (c) 2019-2021 Rudy Matela
-- License     : 3-Clause BSD  (see the file LICENSE)
-- Maintainer  : Rudy Matela <rudy@matela.com.br>
--
-- Defines the 'Name' type class.
module Data.Express.Name
  ( Name (..)
  , names
  , variableNamesFromTemplate
  )
where

import Data.Express.Utils.String

import Data.List

-- for instances:
import Data.Int
import Data.Word
import Data.Ratio
import Data.Complex
import Data.Char

-- |
-- If we were to come up with a variable name for the given type
-- what 'name' would it be?
--
-- An instance for a given type @ Ty @ is simply given by:
--
-- > instance Name Ty where name _ = "x"
--
-- Examples:
--
-- > > name (undefined :: Int)
-- > "x"
--
-- > > name (undefined :: Bool)
-- > "p"
--
-- > > name (undefined :: [Int])
-- > "xs"
--
-- This is then used to generate an infinite list of variable 'names':
--
-- > > names (undefined :: Int)
-- > ["x", "y", "z", "x'", "y'", "z'", "x''", "y''", "z''", ...]
--
-- > > names (undefined :: Bool)
-- > ["p", "q", "r", "p'", "q'", "r'", "p''", "q''", "r''", ...]
--
-- > > names (undefined :: [Int])
-- > ["xs", "ys", "zs", "xs'", "ys'", "zs'", "xs''", "ys''", ...]
class Name a where
  -- | /O(1)./
  --
  -- Returns a name for a variable of the given argument's type.
  --
  -- > > name (undefined :: Int)
  -- > "x"
  --
  -- > > name (undefined :: [Bool])
  -- > "ps"
  --
  -- > > name (undefined :: [Maybe Integer])
  -- > "mxs"
  --
  -- The default definition is:
  --
  -- > name _ = "x"
  name :: a -> String
  name _ = "x"

-- |
-- > name (undefined :: ()) = "u"
-- > names (undefined :: ()) = ["u", "v", "w", "u'", "v'", ...]
instance Name ()        where  name _  =  "u"

-- |
-- > name (undefined :: Bool) = "p"
-- > names (undefined :: Bool) = ["p", "q", "r", "p'", "q'", ...]
instance Name Bool      where  name _  =  "p"

-- |
-- > name (undefined :: Int) = "x"
-- > names (undefined :: Int) = ["x", "y", "z", "x'", "y'", ...]
instance Name Int       where  name _  =  "x"

-- |
-- > name (undefined :: Integer) = "x"
-- > names (undefined :: Integer) = ["x", "y", "z", "x'", ...]
instance Name Integer   where  name _  =  "x"

-- |
-- > name (undefined :: Char) = "c"
-- > names (undefined :: Char) = ["c", "d", "e", "c'", "d'", ...]
instance Name Char      where  name _  =  "c"

-- |
-- > name (undefined :: Ordering) = "o"
-- > names (undefined :: Ordering) = ["o", "p", "q", "o'", ...]
instance Name Ordering  where  name _  =  "o"

-- |
-- > name (undefined :: Rational) = "q"
-- > names (undefined :: Rational) = ["q", "r", "s", "q'", ...]
instance Name (Ratio a) where  name _  =  "q"

-- |
-- > name (undefined :: Complex) = "x"
-- > names (undefined :: Complex) = ["x", "y", "z", "x'", ...]
instance Name (Complex a)  where  name _  =  "x"

-- |
-- > name (undefined :: Float) = "x"
-- > names (undefined :: Float) = ["x", "y", "z", "x'", ...]
instance Name Float     where  name _  =  "x"

-- |
-- > name (undefined :: Double) = "x"
-- > names (undefined :: Double) = ["x", "y", "z", "x'", ...]
instance Name Double    where  name _  =  "x"

-- |
-- > names (undefined :: ()->()) = ["f", "g", "h", "f'", ...]
-- > names (undefined :: Int->Int) = ["f", "g", "h", ...]
instance Name (a -> b)  where  name _  =  "f"

-- |
-- > names (undefined :: Maybe Int) = ["mx", "mx1", "mx2", ...]
-- > nemes (undefined :: Maybe Bool) = ["mp", "mp1", "mp2", ...]
instance Name a => Name (Maybe a) where
  name mx  =  "m" ++ name x
    where
    Just x = mx

-- |
-- > names (undefined :: Either Int Int) = ["exy", "exy1", ...]
-- > names (undefined :: Either Int Bool) = ["exp", "exp1", ...]
instance (Name a, Name b) => Name (Either a b) where
  name exy  =  "e" ++ n ++ m
    where
    Left x  = exy
    Right y = exy
    n = name x
    m = head $ names y \\ [n]

-- |
-- > names (undefined :: (Int,Int)) = ["xy", "zw", "xy'", ...]
-- > names (undefined :: (Bool,Bool)) = ["pq", "rs", "pq'", ...]
instance (Name a, Name b) => Name (a,b) where
  name xy  =  n ++ m
    where
    (x,y)  =  xy
    n  =  name x
    m  =  head $ names y \\ [n]

-- |
-- > names (undefined :: (Int,Int,Int)) = ["xyz","uvw", ...]
-- > names (undefined :: (Int,Bool,Char)) = ["xpc", "xpc1", ...]
instance (Name a, Name b, Name c) => Name (a,b,c) where
  name xyz  =  n ++ m ++ o
    where
    (x,y,z)  =  xyz
    n  =  name x
    m  =  head $ names y \\ [n]
    o  =  head $ names z \\ [n,m]

-- |
-- > names (undefined :: ((),(),(),())) = ["uuuu", "uuuu1", ...]
-- > names (undefined :: (Int,Int,Int,Int)) = ["xxxx", ...]
instance (Name a, Name b, Name c, Name d) => Name (a,b,c,d) where
  name xyzw  =  name x ++ name y ++ name z ++ name w  where  (x,y,z,w)  =  xyzw

-- |
-- > names (undefined :: [Int]) = ["xs", "ys", "zs", "xs'", ...]
-- > names (undefined :: [Bool]) = ["ps", "qs", "rs", "ps'", ...]
instance Name a => Name [a] where
  name xs  =  name (head xs) ++ "s"

-- |
-- Returns na infinite list of variable names from the given type:
-- the result of 'variableNamesFromTemplate' after 'name'.
--
-- > > names (undefined :: Int)
-- > ["x", "y", "z", "x'", "y'", "z'", "x''", "y''", "z''", ...]
--
-- > > names (undefined :: Bool)
-- > ["p", "q", "r", "p'", "q'", "r'", "p''", "q''", "r''", ...]
--
-- > > names (undefined :: [Int])
-- > ["xs", "ys", "zs", "xs'", "ys'", "zs'", "xs''", "ys''", ...]
names :: Name a => a -> [String]
names  =  variableNamesFromTemplate . name

-- instances of further types and arities --

instance Name Word where  name _  =  "x"

instance Name Int8   where  name _  =  "x"
instance Name Int16  where  name _  =  "x"
instance Name Int32  where  name _  =  "x"
instance Name Int64  where  name _  =  "x"
instance Name Word8  where  name _  =  "x"
instance Name Word16 where  name _  =  "x"
instance Name Word32 where  name _  =  "x"
instance Name Word64 where  name _  =  "x"
instance Name GeneralCategory where  name _  =  "c"

instance (Name a, Name b, Name c, Name d, Name e) => Name (a,b,c,d,e) where
  name xyzwv  =  name x ++ name y ++ name z ++ name w ++ name v
    where  (x,y,z,w,v)  =  xyzwv

instance (Name a, Name b, Name c, Name d, Name e, Name f)
      => Name (a,b,c,d,e,f) where
  name xyzwvu  =  name x ++ name y ++ name z ++ name w ++ name v ++ name u
    where  (x,y,z,w,v,u)  =  xyzwvu

instance (Name a, Name b, Name c, Name d, Name e, Name f, Name g)
      => Name (a,b,c,d,e,f,g) where
  name xyzwvut  =  name x ++ name y ++ name z ++ name w
                ++ name v ++ name u ++ name t
    where  (x,y,z,w,v,u,t)  =  xyzwvut

instance (Name a, Name b, Name c, Name d, Name e, Name f, Name g, Name h)
      => Name (a,b,c,d,e,f,g,h) where
  name xyzwvuts  =  name x ++ name y ++ name z ++ name w
                 ++ name v ++ name u ++ name t ++ name s
    where  (x,y,z,w,v,u,t,s)  =  xyzwvuts

instance ( Name a, Name b, Name c, Name d
         , Name e, Name f, Name g, Name h
         , Name i)
      => Name (a,b,c,d,e,f,g,h,i) where
  name xyzwvutsr  =  name x ++ name y ++ name z ++ name w
                  ++ name v ++ name u ++ name t ++ name s
                  ++ name r
    where  (x,y,z,w,v,u,t,s,r)  =  xyzwvutsr

instance ( Name a, Name b, Name c, Name d
         , Name e, Name f, Name g, Name h
         , Name i, Name j )
      => Name (a,b,c,d,e,f,g,h,i,j) where
  name xyzwvutsrq  =  name x ++ name y ++ name z ++ name w
                   ++ name v ++ name u ++ name t ++ name s
                   ++ name r ++ name q
    where  (x,y,z,w,v,u,t,s,r,q)  =  xyzwvutsrq

instance ( Name a, Name b, Name c, Name d
         , Name e, Name f, Name g, Name h
         , Name i, Name j, Name k )
      => Name (a,b,c,d,e,f,g,h,i,j,k) where
  name xyzwvutsrqp  =  name x ++ name y ++ name z ++ name w
                    ++ name v ++ name u ++ name t ++ name s
                    ++ name r ++ name q ++ name p
    where  (x,y,z,w,v,u,t,s,r,q,p)  =  xyzwvutsrqp

instance ( Name a, Name b, Name c, Name d
         , Name e, Name f, Name g, Name h
         , Name i, Name j, Name k, Name l )
      => Name (a,b,c,d,e,f,g,h,i,j,k,l) where
  name xyzwvutsrqpo  =  name x ++ name y ++ name z ++ name w
                     ++ name v ++ name u ++ name t ++ name s
                     ++ name r ++ name q ++ name p ++ name o
    where  (x,y,z,w,v,u,t,s,r,q,p,o)  =  xyzwvutsrqpo


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell   #-}
{-# LANGUAGE CApiFFI           #-}
-- |
-- Copyright: 2018 (C) Tweag I/O Limited.
--
-- inline-c context.
module Foreign.R.Context
  ( rCtx
  , SEXPREC
  , SEXP0
  , Logical(..)
  ) where

import Data.Complex
import qualified Data.Map as Map
import Foreign.C
import Foreign.Ptr
import Foreign.Storable
import Language.C.Types (TypeSpecifier(TypeName))
import Language.C.Inline.Context (Context(..))
import Internal.Error

#include <Rinternals.h>

data {-# CTYPE  "SEXPREC" #-} SEXPREC

type SEXP0 = Ptr SEXPREC

-- | R uses three-valued logic.
data {-# CTYPE "Logical" #-} Logical = FALSE
             | TRUE
             | NA
-- XXX no Enum instance because NA = INT_MIN, not representable as an Int on
-- 32-bit systems.
               deriving (Eq, Ord, Show)

instance Storable Logical where
  sizeOf _       = sizeOf (undefined :: CInt)
  alignment _    = alignment (undefined :: CInt)
  poke ptr FALSE = poke (castPtr ptr) (0 :: CInt)
  poke ptr TRUE  = poke (castPtr ptr) (1 :: CInt)
  -- Currently NA_LOGICAL = INT_MIN.
  poke ptr NA    = poke (castPtr ptr) (#{const INT_MIN} :: CInt)
  peek ptr = do
      x <- peek (castPtr ptr)
      case x :: CInt of
          0 -> return FALSE
          1 -> return TRUE
          #{const INT_MIN} -> return NA
          _ -> failure "Storable Logical peek" "Not a Logical."

rCtx :: Context
rCtx = mempty { ctxTypesTable = Map.fromList tytabs }
  where
    tytabs =
      [ (TypeName "SEXP", [t| Ptr SEXPREC |])
      , (TypeName "Rcomplex", [t| Complex Double |])
      ]


-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Signal.EEG.BDF
-- Copyright   :  (c) Alexander Vivian Hugh McPhail 2010, 2014, 2015
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :
--
-- EEG BDF data file functions
--
-----------------------------------------------------------------------------

module Numeric.Signal.EEG.BDF (
                               BDF(..)
                              , loadBDFRaw
                              , loadBDF
                          ) where

import qualified Data.ByteString as BS

--import Data.Char
--import Data.List
import Data.Word
import Data.Bits
--import Data.Array.Storable
import Numeric.LinearAlgebra.Data

import qualified Numeric.Signal.Multichannel as M

import Control.Monad hiding(join)
import Control.Monad.State hiding(join)

--import System.IO

--import Debug.Trace

{- BDF Header format: http://www.biosemi.com/faq/file_format.htm
Length in bytes  	BDF Header:  	EDF Header:  	Description
8 bytes 	Byte 1: "255" (non ascii) 	Byte 1: "0" (ASCII) 	Identification code
Bytes 2-8 : "BIOSEMI" (ASCII) 	Bytes 2-8 : " "(ASCII)
80 bytes
User text input (ASCII)
	Local subject identification
80 bytes
User text input (ASCII)
	Local recording identification
8 bytes
dd.mm.yy (ASCII)
	Startdate of recording
8 bytes
hh.mm.ss (ASCII)
	Starttime of recording
8 bytes
(ASCII)
	Number of bytes in header record
44 bytes
"24BIT" (ASCII)

"BIOSEMI" (ASCII)
	Version of data format.
8 bytes
(ASCII)
	Number of data records "-1" if unknown
8 bytes
e.g.: "1" (ASCII)
	Duration of a data record, in seconds
4 bytes
e.g.: "257" or "128" (ASCII)
	Number of channels (N) in data record
N x 16 bytes
e.g.: "Fp1", "Fpz", "Fp2", etc (ASCII)
	Labels of the channels
N x 80 bytes
e.g.: "active electrode", "respiration belt" (ASCII)
	Transducer type
N x 8 bytes
e.g.: "uV", "Ohm" (ASCII)
	Physical dimension of channels
N x 8 bytes 	e.g.: "-262144" (ASCII) 	e.g.: "-32768" (ASCII) 	Physical minimum in units of physical dimension
N x 8 bytes 	e.g.: "262143" (ASCII) 	e.g.: "32767" (ASCII) 	Physical maximum in units of physical dimension
N x 8 bytes 	e.g.: "-8388608" (ASCII) 	e.g.: "-32768" (ASCII) 	Digital minimum
N x 8 bytes 	e.g.: "8388607" (ASCII) 	e.g.: "32767" (ASCII) 	Digital maximum
N x 80 bytes 	e.g.: "HP:DC; LP:410" 	e.g.: "HP:0,16; LP:500" 	Prefiltering
N x 8 bytes
For example: "2048" (ASCII)

Number of samples in each data record
(Sample-rate if Duration of data record = "1")
N x 32 bytes
(ASCII)
	Reserved
-}

type BSM a = StateT BS.ByteString IO a

getN :: Int -> BSM BS.ByteString
getN n = do
         bs <- get
         let (v,bs') = BS.splitAt n bs
         put bs'
         return v

data Date = Date { day :: Int, month :: Int, year :: Int }
          deriving(Eq,Ord,Show)

data Time = Time { hour :: Int, minute :: Int, second :: Int }
          deriving(Eq,Ord,Show)

-- | BDF datatype.  NOTE: data has already been converted to Float
data BDF = BDF {
                id_ :: !Word8
                , type_ :: !String
                , subject :: !String
                , recording :: !String
                , date :: !Date
                , time :: !Time
                , head_bytes :: !Int
                , data_version :: !String
                , num_records :: !Int
                , duration :: !Int
                , channels :: !Int
                , chan_labels :: ![String]
                , tran_type :: ![String]
                , dimensions :: ![String]
                , phys_min :: ![Int]
                , phys_max :: ![Int]
                , dig_min :: ![Int]
                , dig_max :: ![Int]
                , prefilter :: ![String]
                , samples :: ![Int]
                , reserved :: ![String]
                , data_ :: ![Vector Float]
               } --deriving(Show)

getString :: Int -> BSM String
getString n = do
              bs <- getN n
              return $ (reverse . dropWhile (== ' ') . tail . reverse . tail . show) bs

getInt :: Int -> BSM Int
getInt n = do
           s <- getString n
           return $ read s

getDate = do
          s <- getString 8
          return $ strToDate' s

strToDate' :: String -> Date
strToDate' (d1:d2:'.':m1:m2:'.':y1:y2:[]) = Date (read [d1,d2]) (read [m1,m2]) (read [y1,y2])
strToDate' _                              = error "strToDate"

getTime = do
          s <- getString 8
          return $ strToTime' s

strToTime' :: String -> Time
strToTime' (h1:h2:'.':m1:m2:'.':s1:s2:[]) = Time (read [h1,h2]) (read [m1,m2]) (read [s1,s2])
strToTime' _                              = error "strToTime"

reverseBits :: Word8 -> Word8
--reverseBits w = foldl setBit 0 $ snd . unzip . filter ((== False) . fst) $ zip (map (testBit w) [1..8]) $ reverse [1..8]
--reverseBits = foldl setBit 0 . snd . unzip . filter ((False ==) . fst) . ($ reverse [0..7]) . zip . flip map [0..7] . testBit
--reverseBits w = foldl setBit 0 . foldl ((++) . \b -> if testBit w b then [9-b] else []) [] [1..8]
reverseBits w = foldr (\b r -> if not $ testBit w b then setBit r (7-b) else r) 0 [0..7]

get24Bit :: (Int -> Float) -> BSM Float
get24Bit f = do
             b1 <- getN 1
             b2 <- getN 1
             b3 <- getN 1
--           return $ (fromIntegral $ BS.head b1) `shiftL` 16 .|. (fromIntegral $ BS.head b2) `shiftL` 8 .|. (fromIntegral $ BS.head b3)
             return $ f $ (to32 b3) `shiftL` 16 .|. (to32 b2) `shiftL` 8 .|. (to32 b1)
    where to32 = fromIntegral {- . reverseBits -} . BS.head

readRecord :: (Int -> Float) -> Int -> BSM (Vector Float)
readRecord f s = do
                 m <- replicateM s $ get24Bit f
                 return $! fromList m

readRecordBlock :: (Int -> Int -> Float) -> [(Int,Int)] -> BSM [Vector Float]
readRecordBlock f ss = mapM (\(i,s) -> readRecord (f i) s) ss

convert :: [Int] -> [Int] -> [Int] -> [Int] -> Int -> Int -> Float
convert p_min p_max d_min d_max i = \x -> (fromIntegral x) - (fromIntegral (d_min !! i))*(fromIntegral ((p_max !! i) - (p_min !! i)))/(fromIntegral ((d_max !! i) - (d_min !! i)))

readData :: (Int -> Int -> Float) -> Int -> [(Int,Int)] -> BSM [Vector Float]
readData f rs ss = do
                   lift $ putStrLn $ "channels: " ++ (show $ length ss) ++ ", records: " ++ (show rs) ++ ", samples per record: " ++ (show $ snd $ head ss)
                   d <- mapM (\x -> do
                                    lift $ putStrLn $ "Record: " ++ show x
                                    readRecordBlock f ss) [1..rs]
                   -- let v = rotate d
                   -- lift $ putStrLn $ "vectors: " ++ (show $ length v)
                   -- lift $ putStrLn $ "slices: " ++ (show $ length $ head v)
                   return $! map vjoin $! rotate_ d
    where rotate_ []            = []
          rotate_ xs@((_:[]):_) = [concat xs]
          rotate_ ((x:xs):xss)  = (x : (map head xss)) : (rotate_ (xs : (map tail xss)))

{-
readData _ 0     _  = error "readData, zeroth record requested"
readData f 1     ss = do
                      lift $ putStrLn $ "data record: 1"
                      readRecordBlock f ss
readData f (n+1) ss = do
                      lift $ putStrLn $ "data record: " ++ show (n+1)
                      bs <- readRecordBlock f ss
                      ds <- readData f n ss
                      let result = zipWith (\x y -> join [x,y]) bs ds
                      return $ result
-}

readBDF :: BSM (Maybe BDF)
readBDF = do
          id_' <- getN 1
          type_' <- getString 7
          if (not $ BS.head id_' == 255 && type_' == "BIOSEMI")
               then do
                    lift $ putStrLn "Error: File not BDF Format"
                    return Nothing
               else do
                    subject' <- getString 80
                    recording' <- getString 80
                    date' <- getDate
                    time' <- getTime
                    head_bytes' <- getInt 8
                    data_version' <- getString 44
                    num_records' <- getInt 8
                    if (num_records' == -1)
                       then do
                            lift $ putStrLn "This file is probably a valid BDF file..."
                            lift $ putStrLn "   but this program cannot grok an unspecified"
                            lift $ putStrLn "   number of records"
                            lift $ putStrLn "So complain to the software author"
                            return Nothing
                       else do
                            duration' <- getInt 8
                            channels' <- getInt 4
                            chan_labels' <- replicateM channels' $ getString 16
                            tran_type' <- replicateM channels' $ getString 80
                            dimensions' <- replicateM channels' $ getString 8
                            phys_min' <- replicateM channels' $ getInt 8
                            phys_max' <- replicateM channels' $ getInt 8
                            dig_min' <- replicateM channels' $ getInt 8
                            dig_max' <- replicateM channels' $ getInt 8
                            prefilter' <- replicateM channels' $ getString 80
                            samples' <- replicateM channels' $ getInt 8
                            reserved' <- replicateM channels' $ getString 32
                            data_' <- readData (convert phys_min' phys_max' dig_min' dig_max') num_records' (zip [0..] samples')
                            return $ Just $ BDF {
                                                 id_ = BS.head id_'
                                                 , type_ = type_'
                                                 , subject = subject'
                                                 , recording = recording'
                                                 , date = date'
                                                 , time = time'
                                                 , head_bytes = head_bytes'
                                                 , data_version = data_version'
                                                 , num_records = num_records'
                                                 , duration = duration'
                                                 , channels = channels'
                                                 , chan_labels = chan_labels'
                                                 , tran_type = tran_type'
                                                 , dimensions = dimensions'
                                                 , phys_min = phys_min'
                                                 , phys_max = phys_max'
                                                 , dig_min = dig_min'
                                                 , dig_max = dig_max'
                                                 , prefilter = prefilter'
                                                 , samples = samples'
                                                 , reserved = reserved'
                                                 , data_ = data_'
                                                }

loadBDF :: FilePath -> IO (Maybe (M.Multichannel Float))
loadBDF fn = do
             bs <- BS.readFile fn
             (bdf,bs') <- runStateT readBDF bs
             m <- case bdf of
                           (Just b) -> do
                                       return $ Just $ M.createMultichannel (head $ samples b) 24 (data_ b)
                           _        -> do
                                       putStrLn "File not read"
                                       return Nothing
             when (not (BS.null bs')) $ do
                  putStrLn "data remaining..."
             return m

-- | Load a BDF data structure
--   NOTE: data has already been converted to Float
loadBDFRaw :: FilePath -> IO (Maybe BDF)
loadBDFRaw fn = do
  bs <- BS.readFile fn
  (bdf,bs') <- runStateT readBDF bs
  m <- case bdf of
        Just b -> return (Just b)
        _      -> do
                  putStrLn "File not read"
                  return Nothing
  when (not (BS.null bs')) $ do
    putStrLn "data remaining..."
  return m

{-
getDPConversion :: BDF -> [Word32 -> Float]
getDPConversion b = let chan = channels b
                        tup4 = zip4 (phys_min b) (phys_max b) (dig_min b) (dig_max b)
                    in map (\(a,b,c,d) -> \x -> ((fromIntegral x) - (fromIntegral c))*(fromIntegral (b-a))/(fromIntegral (d-c))) tup4

writeLine :: [Float] -> String
writeLine = ((++ "\n") . concat . intersperse " " . map show)

main = do
       bs <- BS.readFile "VivianMeditation.bdf"
       (bdf,bs') <- runStateT readBDF bs
       case bdf of
                (Just b) -> do
                            let conv = getDPConversion b
                            let b' = map (zipWith ($) conv) (data_ b)
                            putStrLn $ show $ length $ head b'
                            writeFile "VivianMeditation.txt" $ concat $ map writeLine b'
                _        -> putStrLn "File not read"
       when (not (BS.null bs')) $
            do
            putStrLn "data remaining..."
            return ()
       return ()
-}


module Main where

import Prelude hiding (id, (.))
import Numeric.LinearAlgebra.Array
import Numeric.LinearAlgebra.Array.Util
import System.Random
import Control.Monad
import Control.Lens hiding ((#), para)

import CategoricDefinitions
import Autodiff.D
import Para
import Examples

main :: IO ()
main = run


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ViewPatterns      #-}
{-# LANGUAGE RecordWildCards   #-}
{-# LANGUAGE CPP               #-}
module ArrayFire.Internal.Types where

#include "af/seq.h"
#include "af/complex.h"
#include "af/graphics.h"
#include "af/index.h"

import ArrayFire.Internal.Defines

import Data.Complex
import Data.Proxy
import Data.Word
import Foreign.C.String
import Foreign.C.Types
import Foreign.ForeignPtr
import Foreign.Storable
import GHC.Int

data AFSeq
  = AFSeq
  { afSeqBegin :: {-# UNPACK #-} !Double
  , afSeqEnd   :: {-# UNPACK #-} !Double
  , afSeqStep  :: {-# UNPACK #-} !Double
  } deriving (Show, Eq)

instance Storable AFSeq where
  sizeOf _ = #{size af_seq}
  alignment _ = #{alignment af_seq}
  peek ptr = do
    afSeqBegin <- #{peek af_seq, begin} ptr
    afSeqEnd <- #{peek af_seq, end} ptr
    afSeqStep <- #{peek af_seq, step} ptr
    pure AFSeq {..}
  poke ptr AFSeq{..} = do
    #{poke af_seq, begin} ptr afSeqBegin
    #{poke af_seq, end} ptr afSeqEnd
    #{poke af_seq, step} ptr afSeqStep

-- | Type used for indexing into 'Array'
data AFIndex
  = AFIndex
  { afIdx :: !(Either AFArray AFSeq)
  , afIsSeq :: !Bool
  , afIsBatch :: !Bool
  }
instance Storable AFIndex where
  sizeOf _ = #{size af_index_t}
  alignment _ = #{alignment af_index_t}
  peek ptr = do
    afIsSeq <- #{peek af_index_t, isSeq} ptr
    afIsBatch <- #{peek af_index_t, isBatch} ptr
    afIdx <-
      if afIsSeq
        then Left <$> #{peek af_index_t, idx.arr} ptr
        else Right <$> #{peek af_index_t, idx.seq} ptr
    pure AFIndex{..}
  poke ptr AFIndex{..} = do
    case afIdx of
      Left afarr -> #{poke af_index_t, idx.arr} ptr afarr
      Right afseq -> #{poke af_index_t, idx.seq} ptr afseq
    #{poke af_index_t, isSeq} ptr afIsSeq
    #{poke af_index_t, isBatch} ptr afIsBatch

data AFCFloat
  = AFCFloat
  { afcReal :: {-# UNPACK #-} !Float
  , afcImag :: {-# UNPACK #-} !Float
  } deriving (Eq, Show)

instance Storable AFCFloat where
  sizeOf _ = #{size af_cfloat}
  alignment _ = #{alignment af_cfloat}
  peek ptr = do
    afcReal <- #{peek af_cfloat, real} ptr
    afcImag <- #{peek af_cfloat, imag} ptr
    pure AFCFloat{..}
  poke ptr AFCFloat{..} = do
    #{poke af_cfloat, real} ptr afcReal
    #{poke af_cfloat, imag} ptr afcImag

data AFCell
  = AFCell
  { afCellRow :: {-# UNPACK #-} !Int
  , afCellCol :: {-# UNPACK #-} !Int
  , afCellTitle :: {-# UNPACK #-} !CString
  , afCellColorMap :: {-# UNPACK #-} !AFColorMap
  } deriving (Show, Eq)

instance Storable AFCell where
  sizeOf _ = #{size af_cell}
  alignment _ = #{alignment af_cell}
  peek ptr = do
    afCellRow <- #{peek af_cell, row} ptr
    afCellCol <- #{peek af_cell, col} ptr
    afCellTitle <- #{peek af_cell, title} ptr
    afCellColorMap <- #{peek af_cell, cmap} ptr
    pure AFCell{..}
  poke ptr AFCell{..} = do
    #{poke af_cell, row} ptr afCellRow
    #{poke af_cell, col} ptr afCellCol
    #{poke af_cell, title} ptr afCellTitle
    #{poke af_cell, cmap} ptr afCellColorMap

-- | ArrayFire 'Array'
newtype Array a = Array (ForeignPtr ())

-- | ArrayFire 'Features'
newtype Features = Features (ForeignPtr ())

-- | ArrayFire 'RandomEngine'
newtype RandomEngine = RandomEngine (ForeignPtr ())

-- | ArrayFire 'Window'
newtype Window = Window (ForeignPtr ())

-- | Mapping of Haskell types to ArrayFire types
class Storable a => AFType a where
  afType :: Proxy a -> AFDtype

instance AFType Double where
  afType Proxy = f64

instance AFType Float where
  afType Proxy = f32

instance AFType (Complex Double) where
  afType Proxy = c64

instance AFType (Complex Float) where
  afType Proxy = c32

instance AFType CBool where
  afType Proxy = b8

instance AFType Int32 where
  afType Proxy = s32

instance AFType Word32 where
  afType Proxy = u32

instance AFType Word8 where
  afType Proxy = u8

instance AFType Int64 where
  afType Proxy = s64

instance AFType Int where
  afType Proxy = s64

instance AFType Int16 where
  afType Proxy = s16

instance AFType Word16 where
  afType Proxy = u16

instance AFType Word64 where
  afType Proxy = u64

instance AFType Word where
  afType Proxy = u64

-- | ArrayFire backends
data Backend
  = Default
  | CPU
  | CUDA
  | OpenCL
  deriving (Show, Eq, Ord)

-- | Low-level to high-level Backend conversion
toBackend :: AFBackend -> Backend
toBackend (AFBackend 0) = Default
toBackend (AFBackend 1) = CPU
toBackend (AFBackend 2) = CUDA
toBackend (AFBackend 4) = OpenCL
toBackend (AFBackend x) = error $ "Invalid backend: " <> show x

-- | High-level to low-level Backend conversion
toAFBackend :: Backend -> AFBackend
toAFBackend Default = (AFBackend 0)
toAFBackend CPU = (AFBackend 1)
toAFBackend CUDA = (AFBackend 2)
toAFBackend OpenCL = (AFBackend 4)

-- | Read multiple backends
toBackends :: Int -> [Backend]
toBackends 1 = [CPU]
toBackends 2 = [CUDA]
toBackends 3 = [CPU,CUDA]
toBackends 4 = [OpenCL]
toBackends 5 = [CPU,OpenCL]
toBackends 6 = [CUDA,OpenCL]
toBackends 7 = [CPU,CUDA,OpenCL]
toBackends _ = []

-- | Matrix properties
data MatProp
  = None
  | Trans
  | CTrans
  | Conj
  | Upper
  | Lower
  | DiagUnit
  | Sym
  | PosDef
  | Orthog
  | TriDiag
  | BlockDiag
  deriving (Show, Eq, Ord)

-- | Low-level to High-level 'MatProp' conversion
fromMatProp
  :: AFMatProp
  -> MatProp
fromMatProp (AFMatProp 0) = None
fromMatProp (AFMatProp 1) = Trans
fromMatProp (AFMatProp 2) = CTrans
fromMatProp (AFMatProp 4) = Conj
fromMatProp (AFMatProp 32) = Upper
fromMatProp (AFMatProp 64) = Lower
fromMatProp (AFMatProp 128) = DiagUnit
fromMatProp (AFMatProp 512) = Sym
fromMatProp (AFMatProp 1024) = PosDef
fromMatProp (AFMatProp 2048) = Orthog
fromMatProp (AFMatProp 4096) = TriDiag
fromMatProp (AFMatProp 8192) = BlockDiag
fromMatProp x = error $ "Invalid AFMatProp value: " <> show x

-- | High-level to low-level 'MatProp' conversion
toMatProp
  :: MatProp
  -> AFMatProp
toMatProp None = (AFMatProp 0)
toMatProp Trans = (AFMatProp 1)
toMatProp CTrans = (AFMatProp 2)
toMatProp Conj = (AFMatProp 4)
toMatProp Upper = (AFMatProp 32)
toMatProp Lower = (AFMatProp 64)
toMatProp DiagUnit = (AFMatProp 128)
toMatProp Sym = (AFMatProp 512)
toMatProp PosDef = (AFMatProp 1024)
toMatProp Orthog = (AFMatProp 2048)
toMatProp TriDiag = (AFMatProp 4096)
toMatProp BlockDiag = (AFMatProp 8192)

-- | Binary operation support
data BinaryOp
  = Add
  | Mul
  | Min
  | Max
  deriving (Show, Eq, Ord)

-- | High-level to low-level 'MatProp' conversion
toBinaryOp :: BinaryOp -> AFBinaryOp
toBinaryOp Add = AFBinaryOp 0
toBinaryOp Mul = AFBinaryOp 1
toBinaryOp Min = AFBinaryOp 2
toBinaryOp Max = AFBinaryOp 3

-- | 'BinaryOp' conversion helper
fromBinaryOp :: AFBinaryOp -> BinaryOp
fromBinaryOp (AFBinaryOp 0) = Add
fromBinaryOp (AFBinaryOp 1) = Mul
fromBinaryOp (AFBinaryOp 2) = Min
fromBinaryOp (AFBinaryOp 3) = Max
fromBinaryOp x = error ("Invalid Binary Op: " <> show x)

-- | Storage type used for Sparse arrays
data Storage
  = Dense
  | CSR
  | CSC
  | COO
  deriving (Show, Eq, Ord, Enum)

toStorage :: Storage -> AFStorage
toStorage = AFStorage . fromIntegral . fromEnum

fromStorage :: AFStorage -> Storage
fromStorage (AFStorage (fromIntegral -> x))
  | x `elem` [0..3] = toEnum x
  | otherwise = error $ "Invalid Storage " <> (show x)

-- | Type for different RandomEngines
data RandomEngineType
  = Philox
  | ThreeFry
  | Mersenne
  deriving (Eq, Show)

toRandomEngine :: AFRandomEngineType -> RandomEngineType
toRandomEngine (AFRandomEngineType 100) = Philox
toRandomEngine (AFRandomEngineType 200) = ThreeFry
toRandomEngine (AFRandomEngineType 300) = Mersenne
toRandomEngine (AFRandomEngineType x) =
  error ("Invalid random engine: " <> show x)

fromRandomEngine :: RandomEngineType ->  AFRandomEngineType
fromRandomEngine Philox = (AFRandomEngineType 100)
fromRandomEngine ThreeFry = (AFRandomEngineType 200)
fromRandomEngine Mersenne = (AFRandomEngineType 300)

-- | Interpolation type
data InterpType
  = Nearest
  | Linear
  | Bilinear
  | Cubic
  | LowerInterp
  | LinearCosine
  | BilinearCosine
  | Bicubic
  | CubicSpline
  | BicubicSpline
  deriving (Show, Eq, Ord, Enum)

toInterpType :: AFInterpType -> InterpType
toInterpType (AFInterpType (fromIntegral -> x)) = toEnum x

fromInterpType :: InterpType -> AFInterpType
fromInterpType = AFInterpType . fromIntegral . fromEnum

-- | Border Type
data BorderType
  = PadZero
  | PadSym
  deriving (Show, Ord, Enum, Eq)

toBorderType :: AFBorderType -> BorderType
toBorderType (AFBorderType (fromIntegral -> x)) = toEnum x

fromBorderType :: BorderType -> AFBorderType
fromBorderType = AFBorderType . fromIntegral . fromEnum

-- | Connectivity Type
data Connectivity
  = Conn4
  | Conn8
  deriving (Show, Ord, Enum, Eq)

toConnectivity :: AFConnectivity -> Connectivity
toConnectivity (AFConnectivity 4) = Conn4
toConnectivity (AFConnectivity 8) = Conn4
toConnectivity (AFConnectivity x) = error ("Unknown connectivity option: " <> show x)

fromConnectivity :: Connectivity -> AFConnectivity
fromConnectivity Conn4 = AFConnectivity 4
fromConnectivity Conn8 = AFConnectivity 8

-- | Color Space type
data CSpace
  = Gray
  | RGB
  | HSV
  | YCBCR
  deriving (Show, Eq, Ord, Enum)

toCSpace :: AFCSpace -> CSpace
toCSpace (AFCSpace (fromIntegral -> x)) = toEnum x

fromCSpace :: CSpace -> AFCSpace
fromCSpace = AFCSpace . fromIntegral . fromEnum

-- | YccStd type
data YccStd
  = Ycc601
  | Ycc709
  | Ycc2020
  deriving (Show, Eq, Ord)

toAFYccStd :: AFYccStd -> YccStd
toAFYccStd (AFYccStd 601) = Ycc601
toAFYccStd (AFYccStd 709) = Ycc709
toAFYccStd (AFYccStd 2020) = Ycc2020
toAFYccStd (AFYccStd x) = error ("Unknown AFYccStd option: " <> show x)

fromAFYccStd :: YccStd -> AFYccStd
fromAFYccStd Ycc601 = afYcc601
fromAFYccStd Ycc709 = afYcc709
fromAFYccStd Ycc2020 = afYcc2020

-- | Moment types
data MomentType
  = M00
  | M01
  | M10
  | M11
  | FirstOrder
  deriving (Show, Eq, Ord)

toMomentType :: AFMomentType -> MomentType
toMomentType x
  | x == afMomentM00 = M00
  | x == afMomentM01 = M01
  | x == afMomentM10 = M10
  | x == afMomentM11 = M11
  | x == afMomentFirstOrder = FirstOrder
  | otherwise = error ("Unknown moment type: " <> show x)

fromMomentType :: MomentType -> AFMomentType
fromMomentType M00 = afMomentM00
fromMomentType M01 = afMomentM01
fromMomentType M10 = afMomentM10
fromMomentType M11 = afMomentM11
fromMomentType FirstOrder = afMomentFirstOrder

-- | Canny Theshold type
data CannyThreshold
  = Manual
  | AutoOtsu
  deriving (Show, Eq, Ord, Enum)

toCannyThreshold :: AFCannyThreshold -> CannyThreshold
toCannyThreshold (AFCannyThreshold (fromIntegral -> x)) = toEnum x

fromCannyThreshold :: CannyThreshold -> AFCannyThreshold
fromCannyThreshold = AFCannyThreshold . fromIntegral . fromEnum

-- | Flux function type
data FluxFunction
  = FluxDefault
  | FluxQuadratic
  | FluxExponential
  deriving (Show, Eq, Ord, Enum)

toFluxFunction :: AFFluxFunction -> FluxFunction
toFluxFunction (AFFluxFunction (fromIntegral -> x)) = toEnum x

fromFluxFunction :: FluxFunction -> AFFluxFunction
fromFluxFunction = AFFluxFunction . fromIntegral . fromEnum

-- | Diffusion type
data DiffusionEq
  = DiffusionDefault
  | DiffusionGrad
  | DiffusionMCDE
  deriving (Show, Eq, Ord, Enum)

toDiffusionEq :: AFDiffusionEq -> DiffusionEq
toDiffusionEq (AFDiffusionEq (fromIntegral -> x)) = toEnum x

fromDiffusionEq :: DiffusionEq -> AFDiffusionEq
fromDiffusionEq = AFDiffusionEq . fromIntegral . fromEnum

-- | Iterative deconvolution algo type
data IterativeDeconvAlgo
  = DeconvDefault
  | DeconvLandweber
  | DeconvRichardsonLucy
  deriving (Show, Eq, Ord, Enum)

toIterativeDeconvAlgo :: AFIterativeDeconvAlgo -> IterativeDeconvAlgo
toIterativeDeconvAlgo (AFIterativeDeconvAlgo (fromIntegral -> x)) = toEnum x

fromIterativeDeconvAlgo :: IterativeDeconvAlgo -> AFIterativeDeconvAlgo
fromIterativeDeconvAlgo = AFIterativeDeconvAlgo . fromIntegral . fromEnum

-- | Inverse deconvolution algo type
data InverseDeconvAlgo
  = InverseDeconvDefault
  | InverseDeconvTikhonov
  deriving (Show, Eq, Ord, Enum)

toInverseDeconvAlgo :: AFInverseDeconvAlgo -> InverseDeconvAlgo
toInverseDeconvAlgo (AFInverseDeconvAlgo (fromIntegral -> x)) = toEnum x

fromInverseDeconvAlgo :: InverseDeconvAlgo -> AFInverseDeconvAlgo
fromInverseDeconvAlgo = AFInverseDeconvAlgo . fromIntegral . fromEnum

-- | Cell type, used in Graphics module
data Cell
  = Cell
  { cellRow :: Int
  , cellCol :: Int
  , cellTitle :: String
  , cellColorMap :: ColorMap
  } deriving (Show, Eq)

cellToAFCell :: Cell -> IO AFCell
cellToAFCell Cell {..} =
  withCString cellTitle $ \cstr ->
    pure AFCell { afCellRow = cellRow
                , afCellCol = cellCol
                , afCellTitle = cstr
                , afCellColorMap = fromColorMap cellColorMap
                }

-- | ColorMap type
data ColorMap
  = ColorMapDefault
  | ColorMapSpectrum
  | ColorMapColors
  | ColorMapRed
  | ColorMapMood
  | ColorMapHeat
  | ColorMapBlue
  | ColorMapInferno
  | ColorMapMagma
  | ColorMapPlasma
  | ColorMapViridis
  deriving (Show, Eq, Ord, Enum)

fromColorMap :: ColorMap -> AFColorMap
fromColorMap = AFColorMap . fromIntegral . fromEnum

toColorMap :: AFColorMap -> ColorMap
toColorMap (AFColorMap (fromIntegral -> x)) = toEnum x

-- | Marker type
data MarkerType
  = MarkerTypeNone
  | MarkerTypePoint
  | MarkerTypeCircle
  | MarkerTypeSquare
  | MarkerTypeTriangle
  | MarkerTypeCross
  | MarkerTypePlus
  | MarkerTypeStar
  deriving (Show, Eq, Ord, Enum)

fromMarkerType :: MarkerType -> AFMarkerType
fromMarkerType = AFMarkerType . fromIntegral . fromEnum

toMarkerType :: AFMarkerType -> MarkerType
toMarkerType (AFMarkerType (fromIntegral -> x)) = toEnum x

-- | Match type
data MatchType
  = MatchTypeSAD
  | MatchTypeZSAD
  | MatchTypeLSAD
  | MatchTypeSSD
  | MatchTypeZSSD
  | MatchTypeLSSD
  | MatchTypeNCC
  | MatchTypeZNCC
  | MatchTypeSHD
  deriving (Show, Eq, Ord, Enum)

fromMatchType :: MatchType -> AFMatchType
fromMatchType = AFMatchType . fromIntegral . fromEnum

toMatchType :: AFMatchType -> MatchType
toMatchType (AFMatchType (fromIntegral -> x)) = toEnum x

-- | TopK type
data TopK
  = TopKDefault
  | TopKMin
  | TopKMax
  deriving (Show, Eq, Ord, Enum)

fromTopK :: TopK -> AFTopkFunction
fromTopK = AFTopkFunction . fromIntegral . fromEnum

toTopK :: AFTopkFunction -> TopK
toTopK (AFTopkFunction (fromIntegral -> x)) = toEnum x

-- | Homography Type
data HomographyType
  = RANSAC
  | LMEDS
  deriving (Show, Eq, Ord, Enum)

fromHomographyType :: HomographyType -> AFHomographyType
fromHomographyType = AFHomographyType . fromIntegral . fromEnum

toHomographyType :: AFHomographyType -> HomographyType
toHomographyType (AFHomographyType (fromIntegral -> x)) = toEnum x

-- | Sequence Type
data Seq
  = Seq
  { seqBegin :: !Double
  , seqEnd :: !Double
  , seqStep :: !Double
  } deriving (Show, Eq, Ord)

toAFSeq :: Seq -> AFSeq
toAFSeq (Seq x y z) = (AFSeq x y z)

-- | Index Type
data Index a
  = Index
  { idx :: Either (Array a) Seq
  , isSeq :: !Bool
  , isBatch :: !Bool
  }

seqIdx :: Seq -> Bool -> Index a
seqIdx s = Index (Right s) True

arrIdx :: Array a -> Bool -> Index a
arrIdx a = Index (Left a) False

toAFIndex :: Index a -> IO AFIndex
toAFIndex (Index a b c) = do
  case a of
    Right s -> pure $ AFIndex (Right (toAFSeq s)) b c
    Left (Array fptr) -> do
      withForeignPtr fptr $ \ptr ->
        pure $ AFIndex (Left ptr) b c

-- | Type alias for ArrayFire API version
type Version = (Int,Int,Int)

-- | Norm Type
data NormType
  = NormVectorOne
  -- ^ treats the input as a vector and returns the sum of absolute values
  | NormVectorInf
  -- ^ treats the input as a vector and returns the max of absolute values
  | NormVector2
  -- ^ treats the input as a vector and returns euclidean norm
  | NormVectorP
  -- ^ treats the input as a vector and returns the p-norm
  | NormMatrix1
  -- ^ return the max of column sums
  | NormMatrixInf
  -- ^ return the max of row sums
  | NormMatrix2
  -- ^ returns the max singular value). Currently NOT SUPPORTED
  | NormMatrixLPQ
  -- ^ returns Lpq-norm
  | NormEuclid
  -- ^ The default. Same as AF_NORM_VECTOR_2
  deriving (Show, Eq, Enum)

fromNormType :: NormType -> AFNormType
fromNormType = AFNormType . fromIntegral . fromEnum

toNormType :: AFNormType -> NormType
toNormType (AFNormType (fromIntegral -> x)) = toEnum x

-- | Convolution Domain
data ConvDomain
  = ConvDomainAuto
  -- ^ ArrayFire automatically picks the right convolution algorithm
  | ConvDomainSpatial
  -- ^ Perform convolution in spatial domain
  | ConvDomainFreq
  -- ^ Perform convolution in frequency domain
  deriving (Show, Eq, Enum)

-- | Convolution Mode
data ConvMode
  = ConvDefault
  -- ^ Output of the convolution is the same size as input
  | ConvExpand
  -- ^ Output of the convolution is signal_len + filter_len - 1
  deriving (Show, Eq, Enum)

fromConvDomain :: ConvDomain -> AFConvDomain
fromConvDomain = AFConvDomain . fromIntegral . fromEnum

toConvDomain :: AFConvDomain -> ConvDomain
toConvDomain (AFConvDomain (fromIntegral -> x)) = toEnum x

fromConvMode :: AFConvMode -> ConvMode
fromConvMode (AFConvMode (fromIntegral -> x)) = toEnum x

toConvMode :: ConvMode -> AFConvMode
toConvMode = AFConvMode . fromIntegral . fromEnum

-- | Array Fire types
data AFDType
  = F32
  | C32
  | F64
  | C64
  | B8
  | S32
  | U32
  | U8
  | S64
  | U64
  | S16
  | U16
  deriving (Show, Eq, Enum)

fromAFType :: AFDtype -> AFDType
fromAFType (AFDtype (fromIntegral -> x)) = toEnum x

toAFType :: AFDType -> AFDtype
toAFType = AFDtype . fromIntegral . fromEnum


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
module Data.Array.Accelerate.Utility.Lift.Exp (
   Unlift,
   Unlifted,
   Tuple,
   unlift,
   modify,
   modify2,
   modify3,
   modify4,
   Exp(Exp), expr, atom,
   unliftPair,
   unliftTriple,
   unliftQuadruple,
   asExp,
   mapFst,
   mapSnd,
   fst3,
   snd3,
   thd3,
   indexCons,
   ) where

import qualified Data.Array.Accelerate.Data.Complex as Complex
import qualified Data.Array.Accelerate as A
import Data.Complex (Complex((:+)))
import Data.Array.Accelerate ((:.)((:.)))

import qualified Data.Tuple.HT as Tuple
import Data.Tuple.HT (mapTriple)

{- |
This class simplifies untupling of expressions.
If you have a function

> g :: ((Exp a, Exp b), Exp (c,d)) -> (Exp e, Exp f)

you cannot apply it to an array @arr :: Array sh ((a,b),(c,d))@ using 'A.map'.
Here, the 'modify' function helps:

> modify ((expr,expr),expr) g :: Exp ((a,b),(c,d)) -> Exp (e,f)

The 'expr'-pattern tells, how deep the tuple shall be unlifted.
This way you can write:

> A.map
>    (Exp.modify ((expr,expr),expr) $ \((a,b), cd) -> g ((a,b), cd))
>    arr

'modify' is based on 'unlift'.
In contrast to 'A.unlift' it does not only unlift one level of tupels,
but is guided by an 'expr'-pattern.
In the example I have demonstrated,
how the pair @(a,b)@ is unlifted, but the pair @(c,d)@ is not.
For the result tuple, 'modify' simply calls 'A.lift'.
In contrast to 'A.unlift',
'A.lift' lifts over all tupel levels until it obtains a single 'Exp'.
-}
class
   (A.Elt (Tuple pattern), A.Plain (Unlifted pattern) ~ Tuple pattern) =>
      Unlift pattern where
   type Unlifted pattern
   type Tuple pattern
   unlift :: pattern -> A.Exp (Tuple pattern) -> Unlifted pattern

modify ::
   (A.Lift A.Exp a, Unlift pattern) =>
   pattern ->
   (Unlifted pattern -> a) ->
   A.Exp (Tuple pattern) -> A.Exp (A.Plain a)
modify p f = A.lift . f . unlift p

modify2 ::
   (A.Lift A.Exp a, Unlift patternA, Unlift patternB) =>
   patternA ->
   patternB ->
   (Unlifted patternA -> Unlifted patternB -> a) ->
   A.Exp (Tuple patternA) -> A.Exp (Tuple patternB) -> A.Exp (A.Plain a)
modify2 pa pb f a b = A.lift $ f (unlift pa a) (unlift pb b)

modify3 ::
   (A.Lift A.Exp a, Unlift patternA, Unlift patternB, Unlift patternC) =>
   patternA ->
   patternB ->
   patternC ->
   (Unlifted patternA -> Unlifted patternB -> Unlifted patternC -> a) ->
   A.Exp (Tuple patternA) -> A.Exp (Tuple patternB) ->
   A.Exp (Tuple patternC) -> A.Exp (A.Plain a)
modify3 pa pb pc f a b c =
   A.lift $ f (unlift pa a) (unlift pb b) (unlift pc c)

modify4 ::
   (A.Lift A.Exp a,
    Unlift patternA, Unlift patternB, Unlift patternC, Unlift patternD) =>
   patternA ->
   patternB ->
   patternC ->
   patternD ->
   (Unlifted patternA -> Unlifted patternB ->
    Unlifted patternC -> Unlifted patternD -> a) ->
   A.Exp (Tuple patternA) -> A.Exp (Tuple patternB) ->
   A.Exp (Tuple patternC) -> A.Exp (Tuple patternD) -> A.Exp (A.Plain a)
modify4 pa pb pc pd f a b c d =
   A.lift $ f (unlift pa a) (unlift pb b) (unlift pc c) (unlift pd d)

instance (A.Elt a) => Unlift (Exp a) where
   type Unlifted (Exp a) = A.Exp a
   type Tuple (Exp a) = a
   unlift _ = id

data Exp e = Exp

expr :: Exp e
expr = Exp

{-# DEPRECATED atom "use expr instead" #-}
-- | for compatibility with accelerate-utility-0.0
atom :: Exp e
atom = expr

instance (Unlift pa, Unlift pb) => Unlift (pa,pb) where
   type Unlifted (pa,pb) = (Unlifted pa, Unlifted pb)
   type Tuple (pa,pb) = (Tuple pa, Tuple pb)
   unlift (pa,pb) ab =
      (unlift pa $ A.fst ab, unlift pb $ A.snd ab)

instance
   (Unlift pa, Unlift pb, Unlift pc) =>
      Unlift (pa,pb,pc) where
   type Unlifted (pa,pb,pc) = (Unlifted pa, Unlifted pb, Unlifted pc)
   type Tuple (pa,pb,pc) = (Tuple pa, Tuple pb, Tuple pc)
   unlift (pa,pb,pc) =
      mapTriple (unlift pa, unlift pb, unlift pc) . A.unlift

instance (Unlift pa, A.Slice (Tuple pa), int ~ Exp Int) => Unlift (pa :. int) where
   type Unlifted (pa :. int) = Unlifted pa :. A.Exp Int
   type Tuple (pa :. int) = Tuple pa :. Int
   unlift (pa:.pb) ab =
      (unlift pa $ A.indexTail ab) :. (unlift pb $ A.indexHead ab)

instance (Unlift p) => Unlift (Complex p) where
   type Unlifted (Complex p) = Complex (Unlifted p)
   type Tuple (Complex p) = Complex (Tuple p)
   unlift (preal:+pimag) z =
      unlift preal (Complex.real z)
      :+
      unlift pimag (Complex.imag z)

unliftPair :: (A.Elt a, A.Elt b) => A.Exp (a,b) -> (A.Exp a, A.Exp b)
unliftPair = A.unlift

unliftTriple ::
   (A.Elt a, A.Elt b, A.Elt c) => A.Exp (a,b,c) -> (A.Exp a, A.Exp b, A.Exp c)
unliftTriple = A.unlift

unliftQuadruple ::
   (A.Elt a, A.Elt b, A.Elt c, A.Elt d) =>
   A.Exp (a,b,c,d) -> (A.Exp a, A.Exp b, A.Exp c, A.Exp d)
unliftQuadruple = A.unlift

asExp :: A.Exp a -> A.Exp a
asExp = id

mapFst ::
   (A.Elt a, A.Elt b, A.Elt c) =>
   (A.Exp a -> A.Exp b) -> A.Exp (a,c) -> A.Exp (b,c)
mapFst f = modify (expr,expr) $ \(a,c) -> (f a, c)

mapSnd ::
   (A.Elt a, A.Elt b, A.Elt c) =>
   (A.Exp b -> A.Exp c) -> A.Exp (a,b) -> A.Exp (a,c)
mapSnd f = modify (expr,expr) $ \(a,b) -> (a, f b)

fst3 ::
   (A.Elt a, A.Elt b, A.Elt c) =>
   A.Exp (a,b,c) -> A.Exp a
fst3 = modify (expr,expr,expr) Tuple.fst3

snd3 ::
   (A.Elt a, A.Elt b, A.Elt c) =>
   A.Exp (a,b,c) -> A.Exp b
snd3 = modify (expr,expr,expr) Tuple.snd3

thd3 ::
   (A.Elt a, A.Elt b, A.Elt c) =>
   A.Exp (a,b,c) -> A.Exp c
thd3 = modify (expr,expr,expr) Tuple.thd3

indexCons ::
   (A.Slice ix) => A.Exp ix -> A.Exp Int -> A.Exp (ix :. Int)
indexCons ix n = A.lift $ ix:.n