{-# LANGUAGE Safe #-}
-- | Import all modules from prelude that should be safe
module Main where

import Numeric
import Prelude
-- import Foreign

import Control.Applicative
import Control.Arrow
import Control.Category

-- import Control.ST
-- import Control.ST.Lazy
-- import Control.ST.Strict

import Control.Concurrent
import Control.Concurrent.Chan
import Control.Concurrent.MVar

import Control.Exception
import Control.Exception.Base

import Control.Monad
import Control.Monad.Fix
import Control.Monad.Zip

import Data.Bits
import Data.Bool
import Data.Char
import Data.Complex
import Data.Either
import Data.Eq
import Data.Fixed
import Data.Foldable
import Data.Function
import Data.Functor
import Data.IORef
import Data.Int
import Data.Ix
import Data.List
import Data.Maybe
import Data.Monoid
import Data.Ord
import Data.Ratio
import Data.String
import Data.Traversable
import Data.Tuple
import Data.Typeable
import Data.Unique
import Data.Version
import Data.Word

import Data.STRef
import Data.STRef.Lazy
import Data.STRef.Strict

-- import Debug.Trace

import Foreign.Concurrent
-- import Foreign.ForeignPtr
import Foreign.Ptr
import Foreign.StablePtr
import Foreign.Storable

import Foreign.C
import Foreign.C.Error
import Foreign.C.String
import Foreign.C.Types

-- import Foreign.Marshal
import Foreign.Marshal.Alloc
import Foreign.Marshal.Array
import Foreign.Marshal.Error
import Foreign.Marshal.Pool
import Foreign.Marshal.Utils

import System.CPUTime
import System.Environment
import System.Exit
import System.Info
import System.Mem
import System.Timeout

import System.Console.GetOpt

import System.IO
import System.IO.Error

import System.Mem.StableName
import System.Mem.Weak

import System.Posix.Internals
import System.Posix.Types

import Text.Printf

import Text.Read
import Text.Read.Lex

import Text.Show
import Text.Show.Functions

import Text.ParserCombinators.ReadP
import Text.ParserCombinators.ReadPrec

-- import Unsafe.Coerce

f :: Int
f = 2

main :: IO ()
main = putStrLn $ "X is: " ++ show f


{-# LANGUAGE RecordWildCards #-}

module Arpack.Foreign.Class where

import Data.Vector.Storable.Mutable (IOVector)
import Numeric.LinearAlgebra (Matrix, Vector)

import Arpack.Options

class Arpack t where
  arpack :: Options t -> Int -> (IOVector t -> IOVector t -> IO ())
         -> IO (Vector t, Matrix t)


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Statistics.BBVI.Observed
  ( Obs(..)
  , defaultObs
  )
where

import qualified Data.Vector                   as V
import           Statistics.BBVI.Class
import           Statistics.BBVI.Propagator     ( DistCell(..)
                                                , SampleVector
                                                , SampleDouble
                                                )
import           System.Random.MWC              ( uniformR )

-- | Dummy distribution representing a vector of observations, for use
-- with building "observation" distribution cells.  these can be used
-- to easily/selectively subsample (and resample) observations using
-- existing typeclass methods (useful with stochatic gradient
-- updates).
newtype Obs a = O (V.Vector a) deriving (Show, Eq, Ord, Read)

-- | helper function to build a distribution cell of observations
defaultObs
  :: V.Vector a -- ^ observations
  -> DistCell (Obs a)
defaultObs d = (Node 1 V.empty (O d))

instance DistUtil (Obs Double) where
  nParams _x = 0
  toParamVector _ = V.empty
  fromParamVector _ = O V.empty

instance Dist (Obs Double) SampleDouble where
  resample (O d) gen = return . (d V.!) =<< uniformR (0, (V.length d - 1)) gen
  logProb _d _x = 0
  paramGradOfLogQ _d _x = V.empty

instance DistUtil (Obs SampleVector) where
  nParams _x = 0
  toParamVector _ = V.empty
  fromParamVector _ = O V.empty

instance Dist (Obs SampleVector) SampleVector where
  resample (O d) gen = return . (d V.!) =<< uniformR (0, (V.length d - 1)) gen
  logProb _d _x = 0
  paramGradOfLogQ _d _x = V.empty


{-# LANGUAGE CPP
           , GADTs
           , KindSignatures
           , TypeOperators
           , TypeFamilies
           , EmptyCase
           , DataKinds
           , PolyKinds
           , ExistentialQuantification
           , FlexibleContexts
           , OverloadedStrings
           #-}

{-# OPTIONS_GHC -Wall -fwarn-tabs #-}

module Language.Hakaru.Sample where

import           Numeric.SpecFunctions            (logFactorial)
import qualified Data.Number.LogFloat             as LF
import qualified Math.Combinatorics.Exact.Binomial as EB
-- import qualified Numeric.Integration.TanhSinh     as TS
import qualified System.Random.MWC                as MWC
import qualified System.Random.MWC.CondensedTable as MWC
import qualified System.Random.MWC.Distributions  as MWCD

import qualified Data.Vector                      as V
import           Data.STRef
import           Data.Sequence (Seq)
import qualified Data.Foldable                    as F
import qualified Data.List.NonEmpty               as L
import           Data.List.NonEmpty               (NonEmpty(..))
import           Data.Maybe                       (fromMaybe)

#if __GLASGOW_HASKELL__ < 710
import           Control.Applicative   (Applicative(..), (<$>))
#endif
import           Control.Monad
import           Control.Monad.ST
import           Control.Monad.Identity
import           Control.Monad.Trans.Maybe
import           Control.Monad.State.Strict
import qualified Data.IntMap                      as IM

import Data.Number.Nat     (fromNat)
import Data.Number.Natural (fromNatural, fromNonNegativeRational, Natural, unsafeNatural)
import Language.Hakaru.Types.DataKind
import Language.Hakaru.Types.Coercion
import Language.Hakaru.Types.Sing
import Language.Hakaru.Types.HClasses
import Language.Hakaru.Syntax.IClasses
import Language.Hakaru.Syntax.TypeOf
import Language.Hakaru.Syntax.Value
import Language.Hakaru.Syntax.Reducer
import Language.Hakaru.Syntax.Datum
import Language.Hakaru.Syntax.DatumCase
import Language.Hakaru.Syntax.AST
import Language.Hakaru.Syntax.ABT

data EAssoc =
    forall a. EAssoc {-# UNPACK #-} !(Variable a) !(Value a)

newtype Env = Env (IM.IntMap EAssoc)

emptyEnv :: Env
emptyEnv = Env IM.empty

updateEnv :: EAssoc -> Env -> Env
updateEnv v@(EAssoc x _) (Env xs) =
    Env $ IM.insert (fromNat $ varID x) v xs

updateEnvs
    :: List1 Variable xs
    -> List1 Value xs
    -> Env
    -> Env
updateEnvs Nil1         Nil1         env = env
updateEnvs (Cons1 x xs) (Cons1 y ys) env =
    updateEnvs xs ys (updateEnv (EAssoc x y) env)

lookupVar :: Variable a -> Env -> Maybe (Value a)
lookupVar x (Env env) = do
    EAssoc x' e' <- IM.lookup (fromNat $ varID x) env
    Refl         <- varEq x x'
    return e'

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

-- Makes use of Atkinson's algorithm as described in:
-- Monte Carlo Statistical Methods pg. 55
--
-- Further discussion at:
-- http://www.johndcook.com/blog/2010/06/14/generating-poisson-random-values/
poisson_rng :: Double -> MWC.GenIO -> IO Int
poisson_rng lambda g' = make_poisson g'
    where
    smu   = sqrt lambda
    b     = 0.931 + 2.53*smu
    a     = -0.059 + 0.02483*b
    vr    = 0.9277 - 3.6224/(b - 2)
    arep  = 1.1239 + 1.1368/(b - 3.4)
    lnlam = log lambda

    make_poisson :: MWC.GenIO -> IO Int
    make_poisson g = do
        u <- MWC.uniformR (-0.5,0.5) g
        v <- MWC.uniformR (0,1) g
        let us = 0.5 - abs u
            k = floor $ (2*a / us + b)*u + lambda + 0.43
        case () of
            () | us >= 0.07 && v <= vr -> return k
            () | k < 0                 -> make_poisson g
            () | us <= 0.013 && v > us -> make_poisson g
            () | accept_region us v k  -> return k
            _                          -> make_poisson g

    accept_region :: Double -> Double -> Int -> Bool
    accept_region us v k =
        log (v * arep / (a/(us*us)+b))
        <=
        -lambda + fromIntegral k * lnlam - logFactorial k

normalize :: [Value 'HProb] -> (LF.LogFloat, Double, [Double])
normalize []          = (0, 0, [])
normalize [(VProb x)] = (x, 1, [1])
normalize xs          = (m, y, ys)
    where
    xs' = map (\(VProb x) -> x) xs
    m   = maximum xs'
    ys  = [ LF.fromLogFloat (x/m) | x <- xs' ]
    y   = sum ys

normalizeVector
    :: Value ('HArray 'HProb) -> (LF.LogFloat, Double, V.Vector Double)
normalizeVector (VArray xs) =
    let xs' = V.map (\(VProb x) -> x) xs in
    case V.length xs of
    0 -> (0, 0, V.empty)
    1 -> (V.unsafeHead xs', 1, V.singleton 1)
    _ ->
        let m   = V.maximum xs'
            ys  = V.map (\x -> LF.fromLogFloat (x/m)) xs'
            y   = V.sum ys
        in (m, y, ys)

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

runEvaluate
    :: (ABT Term abt)
    => abt '[] a
    -> Value a
runEvaluate prog = evaluate prog emptyEnv

evaluate
    :: (ABT Term abt)
    => abt '[] a
    -> Env
    -> Value a
evaluate e env = caseVarSyn e (evaluateVar env) (flip evaluateTerm env)

evaluateVar :: Env -> Variable a -> Value a
evaluateVar env v =
    case lookupVar v env of
    Nothing -> error "variable not found!"
    Just a  -> a

evaluateTerm
    :: (ABT Term abt)
    => Term abt a
    -> Env
    -> Value a
evaluateTerm t env =
    case t of
    o :$          es -> evaluateSCon    o es    env
    NaryOp_  o    es -> evaluateNaryOp  o es    env
    Literal_ v       -> evaluateLiteral v
    Empty_   _       -> evaluateEmpty
    Array_   n    es -> evaluateArray   n es    env
    ArrayLiteral_ es -> VArray . V.fromList $ map (flip evaluate env) es
    Bucket b e    rs -> evaluateBucket  b e  rs env
    Datum_   d       -> evaluateDatum   d       env
    Case_    o    es -> evaluateCase    o es    env
    Superpose_    es -> evaluateSuperpose es    env
    Reject_  _       -> VMeasure $ \_ _ -> return Nothing

evaluateSCon
    :: (ABT Term abt)
    => SCon args a
    -> SArgs abt args
    -> Env
    -> Value a
evaluateSCon Lam_ (e1 :* End) env =
    caseBind e1 $ \x e1' ->
        VLam $ \v -> evaluate e1' (updateEnv (EAssoc x v) env)
evaluateSCon App_ (e1 :* e2 :* End) env =
    case evaluate e1 env of
    VLam f -> f (evaluate e2 env)
evaluateSCon Let_ (e1 :* e2 :* End) env =
    let v = evaluate e1 env
    in caseBind e2 $ \x e2' ->
        evaluate e2' (updateEnv (EAssoc x v) env)
evaluateSCon (CoerceTo_   c) (e1 :* End) env =
    coerceTo c $ evaluate e1 env
evaluateSCon (UnsafeFrom_ c) (e1 :* End) env =
    coerceFrom c $ evaluate e1 env
evaluateSCon (PrimOp_ o)     es env = evaluatePrimOp    o es env
evaluateSCon (ArrayOp_ o)    es env = evaluateArrayOp   o es env
evaluateSCon (MeasureOp_  m) es env = evaluateMeasureOp m es env
evaluateSCon Dirac           (e1 :* End) env =
    VMeasure $ \p _ -> return $ Just (evaluate e1 env, p)
evaluateSCon MBind (e1 :* e2 :* End) env =
    case evaluate e1 env of
    VMeasure m1 -> VMeasure $ \ p g -> do
        x <- m1 p g
        case x of
            Nothing -> return Nothing
            Just (a, p') ->
                caseBind e2 $ \x' e2' ->
                    case evaluate e2' (updateEnv (EAssoc x' a) env) of
                    VMeasure y -> y p' g

evaluateSCon Plate (n :* e2 :* End) env =
    case evaluate n env of
    VNat n' -> caseBind e2 $ \x e' ->
        VMeasure $ \(VProb p) g -> runMaybeT $ do
            (v', ps) <- fmap V.unzip . V.mapM (performMaybe g) $
                V.generate (fromInteger $ fromNatural n') $ \v ->
                    evaluate e' $
                    updateEnv (EAssoc x . VNat $ intToNatural v) env
            return
                ( VArray v'
                , VProb $ p * V.product (V.map (\(VProb y) -> y) ps)
                )
    where
    performMaybe
        :: MWC.GenIO
        -> Value ('HMeasure a)
        -> MaybeT IO (Value a, Value 'HProb)
    performMaybe g (VMeasure m) = MaybeT $ m (VProb 1) g

evaluateSCon Chain (n :* s :* e :* End) env =
    case (evaluate n env, evaluate s env) of
    (VNat n', start) ->
        caseBind e $ \x e' ->
            let s' = VLam $ \v -> evaluate e' (updateEnv (EAssoc x v) env) in
            VMeasure (\(VProb p) g -> runMaybeT $ do
                (evaluates, sout) <- runStateT (replicateM (unsafeInt n') $ convert g s') start
                let (v', ps) = unzip evaluates
                    bodyType :: Sing ('HMeasure (HPair a b)) -> Sing ('HArray a)
                    bodyType = SArray . fst . sUnPair . sUnMeasure
                return
                    ( VDatum $ dPair_ (bodyType $ caseBind e (const typeOf)) (typeOf s)
                        (VArray . V.fromList $ v') sout
                    , VProb $ p * product (map (\(VProb y) -> y) ps)
                    ))
    where
    convert
        :: MWC.GenIO
        -> Value (s ':-> 'HMeasure (HPair a s))
        -> StateT (Value s) (MaybeT IO) (Value a, Value 'HProb)
    convert g (VLam f) = StateT $ \s' ->
        case f s' of
        VMeasure f' -> do
            (as'', p') <- MaybeT (f' (VProb 1) g)
            let (a, s'') = unPair as''
            return ((a, p'), s'')

    unPair :: Value (HPair a b) -> (Value a, Value b)
    unPair (VDatum (Datum "pair" _typ
        (Inl (Et (Konst a)
            (Et (Konst b) Done))))) = (a, b)
    unPair x = case x of {}

evaluateSCon (Summate hd hs) (e1 :* e2 :* e3 :* End) env =
    case (evaluate e1 env, evaluate e2 env) of
    (lo, hi) ->
        caseBind e3 $ \x e3' ->
            foldl (\t i ->
                   evalOp (Sum  hs) t $
                     evaluate e3' (updateEnv (EAssoc x i) env))
                  (identityElement $ Sum hs)
                  (enumFromUntilValue hd lo hi)

evaluateSCon (Product hd hs) (e1 :* e2 :* e3 :* End) env =
    case (evaluate e1 env, evaluate e2 env) of
    (lo, hi) ->
        caseBind e3 $ \x e3' ->
            foldl (\t i ->
                   evalOp (Prod hs) t $
                     evaluate e3' (updateEnv (EAssoc x i) env))
                  (identityElement $ Prod hs)
                  (enumFromUntilValue hd lo hi)

evaluateSCon s _ _ = error $ "TODO: evaluateSCon{" ++ show s ++ "}"

evaluatePrimOp
    ::  ( ABT Term abt, typs ~ UnLCs args, args ~ LCs typs)
    => PrimOp typs a
    -> SArgs abt args
    -> Env
    -> Value a
evaluatePrimOp Not (e1 :* End) env =
    case evaluate e1 env of
      VDatum a -> if a == dTrue
                  then VDatum dFalse
                  else VDatum dTrue

evaluatePrimOp Pi  End         _   = VProb . LF.logFloat $ pi
evaluatePrimOp Cos (e1 :* End) env =
    case evaluate e1 env of
      VReal v1 -> VReal . cos $ v1

evaluatePrimOp Sin (e1 :* End) env =
    case evaluate e1 env of
      VReal v1 -> VReal . sin $ v1

evaluatePrimOp Tan (e1 :* End) env =
    case evaluate e1 env of
      VReal v1 -> VReal . tan $ v1

evaluatePrimOp RealPow (e1 :* e2 :* End) env =
    case (evaluate e1 env, evaluate e2 env) of
      (VProb v1, VReal v2) -> VProb $ LF.pow v1 v2

evaluatePrimOp Choose (e1 :* e2 :* End) env =
    case (evaluate e1 env, evaluate e2 env) of
      (VNat v1, VNat v2) -> VNat $ EB.choose v1 v2

evaluatePrimOp Exp (e1 :* End) env =
    case evaluate e1 env of
      VReal v1 -> VProb . LF.logToLogFloat $ v1

evaluatePrimOp Log (e1 :* End) env =
    case evaluate e1 env of
      VProb v1 -> VReal . LF.logFromLogFloat $ v1

evaluatePrimOp (Infinity h) End _ =
    case h of
      HIntegrable_Nat  -> error "Can not evaluate infinity for natural numbers"
      HIntegrable_Prob -> VProb $ LF.logFloat LF.infinity

evaluatePrimOp (Equal _) (e1 :* e2 :* End) env = (VDatum . dBool) $ evaluate e1 env == evaluate e2 env

evaluatePrimOp (Less _) (e1 :* e2 :* End) env =
    case (evaluate e1 env, evaluate e2 env) of
    (VNat  v1, VNat  v2) -> VDatum $ if v1 < v2 then dTrue else dFalse
    (VInt  v1, VInt  v2) -> VDatum $ if v1 < v2 then dTrue else dFalse
    (VProb v1, VProb v2) -> VDatum $ if v1 < v2 then dTrue else dFalse
    (VReal v1, VReal v2) -> VDatum $ if v1 < v2 then dTrue else dFalse
    _                    -> error "TODO: evaluatePrimOp{Less}"
evaluatePrimOp (NatPow _) (e1 :* e2 :* End) env =
    case evaluate e2 env of
    VNat  v2 ->
        let v2' = fromNatural v2 in
        case evaluate e1 env of
          VNat  v1 -> VNat  (v1 ^ v2')
          VInt  v1 -> VInt  (v1 ^ v2')
          VProb v1 -> VProb (v1 ^ v2')
          VReal v1 -> VReal (v1 ^ v2')
          _        -> error "NatPow should always return some kind of number"
evaluatePrimOp (Negate _) (e1 :* End) env =
    case evaluate e1 env of
    VInt  v -> VInt  (negate v)
    VReal v -> VReal (negate v)
    v       -> case v of {}
evaluatePrimOp (Abs   _) (e1 :* End) env =
    case evaluate e1 env of
    VInt  v -> VNat  . unsafeNatural   $ abs v
    VReal v -> VProb . LF.logFloat $ abs v
    v       -> case v of {}
evaluatePrimOp (Recip _) (e1 :* End) env =
    case evaluate e1 env of
    VProb v -> VProb (recip v)
    VReal v -> VReal (recip v)
    v       -> case v of {}
evaluatePrimOp (NatRoot _) (e1 :* e2 :* End) env =
    case (evaluate e1 env, evaluate e2 env) of
    (VProb v1, VNat v2) -> VProb $ LF.pow v1 (recip . fromIntegral $ v2)
    v                   -> case v of {}

evaluatePrimOp (Floor) (e1 :* End) env =
    case (evaluate e1 env) of
    VProb v1 -> VNat (floor (LF.fromLogFloat v1))

evaluatePrimOp prim _ _ =
    error ("TODO: evaluatePrimOp{" ++ show prim ++ "}")

evaluateArrayOp
    :: ( ABT Term abt
       , typs ~ UnLCs args
       , args ~ LCs typs)
    => ArrayOp typs a
    -> SArgs abt args
    -> Env
    -> Value a
evaluateArrayOp (Index _) = \(e1 :* e2 :* End) env ->
    case (evaluate e1 env, evaluate e2 env) of
    (VArray v, VNat n) -> v V.! unsafeInt n

evaluateArrayOp (Size _) = \(e1 :* End) env ->
    case evaluate e1 env of
    VArray v -> VNat . intToNatural $ V.length v

evaluateArrayOp (Reduce _) = \(e1 :* e2 :* e3 :* End) env ->
    case ( evaluate e1 env
         , evaluate e2 env
         , evaluate e3 env) of
    (f, a, VArray v) -> V.foldl' (lam2 f) a v

evaluateMeasureOp
    :: ( ABT Term abt
       , typs ~ UnLCs args
       , args ~ LCs typs)
    => MeasureOp typs a
    -> SArgs abt args
    -> Env
    -> Value ('HMeasure a)

evaluateMeasureOp Lebesgue = \(e1 :* e2 :* End) env ->
  case (evaluate e1 env, evaluate e2 env) of
    (VReal v1, VReal v2) | v1 < v2 ->
      VMeasure $ \(VProb p) g ->
        case (isInfinite v1, isInfinite v2) of
          (False, False) -> do
            x <- MWC.uniformR (v1, v2) g
            return $ Just (VReal $ x,
                           VProb $ p * LF.logFloat (v2 - v1))
          (False, True) -> do
            u <- MWC.uniform g
            let l = log u
            let n = -l
            return $ Just (VReal $ v1 + n,
                           VProb $ p * LF.logToLogFloat n)
          (True, False) -> do
            u <- MWC.uniform g
            let l = log u
            let n = -l
            return $ Just (VReal $ v2 - n,
                           VProb $ p * LF.logToLogFloat n)
          (True, True) -> do
            (u,b) <- MWC.uniform g
            let l = log u
            let n = -l
            return $ Just (VReal $ if b then n else l,
                           VProb $ p * 2 * LF.logToLogFloat n)
    (VReal _, VReal _) -> error "Lebesgue with length 0 or flipped endpoints"

evaluateMeasureOp Counting = \End _ ->
    VMeasure $ \(VProb p) g -> do
        let success = LF.logToLogFloat (-3 :: Double)
        let pow x y = LF.logToLogFloat (LF.logFromLogFloat x *
                                       (fromIntegral y :: Double))
        u' <- MWCD.geometric0 (LF.fromLogFloat success) g
        let u = toInteger u'
        b <- MWC.uniform g
        return $ Just
            ( VInt  $ if b then -1-u else u
            , VProb $ p * 2 / pow (1-success) u / success)

evaluateMeasureOp Categorical = \(e1 :* End) env ->
    VMeasure $ \p g -> do
        let (_,y,ys) = normalizeVector (evaluate e1 env)
        if not (y > (0::Double)) -- TODO: why not use @y <= 0@ ??
        then error "Categorical needs positive weights"
        else do
            u <- MWC.uniformR (0, y) g
            return $ Just
                ( VNat
                . intToNatural
                . fromMaybe 0
                . V.findIndex (u <=)
                . V.scanl1' (+)
                $ ys
                , p)

evaluateMeasureOp Uniform = \(e1 :* e2 :* End) env ->
    case (evaluate e1 env, evaluate e2 env) of
    (VReal v1, VReal v2) -> VMeasure $ \p g -> do
        x <- MWC.uniformR (v1, v2) g
        return $ Just (VReal x, p)

evaluateMeasureOp Normal = \(e1 :* e2 :* End) env ->
    case (evaluate e1 env, evaluate e2 env) of
    (VReal v1, VProb v2) -> VMeasure $ \ p g -> do
        x <- MWCD.normal v1 (LF.fromLogFloat v2) g
        return $ Just (VReal x, p)

evaluateMeasureOp Poisson = \(e1 :* End) env ->
    case evaluate e1 env of
    VProb v1 -> VMeasure $ \ p g -> do
        x <- MWC.genFromTable (MWC.tablePoisson (LF.fromLogFloat v1)) g
        return $ Just (VNat $ intToNatural x, p)

evaluateMeasureOp Gamma = \(e1 :* e2 :* End) env ->
    case (evaluate e1 env, evaluate e2 env) of
    (VProb v1, VProb v2) -> VMeasure $ \ p g -> do
        x <- MWCD.gamma (LF.fromLogFloat v1) (LF.fromLogFloat v2) g
        return $ Just (VProb $ LF.logFloat x, p)

evaluateMeasureOp Beta = \(e1 :* e2 :* End) env ->
    case (evaluate e1 env, evaluate e2 env) of
    (VProb v1, VProb v2) -> VMeasure $ \ p g -> do
        x <- MWCD.beta (LF.fromLogFloat v1) (LF.fromLogFloat v2) g
        return $ Just (VProb $ LF.logFloat x, p)

evaluateNaryOp
    :: (ABT Term abt)
    => NaryOp a -> Seq (abt '[] a) -> Env -> Value a
evaluateNaryOp s es =
    F.foldr (evalOp s) (identityElement s) . mapEvaluate es

identityElement :: NaryOp a -> Value a
identityElement And                   = VDatum dTrue
identityElement (Sum HSemiring_Nat)   = VNat  0
identityElement (Sum HSemiring_Int)   = VInt  0
identityElement (Sum HSemiring_Prob)  = VProb 0
identityElement (Sum HSemiring_Real)  = VReal 0
identityElement (Prod HSemiring_Nat)  = VNat  1
identityElement (Prod HSemiring_Int)  = VInt  1
identityElement (Prod HSemiring_Prob) = VProb 1
identityElement (Prod HSemiring_Real) = VReal 1
identityElement (Max  HOrd_Prob)      = VProb 0
identityElement (Max  HOrd_Real)      = VReal LF.negativeInfinity
identityElement (Min  HOrd_Prob)      = VProb (LF.logFloat LF.infinity)
identityElement (Min  HOrd_Real)      = VReal LF.infinity
identityElement _                     = error "Missing identity elements?"

evalOp
    :: NaryOp a -> Value a -> Value a -> Value a
evalOp And (VDatum a) (VDatum b)
    | a == dTrue && b == dTrue = VDatum dTrue
    | otherwise = VDatum dFalse
evalOp (Sum  HSemiring_Nat)  (VNat  a) (VNat  b) = VNat  (a + b)
evalOp (Sum  HSemiring_Int)  (VInt  a) (VInt  b) = VInt  (a + b)
evalOp (Sum  HSemiring_Prob) (VProb a) (VProb b) = VProb (a + b)
evalOp (Sum  HSemiring_Real) (VReal a) (VReal b) = VReal (a + b)
evalOp (Prod HSemiring_Nat)  (VNat  a) (VNat  b) = VNat  (a * b)
evalOp (Prod HSemiring_Int)  (VInt  a) (VInt  b) = VInt  (a * b)
evalOp (Prod HSemiring_Prob) (VProb a) (VProb b) = VProb (a * b)
evalOp (Prod HSemiring_Real) (VReal a) (VReal b) = VReal (a * b)
evalOp (Max  HOrd_Prob)      (VProb a) (VProb b) = VProb (max a b)
evalOp (Max  HOrd_Real)      (VReal a) (VReal b) = VReal (max a b)
evalOp (Min  HOrd_Prob)      (VProb a) (VProb b) = VProb (min a b)
evalOp (Min  HOrd_Real)      (VReal a) (VReal b) = VReal (min a b)

evalOp op                    _          _        =
    error ("TODO: evalOp{" ++ show op ++ "}")

mapEvaluate
    :: (ABT Term abt)
    => Seq (abt '[] a) -> Env -> Seq (Value a)
mapEvaluate es env = fmap (flip evaluate env) es

evaluateLiteral :: Literal a -> Value a
evaluateLiteral (LNat  n) = VNat  . fromInteger $ fromNatural n -- TODO: catch overflow errors
evaluateLiteral (LInt  n) = VInt  $ fromInteger n -- TODO: catch overflow errors
evaluateLiteral (LProb n) = VProb . fromRational $ fromNonNegativeRational n
evaluateLiteral (LReal n) = VReal $ fromRational n

evaluateEmpty :: Value ('HArray a)
evaluateEmpty = VArray V.empty

evaluateArray
    :: (ABT Term abt)
    => (abt '[] 'HNat)
    -> (abt '[ 'HNat ] a)
    -> Env
    -> Value ('HArray a)
evaluateArray n e env =
    case evaluate n env of
    VNat n' -> caseBind e $ \x e' ->
        VArray $ V.generate (unsafeInt n') $ \v ->
            let v' = VNat $ intToNatural v in
            evaluate e' (updateEnv (EAssoc x v') env)

evaluateBucket
    :: (ABT Term abt)
    => abt '[] 'HNat
    -> abt '[] 'HNat
    -> Reducer abt '[] a
    -> Env
    -> Value a
evaluateBucket b e rs env =
    case (evaluate b env, evaluate e env) of
      (VNat b', VNat e') -> runST $ do
          s' <- init Nil1 rs env
          mapM_ (\i -> accum (VNat i) Nil1 rs s' env) [b' .. e' - 1]
          done s'
    where init :: (ABT Term abt)
               => List1 Value xs
               -> Reducer abt xs a
               -> Env
               -> ST s (VReducer s a)
          init ix (Red_Fanout r1 r2)    env  =
              VRed_Pair (type_ r1) (type_ r2) <$> init ix r1 env <*> init ix r2 env
          init ix (Red_Index  n  _  mr) env' =
              let (vars, n') = caseBinds n in
              case evaluate n' (updateEnvs vars ix env') of
                VNat n'' -> VRed_Array <$> V.generateM (fromIntegral n'')
                            (\bb -> init (Cons1 (vnat bb) ix) mr env')
          init ix (Red_Split _ r1 r2)   env' =
              VRed_Pair (type_ r1) (type_ r2) <$> init ix r1 env <*> init ix r2 env'
          init _  Red_Nop               _    = return VRed_Unit
          init _  (Red_Add h _) _ = VRed_Num <$> newSTRef (identityElement (Sum h))

          type_ = typeOfReducer

          vnat :: Int -> Value 'HNat
          vnat  = VNat . fromIntegral

          accum :: (ABT Term abt)
                => Value 'HNat
                -> List1 Value xs
                -> Reducer abt xs a
                -> VReducer s a
                -> Env
                -> ST s ()
          accum n ix (Red_Fanout r1 r2)   (VRed_Pair _ _ v1 v2) env' =
              accum n ix r1 v1 env >> accum n ix r2 v2 env'
          accum n ix (Red_Index n' a1 r2) (VRed_Array v)          env' =
              caseBind a1 $ \i a1' ->
              let (vars, a1'') = caseBinds a1'
                  VNat ov = evaluate a1''
                            (updateEnv (EAssoc i n) (updateEnvs vars ix env'))
                  ov' = fromIntegral ov in
              accum n (Cons1 (VNat ov) ix) r2 (v V.! ov') env
          accum n ix (Red_Split bb r1 r2) (VRed_Pair _ _ v1 v2) env' =
              caseBind bb $ \i b' ->
                  let (vars, b'') = caseBinds b' in
                  case evaluate b''
                       (updateEnv (EAssoc i n) (updateEnvs vars ix env')) of
                  VDatum bb -> if bb == dTrue then
                                   accum n ix r1 v1 env'
                               else
                                   accum n ix r2 v2 env'
          accum n ix (Red_Add h ee) (VRed_Num s) env' =
              caseBind ee $ \i e' ->
                  let (vars, e'') = caseBinds e'
                      v = evaluate e''
                          (updateEnv (EAssoc i n) (updateEnvs vars ix env')) in
                  modifySTRef' s (evalOp (Sum h) v)
          accum _ _ Red_Nop _ _ = return ()
          accum _ _ _ _ _ = error "Some impossible combinations happened?"

          done :: VReducer s a -> ST s (Value a)
          done (VRed_Num s)            = readSTRef s
          done VRed_Unit               = return (VDatum dUnit)
          done (VRed_Pair s1 s2 v1 v2) = do
            v1' <- done v1
            v2' <- done v2
            return (VDatum $ dPair_ s1 s2 v1' v2')
          done (VRed_Array v)          = VArray <$> V.sequence (V.map done v)

evaluateDatum
    :: (ABT Term abt)
    => Datum (abt '[]) (HData' a)
    -> Env
    -> Value (HData' a)
evaluateDatum d env = VDatum (fmap11 (flip evaluate env) d)

evaluateCase
    :: forall abt a b
    .  (ABT Term abt)
    => abt '[] a
    -> [Branch a abt b]
    -> Env
    -> Value b
evaluateCase o es env =
    case runIdentity $ matchBranches evaluateDatum' (evaluate o env) es of
    Just (Matched rho b) ->
        evaluate b (extendFromMatch (fromAssocs rho) env)
    _ -> error "Missing cases in match expression"
    where
    extendFromMatch :: [Assoc Value] -> Env -> Env
    extendFromMatch []                env' = env'
    extendFromMatch (Assoc x v : xvs) env' =
        extendFromMatch xvs (updateEnv (EAssoc x v) env')

    evaluateDatum' :: DatumEvaluator Value Identity
    evaluateDatum' = return . Just . getVDatum

    getVDatum :: Value (HData' a) -> Datum Value (HData' a)
    getVDatum (VDatum a) = a

evaluateSuperpose
    :: (ABT Term abt)
    => NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
    -> Env
    -> Value ('HMeasure a)
evaluateSuperpose ((q, m) :| []) env =
    case evaluate m env of
    VMeasure m' ->
        let VProb q' = evaluate q env
        in  VMeasure (\(VProb p) g -> m' (VProb $ p * q') g)

evaluateSuperpose pms@((_, m) :| _) env =
    case evaluate m env of
    VMeasure m' ->
        let pms'     = L.toList pms
            weights  = map ((flip evaluate env) . fst) pms'
            (x,y,ys) = normalize weights
        in VMeasure $ \(VProb p) g ->
            if not (y > (0::Double)) then return Nothing else do
            u <- MWC.uniformR (0, y) g
            case [ m1 | (v,(_,m1)) <- zip (scanl1 (+) ys) pms', u <= v ] of
                m2 : _ ->
                    case evaluate m2 env of
                    VMeasure m2' -> m2' (VProb $ p * x * LF.logFloat y) g
                []     -> m' (VProb $ p * x * LF.logFloat y) g

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

-- Useful 'short-hand'
intToNatural :: Int -> Natural
intToNatural = unsafeNatural . toInteger

unsafeInt :: Natural -> Int
unsafeInt = fromInteger . fromNatural
----------------------------------------------------------- fin.


-- Copyright (c) 1982-1999 Lennart Augustsson, Thomas Johnsson
-- See LICENSE for the full license.
--
module Hash(Hash(..), combineHash, emptyHash, hashToInt, hashToMax, Hashable(..)) where

--
-- Hash a value.  Hashing produces an Int of
-- unspecified range.
--
import Data.Array
import Data.Complex
import Data.Ratio

-- #define HASH(x) (case hash (x) of H h -> h)
hASH x = case hash (x) of H h -> h

newtype Hash = H Int deriving (Eq)

instance Show Hash where
    --showsType _ = showString "Hash"
    show _ = "<<Hash>>"

combineHash :: Hash -> Hash -> Hash
combineHash (H x) (H y) = H (x+y)

emptyHash :: Hash
emptyHash = H 0

class Hashable a where
    hash :: a -> Hash

instance Hashable Char where
    hash x = H $ fromEnum x

instance Hashable Int where
    hash x = H $ x

instance Hashable Integer where
    hash x = H $ fromInteger x

instance Hashable Float where
    hash x = H $ truncate x

instance Hashable Double where
    hash x = H $ truncate x

instance Hashable (IO a) where
    hash x = H 0

instance Hashable () where
    hash x = H 0

instance Hashable (a -> b) where
    hash x = H 0

instance (Hashable a) => Hashable (Maybe a) where
    hash Nothing = H 0
    hash (Just x) = hash x

instance (Hashable a, Hashable b) => Hashable (Either a b) where
    hash (Left x) = hash x
    hash (Right y) = hash y

instance Hashable a => Hashable [a] where
    hash l = H $ f l 0
        where f :: (Hashable a') => [a'] -> Int -> Int
              f [] r = r
              f (c:cs) r = f cs (3*r + hASH(c))

{-
#ifndef __HBCC__
instance Hashable [Char] where
    hash l = H $ f l 0
        where f :: String -> Int -> Int
              f [] r = r
              f (c:cs) r = f cs (3*r + fromEnum c)
#endif
-}

instance (Hashable a, Hashable b) => Hashable (a,b) where
    hash (a,b) = H $ hASH(a) + 3 * hASH(b)

instance (Hashable a, Hashable b, Hashable c) => Hashable (a,b,c) where
    hash (a,b,c) = H $ hASH(a) + 3 * hASH(b) + 5 * hASH(c)

instance (Hashable a, Hashable b, Hashable c, Hashable d) => Hashable (a,b,c,d) where
    hash (a,b,c,d) = H $ hASH(a) + 3 * hASH(b) + 5 * hASH(c) + 7 * hASH(d)

instance (Hashable a, Hashable b, Hashable c, Hashable d, Hashable e) => Hashable (a,b,c,d,e) where
    hash (a,b,c,d,e) = H $ hASH(a) + 3 * hASH(b) + 5 * hASH(c) + 7 * hASH(d) + 9 * hASH(e)

instance Hashable Bool where
    hash False = H 0
    hash True = H 1

instance (Integral a, Hashable a) => Hashable (Ratio a) where
    hash x = H $ hASH (denominator x) + hASH (numerator x)

instance (RealFloat a, Hashable a) => Hashable (Complex a) where
    hash (x :+ y) = H $ hASH(x) + hASH(y)

instance (Ix a) => Hashable (Array a b) where
    hash x = H $ 0 -- !!!

hashToInt :: Int -> Hash -> Int
hashToInt maxhash x =
    case x of
    H h ->
        if h < 0 then
            if -h < 0 then 0
            else (-h) `rem` maxhash
        else h `rem` maxhash

hashToMax maxhash x =
    case hash x of
    H h ->
        if h < 0 then
            if -h < 0 then 0
            else (-h) `rem` maxhash
        else h `rem` maxhash


{-# LANGUAGE ForeignFunctionInterface #-}
module Grenade.Layers.Internal.Pad (
    pad
  , crop
  ) where

import qualified Data.Vector.Storable        as U (unsafeFromForeignPtr0,
                                                   unsafeToForeignPtr0)

import           Foreign                     (mallocForeignPtrArray, withForeignPtr)
import           Foreign.Ptr                 (Ptr)
import           Numeric.LinearAlgebra       (Matrix, flatten)
import qualified Numeric.LinearAlgebra.Devel as U
import           System.IO.Unsafe            (unsafePerformIO)

import           Grenade.Types

pad :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
pad channels padLeft padTop padRight padBottom rows cols rows' cols' m
 = let outMatSize      = rows' * cols' * channels
       vec             = flatten m
   in unsafePerformIO $ do
     outPtr        <- mallocForeignPtrArray outMatSize
     let (inPtr, _) = U.unsafeToForeignPtr0 vec

     withForeignPtr inPtr $ \inPtr' ->
       withForeignPtr outPtr $ \outPtr' ->
         pad_cpu inPtr' channels rows cols padLeft padTop padRight padBottom outPtr'

     let matVec = U.unsafeFromForeignPtr0 outPtr outMatSize
     return (U.matrixFromVector U.RowMajor (rows' * channels) cols' matVec)
{-# INLINE pad #-}

foreign import ccall unsafe
    pad_cpu
      :: Ptr RealNum -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Ptr RealNum -> IO ()

crop :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
crop channels padLeft padTop padRight padBottom rows cols _ _ m
 = let outMatSize      = rows * cols * channels
       vec             = flatten m
   in unsafePerformIO $ do
     outPtr        <- mallocForeignPtrArray outMatSize
     let (inPtr, _) = U.unsafeToForeignPtr0 vec

     withForeignPtr inPtr $ \inPtr' ->
       withForeignPtr outPtr $ \outPtr' ->
         crop_cpu inPtr' channels rows cols padLeft padTop padRight padBottom outPtr'

     let matVec = U.unsafeFromForeignPtr0 outPtr outMatSize
     return (U.matrixFromVector U.RowMajor (rows * channels) cols matVec)

foreign import ccall unsafe
    crop_cpu
      :: Ptr RealNum -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Ptr RealNum -> IO ()


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}

{-|
Module      : Data.Picture
Description : picture manipulation functions
Copyright   : (c) Mahdi Dibaiee, 2016
License     : GPL-3
Maintainer  : mdibaiee@aol.com
Stability   : experimental
Portability : POSIX
-}
module Data.Picture ( Picture
                    -- * Manipulation functions
                    , grayscale
                    , fade
                    , rotate
                    , contrast
                    , brightness
                    , gamma
                    , invert
                    , compress
                    , embed
                    , resize
                    , Data.Picture.scale
                    -- * Converting between Image and Picture
                    , fromImage
                    , toImage
                    -- * IO operations
                    , readPicture
                    , writePicturePng
                    )
  where

    import Codec.Picture
    import Numeric.LinearAlgebra
    import qualified Data.Vector.Storable as V
    import System.IO
    import Data.Maybe
    import Debug.Trace
    import Data.List (zipWith4)

    -- | (R, G, B, A) color channels
    type Picture = (Matrix Double, Matrix Double, Matrix Double, Matrix Double)

    -- | Converts a JuicyPixel 'Image PixelRGBA8' to 'Picture'
    fromImage :: Image PixelRGBA8 -> Picture
    fromImage Image { imageWidth = w, imageHeight = h, imageData = vec } =
     let [r, g, b, a] = map (reshape w . V.fromList . reverse) (snd $ V.foldl' gp (0, [[],[],[],[]]) (V.map fromIntegral vec))
     in (r, g, b, a)
     where
       gp acc x =
         case acc of
           (0, [r, g, b, a]) -> (1, [x:r, g, b, a])
           (1, [r, g, b, a]) -> (2, [r, x:g, b, a])
           (2, [r, g, b, a]) -> (3, [r, g, x:b, a])
           (3, [r, g, b, a]) -> (0, [r, g, b, x:a])

    -- | Converts a 'Picture' to JuicyPixel 'Image PixelRGBA8'
    toImage :: Picture -> Image PixelRGBA8
    toImage (r, g, b, a) =
      let (fr, fg, fb, fa) = (toList $ flatten r, toList $ flatten g, toList $ flatten b, toList $ flatten a)
          img = V.map (fromIntegral . floor) . V.concat $ zipWith4 (\a b c d -> vector [a, b, c, d]) fr fg fb fa
      in Image { imageWidth = cols r, imageHeight = rows r, imageData = img }

    -- | Reads a 'Picture' from specified path
    readPicture :: FilePath -> IO (Either String Picture)
    readPicture path = do
      img <- readImage path
      return $ case img of
        Left err -> Left err
        Right im -> Right $ fromImage (convertRGBA8 im)

    -- | Write the specified 'Picture' to a PNG file
    writePicturePng :: FilePath -> Picture -> IO ()
    writePicturePng path pic = writePng path (toImage pic)

    -- | Turn the 'Picture' grayscale
    grayscale :: Picture -> Picture
    grayscale (r, g, b, a) =
      let (fr, fg, fb) = (flatten r, flatten g, flatten b)
          mean = reshape (cols r) $ V.map (/ 3) (fr + fg + fb)
      in (mean, mean, mean, a)

    -- | Fade the 'Picture' by a number between 0 and 1
    fade :: Double -> Picture -> Picture
    fade opacity (r, g, b, a) = (r, g, b, f a)
      where
        f = cmap (*opacity)

    -- | Set contrast level of 'Picture', a number between -255 and 255
    contrast :: Double -> Picture -> Picture
    contrast level (r, g, b, a) = (f r, f g, f b, a)
      where
        cfactor = (259 * (level + 255)) / (255 * (259 - level))
        f = cmap (\x -> pixelBound $ cfactor * (x - 128) + 128)

    -- | Set brightness level of 'Picture', a number between -255 and 255
    brightness :: Double -> Picture -> Picture
    brightness level (r, g, b, a) = (f r, f g, f b, a)
      where
        f = cmap (pixelBound . (+level))

    -- | Set gamma level of 'Picture'
    gamma :: Int -> Picture -> Picture
    gamma level (r, g, b, a) = (f r, f g, f b, a)
      where
        f = cmap (\x -> pixelBound $ 255 * (x / 255) ^ level)

    -- | Inverts the 'Picture'
    invert :: Picture -> Picture
    invert (r, g, b, a) = (f r, f g, f b, a)
      where
        f = cmap (`subtract` 255)

    -- | Rotate 'Picture' for the specified degrees, around the specified origin.
    -- If the origin is `Nothing`, rotates around the center
    rotate :: Double -> Maybe (Int, Int) -> Picture -> Picture
    rotate deg orig (r, g, b, a) = (f r, f g, f b, f a)
      where
        -- rotation in radians
        rad = deg * pi / 180

        -- rotation matrix (clockwise)
        rm = fromLists [[cos rad, sin rad],
                        [negate $ sin rad, cos rad]]

        -- origin of rotation
        (originX, originY) = if isJust orig then fromJust orig else (cols r `div` 2, rows r `div` 2)

        -- all index pairs
        indices = [vector [fromIntegral x - fromIntegral originX, fromIntegral y - fromIntegral originY] | y <- [0..rows r - 1], x <- [0..cols r - 1]]

        -- rotate them using rotation matrix
        rotatedIndices :: [[Int]] = map (\r -> toList . V.map (fromIntegral . round) $ rm #> r) indices

        -- move them back to the origin
        movedIndices = map (\[x, y] -> [x + originX, y + originY]) rotatedIndices

        f m = reshape (cols m) $ fromList $ map (\[x, y] -> if y < 0 || y >= rows r || x < 0 || x >= cols r then 0 else m `atIndex` (y, x)) movedIndices

    -- | Compress the image using SVD
    -- note: this is not size compression, it's just a k-rank approximation of the image
    compress :: Int -> Picture -> Picture
    compress rate (r, g, b, a) = (f r, f g, f b, a)
      where
        k = cols r - rate
        f m =
          let (u, s, v) = svd m
              si = diagRect 0 s (rows m) (cols m)
              (mu, ms, mv) = (u ?? (All, Take k), si ?? (Take k, Take k), (tr v) ?? (Take k, All))
          in mu <> ms <> mv

    -- | Embed a 'Picture' into another one, in the specified position-}
    embed :: Picture -> (Int, Int) -> Picture -> Picture
    embed (br, bg, bb, ba) (x, y) (lr, lg, lb, la) = (f br lmr, f bg lmg, f bb lmb, maxAlpha)
      where
        (lmr, lmg, lmb, lma) = (fit lr, fit lg, fit lb, fit la)
        scaledAlpha = cmap (/255) lma

        fit m =
          let distance = y * cols br
              total = rows br * cols br
              xPush = (rows m><x) (repeat 0)
              xPast = (rows m><(cols br - x - cols m)) (repeat 0)
              positioned = xPush ||| m ||| xPast
              flat = toList $ flatten positioned
          in (rows br><cols br) $ replicate distance 0 ++ flat ++ repeat 0

        f b lm = (b * (cmap (1-) scaledAlpha)) + (lm * scaledAlpha)
        maxAlpha = (rows ba><cols ba) $ zipWith max (toList . flatten $ ba) (toList . flatten . fit $ la)

    -- | Resize an image using nearest-neighbor interpolation
    resize :: (Int, Int) -> Picture -> Picture
    resize (sWidth, sHeight) (r, g, b, a) = (f r, f g, f b, f a)
      where
        initial = vector [0..fromIntegral sWidth * fromIntegral sHeight - 1]
        (width, height) = (cols r, rows r)
        factor = 2 ^ 16
        (xRatio, yRatio) = (width * factor `div` sWidth + 1, height * factor `div` sHeight + 1)
        f m = reshape sWidth $ V.map replace initial
          where
            v = flatten m
            replace index =
              let (x, y) = (floor index `mod` sWidth, floor index `div` sWidth)
                  (px, py) = (x * xRatio `div` factor, y * yRatio `div` factor)
              in v ! (py * width + px)

    -- | Scale an image using the resize function
    scale :: Double -> Picture -> Picture
    scale 1 p = p
    scale s (r, g, b, a) = resize (floor $ s * width, floor $ s * height) (r, g, b, a)
      where
        (width, height) = (fromIntegral $ cols r, fromIntegral $ rows r)

    bound (l, u) x = max l $ min u x
    pixelBound = bound (0, 255)


{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE Rank2Types            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
-- |
-- This module provides function for working with product types and
-- comes in two variants. First works with plain product, types like
-- @(a,b)@ or @data Prod = Prod A B@, etc. Second one is for
-- parameterized products (it seems there's no standard name for
-- them), that is types like: @data ProdF f = ProdF (f Int) (f Char)@.
--
-- Most examples in this module use tuple but library is not limited
-- to them in any way. They're just in base and convenient to work
-- with.
module Data.Vector.HFixed (
    -- * HVector type classes
    HVector(..)
  , tupleSize
  , HVectorF(..)
  , tupleSizeF
  , ContVec
  , ContVecF(..)
  , asCVec
  , asCVecF
    -- * Plain product types
    -- ** Construction
    -- *** Simple constructor
    -- $construction
  , mk0
  , mk1
  , mk2
  , mk3
  , mk4
  , mk5
    -- *** Unfoldr & replicate
  , unfoldr
  , replicate
  , replicateM
  -- ** Position based functions
  , convert
  , head
  , tail
  , cons
  , concat
    -- *** Indexing
  , ValueAt
  , Index
  , index
  , set
  , element
  , elementCh
  , tyLookup
  , tyLookupF
    -- ** Folds & unfolds
  , foldr
  , foldl
  , foldMap
  , mapM_
    -- ** Zips
  , zipWith
  , zipFold
    -- ** Specializations
  , eq
  , compare
  , rnf
    -- * Parametrized products
    -- ** Construction
    -- *** Simple constructors
    -- $construction_F
  , mk0F
  , mk1F
  , mk2F
  , mk3F
  , mk4F
  , mk5F
    -- *** Unfoldr & replicate
  , unfoldrF
  , replicateF
  , replicateNatF
    -- ** Conversion to\/from products
  , wrap
  , unwrap
  , monomorphize
  , monomorphizeF
    -- ** Functor\/Applicative like
  , map
  , mapNat
  , sequence
  , sequence_
  , sequenceF
  , distribute
  , distributeF
    -- ** Folds and unfolds
  , foldrF
  , foldlF
  , foldMapF
  , foldrNatF
  , foldlNatF
  , foldMapNatF
    -- ** Zips
  , zipWithF
  , zipWithNatF
  , zipFoldF
    -- ** Reexports
  , Arity
  , ArityC
  , Proxy(..)
  ) where

import Control.Applicative  (Applicative(..),(<$>))
import qualified Control.DeepSeq as NF

import Data.Coerce           (coerce)
import Data.Functor.Compose  (Compose(..))
import Data.Functor.Identity (Identity(..))
import Data.Monoid           (Monoid,All(..))
import Prelude ( Functor(..),Eq(..),Ord,Bool,Ordering
               , id,(.),($),seq)
import qualified Prelude

import           Data.Vector.HFixed.Class hiding (cons,consF)
import           Data.Vector.Fixed.Cont       (Peano)
import qualified Data.Vector.Fixed          as F
import qualified Data.Vector.HFixed.Cont    as C

----------------------------------------------------------------
-- Generic API
----------------------------------------------------------------

-- | Restrict type of vector to 'ContVec'. This function is useful for
--   resolving type ambiguity when composing functions. For example
--   following code would not compile because intermediate type is
--   ambiguous:
--
-- > cons 'a' . tail
--
--   GHC cannot guess what type should be produced by @tail@. However
--   we can fix type of intermediate vector with @asCVec@, so code
--   below will work just fine:
--
-- > cons 'a' . asCVec . tail
asCVec :: ContVec xs -> ContVec xs
asCVec = id

asCVecF :: ContVecF f xs -> ContVecF f xs
asCVecF = id

-- | We can convert between any two vector which have same
--   structure but different representations.
--
-- >>> convert (1 :+ 2) :: (Double,Double)
-- (1.0,2.0)
convert :: (HVector v, HVector w, Elems v ~ Elems w)
        => v -> w
{-# INLINE convert #-}
convert v = inspect v construct

-- | Tail of the vector. Note that in the example we only tell GHC
--   that resulting value is 2-tuple via pattern matching and let
--   typechecker figure out the rest.
--
-- >>> case tail ('a',"aa",()) of x@(_,_) -> x
-- ("aa",())
tail :: (HVector v, HVector w, (a : Elems w) ~ Elems v)
     => v -> w
{-# INLINE tail #-}
tail = C.vector . C.tail . C.cvec

-- | Head of the vector
--
-- >>> head ('a',"ABC")
-- 'a'
head :: (HVector v, Elems v ~ (a : as), Arity as)
     => v -> a
{-# INLINE head #-}
head = C.head . C.cvec

-- | Prepend element to the product.
--
-- >>> cons 'c' ('d','e') :: (Char,Char,Char)
-- ('c','d','e')
cons :: (HVector v, HVector w, Elems w ~ (a : Elems v))
     => a -> v -> w
{-# INLINE cons #-}
cons a = C.vector . C.cons a . C.cvec

-- | Concatenate two vectors
--
-- >>> concat ('c','d') ('e','f') :: (Char,Char,Char,Char)
-- ('c','d','e','f')
concat :: ( HVector v, HVector u, HVector w
          , Elems w ~ (Elems v ++ Elems u)
          )
       => v -> u -> w
concat v u = C.vector $ C.concat (C.cvec v) (C.cvec u)
{-# INLINE concat #-}

----------------------------------------------------------------
-- Indexing
----------------------------------------------------------------

-- | Index heterogeneous vector.
--
-- >>> index (Proxy @0) ('c',"str")
-- 'c'
-- >>> index (Proxy @1) ('c',"str")
-- "str"
index
  :: forall n v proxy. (Index (Peano n) (Elems v), HVector v)
  => proxy n                     -- ^ Type level index
  -> v                           -- ^ Vector to index
  -> ValueAt (Peano n) (Elems v)
{-# INLINE index #-}
index _ v = C.index (C.cvec v) (Proxy @(Peano n))

-- | Set element in the vector
--
-- >>> set (Proxy @0) 'X' ('_',"str")
-- ('X',"str")
set :: forall n v proxy. (Index (Peano n) (Elems v), HVector v)
    => proxy n                     -- ^ Type level index
    -> ValueAt (Peano n) (Elems v) -- ^ New value at index
    -> v
    -> v
{-# INLINE set #-}
set _ x = C.vector
        . C.set (Proxy @(Peano n)) x
        . C.cvec

-- | Twan van Laarhoven's lens for i'th element.
element :: forall n v proxy.
           ( Index (Peano n) (Elems v)
           , HVector v
           )
        => proxy n              -- ^ Type level index
        -> Lens' v (ValueAt (Peano n) (Elems v))
{-# INLINE element #-}
element _ f v = inspect v
              $ lensF (Proxy @(Peano n)) f construct

-- | Type changing Twan van Laarhoven's lens for i'th element.
elementCh :: forall n v w a b proxy.
             ( Index   (Peano n) (Elems v)
             , ValueAt (Peano n) (Elems v) ~ a
             , HVector v
             , HVector w
             , Elems w ~ NewElems (Peano n) (Elems v) b
             )
          => proxy n            -- ^ Type level index
          -> Lens v w a b
{-# INLINE elementCh #-}
elementCh _ f v = inspect v
                $ lensChF (Proxy @(Peano n)) f construct

-- | Lookup field from product by its type. Product must contain one
--   and only one field of type @a@
--
-- >>> tyLookup ('c',"str") :: Char
-- 'c'
--
-- >>> tyLookup ('c',"str") :: Int
-- ...
--     • Cannot find type:
--       Int
--     • In the expression: tyLookup ('c', "str") :: Int
--       In an equation for ‘it’: it = tyLookup ('c', "str") :: Int
--
-- >>> tyLookup ('c','c') :: Char
-- ...
--     • Duplicate type found:
--       Char
--     • In the expression: tyLookup ('c', 'c') :: Char
--       In an equation for ‘it’: it = tyLookup ('c', 'c') :: Char
tyLookup :: (HVector v, TyLookup a (Elems v)) => v -> a
tyLookup = C.tyLookup . C.cvec
{-# INLINE tyLookup #-}

-- | Analog of 'tyLookup' for @HVectorF@
tyLookupF :: (HVectorF v, TyLookup a (ElemsF v)) => v f -> f a
tyLookupF = C.tyLookupF . C.cvecF
{-# INLINE tyLookupF #-}

----------------------------------------------------------------
-- Folds over vector
----------------------------------------------------------------

-- | Right fold over heterogeneous vector
--
-- >>> foldr (Proxy @Show) (\x str -> show x : str) [] (12,'c')
-- ["12","'c'"]
foldr :: (HVector v, ArityC c (Elems v))
      => Proxy c -> (forall a. c a => a -> b -> b) -> b -> v -> b
{-# INLINE foldr #-}
foldr c f b0 = C.foldrF c (\(Identity a) b -> f a b) b0 . C.cvec

-- | Left fold over heterogeneous vector
foldl :: (HVector v, ArityC c (Elems v))
      => Proxy c -> (forall a. c a => b -> a -> b) -> b -> v -> b
{-# INLINE foldl #-}
foldl c f b0 = C.foldlF c (\b (Identity a) -> f b a) b0 . C.cvec

-- | Monoidal fold over heterogeneuous vector
--
-- >>> foldMap (Proxy @Show) show (12,'c',"str")
-- "12'c'\"str\""
foldMap
  :: (HVector v, ArityC c (Elems v), Monoid m)
  => Proxy c -> (forall a. c a => a -> m) -> v -> m
{-# INLINE foldMap #-}
foldMap c f = C.foldMapF c (\(Identity a) -> f a) . C.cvec

-- | Right fold over heterogeneous vector
foldrF :: (HVectorF v, ArityC c (ElemsF v))
       => Proxy c -> (forall a. c a => f a -> b -> b) -> b -> v f -> b
{-# INLINE foldrF #-}
foldrF c f b0 = C.foldrF c f b0 . C.cvecF

-- | Left fold over heterogeneous vector
foldlF :: (HVectorF v, ArityC c (ElemsF v))
       => Proxy c -> (forall a. c a => b -> f a -> b) -> b -> v f -> b
{-# INLINE foldlF #-}
foldlF c f b0 = C.foldlF c f b0 . C.cvecF

-- | Monoidal fold over heterogeneous vector
--
-- >>> foldMapF (Proxy @Show) show (mk2F (Just 1) Nothing :: HVecF '[Int,Char] Maybe)
-- "Just 1Nothing"
foldMapF :: (HVectorF v, ArityC c (ElemsF v), Monoid m)
         => Proxy c -> (forall a. c a => f a -> m) -> v f -> m
{-# INLINE foldMapF #-}
foldMapF c f = C.foldMapF c f . C.cvecF

-- | Right fold over heterogeneous vector
foldrNatF :: (HVectorF v)
          => (forall a. f a -> b -> b) -> b -> v f -> b
{-# INLINE foldrNatF #-}
foldrNatF f b0 = C.foldrNatF f b0 . C.cvecF

-- | Left fold over heterogeneous vector
foldlNatF :: (HVectorF v)
          => (forall a. b -> f a -> b) -> b -> v f -> b
{-# INLINE foldlNatF #-}
foldlNatF f b0 = C.foldlNatF f b0 . C.cvecF

-- | Monoidal fold over heterogeneous vector
--
-- >>> foldMapNatF (Sum . getConst) (mk2F (Const 1) (Const 2) :: HVecF '[Char,String] (Const Int))
-- Sum {getSum = 3}
foldMapNatF :: (HVectorF v, Monoid m)
            => (forall a. f a -> m) -> v f -> m
{-# INLINE foldMapNatF #-}
foldMapNatF f = C.foldMapNatF f . C.cvecF

-- | Apply monadic action to every element in the vector
mapM_ :: (HVector v, ArityC c (Elems v), Applicative f)
      => Proxy c -> (forall a. c a => a -> f ()) -> v -> f ()
{-# INLINE mapM_ #-}
mapM_ c f = foldl c (\m a -> m *> f a) (pure ())

-- | Unfold vector.
unfoldr :: (HVector v, ArityC c (Elems v))
        => Proxy c -> (forall a. c a => b -> (a,b)) -> b -> v
{-# INLINE unfoldr #-}
unfoldr c f = C.vector . C.unfoldrF c (\b -> let (a,b') = f b in (Identity a, b'))

-- | Unfold vector.
unfoldrF :: (HVectorF v, ArityC c (ElemsF v))
        => Proxy c -> (forall a. c a => b -> (f a,b)) -> b -> v f
{-# INLINE unfoldrF #-}
unfoldrF c f = C.vectorF . C.unfoldrF c f

----------------------------------------------------------------
-- Constructors
----------------------------------------------------------------

-- $construction
--
-- Functions below allow to construct products up to 5 elements. Here
-- are example for product types from base:
--
-- >>> mk0 :: ()
-- ()
--
-- >>> mk3 12 'x' "xyz" :: (Int,Char,String)
-- (12,'x',"xyz")
--
-- >>> mk2 0 1 :: Complex Double
-- 0.0 :+ 1.0

mk0 :: forall v. (HVector v, Elems v ~ '[]) => v
mk0 = coerce (construct :: Fun '[] v)
{-# INLINE mk0 #-}

mk1 :: forall v a. (HVector v, Elems v ~ '[a])
    => a -> v
mk1 = coerce (construct :: Fun '[a] v)
{-# INLINE mk1 #-}

mk2 :: forall v a b. (HVector v, Elems v ~ '[a,b])
    => a -> b -> v
mk2 = coerce (construct :: Fun '[a,b] v)
{-# INLINE mk2 #-}

mk3 :: forall v a b c. (HVector v, Elems v ~ '[a,b,c])
    => a -> b -> c -> v
mk3 = coerce (construct :: Fun '[a,b,c] v)
{-# INLINE mk3 #-}

mk4 :: forall v a b c d. (HVector v, Elems v ~ '[a,b,c,d])
    => a -> b -> c -> d -> v
mk4 = coerce (construct :: Fun '[a,b,c,d] v)
{-# INLINE mk4 #-}

mk5 :: forall v a b c d e. (HVector v, Elems v ~ '[a,b,c,d,e])
    => a -> b -> c -> d -> e -> v
mk5 = coerce (construct :: Fun '[a,b,c,d,e] v)
{-# INLINE mk5 #-}

-- $construction_F
--
-- Construction function for parametrized products are fully
-- analogous to plain products:
--
-- >>>mk2F (Identity 'c') (Identity 1) :: HVecF '[Char, Int] Identity
-- [Identity 'c',Identity 1]
--
-- >>>mk2F (Nothing) (Just 1) :: HVecF '[Char, Int] Maybe
-- [Nothing,Just 1]

mk0F :: forall f v. (HVectorF v, ElemsF v ~ '[]) => v f
mk0F = coerce (constructF :: TFun f '[] (v f))
{-# INLINE mk0F #-}

mk1F :: forall f v a. (HVectorF v, ElemsF v ~ '[a])
     => f a -> v f
mk1F = coerce (constructF :: TFun f '[a] (v f))
{-# INLINE mk1F #-}

mk2F :: forall f v a b. (HVectorF v, ElemsF v ~ '[a,b])
     => f a -> f b -> v f
mk2F = coerce (constructF :: TFun f '[a,b] (v f))
{-# INLINE mk2F #-}

mk3F :: forall f v a b c. (HVectorF v, ElemsF v ~ '[a,b,c])
     => f a -> f b -> f c -> v f
mk3F = coerce (constructF :: TFun f '[a,b,c] (v f))
{-# INLINE mk3F #-}

mk4F :: forall f v a b c d. (HVectorF v, ElemsF v ~ '[a,b,c,d])
     => f a -> f b -> f c -> f d -> v f
mk4F = coerce (constructF :: TFun f '[a,b,c,d] (v f))
{-# INLINE mk4F #-}

mk5F :: forall f v a b c d e. (HVectorF v, ElemsF v ~ '[a,b,c,d,e])
     => f a -> f b -> f c -> f d -> f e -> v f
mk5F = coerce (constructF :: TFun f '[a,b,c,d,e] (v f))
{-# INLINE mk5F #-}

----------------------------------------------------------------
-- Collective operations
----------------------------------------------------------------

-- | Apply function to every value of parametrized product.
--
-- >>> map (Proxy @Num) (Identity . fromMaybe 0) (mk2F (Just 12) Nothing :: HVecF '[Double, Int] Maybe)
-- [Identity 12.0,Identity 0]
map :: (HVectorF v, ArityC c (ElemsF v))
    => Proxy c -> (forall a. c a => f a -> g a) -> v f -> v g
{-# INLINE map #-}
map cls f = C.vectorF . C.map cls f . C.cvecF

-- | Apply natural transformation to every element of the tuple.
--
-- >>> mapNat (Just . runIdentity) (mk2F (pure 'c') (pure 1) :: HVecF '[Char, Int] Identity)
-- [Just 'c',Just 1]
mapNat :: (HVectorF v)
       => (forall a. f a -> g a) -> v f -> v g
{-# INLINE mapNat #-}
mapNat f = C.vectorF . C.mapNat f . C.cvecF

-- | Sequence effects for every element in the vector
--
-- >>> sequence (mk2F [1,2] "ab" :: HVecF '[Int,Char] []) :: [(Int,Char)]
-- [(1,'a'),(1,'b'),(2,'a'),(2,'b')]
sequence
  :: ( Applicative f, HVectorF v, HVector w, ElemsF v ~ Elems w )
  => v f -> f w
{-# INLINE sequence #-}
sequence
  = fmap C.vector
  . C.sequenceF
  . C.mapNat (Compose . fmap Identity)
  . C.cvecF

-- | Sequence effects for every element in the vector
sequence_ :: (Applicative f, HVectorF v) => v f -> f ()
{-# INLINE sequence_ #-}
sequence_ = foldlNatF (\m a -> m <* a) (pure ())

-- | Sequence effects for every element in the vector
sequenceF :: ( Applicative f, HVectorF v) => v (f `Compose` g) -> f (v g)
{-# INLINE sequenceF #-}
sequenceF v = C.vectorF <$> C.sequenceF (C.cvecF v)

-- | Wrap every value in the vector into type constructor.
wrap :: ( HVector v, HVectorF w, Elems v ~ ElemsF w )
     => (forall a. a -> f a) -> v -> w f
{-# INLINE wrap #-}
wrap f = C.vectorF . C.mapNat (f . runIdentity) . C.cvec

-- | Unwrap every value in the vector from the type constructor.
unwrap :: ( HVectorF v, HVector w, ElemsF v ~ Elems w )
       => (forall a. f a -> a) -> v f -> w
{-# INLINE unwrap #-}
unwrap  f = C.vector . C.mapNat (Identity . f) . C.cvecF

-- | Analog of /distribute/ from /Distributive/ type class.
distribute
  :: ( Functor f, HVector v, HVectorF w,  Elems v ~ ElemsF w )
  => f v -> w f
{-# INLINE distribute #-}
distribute
  = C.vectorF
  . mapNat (fmap runIdentity . getCompose)
  . C.distributeF
  . fmap C.cvec

-- | Analog of /distribute/ from /Distributive/ type class.
distributeF
  :: ( Functor f, HVectorF v)
  => f (v g) -> v (f `Compose` g)
{-# INLINE distributeF #-}
distributeF = C.vectorF . C.distributeF . fmap C.cvecF

----------------------------------------------------------------
-- Type class based ops
----------------------------------------------------------------

-- | Replicate polymorphic value n times. Concrete instance for every
--   element is determined by their respective types.
--
-- >>> replicate (Proxy :: Proxy Monoid) mempty :: ((),String)
-- ((),"")
--
-- Or a bit contrived example which illustrate what how to call
-- function that require multiple type class constraints:
--
-- >>> replicate (Proxy @(Monoid :&&: Num)) (mempty * 10) :: (Product Int, Sum Int)
-- (Product {getProduct = 10},Sum {getSum = 0})
replicate :: (HVector v, ArityC c (Elems v))
          => Proxy c -> (forall x. c x => x) -> v
{-# INLINE replicate #-}
replicate c x = C.vector $ C.replicateF c (Identity x)

-- | Replicate monadic action n times. Example below is a bit awkward does convey what's
--
-- >>> :{
--   Prelude.mapM_ print
--     (replicateM (Proxy @(Monoid :&&: Num)) [mempty+1, mempty * 10] :: [(Product Int, Sum Int)])
-- :}
-- (Product {getProduct = 2},Sum {getSum = 1})
-- (Product {getProduct = 2},Sum {getSum = 0})
-- (Product {getProduct = 10},Sum {getSum = 1})
-- (Product {getProduct = 10},Sum {getSum = 0})
replicateM :: (HVector v, Applicative f, ArityC c (Elems v))
           => Proxy c -> (forall a. c a => f a) -> f v
{-# INLINE replicateM #-}
replicateM c x
  = fmap C.vector
  $ C.sequenceF
  $ C.replicateF c (Compose $ fmap Identity x)

-- | Replicate value @f a@ which is valid for every type a n times.
--
-- >>> replicateNatF Nothing :: HVecF '[Char,Int] Maybe
-- [Nothing,Nothing]
replicateNatF :: (HVectorF v, Arity (ElemsF v))
           => (forall a. f a) -> v f
{-# INLINE replicateNatF #-}
replicateNatF x = C.vectorF $ C.replicateNatF x

-- | Replicate polymorphic value n times:
--
-- >>> replicateF (Proxy @Num) (Just 0) :: HVecF '[Double,Int] Maybe
-- [Just 0.0,Just 0]
replicateF :: (HVectorF v, ArityC c (ElemsF v))
            => Proxy c -> (forall a. c a => f a) -> v f
{-# INLINE replicateF #-}
replicateF c x = C.vectorF $ C.replicateF c x

----------------------------------------------------------------
-- Zipping of vectors
----------------------------------------------------------------

-- | Zip two heterogeneous vectors
--
-- >>> zipWith (Proxy @Num) (+) (0, 1.2) (1, 10) :: (Int,Double)
-- (1,11.2)
zipWith :: (HVector v, ArityC c (Elems v))
        => Proxy c -> (forall a. c a => a -> a -> a) -> v -> v -> v
{-# INLINE zipWith #-}
zipWith c f v u
  = C.vector
  $ C.zipWithF c (\(Identity a) (Identity b) -> Identity (f a b)) (C.cvec v) (C.cvec u)

-- | Zip two heterogeneous vectors
zipWithF :: (HVectorF v, ArityC c (ElemsF v))
         => Proxy c -> (forall a. c a => f a -> g a -> h a) -> v f -> v g -> v h
{-# INLINE zipWithF #-}
zipWithF c f v u
  = C.vectorF $ C.zipWithF c f (C.cvecF v) (C.cvecF u)

-- | Zip two heterogeneous vectors
zipWithNatF :: (HVectorF v)
        => (forall a. f a -> g a -> h a) -> v f -> v g -> v h
{-# INLINE zipWithNatF #-}
zipWithNatF f v u
  = C.vectorF $ C.zipWithNatF f (C.cvecF v) (C.cvecF u)

-- | Zip two heterogeneous vectors and immediately fold resulting
--   value.
--
-- >>> zipFold (Proxy @Show) (\a b -> show (a,b)) ((),'c',10) ((),'D',1)
-- "((),())('c','D')(10,1)"
zipFold :: (HVector v, ArityC c (Elems v), Monoid m)
        => Proxy c -> (forall a. c a => a -> a -> m) -> v -> v -> m
{-# INLINE zipFold #-}
zipFold c f v u
  = C.zipFoldF c (\(Identity a) (Identity b) -> f a b) (C.cvec v) (C.cvec u)

zipFoldF :: (HVectorF v, ArityC c (ElemsF v), Monoid m)
        => Proxy c -> (forall a. c a => f a -> f a -> m) -> v f -> v f -> m
{-# INLINE zipFoldF #-}
zipFoldF c f v u
  = C.zipFoldF c f (C.cvecF v) (C.cvecF u)

-- | Convert heterogeneous vector to homogeneous
monomorphize :: ( HVector v
                , Peano n ~ Len (Elems v)
                , ArityC c (Elems v))
             => Proxy c -> (forall a. c a => a -> x)
             -> v -> F.ContVec n x
{-# INLINE monomorphize #-}
monomorphize c f = C.monomorphizeF c (f . runIdentity) . C.cvec

-- | Convert heterogeneous vector to homogeneous
monomorphizeF :: ( HVectorF v
                 , Peano n ~ Len (ElemsF v)
                 , ArityC c (ElemsF v)
                 )
             => Proxy c -> (forall a. c a => f a -> x)
             -> v f -> F.ContVec n x
{-# INLINE monomorphizeF #-}
monomorphizeF c f = C.monomorphizeF c f . C.cvecF

-- | Generic equality for heterogeneous vectors
--
-- >>> data A = A Int Char deriving Generic
-- >>> instance HVector A
-- >>> eq (A 1 'c') (A 2 'c')
-- False
eq :: (HVector v, ArityC Eq (Elems v)) => v -> v -> Bool
eq v u = getAll $ zipFold (Proxy :: Proxy Eq) (\x y -> All (x == y)) v u
{-# INLINE eq #-}

-- | Generic comparison for heterogeneous vectors. It works same way
--   as Ord instance for tuples.
--
-- >>> data A = A Int Char deriving Generic
-- >>> instance HVector A
-- >>> compare (A 1 'c') (A 2 'c')
-- LT
compare :: (HVector v, ArityC Ord (Elems v)) => v -> v -> Ordering
compare = zipFold (Proxy :: Proxy Ord) Prelude.compare
{-# INLINE compare #-}

-- | Reduce vector to normal form
rnf :: (HVector v, ArityC NF.NFData (Elems v)) => v -> ()
rnf = foldl (Proxy :: Proxy NF.NFData) (\r a -> NF.rnf a `seq` r) ()
{-# INLINE rnf #-}

----------------------------------------------------------------
-- Doctest
----------------------------------------------------------------

-- $setup
--
-- >>> :set -XDeriveGeneric
-- >>> :set -XTypeApplications
-- >>> :set -XTypeOperators
-- >>> :set -XDataKinds
-- >>> import Prelude (Int,Double,String,Char,IO,(++),Maybe(..))
-- >>> import Prelude (Show(..),Read(..),read,Num(..),Monoid(..))
-- >>> import Prelude (print)
-- >>> import Control.Applicative     (Const(..))
-- >>> import Data.Complex            (Complex(..))
-- >>> import Data.Monoid             (Sum(..),Product(..))
-- >>> import Data.Maybe              (fromMaybe)
-- >>> import Data.Vector.HFixed.HVec (HVec,HVecF)
-- >>> import GHC.Generics            (Generic)


{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ConstraintKinds #-}
module CodeGen.Type where

import           Data.Bits
import           Data.Typeable
import           Data.List
import           Foreign.Storable               ( Storable )
import           Data.Complex

type Serialise a = (Repr a, Show a, Read a)
type CVal a = Serialise a

type Complex a = (a, a)

-- need to find a way to represent recursive single type
data SingleType a where
    NumSingleType :: NumType a -> SingleType a
    LabelSingleType :: SingleType Label
    SumSingleType :: SingleType a -> SingleType b -> SingleType (Either a b)
    UnitSingleType :: SingleType ()
    ProductSingleType :: SingleType a -> SingleType b -> SingleType (a, b)
    ListSingleType :: SingleType a -> SingleType [a]
    FuncSingleType :: SingleType a -> SingleType b -> SingleType (a -> b)

data ASingleType where
    ASingleType :: forall a. SingleType a -> ASingleType

equal :: SingleType a -> SingleType b -> Bool
equal (NumSingleType (IntegralNumType _)) (NumSingleType (IntegralNumType _)) =
    True
equal (NumSingleType (FloatingNumType _)) (NumSingleType (FloatingNumType _)) =
    True
equal LabelSingleType     LabelSingleType       = True
equal (SumSingleType a b) (SumSingleType a' b') = equal a a' && equal b b'
equal (ProductSingleType a b) (ProductSingleType a' b') =
    equal a a' && equal b b'
equal UnitSingleType       UnitSingleType         = True
equal (ListSingleType a  ) (ListSingleType b    ) = a `equal` b
equal (FuncSingleType a b) (FuncSingleType a' b') = equal a a' && equal b b'
equal _                    _                      = False

sTypeHeight :: SingleType a -> Int
sTypeHeight (NumSingleType _)       = 0
sTypeHeight LabelSingleType         = 0
sTypeHeight (SumSingleType     a b) = 1 + max (sTypeHeight a) (sTypeHeight b)
sTypeHeight (ProductSingleType a b) = 1 + max (sTypeHeight a) (sTypeHeight b)
sTypeHeight (UnitSingleType       ) = 0
sTypeHeight (ListSingleType a     ) = 1 + sTypeHeight a
sTypeHeight (FuncSingleType _ _   ) = error "Func not height"

compareSingleType :: SingleType a -> SingleType b -> Ordering
compareSingleType a b
    | a `equal` b = EQ
    | otherwise = case (sTypeHeight a) `compare` (sTypeHeight b) of
        EQ -> LT
        x  -> x

instance Show (SingleType a) where
    show (NumSingleType (IntegralNumType _)) = "int"
    show (NumSingleType (FloatingNumType _)) = "float"
    show UnitSingleType                      = "unit"
    show LabelSingleType                     = "Label"
    show (SumSingleType a b) = intercalate "_" ["Sum", show a, show b]
    show (ProductSingleType a b) = intercalate "_" ["Prod", show a, show b]
    show (ListSingleType a     )             = intercalate "_" ["List", show a]

instance Show ASingleType where
    show (ASingleType s) = show s

toASingleType :: SingleType a -> ASingleType
toASingleType stype = ASingleType stype

instance Eq (SingleType a) where
    (==) = equal

instance Ord (SingleType a) where
    compare a b = compareSingleType a b

instance Eq (ASingleType) where
    (ASingleType left) == (ASingleType right) = equal left right

instance Ord ASingleType where
    compare (ASingleType left) (ASingleType right) =
        compareSingleType left right

data NumType a where
    IntegralNumType :: IntegralType a -> NumType a
    FloatingNumType :: FloatingType a -> NumType a

data IntegralType a where
    TypeInt     :: IntegralDict Int     -> IntegralType Int

data FloatingType a where
    TypeFloat   :: FloatingDict Float   -> FloatingType Float

data IntegralDict a where
    IntegralDict :: ( Typeable a, Bounded a, Eq a, Ord a, Show a
                    , Bits a, FiniteBits a, Integral a, Num a, Real a, Storable a )
                    => IntegralDict a

data FloatingDict a where
    FloatingDict :: ( Typeable a, Eq a, Ord a, Show a
                    , Floating a, Fractional a, Num a, Real a, RealFrac a
                    , RealFloat a, Storable a )
                    => FloatingDict a

data Label = Le | Ri
    deriving (Eq, Show, Typeable)

typeInt :: IntegralType Int
typeInt = TypeInt (IntegralDict @Int)

typeFloat :: FloatingType Float
typeFloat = TypeFloat (FloatingDict @Float)

numTypeInt :: NumType Int
numTypeInt = IntegralNumType typeInt

numTypeFloat :: NumType Float
numTypeFloat = FloatingNumType typeFloat

singleTypeInt :: SingleType Int
singleTypeInt = NumSingleType numTypeInt

singleTypeLabel :: SingleType Label
singleTypeLabel = LabelSingleType

singleTypeUnionInt :: SingleType (Either Int Int)
singleTypeUnionInt = SumSingleType singleTypeInt singleTypeInt

class Typeable a => Repr a where
    singleType :: SingleType a

instance Repr () where
    singleType = UnitSingleType

instance Repr Int where
    singleType = singleTypeInt

instance Repr Label where
    singleType = LabelSingleType

instance Repr Float where
    singleType = NumSingleType numTypeFloat

instance (Repr a, Repr b) => Repr (Either a b) where
    singleType = SumSingleType singleType singleType

instance (Repr a, Repr b) => Repr (a, b) where
    singleType = ProductSingleType singleType singleType

instance Repr a => Repr [a] where
    singleType = ListSingleType singleType

instance (Repr a, Repr b) => Repr (a -> b) where
    singleType = FuncSingleType singleType singleType


import Graphics.UI.SDL.General (withInit, InitFlag(InitVideo, InitEventthread))
import Graphics.UI.SDL.Types (Surface, surfaceGetPixels, surfaceGetPixelFormat, SurfaceFlag(SWSurface))
import Graphics.UI.SDL.Video (setVideoMode, mapRGB)
import qualified Graphics.UI.SDL.Video as V (flip)
import Graphics.UI.SDL.WindowManagement (setCaption)
import Graphics.UI.SDL.Events (Event(Quit, KeyDown), waitEvent)
import Graphics.UI.SDL.Keysym (Keysym(Keysym), SDLKey(SDLK_SPACE))
import Foreign.Ptr (castPtr)
import Data.Array (elems)
import Foreign.Marshal.Array (pokeArray)
import Control.Monad (liftM)

import Control.Applicative ((<*>))

import Data.Maybe (fromJust, isJust)
import Data.List (intersperse)

import Numeric.LinearAlgebra.Data ((|>))

import Raytracer.Geometry (cube)
import Raytracer.Camera (Camera(Camera), calculate_rays, fire_rays)

wres = 400
hres = 300
color_depth = 32

test_camera = Camera 4 3 wres hres (Just 1)
test_positions = map (3 |>) [[1, 1, -1], [2, 1, -1], [2, 2, -1], [2, 2, -2]]
test_directions = map (3 |>) [[0,0,1], [-1, 0, 1], [0,1,1], [0,-1,1]]

test_cube = cube (3 |> [2, 0, 0]) (3 |> [0, 2, 0]) (3 |> [0, 0, 2]) (3 |> [0, 0, 0])

transpose (a, b) (c, d) v = (((v-c) * (b-a)) / (d-c)) + a

compute_pixels surface array = mapM (pixel (minimum just_list, maximum just_list)) list
	where
	list = elems array
	just_list = map fromJust $ filter isJust list
	pixel_format = surfaceGetPixelFormat surface
	pixel range (Just n) = greyscale $ round $ transpose (200, 100) range n
	pixel _ Nothing = greyscale 0
	greyscale a = mapRGB pixel_format a a a

array_to_surface surface array = array_ptr >>= ((flip pokeArray) array)
	where
	array_ptr = liftM castPtr $ surfaceGetPixels surface

wait_to_quit _ = waitEvent >>= handle
	where
	handle Quit = return ()
	handle _ = wait_to_quit ()

wait_to_space _ = waitEvent >>= handle
	where
	handle (KeyDown (Keysym SDLK_SPACE _ _)) = return ()
	handle _ = wait_to_space ()

render _ = do
	screen <- setVideoMode wres hres color_depth [SWSurface]
	pixels <- mapM (compute_pixels screen) $ fmap ((flip fire_rays) test_cube) $ fmap calculate_rays $ fmap test_camera test_positions <*> test_directions
	putStrLn "Images Rendered"
	sequence $ intersperse (V.flip screen >>= wait_to_space) $ map (array_to_surface screen) pixels
	V.flip screen

main = withInit [InitVideo, InitEventthread] $ setCaption "SDL Test" "" >>= render >> putStrLn "Ready To Quit" >>= wait_to_quit


{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}
module Myelin.Neuron where

import Control.Lens hiding ((.=), (*~))
import Data.Aeson
import Numeric.LinearAlgebra

{--| NeuronType corresponds to the different point neuron
types that are constructible in SNN. Many of them
are back end specific.

Things that can be improved:

- Currently none of the parameters have any units,
  instead their units are indicated by their names

  - tau_* are time constants
  - v_* are voltages
  - i_* are currents
  - cm is a capacitance

- There is no restriction to "biological" ranges of
  parameters in place and no check of consistency

- In the case of the Heidelberg Hardware system additionally
  every neuron on the chip has analog parameter variation,
  and finite (much coarser than floating point) precision
  in parameter adjustment.

TODO:

HH_cond_exp
GIF_cond_exp

SpikeSourceGamma
SpikeSourceArray
SpikeSourcePoisson
SpikeSourcePoissonRefractory
--}
data NeuronType =
    IFCurrentAlpha {
        _tau_m :: Float, -- ^ membrane time constant
        _tau_refrac :: Float, -- ^ refractory time
        _v_thresh :: Float, -- ^ threshhold potential
        _tau_syn_E :: Float, -- ^ excitatory synaptic time constant
        _v_rest :: Float, -- ^ resting potential
        _cm :: Float, -- ^ membrane capactitance
        _v_reset :: Float, -- ^ reset potential
        _tau_syn_I :: Float, -- ^ inhibitory synaptic time constant
        _i_offset :: Float -- ^ offset current
    }
    | IFCondAlpha {
        _v_rest :: Float, -- ^ resting membrane potential
        _cm :: Float, -- ^ capacity of the membrane
        _tau_m :: Float, -- ^ membrane time constant
        _tau_refrac :: Float, -- ^ duration of refractory period
        _tau_syn_E :: Float, -- ^ rise time of the excitatory synaptic alpha function
        _tau_syn_I :: Float, -- ^ rise time of the inhibitory synaptic alpha function
        _e_rev_E :: Float, -- ^ reversal potential for excitatory input
        _e_rev_I :: Float, -- ^ reversal potential for inhibitory input
        _v_thresh :: Float, -- ^ spike threshold
        _v_reset :: Float, -- ^ reset potential after a spike
        _i_offset :: Float -- ^ offset current
    }
    | IFSpikey {
        _e_rev_I :: Float, -- ^ excitatory reversal current
        _g_leak :: Float, -- ^ leak conductance
        _tau_refrac :: Float, -- ^ refractory time
        _v_reset :: Float, -- ^ reset potential
        _v_rest :: Float, -- ^ resting potential
        _v_thresh :: Float -- ^ threshhold potential
    }
    | IFCurrExp {
        _cm :: Float, -- ^ membrane capacitance
        _tau_m :: Float, -- ^ membrane time constant
        _tau_syn_E :: Float, -- ^ excitatory synaptic time constant
        _tau_syn_I :: Float, -- ^ inhibitory synaptic time constant
        _tau_refrac :: Float, -- ^ refractory time
        _v_thresh :: Float, -- ^ threshhold voltage
        _v_rest :: Float, -- ^ resting potential
        _v_reset :: Float, -- ^ reset potential
        _i_offset :: Float -- ^ offset current
    }
    | IFCondExp {
        _v_rest :: Float, -- ^ resting potential
        _cm :: Float, -- ^ membrane capacitance
        _tau_m :: Float, -- ^ membrane time constant
        _tau_refrac :: Float, -- ^ refractory time
        _tau_syn_E :: Float, -- ^ excitatory synaptic time constant
        _tau_syn_I :: Float, -- ^ inhibitory synaptic time constant
        _e_rev_E :: Float, -- ^ excitatory reversal potential
        _e_rev_I :: Float, -- ^ inhibitory reversal potential
        _v_thresh :: Float, -- ^ spike initiation threshold
        _v_reset :: Float, -- ^ reset value for membrane potential after a spike
        _i_offset :: Float -- ^ offset current
    }
    | Izhikevich {
        _a :: Float,
        _i_offset :: Float,
        _c :: Float,
        _d :: Float,
        _b :: Float
    }
    | EIFCondExp {
        _v_reset :: Float,
        _i_offset :: Float,
        _tau_w :: Float,
        _tau_syn_I :: Float,
        _e_rev_E :: Float,
        _v_rest :: Float,
        _cm :: Float,
        _tau_syn_E :: Float,
        _tau_m :: Float,
        _a :: Float,
        _delta_T :: Float,
        _v_thresh :: Float,
        _b :: Float,
        _v_spike :: Float,
        _e_rev_I :: Float,
        _tau_refrac :: Float
    }
    | EIFCondAlpha {
        _cm :: Float, -- ^ capacity of the membrane
        _tau_refrac :: Float, -- ^ duration of the refractory period
        _v_spike :: Float, -- ^ spike detection threshold
        _v_reset :: Float, -- ^ reset value for membrane potential after a spike
        _v_rest :: Float, -- ^ resting membrane potential (Leak reversal potential)
        _tau_m :: Float, -- ^ membrane time constant
        _i_offset :: Float, -- ^ offset current
        _a :: Float, -- ^ subthreshold adaptation conductance
        _b :: Float, -- ^ spike-triggered adaptation
        _delta_T :: Float, -- ^ slope factor
        _tau_w :: Float, -- ^ adaptation time constant
        _v_thresh :: Float, -- ^ spike initiation threshold
        _e_rev_E :: Float, -- ^ excitatory reversal potential
        _tau_syn_E :: Float, -- ^ rise time of excitatory synaptic conductance (alpha function)
        _e_rev_I :: Float, -- ^ inhibitory reversal potential
        _tau_syn_I :: Float -- ^ rise time of the inhibitory synaptic conductance (alpha function)
    }
    | HHCondExp {
        _cm :: Float, -- ^ capacity of the membrane
        _e_rev_E :: Float, -- ^ excitatory reversal potential
        _e_rev_I :: Float, -- ^ inhibitory reversal potential
        _e_rev_K :: Float,
        _e_rev_Na :: Float,
        _e_rev_leak :: Float,
        _g_leak :: Float,
        _gbar_K :: Float,
        _gbar_Na :: Float,
        _i_offset :: Float,
        _tau_syn_E :: Float,
        _tau_syn_I :: Float,
        _v_offset :: Float
    } deriving (Eq, Show)

makeLenses ''NeuronType
makePrisms ''NeuronType

-- ^  Defaults taken from http://neuralensemble.org/docs/PyNN/standardmodels.html
if_cond_exp :: NeuronType
if_cond_exp = IFCondExp {
    _v_rest = -65.0,
    _cm = 1.0,
    _tau_m = 20.0,
    _tau_refrac = 0.0,
    _tau_syn_E = 5.0,
    _tau_syn_I = 5.0,
    _e_rev_E = 0.0,
    _e_rev_I = -70.0,
    _v_thresh = -50.0,
    _v_reset = -65.0,
    _i_offset = 0.0
}

if_current_alpha :: NeuronType
if_current_alpha = IFCurrentAlpha {
    _tau_m = 20.0,
    _tau_refrac = 0.1,
    _v_thresh = -50.0,
    _tau_syn_E =  0.5,
    _tau_syn_I = 0.5,
    _v_rest = -65.0,
    _cm = 1.0,
    _v_reset = -65.0,
    _i_offset = 0.0
}

if_cond_alpha :: NeuronType
if_cond_alpha = IFCondAlpha {
    _v_rest = -65.0,
    _cm = 1.0,
    _tau_m = 20.0,
    _tau_refrac = 0.0,
    _tau_syn_E = 5.0,
    _tau_syn_I = 5.0,
    _e_rev_E = 0.0,
    _e_rev_I = -70.0,
    _v_thresh = -50.0,
    _v_reset = -65.0,
    _i_offset = 0.0
}

if_spikey :: NeuronType
if_spikey = IFSpikey {
    _e_rev_I = -80.0,
    _g_leak = 20.0,
    _tau_refrac = 1.0,
    _v_reset = -80.0,
    _v_rest = -75.0,
    _v_thresh = -55.0
}

if_current_exponential :: NeuronType
if_current_exponential = IFCurrExp {
    _cm = 1.0,
    _tau_m = 20.0,
    _tau_syn_E = 5.0,
    _tau_syn_I = 5.0,
    _tau_refrac = 0.1,
    _v_thresh = -50.0,
    _v_rest = -65.0,
    _v_reset = -65.0,
    _i_offset = 0.0
}

izhikevich :: NeuronType
izhikevich = Izhikevich {
    _a = 0.02,
    _i_offset = 0.0,
    _c = -65.0,
    _d = 2.0,
    _b = 0.2
}

eif_cond_exp :: NeuronType
eif_cond_exp = EIFCondExp {
    _v_reset = -70.6,
    _i_offset = 0.0,
    _tau_w = 144.0,
    _tau_syn_I = 5.0,
    _e_rev_E = 0.0,
    _v_rest = -70.6,
    _cm = 0.281,
    _tau_syn_E = 5.0,
    _tau_m = 9.3667,
    _a = 4.0,
    _delta_T = 2.0,
    _v_thresh = -50.4,
    _b = 0.0805,
    _v_spike = -40.0,
    _e_rev_I = -80.0,
    _tau_refrac = 0.1
}

eif_cond_alpha :: NeuronType
eif_cond_alpha = EIFCondAlpha {
    _v_reset = -70.6,
    _i_offset = 0.0,
    _tau_w = 144.0,
    _tau_syn_I = 5.0,
    _e_rev_E = 0.0,
    _v_rest = -70.6,
    _cm = 0.281,
    _tau_syn_E = 5.0,
    _tau_m = 9.3667,
    _a = 4.0,
    _delta_T = 2.0,
    _v_thresh = -50.4,
    _b = 0.0805,
    _v_spike = -40.0,
    _e_rev_I = -80.0,
    _tau_refrac = 0.1
}

instance ToJSON NeuronType where
    toJSON IFCondExp {..} = object [
            "type" .= ("IFCondExp" :: String),
            "tau_m" .= _tau_m,
            "tau_refrac" .= _tau_refrac,
            "v_thresh" .= _v_thresh,
            "tau_syn_E" .= _tau_syn_E,
            "v_rest" .= _v_rest,
            "cm" .= _cm,
            "v_reset" .= _v_reset,
            "tau_syn_I" .= _tau_syn_I,
            "i_offset" .= _i_offset,
            "e_rev_E" .= _e_rev_E,
            "e_rev_I" .= _e_rev_I
        ]
    toJSON IFCurrentAlpha {..} = object [
            "type" .= ("IFCurrentAlpha" :: String),
            "tau_m" .= _tau_m,
            "tau_refrac" .= _tau_refrac,
            "v_thresh" .= _v_thresh,
            "tau_syn_E" .= _tau_syn_E,
            "v_rest" .= _v_rest,
            "cm" .= _cm,
            "v_reset" .= _v_reset,
            "tau_syn_I" .= _tau_syn_I,
            "i_offset" .= _i_offset
        ]
    toJSON IFCondAlpha {..} = object [
            "type" .= ("IFCondAlpha" :: String),
            "v_rest" .= _v_rest,
            "cm" .= _cm,
            "tau_m" .= _tau_m,
            "tau_refrac" .= _tau_refrac,
            "tau_syn_E" .= _tau_syn_E,
            "tau_syn_I" .= _tau_syn_I,
            "e_rev_E" .= _e_rev_E,
            "e_rev_I" .= _e_rev_I,
            "v_thresh" .= _v_thresh,
            "v_reset" .= _v_reset,
            "i_offset" .= _i_offset
        ]
    toJSON IFSpikey {..} = object [
            "type" .= ("IFSpikey" :: String),
            "e_rev_I" .= _e_rev_I,
            "g_leak" .= _g_leak,
            "tau_refrac" .= _tau_refrac,
            "v_reset" .= _v_reset,
            "v_rest" .= _v_rest,
            "v_thresh" .= _v_thresh
        ]
    toJSON IFCurrExp {..} = object [
            "type" .= ("IFCurrentExp" :: String),
            "cm" .= _cm,
            "tau_m" .= _tau_m,
            "tau_syn_E" .= _tau_syn_E,
            "tau_syn_I" .= _tau_syn_I,
            "tau_refrac" .= _tau_refrac,
            "v_thresh" .= _v_thresh,
            "v_rest" .= _v_rest,
            "v_reset" .= _v_reset,
            "i_offset" .= _i_offset
        ]
    toJSON Izhikevich {..} = object [
            "type" .= ("Izhikevich" :: String),
            "a" .= _a,
            "i_offset" .= _i_offset,
            "c" .= _c,
            "d" .= _d,
            "b" .= _b
        ]
    toJSON EIFCondExp {..} = object [
            "type" .= ("EIFCondExp" :: String),
            "v_reset" .= _v_reset,
            "i_offset" .= _i_offset,
            "tau_w" .= _tau_w,
            "tau_syn_I" .= _tau_syn_I,
            "e_rev_E" .= _e_rev_E,
            "v_rest" .= _v_rest,
            "cm" .= _cm,
            "tau_syn_E" .= _tau_syn_E,
            "tau_m" .= _tau_m,
            "a" .= _a,
            "delta_T" .= _delta_T,
            "v_thresh" .= _v_thresh,
            "b" .= _b,
            "v_spike" .= _v_spike,
            "e_rev_I" .= _e_rev_I,
            "tau_refrac" .= _tau_refrac
        ]
    toJSON EIFCondAlpha {..} = object [
            "type" .= ("EIFCondAlpha" :: String),
            "cm" .= _cm,
            "tau_refrac" .= _tau_refrac,
            "v_spike" .= _v_spike,
            "v_reset" .= _v_reset,
            "v_rest" .= _v_rest,
            "tau_m" .= _tau_m,
            "i_offset" .= _i_offset,
            "a" .= _a,
            "b" .= _b,
            "delta_T" .= _delta_T,
            "tau_w" .= _tau_w,
            "v_thresh" .= _v_thresh,
            "e_rev_E" .= _e_rev_E,
            "tau_syn_E" .= _tau_syn_E,
            "e_rev_I" .= _e_rev_I,
            "tau_syn_I" .= _tau_syn_I
        ]

instance FromJSON NeuronType where
    parseJSON = withObject "neuron" $ \o -> do
        typ :: String <- o .: "type"
        case typ of
            "IFCondExp" -> IFCondExp <$>
                o .: "v_rest" <*>
                o .: "cm" <*>
                o .: "tau_m" <*>
                o .: "tau_refrac" <*>
                o .: "tau_syn_E" <*>
                o .: "tau_syn_I" <*>
                o .: "e_rev_E" <*>
                o .: "e_rev_I" <*>
                o .: "v_thresh" <*>
                o .: "v_reset" <*>
                o .: "i_offset"
            "IFCurrentAlpha" -> IFCurrentAlpha <$>
                o .: "tau_m" <*>
                o .: "tau_refrac" <*>
                o .: "v_thresh" <*>
                o .: "tau_syn_E" <*>
                o .: "v_rest" <*>
                o .: "cm" <*>
                o .: "v_reset" <*>
                o .: "tau_syn_I" <*>
                o .: "i_offset"
            "IFCondAlpha" -> IFCondAlpha <$>
                o .: "v_rest" <*>
                o .: "cm" <*>
                o .: "tau_m" <*>
                o .: "tau_refrac" <*>
                o .: "tau_syn_E" <*>
                o .: "tau_syn_I" <*>
                o .: "e_rev_E" <*>
                o .: "e_rev_I" <*>
                o .: "v_thresh" <*>
                o .: "v_reset" <*>
                o .: "i_offset"
            "IFSpikey" -> IFSpikey <$>
                o .: "e_rev_I" <*>
                o .: "g_leak" <*>
                o .: "tau_refrac" <*>
                o .: "v_reset" <*>
                o .: "v_rest" <*>
                o .: "v_thresh"
            "IFCurrExp" -> IFCurrExp <$>
                o .: "cm" <*>
                o .: "tau_m" <*>
                o .: "tau_syn_E" <*>
                o .: "tau_syn_I" <*>
                o .: "tau_refrac" <*>
                o .: "v_thresh" <*>
                o .: "v_rest" <*>
                o .: "v_reset" <*>
                o .: "i_offset"
            "Izhikevich" -> Izhikevich <$>
                o .: "a" <*>
                o .: "i_offset"<*>
                o .: "c" <*>
                o .: "d" <*>
                o .: "b"
            "EIFCondExp" -> EIFCondExp <$>
                o .: "v_reset" <*>
                o .: "i_offset" <*>
                o .: "tau_w" <*>
                o .: "tau_syn_I" <*>
                o .: "e_rev_E" <*>
                o .: "v_rest" <*>
                o .: "cm" <*>
                o .: "tau_syn_E" <*>
                o .: "tau_m" <*>
                o .: "a" <*>
                o .: "delta_T" <*>
                o .: "v_thresh" <*>
                o .: "b" <*>
                o .: "v_spike" <*>
                o .: "e_rev_I" <*>
                o .: "tau_refrac"
            "EIFCondAlpha" -> EIFCondAlpha <$>
                o .: "cm" <*>
                o .: "tau_refrac" <*>
                o .: "v_spike" <*>
                o .: "v_reset" <*>
                o .: "v_rest" <*>
                o .: "tau_m" <*>
                o .: "i_offset" <*>
                o .: "a" <*>
                o .: "b" <*>
                o .: "delta_T" <*>
                o .: "tau_w" <*>
                o .: "v_thresh" <*>
                o .: "e_rev_E" <*>
                o .: "tau_syn_E" <*>
                o .: "e_rev_I" <*>
                o .: "tau_syn_I"


{-# LANGUAGE TemplateHaskell, ScopedTypeVariables, OverloadedLists, BangPatterns, Strict #-}
module Data.Mesh where

import qualified Graphics.Rendering.OpenGL.GL as GL
import           Graphics.Rendering.OpenGL.GL (($=))
import           Control.Lens      hiding (indices)
import qualified Data.Vector.Storable as VS
import           Number.Quaternion (T,real,imag)
import           Numeric.LinearAlgebra (Vector)
import qualified SDL
import qualified SDL.Image as SDLImage
import           Foreign (sizeOf)
import           Foreign.Ptr (intPtrToPtr)
import           Data.Shaders
import           Utils.Quaternions (qZero)
import           Control.Arrow
import           Control.Monad (unless)
import Control.DeepSeq
import Graphics.Rendering.OpenGL.GLU.Errors

instance NFData a => NFData (T a) where
  rnf q = let !_ = real q
              (!_,!_,!_) = imag q
          in ()

instance NFData Mesh where
  rnf m = let !_ = _quaternion m
          in ()

data Mesh = Mesh { _vertices     :: [Vector Float],
                   _indices      :: Vector GL.BaseInstance,
                   _vao          :: GL.VertexArrayObject,
                   _vbo          :: GL.BufferObject,
                   _ebo          :: GL.BufferObject,
                   _shaders      :: Shaders,
                   _quaternion   :: T Float,
                   _texture      :: Maybe GL.TextureObject,
                   _normalMap    :: Maybe GL.TextureObject }

makeLenses ''Mesh

-- loadAndBindImage :: FilePath -> IO ()
loadAndBindImage texType fn = do
  sdlSurf <- SDLImage.load fn
  SDL.V2 w h <- fmap fromIntegral <$> SDL.surfaceDimensions sdlSurf
  SDL.lockSurface sdlSurf
  texturePtr <- SDL.surfacePixels sdlSurf
  let textureSize = GL.TextureSize2D w h
      textureData = GL.PixelData GL.RGB GL.UnsignedByte texturePtr
  GL.texImage2D texType GL.NoProxy 0 GL.RGB' textureSize 0 textureData
  SDL.unlockSurface sdlSurf

loadTexture :: FilePath -> IO GL.TextureObject
loadTexture fn = do
  texName <- GL.genObjectName
  GL.textureBinding GL.Texture2D $= Just texName
  loadAndBindImage GL.Texture2D fn
  GL.textureFilter GL.Texture2D $=
    ((GL.Linear', Nothing),GL.Linear')
  return texName

createMesh (vdata, indices, verts) shad mtex norm = do
    vao <- GL.genObjectName
    GL.bindVertexArrayObject $= Just vao
    ebo <- createAndBindArrayBuffer GL.ElementArrayBuffer $ indices
    vbo <- createAndBindArrayBuffer GL.ArrayBuffer $ vdata

    setAttrib 17 3 "vertexPosition_modelspace" 0  -- vertex location
    setAttrib 17 3 "vertexNormal" 3  -- vertex normal
    setAttrib 17 2 "vertexUV" 6  -- vertex UV
    setAttrib 17 3 "vertex_col" 8  -- vertex color
    setAttrib 17 3 "vertex_tangent" 11 -- vertex tangent
    setAttrib 17 3 "vertex_bitangent" 14 -- vertex bitangent

    GL.bindVertexArrayObject $= Nothing
    return $ Mesh verts indices vao vbo ebo shad qZero mtex norm
  where
    setAttrib size attrSize name offset = do
      let glFloatSize = sizeOf (undefined :: GL.GLfloat)
          glLocation  = fromIntegral $ size * glFloatSize
          glOffset    = intPtrToPtr $ fromIntegral(offset * glFloatSize)
      loc@(GL.AttribLocation n) <- GL.get $ GL.attribLocation (shad^.program) name
      unless (n > 1000000) $ do
        GL.vertexAttribPointer loc $=
          (GL.ToFloat,
           GL.VertexArrayDescriptor attrSize GL.Float glLocation glOffset)
        GL.vertexAttribArray loc $= GL.Enabled

sendData :: forall a. VS.Storable a => GL.BufferTarget -> VS.Vector a -> IO ()
sendData t d =
  VS.unsafeWith d $
    \vptr -> GL.bufferData t $=
               (fromIntegral $ VS.length d * sizeOf (undefined :: a),
                vptr,
                GL.StaticDraw)

createAndBindArrayBuffer
  :: forall a. VS.Storable a =>
  GL.BufferTarget
  -> VS.Vector a
  -> IO GL.BufferObject
createAndBindArrayBuffer arrayType arrayData = do
    bufferObject <- GL.genObjectName
    GL.bindBuffer arrayType $= Just bufferObject
    sendData arrayType arrayData
    return bufferObject

vertexBufferData :: (VS.Vector Float, VS.Vector GL.BaseInstance, [VS.Vector Float])
vertexBufferData =
  let v = --position        normal         UV        color          tangents
          [-2.0, 0.0, -2.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.1, 0.0, 0.0, 0.0, 0.0, 1.0,
            2.0, 0.0, -2.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.1, 0.0, 0.0, 0.0, 0.0, 1.0,
           -2.0, 0.0,  2.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.1, 0.0, 0.0, 0.0, 0.0, 1.0,
            2.0, 0.0,  2.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.1, 0.0, 0.0, 0.0, 0.0, 1.0]
      v3 = [[-2.0, 0.0, -2.0],
            [ 2.0, 0.0, -2.0],
            [-2.0, 0.0,  2.0],
            [ 2.0, 0.0,  2.0]]
      i = [0,1,2,1,3,2]
  in (v,i,v3)

flatQuad shad tex norm =
  createMesh vertexBufferData shad (Just tex) (Just norm)

update :: (Float -> Float -> Bool) -> Int -> Float -> Vector Float -> Vector Float
update pred ix f v = if pred f (v VS.! ix) then VS.update_ v [ix] [f] else v

updatePair pred ix f = update pred ix f *** update ((not .) . pred) ix f

-- getAABB :: Mesh -> (Vector Float, Vector Float)
-- getAABB = VS.ifoldr inserter ([0,0,0],[0,0,0]) . view vertices
--   where inserter ix = case mod ix 17 of
--           1 -> updatePair (<) 0
--           2 -> updatePair (<) 1
--           3 -> updatePair (<) 2
--           _ -> flip const