{-# LANGUAGE DataKinds, NamedFieldPuns, GADTs, KindSignatures, TypeFamilies, TypeOperators, TypeApplications, ScopedTypeVariables, MultiParamTypeClasses, OverloadedLabels, InstanceSigs, FlexibleContexts, AllowAmbiguousTypes, FlexibleInstances, RankNTypes, UndecidableInstances, ConstraintKinds, OverlappingInstances, TypeFamilyDependencies #-}

module TraceTypes where

import           Data.Row.Records        hiding ( Disjoint
                                                , sequence
                                                , map
                                                , zip
                                                )
import           Data.Row.Internal              ( Subset
                                                , type (.\\)
                                                , type (.\/)
                                                )
import           Data.Default
import           Numeric.Log                   as Log
import qualified Data.Vector.Sized             as V
import           Data.Vector.Sized              ( Vector )
import           GHC.OverloadedLabels
import           Data.Finite
import           GHC.TypeLits
import           Control.Monad
import           Control.Monad.Bayes.Class     as Bayes
import Statistics.Distribution (logDensity, complCumulative)
import Statistics.Distribution.Gamma
import Statistics.Distribution.Beta
import Statistics.Distribution.Normal

-- The "Constraint" type is intended only
-- for internal use, and is not user-facing.
-- A Constraint of a certain record type is
-- a record where each field type `a` has been
-- replaced with the type `Maybe a`; a value of
-- `Nothing` indicates an unconstrained value,
-- whereas `Just x` indicates an observed (i.e.
-- constrained) value.
type Constraint r = Rec (Map Maybe r)

-- (ConcatTo l r s) means that l, r, and s are all
-- valid record types, and that l .+ r = r .+ l = s.
type ConcatTo l r s
  = ( WellBehaved l
    , WellBehaved r
    , WellBehaved s
    , Subset l s
    , Subset r s
    , (l .+ r) ≈ s
    , (r .+ l) ≈ s
    , s .\\ l ≈ r
    , s .\\ r ≈ l
    )

-- Are the label sets of l and r disjoint?
type Disjoint l r
  = ( ConcatTo l r (l .+ r)
    , ConcatTo (Map Maybe l) (Map Maybe r) (Map Maybe (l .+ r))
    , Forall (Map Maybe l) Default
    , Forall (Map Maybe r) Default
    )

-- Three main types: probability distributions (D), unnormalized density-carrying measures (U),
-- and probabilistic programs (P).
data U m a = U
  { usampler :: m a
  , udensity :: a -> Log Double
  }
data D m a = D
  { dsampler :: m a
  , ddensity :: a -> Log Double
  }
data P m t a = P
  { traced    :: D m (Rec t)
  , retval    :: Rec t -> a
  , constrain :: Constraint t -> U m (Rec t)
  }

class DensityCarrying f m where
  sampler :: f m a -> m a
  density :: f m a -> a -> Log Double

instance DensityCarrying U m where
  sampler = usampler
  density = udensity

instance DensityCarrying D m where
  sampler = dsampler
  density = ddensity

dist sampler density = D { dsampler = sampler, ddensity = density }
udist sampler density = U { usampler = sampler, udensity = density }

-- DEFAULT INFERENCE
observe
  :: forall m a t s
   . (Disjoint t s, Monad m)
  => P m (t .+ s) a
  -> Rec t
  -> U m (Rec s)
observe p obs =
  let
    -- Turn the observation record `obs` of type `t` into a Constraint of
    -- type `t .+ s`, by filling in the `s` piece of the record with `Nothing`.
      constraints :: Constraint (t .+ s)
      constraints = map' Just obs .+ (default' @Default def :: Constraint s)

      -- Obtain a constrained distribution over traces.
      u :: U m (Rec (t .+ s))
      u = constrain p constraints
  in
    -- Modify the distribution u so that it is over only the unobserved values
    -- (Rec s), not full-but-constrained traces (Rec (t .+ s)).
      udist (fmap restrict $ sampler u) (density u . ((.+) obs))

traceSampler = sampler . traced
traceDensity = density . traced

-- UNIT INTERVAL TYPE

newtype Unit = Sigmoid {logit :: Double}

fromUnit :: Unit -> Double
fromUnit (Sigmoid x) = 1.0 / (1.0 + exp (-x))

mkUnit :: Double -> Unit
mkUnit u = if 0 < u && u < 1
  then Sigmoid (log (u / (1.0 - u)))
  else error "Number must be between 0 and 1."

instance Show Unit where
  show x = show (fromUnit x)

------------------------------
-- STPL LANGUAGE CONSTRUCTS --
------------------------------

-- RETURN

pReturn :: Monad m => a -> P m Empty a
pReturn x =
  let sampler = return Data.Row.Records.empty
      density _ = 1
      traced = dist sampler density
      constrain _ = udist sampler density
      retval _ = x
  in  P { traced, retval, constrain }

-- BIND

pBind :: (Disjoint t s, Monad m) => P m t a -> (a -> P m s b) -> P m (t .+ s) b
pBind f g =
  let
    bindSampler = do
      fTrace <- sampler $ traced f
      let x = retval f fTrace
      gTrace <- sampler $ traced (g x)
      return (fTrace .+ gTrace)

    bindDensity tr =
      let fDensity = density (traced f) (restrict tr)
          x        = retval f (restrict tr)
          gDensity = density (traced $ g x) (restrict tr)
      in  fDensity * gDensity

    bindTraced = dist bindSampler bindDensity

    bindRet tr = let x = retval f (restrict tr) in retval (g x) (restrict tr)

    constrainedSampler c = do
      fTrace <- sampler $ constrain f (restrict c)
      let x = retval f fTrace
      gTrace <- sampler $ constrain (g x) (restrict c)
      return (fTrace .+ gTrace)

    constrainedDensity c tr =
      let fDensity = density (constrain f (restrict c)) (restrict tr)
          x        = retval f (restrict tr)
          gDensity = density (constrain (g x) (restrict c)) (restrict tr)
      in  fDensity * gDensity

    bindConstrain c = udist (constrainedSampler c) (constrainedDensity c)
  in
    P { traced = bindTraced, retval = bindRet, constrain = bindConstrain }

-- SAMPLE

singleLabelConstrain
  :: (KnownSymbol l, MonadInfer m)
  => Label l
  -> D m (Rec (l .== a))
  -> Constraint (l .== a)
  -> U m (Rec (l .== a))
singleLabelConstrain lbl traced c = case c .! lbl of
  Just x ->
    let t                  = lbl .== x
        constrainedSampler = do
          score $ density traced t
          return t
        constrainedDensity _ = density traced t
    in  udist constrainedSampler constrainedDensity

  Nothing -> udist (sampler traced) (density traced)

sample :: (KnownSymbol l, MonadInfer m) => Label l -> D m a -> P m (l .== a) a
sample lbl d =
  let sampleTraced = dist
        (do
          x <- sampler d
          return (lbl .== x)
        )
        (\t -> density d (t .! lbl))

      sampleRet t = t .! lbl

      sampleConstrain = singleLabelConstrain lbl sampleTraced
  in  P { traced    = sampleTraced
        , retval    = sampleRet
        , constrain = sampleConstrain
        }

-- WITH PROBABILITY

withProbability
  :: (KnownSymbol l, MonadInfer m)
  => Label l
  -> Unit
  -> P m t a
  -> P m s a
  -> P m (l .== Either (Rec t) (Rec s)) a
withProbability lbl logit f g =
  let p         = fromUnit logit

      wpSampler = do
        r <- random
        if r < p
          then do
            t <- sampler $ traced f
            return (lbl .== Left t)
          else do
            t <- sampler $ traced g
            return (lbl .== Right t)

      wpDensity tr = case tr .! lbl of
        Left  t -> Exp (log p) * density (traced f) t
        Right t -> Exp (log (1.0 - p)) * density (traced g) t

      wpTraced = dist wpSampler wpDensity

      wpRet t = case t .! lbl of
        Left  t -> retval f t
        Right t -> retval g t

      wpConstrain = singleLabelConstrain lbl wpTraced
  in  P { traced = wpTraced, retval = wpRet, constrain = wpConstrain }

-- FOR EACH

forEach
  :: (KnownSymbol l, KnownNat n, MonadInfer m)
  => Label l
  -> Vector n a
  -> (a -> P m t b)
  -> P m (l .== Vector n (Rec t)) (Vector n b)
forEach lbl xs body =
  let
    forEachSampler = do
      traces <- V.mapM (sampler . traced . body) xs
      return (lbl .== traces)

    forEachDensity t = V.product
      (V.map (\(x, tr) -> density (traced $ body x) tr) (V.zip xs (t .! lbl)))

    forEachTraced = dist forEachSampler forEachDensity

    forEachRet t = V.map (\(x, tr) -> retval (body x) tr) (V.zip xs (t .! lbl))

    forEachConstrain = singleLabelConstrain lbl forEachTraced
  in
    P { traced    = forEachTraced
      , retval    = forEachRet
      , constrain = forEachConstrain
      }

-- FOR ... IN RANDOM RANGE ...

forRandomRange
  :: (KnownSymbol l, MonadInfer m)
  => Label l
  -> D m Int
  -> (Int -> P m t a)
  -> P m (l .==  [Rec t]) [a]
forRandomRange lbl d body =
  let
    -- Sample a trace of the for loop.
    forSampler = do
      n      <- sampler d
      traces <- sequence $ map (sampler . traced . body) [0 .. n - 1]
      return (lbl .== traces)

    -- Compute the density of a trace of the for loop.
    forDensity t =
      density d (length $ t .! lbl)
        * product
            (map (\(tr, i) -> density (traced $ body i) tr)
                 (zip (t .! lbl) [0 ..])
            )

    -- The for loop's density-carrying distribution over traces
    forTraced = dist forSampler forDensity

    -- Compute the return value of the for loop, given a trace of an execution.
    -- The return value is just a list of the return values from each iteration.
    forRet t = map (\(tr, i) -> retval (body i) tr) (zip (t .! lbl) [0 ..])

    -- Constrain the random choices made inside the for loop.
    forConstrain = singleLabelConstrain lbl forTraced
  in
    P { traced = forTraced, retval = forRet, constrain = forConstrain }

{- WHILE:

    A stochastic while loop.
    Takes as input:
    - lbl    : The label at which to trace the while loop's random choices.
    - init   : The initial state of the while loop.
    - pi     : A function computing a probability of continuing, based on the current state.
    - pi_max : An upper bound beyond which the probability of continuing is truncated.
    - body   : The body of the while loop, which simulates a new state based on the current state.

    The resulting probabilistic program has type P m (l .== [Rec t]) a:
    - The trace contains, at the specified label, a list of subtraces, one for each iteration.
    - The return value is the final state at the finish of the while loop.

 -}

while
  :: (KnownSymbol l, MonadInfer m)
  => Label l
  -> a
  -> (a -> Unit)
  -> Unit
  -> (a -> P m t a)
  -> P m (l .== [Rec t]) a
while lbl init pi pi_max body =
  let
    -- Sample a trace of the while loop.
    whileSampler state = do
      shouldContinue <- bernoulli $ min (fromUnit (pi state)) (fromUnit pi_max)
      if shouldContinue
        then do
          nextTrace <- sampler (traced $ body state)
          let nextState = retval (body state) nextTrace
          restOfTraces <- whileSampler nextState
          return (lbl .== (nextTrace : (restOfTraces .! lbl)))
        else return (lbl .== [])

    -- Given traces from each iteration, and an initial state,
    -- compute the final state.
    retvalFromTraces init ts = case ts of
      []     -> init
      t : ts -> retvalFromTraces (retval (body init) t) ts

    -- Given a trace of the entire while loop, compute the final
    -- state.
    whileRet t = retvalFromTraces init (t .! lbl)

    -- Given an initial state and traces from each iteration,
    -- compute the density.
    densityFromTraces init ts = case ts of
        [] -> Exp (log (1.0 - (fromUnit (pi init))))
        t : ts ->
          Exp (log (fromUnit (pi init)))
            * (density (traced $ body init) t)
            * densityFromTraces (retval (body init) t) ts

    -- Given a trace of the entire while loop, compute the density.
    whileDensity t = densityFromTraces init (t .! lbl)

    -- The while loop's density-carrying distribution over traces.
    whileTraced    = dist (whileSampler init) whileDensity

    -- Constrain the choices made inside the loop.
    whileConstrain = singleLabelConstrain lbl whileTraced

  in  P { traced = whileTraced, retval = whileRet, constrain = whileConstrain }

-------------------
-- DISTRIBUTIONS --
-------------------

nrm :: MonadInfer m => Double -> Log Double -> D m Double
nrm mu (Exp sigln) = dist (normal mu (exp sigln)) (normalPdf mu (exp sigln))

brn :: MonadSample m => Unit -> D m Bool
brn logit =
  let p = fromUnit logit
  in  dist (bernoulli p) (\b -> Exp . log $ if b then p else (1.0 - p))

gmma :: MonadSample m => Log Double -> Log Double -> D m (Log Double)
gmma shape scale =
  let toDouble (Exp l) = exp l
      toLogDouble d = Exp (log d)
  in
    dist (fmap toLogDouble (gamma (toDouble shape) (toDouble scale)))
         (Exp . (logDensity $ Statistics.Distribution.Gamma.gammaDistr
                                (toDouble shape) (toDouble scale))
              . toDouble)

geom :: MonadSample m => Unit -> D m Int
geom p' =
  let p = fromUnit p'
  in dist (geometric p)
          (\n -> Exp (fromIntegral n * (log p) + (log (1.0 - p))))

pois :: MonadSample m => Int -> D m Int
pois = undefined

bta :: MonadSample m => Log Double -> Log Double -> D m Unit
bta (Exp lna) (Exp lnb) =
  let
    a = exp lna
    b = exp lnb
  in
    dist (fmap mkUnit (beta a b))
         (Exp . (logDensity $ Statistics.Distribution.Beta.betaDistr a b)
              . fromUnit)

unif :: MonadSample m => D m Unit
unif = dist (fmap mkUnit Bayes.random) (\x -> 1.0 :: Log Double)

cat :: (KnownNat n, MonadSample m) => Vector n Unit -> D m (Finite n)
cat probs = undefined

lognormal :: forall m. MonadInfer m => Log Double -> Log Double -> D m (Log Double)
lognormal (Exp logmu) sigma =
  dist (fmap Exp $ sampler $ nrm logmu sigma)
       (\(Exp logx) -> density (nrm logmu sigma :: D m Double) logx / (Exp logx))

truncnrm :: MonadInfer m => Log Double -> Log Double -> D m (Log Double)
truncnrm (Exp logmu) (Exp logsigma) =
  let mu = exp logmu
      sigma = exp logsigma
      unbounded = nrm mu (Exp logsigma)

      truncNormSampler = do
        x <- sampler unbounded
        if x < 0 then truncNormSampler else return (Exp $ log x)

      truncNormPdf (Exp logy) =
        let y = exp logy
            z = complCumulative (normalDistr mu sigma) 0
        in (density unbounded y / (Exp (log z)))
  in dist truncNormSampler truncNormPdf

----------------------------
-- PROGRAMMABLE INFERENCE --
----------------------------

-- IMPORTANCE SAMPLING

importance
  :: (Disjoint t s, MonadInfer m)
  => P m (t .+ s) a
  -> Rec t
  -> P m s b
  -> m (Rec s)
importance p t q =
  let target = observe p t
  in  do
        tr <- sampler $ traced q
        score ((density target tr) / (density (traced q) tr))
        return tr

-- SMC

unroll
  :: (MonadInfer m, Disjoint t s)
  => P m i a
  -> (a -> P m (t .+ s) a)
  -> [Rec t]
  -> m (Rec i, [Rec s])
unroll init next steps = do
  tInit <- sampler $ traced init
  let processStep soFar t = do {
    (s, results) <- soFar;
    tNext <- sampler $ observe (next s) t;
    return (retval (next s) (tNext .+ t), tNext : results)
  }
  (s, l) <- foldl processStep (return (retval init tInit, [])) steps
  return (tInit, reverse l)

particleFilter
  :: (MonadInfer m, Disjoint t s,
      Disjoint Empty i, (Empty .+ i) ≈ i)
  => P m i a
  -> P m i b
  -> (a -> P m (t .+ s) a)
  -> [(Rec t, a -> P m s c)]
  -> m (Rec i, [Rec s])
particleFilter init qInit next steps = do
  tInit <- importance init Data.Row.Records.empty qInit
  let processStep soFar (t, q) = do {
    (s, results) <- soFar;
    tNext <- importance (next s) t (q s);
    return (retval (next s) (tNext .+ t), tNext : results)
  }
  (s, l) <- foldl processStep (return (retval init tInit, [])) steps
  return (tInit, reverse l)

-- MCMC

type K m t (s :: Row *) = U m (Rec t) -> Rec t -> m (Rec t)

mh
  :: forall s t m a
   . (Disjoint t s, MonadInfer m)
  => (Rec (t .+ s) -> P m t a)
  -> K m (t .+ s) t
mh q p old =
  let proposal = do
        new <- sampler (traced $ q old)
        return $ new .+ (restrict old :: Rec s)
      rho (x, y) =
          ((density p y) * (density (traced $ q y) (restrict x)))
            / ((density p x) * (density (traced $ q x) (restrict y)))
  in  do
        proposed <- proposal
        r        <- Bayes.random
        if Exp (log r) < (min 1 (rho (old, proposed)))
          then return proposed
          else return old

seqK :: Monad m => K m t s -> K m t r -> K m t (s .\/ r)
seqK k1 k2 p old = do
  t' <- k1 p old
  k2 p t'

mixK :: MonadSample m => Double -> K m t s -> K m t r -> K m t (s .\/ r)
mixK p k1 k2 model old = do
  r <- Bayes.random
  if r < p then k1 model old else k2 model old

repeatK :: Monad m => Int -> K m t s -> K m t s
repeatK n k p old = if n == 0
  then return old
  else do
    t <- k p old
    repeatK (n - 1) k p t

ifK
  :: (Disjoint t s, Monad m)
  => (Rec t -> Bool)
  -> K m (t .+ s) s
  -> K m (t .+ s) s
ifK b k p old = if (b (restrict old)) then k p old else return old

-- VARIATIONAL (STUBBED)

svi
  :: (Disjoint t s, Monad m)
  => P m (t .+ s) a
  -> Rec t
  -> (theta -> P m s b)
  -> theta
  -> m theta
svi p t q theta = undefined
-- do tr <- traceSampler $ q theta
--    let sco  = density (observe p t) tr
--    let gradient = grad (\th -> traceDensity (q th) tr) theta
--    return theta + stepSize * sco * gradient

trainAmortized
  :: (Disjoint t s, Monad m)
  => P m (t .+ s) a
  -> (theta -> Rec t -> P m s b)
  -> theta
  -> m theta
trainAmortized p q theta = undefined
-- do tr <- traceSampler p
--    let gradient = grad (\th -> traceDensity (q th (restrict tr)) (restrict tr)) theta
--    return theta + step * gradient

ifThenElse c t f = if c then t else f


{-# LANGUAGE BangPatterns #-}
module Main where
import           Control.Monad
import           Network
import           Numeric.LinearAlgebra
import           System.Directory      (getCurrentDirectory)
import           System.IO

netConfig :: IO NeuralNet
netConfig = initWB NN
    { weights   = []
    , biases    = []
    , eta       = 3.0
    , epochs    = 30
    , layers    = 2
    , layerSize = 30
    , batchSize = 10  -- has to be a divisor of 50000
    , trainData = []
    , testData  = []
    }

main :: IO ()
main = do
    dir <- getCurrentDirectory

    trainingData <- loadData
    testingData <- loadTestData
    net <- netConfig
    model <- train net {trainData = trainingData, testData = testingData}

    forM_ (enum $ weights model) $ \(i, m) -> do
        let path = (dir ++ "/data/result/w" ++ (show i) ++ ".txt") :: FilePath
        saveMatrix path "%g" m
        return ()
    forM_ (enum $ biases model) $ \(i, b) -> do
        let path = (dir ++ "/data/result/b" ++ (show i) ++ ".txt") :: FilePath
        saveMatrix path "%g" $ asColumn b
        return ()
    where
      enum :: [a] -> [(Int, a)]
      enum a = zip (enumFrom 0) a


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE NoImplicitPrelude #-}

module Sound.StreamVector
(
    vfroml
    , vfromsl
    , vfromsrl
    , vgenerate
    , vmap
    , Vslice(..)
    , mkVslice
    , vsvec
    , vshead
    , vsUnsafeIndex
    , lfromv
    , rlfromrv
    , lvappend
    , rlvappend
    , vfromlist
    , vcfromlist
    , lcyclev
)
where

import qualified Data.Complex as Cp

import qualified Data.Vector.Generic as Vg
import qualified Data.Vector.Unboxed as Vu

import Sound.Class
import Sound.InfList
import Sound.Time

lcyclev :: (Vu.Unbox a) => Int -> a -> Vu.Vector a -> L a
lcyclev cylen z =
    lcycle cylen . lfromv z

vfroml :: (Vg.Vector v a) => Int -> L a -> v a
vfroml n s = Vg.fromList $ take n s

vfromsl :: (Vg.Vector v a) => SL Int a -> v a
vfromsl s = Vg.fromList (take (_slcount s) (_unslice s))

vfromsrl :: (Vg.Vector v a) => SRL Int a -> v a
vfromsrl = vfromsl . smap unrated

vgenerate :: (Vg.Vector v a) => Int -> (Int -> a) -> v a
vgenerate = Vg.generate

vmap :: (Vu.Unbox a, Vu.Unbox b) => (a -> b) -> Vu.Vector a -> Vu.Vector b
vmap = Vu.map

data Vslice a
    = MkVslice
    {
        _vsoffset :: !Int
        , _vsstride :: !Int
        , _vslength :: !Int
        , _vsvector :: Vu.Vector a
    }
    deriving (Read, Show)
mkVslice :: Int -> Int -> Vslice a -> Vslice a
mkVslice offset stride (MkVslice o s l v) =
    MkVslice (o + offset * s) (stride * s) (l `div` stride) v

vsvec :: (Vu.Unbox a) => Vu.Vector a -> Vslice a
vsvec x = MkVslice 0 1 (Vu.length x) x

vshead :: (Vu.Unbox a) => Vslice a -> a
vshead x = Vu.unsafeIndex (_vsvector x) (_vsoffset x)

vsUnsafeIndex :: (Vu.Unbox a) => Vslice a -> Int -> a
vsUnsafeIndex s k =
    Vu.unsafeIndex (_vsvector s) (_vsoffset s + _vsstride s * k)

{-# DEPRECATED lfromv "use 'lvappend' and 'point'" #-}
lfromv :: (Vu.Unbox a) => a -> Vu.Vector a -> L a
lfromv pad vec = lvappend vec (point pad)
{-
    loop vec
    where
        loop v =
            if Vu.null v
                then point pad
                else MkL (Vu.head v) (loop (Vu.tail v))
-}

rlfromrv :: (Vu.Unbox a) => a -> Rated (Vu.Vector a) -> Rated (L a)
rlfromrv pad = rmap (lfromv pad)

lvappend :: (Vu.Unbox a) => Vu.Vector a -> L a -> L a
lvappend vec rest =
    loop vec
    where
        loop v =
            if Vu.null v
                then rest
                else MkL (Vu.head v) (loop (Vu.tail v))

rlvappend :: (Vu.Unbox a) => Vu.Vector a -> RL a -> RL a
rlvappend vec = rmap (lvappend vec)

vfromlist :: (Vg.Vector v a) => [a] -> v a
vfromlist = Vg.fromList

vcfromlist :: (RealFloat a, Vu.Unbox a) => [a] -> Vu.Vector (Cp.Complex a)
vcfromlist = vfromlist . map (Cp.:+ 0)


module Scenes.Hexagons (getScene) where

import Numeric.LinearAlgebra (vector)

import Arc
import Group
import Helpers (degToRad, vectorToFloats)
import Hexagon
import Layer (style)
import Line
import Point
import Scene (Scene, mkSceneWithStyle)
import Shape
import Style

r :: Float
r = 0.35

--rotateArc :: Arc -> Arc
--rotateArc (Arc c ra s e) = Arc c ra (s - (degToRad 90)) (e - (degToRad 90))

connectEndpoints :: (Arc, Arc) -> Line
connectEndpoints (a, b) = Line (snd $ arcCoords a) (snd $ arcCoords b)

duplicate :: (Show a, SvgShape a, Transformable a, Mergable a) => a -> Group a
duplicate = let pts = map (\x -> Point x 0) (vectorToFloats $ vector [0..5] - 6 / 2 + 0.25) in
  Group . translatePoints pts

offsetGroup :: (Show a, SvgShape a, Transformable a, Mergable a) => Group a -> Group a
offsetGroup = translateGroupOverPoints [Point 0.5 (r * 2.1), Point 0.5 (-r * 2.1)]

alignedGroup :: (Show a, SvgShape a, Transformable a, Mergable a) => Group a -> Group a
alignedGroup = translateGroupOverPoints [Point 0 (r * 4.1), Point 0 (-r * 4.1)]

getScene :: IO Scene
getScene = pure $ mkSceneWithStyle 7 5 sceneStyle [hexLayer, arcLayer, lineLayer]
  where
      orange = "#ef6101"
      teal = "#045658"

      hexes = duplicate (Hexagon (Point 0 0) r)
      offsetHexes = offsetGroup hexes
      alignedHexes = alignedGroup hexes
      hexLayer = (toLayer "hexagon" (hexes <> offsetHexes <> alignedHexes)) {style=Just $ withStrokeColor orange}

      leftArcs = duplicate $ Arc (Point 0 0) (r * 1.15) (degToRad 90) (degToRad 180) -- <>
      leftOffsetArcs = offsetGroup leftArcs
      leftAlignedArcs = alignedGroup leftArcs

      rightArcs = duplicate $  Arc (Point 0 0) (r * 1.2) (degToRad 270) (degToRad 0)
      rightOffsetArcs = offsetGroup rightArcs
      rightAlignedArcs = alignedGroup rightArcs

      arcLayer = (toLayer "arcs"
        (leftArcs <> leftOffsetArcs <> leftAlignedArcs <> rightArcs <> rightOffsetArcs <> rightAlignedArcs)) {style=Just $ withStrokeColor teal}

      connectingLines = Group $ zipWith (curry connectEndpoints) (toList rightArcs) (tail (toList leftArcs))
      offsetLines = translate (Point (-0.01) 0) $ offsetGroup connectingLines
      alignedLines = alignedGroup connectingLines
      lineLayer = (toLayer "lines" (connectingLines <> offsetLines <> alignedLines)) {style=Just $ withStrokeColor teal}

      sceneStyle = defaultStyleAttrs { strokeColor=Just "#000000", strokeWidth=Just 0.015 }


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLabels      #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-|
Module      : Grenade.Layers.Add
Description : Layer to add some constant, with shape broadcasting
Maintainer  : Theo Charalambous
License     : BSD2
Stability   : experimental
-}

module Grenade.Layers.Add
  (
  -- * Layer instances
    Add (..)

  -- * Helper functions
  , initAdd
  ) where

import           Control.DeepSeq              (NFData (..))

import           Data.Kind                    (Type)
import           Data.Maybe                   (fromJust)
import           Data.Proxy
import           Data.Serialize
import           GHC.TypeLits

import qualified Numeric.LinearAlgebra        as LA
import           Numeric.LinearAlgebra.Static (R)
import qualified Numeric.LinearAlgebra.Static as H

import           Lens.Micro                   ((^.))

import           Grenade.Core
import           Grenade.Layers.Internal.Add
import           Grenade.Onnx

data Add :: Nat -- The number of channels of the bias
         -> Nat -- The number of rows of the bias
         -> Nat -- The number of columns of the bias
         -> Type where
  Add  :: ( KnownNat channels
          , KnownNat rows
          , KnownNat columns )
          => R (channels * rows * columns)
          -> Add channels rows columns

instance UpdateLayer (Add c h w) where
  type Gradient (Add c h w) = ()
  runUpdate _ x _  = x
  reduceGradient _ = ()

instance (KnownNat c, KnownNat h, KnownNat w ) => RandomLayer (Add c h w) where
  createRandomWith _ _ = pure initAdd

-- | Initialize an Add layer with 0 bias
initAdd :: forall c h w. ( KnownNat c, KnownNat h, KnownNat w )
        => Add c h w
initAdd =
  let c'      = fromIntegral $ natVal (Proxy :: Proxy c)
      h'      = fromIntegral $ natVal (Proxy :: Proxy h)
      w'      = fromIntegral $ natVal (Proxy :: Proxy w)
      zeroes  = replicate (c' * h' * w') 0
      bias    = H.fromList zeroes   :: R (c * h * w)
  in Add bias

instance ( KnownNat c, KnownNat h, KnownNat w ) => Serialize (Add c h w) where
  put (Add bias) = putListOf put . LA.toList . H.extract $ bias
  get            = do
    bias <- maybe (fail "Vector of incorrect size") return . H.create . LA.fromList =<< getListOf get
    return $ Add bias

-- | A three dimensions image can have a vector added to it with size equal to the number of channels
--   of the image, this corresponds to adding the nth element of the bias to every pixel of the nth
--   channel of the input image
instance ( KnownNat i, KnownNat j, KnownNat k ) => Layer (Add k 1 1) ('D3 i j k) ('D3 i j k) where
  type Tape (Add k 1 1) ('D3 i j k) ('D3 i j k) = ()

  runForwards (Add b) (S3D m)
    = let c  = fromIntegral $ natVal (Proxy :: Proxy k)
          h  = fromIntegral $ natVal (Proxy :: Proxy i)
          w  = fromIntegral $ natVal (Proxy :: Proxy j)
          m' = H.extract m
          b' = H.extract b
          r  = addPerChannel c h w m' b'
      in  ((), S3D . fromJust . H.create $ r)

  runBackwards = undefined

instance OnnxOperator (Add c h w) where
  onnxOpTypeNames _ = ["Add"]

instance (KnownNat c, KnownNat h, KnownNat w) => OnnxLoadable (Add c h w) where
  loadOnnxNode inits node = case node ^. #input of
    [bias, _] -> do
      loadedBias <- readInitializerTensorIntoVector inits bias
      return $ Add loadedBias
    _ -> onnxIncorrectNumberOfInputs

instance NFData (Add c h w) where
  rnf (Add bias) = rnf bias


module Lib.Practice.P6
    (
        f6
    ) where

import Lib.Utility.Mnist
import Control.Monad.Trans.Maybe
import Control.Monad.IO.Class
import Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra as LA

f6' :: MaybeT IO ()
f6' = do
    ((trD, trL), _) <- getMnistData
    MaybeT $ fmap Just $ print $ size trD
    let (l, d) = (extract trL LA.! 0, extract trD LA.! 0)
    liftIO $ print d
    liftIO $ print l
    -- TODO 画像の表示

f6 :: IO ()
f6 = do
    let (MaybeT result) = f6'
    result' <- result
    case result' of
        Nothing -> print "失敗"
        Just x -> return x


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RankNTypes   #-}

module Criterion.Collection.Sample
  ( Benchmark(..)
  , SampleData(..)
  , MeasurementMode(..)
  , WorkloadMode(..)
  , computeMeanAndStddev
  , compute95thPercentile
  , computeMax
  , takeSample
  , takeSamples
  ) where

import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.Trans
import           Criterion                           hiding (Benchmark)
import           Criterion.Collection.Internal.Types
import           Criterion.Config
import           Criterion.Environment
import           Criterion.IO.Printf
import           Criterion.Measurement
import           Criterion.Monad
import           Data.IORef
import           Data.List                           (foldl')
import           Data.Monoid
import qualified Data.Vector                         as V
import qualified Data.Vector.Unboxed                 as U
import           Statistics.Quantile                 (cadpw, continuousBy)
import           Statistics.Sample
import           System.Mem                          (performGC)
import           System.Random.MWC
import           Text.Printf                         (printf)

------------------------------------------------------------------------------
data MeasurementMode = PerBatch | PerOperation
data WorkloadMode = Pure | Mutating

------------------------------------------------------------------------------
data SampleData = SampleData {
      sdInputSize :: !Int       -- ^ what was the size of the input for this
                                -- sample?
    , sdNumOps    :: !Int       -- ^ how many operations are covered by this
                                -- sample? For a per-operation measurement,
                                -- this value would be \"1\", and for a batch
                                -- measurement this value would be the number
                                -- of items in the batch.
    , sdData      :: !Sample    -- ^ sample data.
    }

instance Show SampleData where
    show (SampleData is nop da) = "<SampleData inputSize=" ++ show is
                                  ++ ", nops=" ++ show nop
                                  ++ ", sample size=" ++ show (U.length da)
                                  ++ ">"

------------------------------------------------------------------------------
data Benchmark op = Benchmark {
      benchmarkName     :: String
    , dataStructures    :: [(String, DataStructure op)]
    , inputSizes        :: [Int]
    , workloadGenerator :: WorkloadGenerator op
}

------------------------------------------------------------------------------
-- | Given some sample data, compute the mean time per operation (in seconds)
-- and standard deviation
computeMeanAndStddev :: SampleData -> (Double, Double)
computeMeanAndStddev (SampleData _ nops sample) = (v,s)
  where
    nopsReal         = fromIntegral nops
    (meanValue, var) = meanVarianceUnb sample
    stddev           = sqrt $ abs var
    !v               = meanValue / nopsReal
    !s               = stddev / nopsReal

------------------------------------------------------------------------------
-- | Given some sample data, compute the 95th percentile.
compute95thPercentile :: SampleData -> Double
compute95thPercentile (SampleData _ nops sample) = v
  where
    nopsReal         = fromIntegral nops
    quantile         = continuousBy cadpw 19 20 sample
    v                = quantile / nopsReal

------------------------------------------------------------------------------
-- | Given some sample data, compute the maximum value
computeMax :: SampleData -> Double
computeMax (SampleData _ nops sample) = v
  where
    nopsReal         = fromIntegral nops
    maxval           = U.foldl' max 0 sample
    v                = maxval / nopsReal

------------------------------------------------------------------------------
takeSample :: (NFData op) =>
              MeasurementMode
           -> WorkloadMode
           -> Benchmark op
           -> Environment
           -> GenIO
           -> Int
           -> Criterion [SampleData]
takeSample !mMode !wMode !benchmark !env !rng !inputSize = do
    workload <- liftIO $ runWorkloadMonad (workGen inputSize) rng
    let setupOperations = setupWork workload
    let genWorkData     = genWorkload workload

    case mMode of
      PerBatch     -> batch setupOperations genWorkData
      PerOperation -> perOp setupOperations genWorkData

  where
    --------------------------------------------------------------------------
    dss       = dataStructures benchmark
    workGen   = workloadGenerator benchmark

    --------------------------------------------------------------------------
    batch setupOperations genWorkData = do
        workData <- liftIO $ runWorkloadMonad (genWorkData $ inputSize `div` 2)
                                              rng
        let nOps = V.length workData
        mapM (batchOne setupOperations workData nOps) dss

    --------------------------------------------------------------------------
    mkRunOp runOpMutating =
        let runOpPure = \m op -> do
                            m' <- runOpMutating m op
                            return $! m' `seq` m
        in case wMode of
             Pure     -> runOpPure
             Mutating -> runOpMutating

    --------------------------------------------------------------------------
    runWorkData workData chunkSize runOp start i val = go i val
      where
        go !i !val
            | i >= chunkSize = return val
            | otherwise = do
                  let op = V.unsafeIndex workData (start+i)
                  !val' <- runOp val op
                  go (i+1) val'

    --------------------------------------------------------------------------
    batchOne setupOperations workData nOps
             (name, (DataStructure emptyValue runOpMutating)) = do
        note $ "running batch benchmark on " ++ name ++ "\n"
        let minTime = envClockResolution env * 1000
        cfg <- getConfig
        let proc = V.foldM' runOpMutating
        let mkStartValue = emptyValue inputSize >>= flip proc setupOperations
        startValue1 <- liftIO mkStartValue

        liftIO performGC
        let tProc = runWorkData workData nOps runOpMutating 0 0

        prolix $ "running test batch with " ++ show nOps
                 ++ " work items\n"

        (tm,_) <- liftIO $ time (tProc startValue1)

        prolix $ "running initial timing on " ++ show nOps
                 ++ " work items\n"

        let iters = max 5 (ceiling (minTime / tm))

        prolix $ "running benchmark on " ++ show nOps
                 ++ " work items, " ++ show iters ++ " iterations\n"

        sample <- liftIO $ U.generateM iters $ \_ -> do
                      sv <- mkStartValue
                      performGC
                      (!tm,_) <- time (tProc sv)
                      return $ tm - clockCost

        prolix $ "finished batch benchmark on " ++ name ++ "\n"
        return (SampleData inputSize nOps sample)

    --------------------------------------------------------------------------
    perOp setupOperations genWorkData = do
        -- FIXME: lifted this code from criterion, is there some way to merge
        -- them?
        _ <- prolix "generating seed workload"
        seedData <- liftIO $ runWorkloadMonad (genWorkData 1000) rng
        _ <- prolix "seed workload generated"

        workData <- liftIO $ runWorkloadMonad (genWorkData inputSize) rng
        mapM (perOpOne setupOperations workData seedData) dss

    --------------------------------------------------------------------------
    perOpOne setupOperations workData seedData
             (name, (DataStructure emptyValue runOpMutating)) = do
        let runOp = mkRunOp runOpMutating
        let proc = V.foldM' runOpMutating

        note $ "running per-op benchmark on " ++ name ++ "\n"

        startValue <- liftIO (emptyValue inputSize >>=
                              flip proc setupOperations)
        liftIO performGC

        -- warm up clock
        _ <- liftIO $ runForAtLeast 0.1 10000 (`replicateM_` getTime)
        let minTime = envClockResolution env * 1000

        (testTime, testIters, startValue') <-
            liftIO $ timeSeed (min minTime 0.1) seedData runOp startValue
        _   <- note "ran %d iterations in %s\n" testIters (secs testTime)
        cfg <- getConfig

        let testItersD       = fromIntegral testIters
        let sampleCount      = fromLJ cfgSamples cfg
        let timePer          = (testTime - testItersD * clockCost) / testItersD
        let chunkSizeD       = minTime / timePer
        let chunkSize        = min (V.length workData) (ceiling chunkSizeD)
        let nSamples1        = min (chunkSize * sampleCount) (V.length workData)
        let numItersD        = fromIntegral nSamples1 / fromIntegral chunkSize
        let nSamples = max 1 (floor numItersD * chunkSize)

        _ <- note "collecting %d samples (in chunks of %d) in estimated %s\n"
                  nSamples chunkSize
                  (secs ((chunkSizeD * timePer + clockCost)*numItersD))

        (sample,_) <- mkSample chunkSize nSamples workData startValue' runOp
        liftIO performGC
        return (SampleData inputSize chunkSize sample)

    --------------------------------------------------------------------------
    mkSample chunkSize nSamples workData startValue runOp = liftIO $ do
        valRef <- newIORef startValue

        let numItersD = fromIntegral nSamples / fromIntegral chunkSize
        -- make sure nSamples is an integral multiple of chunkSize
        let numIters = max 1 (floor (numItersD :: Double))

        sample <- U.generateM numIters $ \chunk -> do
            !val <- readIORef valRef
            (!tm, val') <- time (runWorkData workData chunkSize runOp
                                             (chunk*chunkSize) 0 val)
            writeIORef valRef val'
            return $ tm - clockCost

        val <- readIORef valRef
        return (sample :: U.Vector Double, val)

    --------------------------------------------------------------------------
    clockCost = envClockCost env

    --------------------------------------------------------------------------
    timeSeed howLong seedData runOp startValue =
        loop startValue seedData (0::Int) =<< getTime
      where
        loop sv seed iters initTime = do
            now <- getTime
            let n = V.length seed
            when (now - initTime > howLong * 10) $
                 fail (printf "took too long to run: seed %d, iters %d"
                              (V.length seed) iters)
            (elapsed, (_,sv')) <- time (mkSample 1 n seed sv runOp)

            if elapsed < howLong
              then loop sv' (seed `mappend` seed) (iters+1) initTime
              else return (elapsed, n, sv')

------------------------------------------------------------------------------
takeSamples :: (NFData op) =>
               MeasurementMode
            -> WorkloadMode
            -> Benchmark op
            -> Environment
            -> GenIO
            -> Criterion [(String, [SampleData])]
takeSamples !mMode !wMode !benchmark !env !rng = do
    let szs = inputSizes benchmark
    when (null szs) $ fail "No input sizes defined"

    ssamples <- mapM (takeSample mMode wMode benchmark env rng) szs
    let names = map fst $ dataStructures benchmark
    let inputs = foldl' combine (map (const []) names) (reverse ssamples)

    return $ names `zip` inputs

  where
    combine :: [[SampleData]] -> [SampleData] -> [[SampleData]]
    combine int samples = map (uncurry (flip (:))) (int `zip` samples)


module Main where
import Data.Word
import Data.List.Split
import Data.Scientific as Scientific
import Data.Complex
import Control.Monad
import Codec.Picture
import Data.List
import Data.Time.Clock
import Data.Time.Calendar
import Graphics.X11.Xlib
import System.Exit (exitWith, ExitCode(..))
import Control.Concurrent
import Data.Bits
import GHC.Conc (numCapabilities)
import Graphical
import Mand
import GenerateSet
import GenerateInParallel

main :: IO ()
main = do
  taking_input <- newMVar True
 -- writePng ("benchmark.png") $ generateImage (\x y -> generatingFunction  0.0005 (-2) (2) (2) x y) 8000 8000
  generateInParallel
 -- forever $ loop taking_input

loop taking_input = do
  --print $ head $ reverse $ general_list_R mand_iterationA ((scientific 24 (-2)) :+ scientific 2 (-1))
  --print $ head $ reverse $ general_list_D mand_iteration (0.24 :+ (0.2))

  --writePng ("benchmark.png") $ generateImage (\x y -> f 0.0005 (-2) (2) (2) x y) 8000 8000
  input <- promptForImput taking_input "set for generating mandlebrot set, mov for the movement of the mandelbrot, grp for graphical, rang for range in a gif, sci for using scientific data type "
  case input of
    "set" -> generateSet taking_input
    "mov" -> generateMovement
    "grp" -> forkIO (startGraphical taking_input)>> return ()
    "rang" -> range
    "sci" -> generateSciSet
    "series" -> generateSeriesOfSets
    "test" -> tester
    "acc" -> generateAccurateSet
    "par" -> generateInParallel

generateAccurateSet = print "yet again not finishe"

tester = do
  writePng ("benchmark.png") $ generateImage (\x y -> generatingFunction  0.0005 (-2) (2) (2) x y) 8000 8000
--stepf srf sif realend x y
generateSeriesOfSets = print "in complete"

generateSciSet = do
  print "nothign yet"
-- print $ mi (makeSciComplex "1e-1" "2e-1") $ makeSciComplex "-9e-1" "4e-2"

--mi :: CS -> CS -> CS
--mi c z = c + z*z

makeSciComplex :: String -> String -> CS
makeSciComplex x y = (read x :: S) :+ (read x :: S)

range = do
  print "in complete"

generateMovement = do
  let list = [fromIntegral x/10:+ fromIntegral y/10 |x <-[(-400)..399], y <- [(-400)..399]]
  let zlist = foldl (\acc x -> (0:+0):acc) [] [0..639999]
  let list' = map' mand_iteration list list
  --print list
  --writePng "movv5_i1.png" $ generateImage (genMovImg $ map roundComplex2 list) 399 399
  --writePng "movv5_i2.png" $ generateImage (genMovImg $ map roundComplex2 list') 399 399
  --let list'1 = filterC $ map' mand_iteration list list'
  --writePng "movv5_i3.png" $ generateImage (genMovImg $ map roundComplex2 list'1) 399 399
  --let list'2 = filterC $ map' mand_iteration list list'1
  --writePng "movv5_i4.png" $ generateImage (genMovImg $ map roundComplex2 list'2) 399 399
  --let list'3 = filterC $ map' mand_iteration list list'2
  --writePng "movv5_i5.png" $ generateImage (genMovImg $ map roundComplex2 list'3) 399 399
  --let list'4 = filterC $ map' mand_iteration list list'3

  mapM_ (curry' writePng) $ reverse $ snd $ mapAccumR nextMov (list, zlist) [0..10]
  print "complete"

curry' :: (a -> b -> c) -> (a, b) ->  c
curry' f (a,b) = f a b

nextMov (last, original) x = ((nextl, nextOrig), ("movv10_i"++(show x)++".png",generateImage (genMovImg $ map roundComplex2 nextl) 3999 3999))
  where
    (nextl, nextOrig) = filter2 (map' mand_iteration original last ) original removeLargeC

removeLargeC :: (Num a, Ord a )=> Complex a -> Bool
removeLargeC (a :+ b) = if abs a >2 || abs b > 2 then False else True

filterC = filter removeLargeC

filter2 :: [a] -> [b] -> (a -> Bool) -> ([a],[b])
filter2 x y f = (fst ziped,snd ziped)
  where
    ziped = foldr (\(a,b) (ac,bc) -> if f a then (a:ac,b:bc) else (ac,bc)) ([],[]) $ zip x y

genMovImg arr x y = if elem (fromIntegral (x-2000) /1000 :+ fromIntegral (y-2000)/1000) arr then PixelRGB8 0 0 0 else PixelRGB8 255 255 255

map' :: (a -> a -> b) -> [a] -> [a] -> [b]
map' _ [] [] = []
map' f (x:xs) (y:ys) = (f x y):(map' f  xs ys)

g :: [Bool]
g = foldr (\x acc -> if mod x 10 == 0 then True:acc else False:acc) [] [0..399]

h :: [[Bool]]
h = foldr (\x acc -> if mod x 10 == 0 then g:acc else i:acc) [] [0..399]

i :: [Bool]
i = foldr (\_ acc -> False:acc) [] [0..399]

falseArr = foldr (\_ acc -> i:acc) [] [0..399]

toComplex :: RealFloat a => [[Bool]] -> [Complex a]
toComplex arr = foldr (\x acc1 -> (foldr (\y acc2 -> if arr !! x !! y then (fromIntegral (x-200)/100 :+ fromIntegral (y-200)/100):acc2 else acc2) [] [0..((length $ arr !! 0)-1)])++acc1) [] [0..((length arr) -1)]

toArr :: RealFloat a => [Complex a] -> [[Bool]]
toArr arr = f arr falseArr
  where
    rarr = map roundComplex2 arr
    f :: RealFloat a=> [Complex a] -> [[Bool]]-> [[Bool]]
    f [] fa  = fa
    f (x:xs) fa = f xs $ modifyarr fa  newx newy
      where
        newx = round $ 1000*(realPart x) +2000
        newy = round $ 1000*(imagPart x) +2000

modifyarr :: [[Bool]] -> Int -> Int -> [[Bool]]
modifyarr arr x y = foldr (\a acc -> (foldr (\b acc2 -> if x == a && b == y then True:acc2 else (arr !! x !! y):acc2) [] [0..((length $ arr !! 0)-1)]):acc) [] [0..((length arr)-1)]

roundComplex2 :: RealFloat a => Complex a -> Complex a
roundComplex2 (a :+ b) = (fromIntegral (round $ a*1000)/1000 :+ fromIntegral (round $ b*1000)/1000)


{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Feldspar.Algorithm.FFT.Push
  ( fft
  , butterfly
  , twid
  , itwids
  , twids
  ) where

import qualified Prelude as P
import qualified Data.Complex

import Feldspar
import Feldspar.Mutable
import Feldspar.Vector

import Feldspar.Algorithm.FFT.Utils
import Feldspar.Algorithm.FFT.Twids

import Test.QuickCheck

-- | Utilities that should go into Feldspar.Vector
chnk :: forall arr1 a b . (Pushy arr1, VecShape arr1 ~ DIM1, Syntax b)
      => Data Length          -- ^ Number of chunks
      -> Data Length          -- ^ Size of the chunks
      -> (Pull DIM1 a -> arr1 b) -- ^ Applied to every chunk
      -> Pull DIM1 a
      -> Push DIM1 b
chnk r c f v = Push loop $ extent v
  where loop :: PushK DIM1 b
        loop func = parFor r $ \i ->
                      let (Push k _) = toPush $ f (take c (drop (c*i) v))
                      in k (\(Z:.j) a -> func (Z:.(c*i + j)) a)

unhalve :: (Syntax a)
        => Pull DIM1 (a,a) -> Push DIM1 a
unhalve xs = unpairWith id (\(Z:.i) -> Z:.(i + length xs)) xs

stride :: Data Length -> Data Length
       -> (Data Index -> a -> M b)
       -> Data Index -> (a,a) -> M b
stride n k f ix (a1,a2) = f (n*ix) a1 >> f (n*ix+k) a2

-- | DFT2 for Decimation-In-Frequency
dft2 :: Num a => a -> (a, a) -> (a,a)
dft2 w (x0,x1) = (x0+x1, (x0-x1)*w)

butterfly :: (Syntax a, Num a)
          => (a -> (a,a) -> (a,a))
          -> Pull DIM1 a -> Pull DIM1 a -> Push DIM1 a
butterfly f ws = unhalve . zipWith f ws . uncurry zip . halve

-- | Cooley-Tukey Radix-2 Decimation In Frequency Fast Fourier Transfrom
fft :: (Syntax a, Num a)
    => Pull DIM1 a -> Pull DIM1 a -> Pull DIM1 a
fft ws vs = forLoop (ilog2 len) vs stage
  where
    len = length vs
    stage s = withLen len
            $ toPull
            . store
            . chnk (1 .<<. s) (len .>>. s) (butterfly dft2 (ixmap (.<<. s) ws))


{-# LANGUAGE GADTs, FlexibleContexts, RankNTypes #-}

module Numeric.Trainee.Neural (
	Net, net, netSeed,
	(⭃),
	ndup, npar, ndupWith, nsum, nfold,
	pad, pad2,
	maxPool, maxPool2,
	fc, conv, conv2,
	dconv, dconv2,

	sigma, relu, softplus,
	summator, biaser, activator, convolver, convolver2,
	normVar,

	module Numeric.Trainee.Types,
	module Numeric.Trainee.Learnee
	) where

import Prelude hiding ((.), id)
import Prelude.Unicode

import Control.Monad (replicateM, liftM2)
import Data.Random
import qualified Data.Vector as V
import Numeric.LinearAlgebra hiding (conv, conv2)
import System.Random (newStdGen, StdGen)

import Numeric.Trainee.Types
import Numeric.Trainee.Gradee
import Numeric.Trainee.Learnee

type Net a = Learnee (Vector a) (Vector a)

net ∷ RVar a → IO a
net n = netSeed <$> newStdGen <*> pure n

netSeed ∷ StdGen → RVar a → a
netSeed g n = fst $ sampleState n g

(⭃) ∷ RVar (Learnee a b) → RVar (Learnee b c) → RVar (Learnee a c)
n ⭃ l = liftM2 (⇉) n l

-- | Split layers, just wraps @vdup@
ndup ∷ Num a ⇒ Int → RVar (Learnee a (V.Vector a))
ndup = return ∘ computee ∘ vdup

-- | Process vector of layers
npar ∷ Int → RVar (Learnee a b) → RVar (Learnee (V.Vector a) (V.Vector b))
npar n = fmap (parallel ∘ V.fromList) ∘ replicateM n

-- | Split layers, justs wraps @vdupWith@
ndupWith ∷ (a → a → a) → Int → RVar (Learnee a (V.Vector a))
ndupWith fn = return ∘ computee ∘ vdupWith fn

-- | Join layers with sum
nsum ∷ Num a ⇒ RVar (Learnee (V.Vector a) a)
nsum = return (computee vsum)

-- | Join layers with custom sum
nfold ∷ (a → a → a) → RVar (Learnee (V.Vector a) a)
nfold = return ∘ computee ∘ vfold

-- | Padding
pad ∷ Numeric a ⇒ Int → RVar (Learnee (Vector a) (Vector a))
pad = return ∘ computee ∘ padVec

-- | 2d padding
pad2 ∷ Numeric a ⇒ Int → Int → RVar (Learnee (Matrix a) (Matrix a))
pad2 w h = return ∘ computee $ padMat w h

-- | Max pool layer
maxPool ∷ (RealFrac a, Numeric a) ⇒ Int → RVar (Learnee (Vector a) (Vector a))
maxPool = return ∘ computee ∘ maxPoolVec

-- | 2d max pool layer
maxPool2 ∷ (RealFrac a, Numeric a) ⇒ Int → Int → RVar (Learnee (Matrix a) (Matrix a))
maxPool2 w h = return ∘ computee $ maxPoolMat w h

-- | Fully connected layer
fc ∷ (Distribution Normal a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Unary (Vector a) → Int → Int → RVar (Net a)
fc f inputs outputs = do
	s ← summator inputs outputs
	b ← biaser outputs
	return $ s ⇉ b ⇉ activator f

-- | Convolution layer 1-d
conv ∷ (Distribution Normal a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Unary (Vector a) → Int → RVar (Net a)
conv f w = do
	c ← convolver w
	b ← do
		bs ← normVar
		return $ learnee biasVec bs
	return $ c ⇉ b ⇉ activator f

-- | Convolution layer 2-d
conv2 ∷ (Distribution Normal a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Unary (Matrix a) → (Int, Int) → RVar (Learnee (Matrix a) (Matrix a))
conv2 f (w, h) = do
	c ← convolver2 w h
	b ← do
		bs ← normVar
		return $ learnee biasMat bs
	return $ c ⇉ b ⇉ activator f

-- | Depth (with several input and output channels) convolution 1-d
dconv ∷ (Distribution Normal a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Unary (Vector a) → Int → Int → Int → RVar (Learnee (V.Vector (Vector a)) (V.Vector (Vector a)))
dconv f inputs outputs width =
	ndupWith (V.zipWith (+)) outputs ⭃
	npar outputs (npar inputs (conv f width) ⭃ nsum)

-- | Depth (with several input and output channels) convolution 2-d
dconv2 ∷ (Distribution Normal a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Unary (Matrix a) → Int → Int → (Int, Int) → RVar (Learnee (V.Vector (Matrix a)) (V.Vector (Matrix a)))
dconv2 f inputs outputs (w, h) =
	ndupWith (V.zipWith (+)) outputs ⭃
	npar outputs (npar inputs (conv2 f (w, h)) ⭃ nsum)

sigma ∷ Floating a ⇒ a → a
sigma t = 1 / (1 + exp (negate t))

relu ∷ Fractional a ⇒ a → a
relu t = 0.5 * (1 + signum t) * t

softplus ∷ Floating a ⇒ a → a
softplus t = log (1 + exp t)

summator ∷ (Distribution Normal a, Fractional a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Int → Int → RVar (Net a)
summator inputs outputs = do
	ws ← replicateM (inputs * outputs) normVar
	return $ learnee matVec ((outputs >< inputs) ws)

biaser ∷ (Distribution Normal a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Int → RVar (Net a)
biaser sz = do
	bs ← replicateM sz normVar
	return $ learnee odot (fromList bs)

activator ∷ Num a ⇒ Unary a → Learnee a a
activator f = computee (unary f)

convolver ∷ (Distribution Normal a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Int → RVar (Net a)
convolver w = do
	ws ← replicateM w normVar
	return $ learnee corrVec (fromList ws)

convolver2 ∷ (Distribution Normal a, Fractional a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Int → Int → RVar (Learnee (Matrix a) (Matrix a))
convolver2 w h = do
	ws ← replicateM (w * h) normVar
	return $ learnee corrMat ((w >< h) ws)

normVar ∷ (Distribution Normal a, Fractional a) ⇒ RVar a
normVar = normal 0.0 0.25


-- | An ad-hoc representation of matrices that saves allocations
-- and speeds up computing gradients. The major improvement comes
-- from avoiding, often, the construction of some matrices,
-- e.g., outer products or repeated columns that naturally occur
-- in delta-expressions. Minor improvements come from better fusion
-- (via vector or hmatrix/blas/lapack) and so less allocation
-- of intermediate matrix results of arithmetic and other operations.
--
-- The latter improvements are very likely tied to the vagaries
-- of hmatrix (or, less likely, vector or blas/lapack) that work underneath
-- and apparently conspire to fuse some matrix operations but not others.
-- That would explain why vector-based MNIST nn allocates not much less
-- than matrix-based MNIST, despite producing much larger delta-expressions.
-- Matrix operations probably fuse worse, even though they are really
-- vector operations on the underlying vector representation,
-- because the dimensions book-keeping of the matrix comes in the way
-- and also because some operations work on columns, against the grain
-- of the representation.
--
-- For fully connected MNIST nns 10 times larger than usual, vector-based
-- implementation allocates similarly and runs faster than matrix-based,
-- despite the latter spending 8 times less in GC. probably because
-- matrices don't fit in cache at this size and vectors still do.
module HordeAd.Internal.MatrixOuter
  ( MatrixOuter (..)
  , emptyMatrixOuter, nullMatrixOuter, rows, cols
  , convertMatrixOuter, convertMatrixOuterOrNull
  , toRows, toColumns, plus, multiplyWithOuter
  , konst, asRow, asColumn, rowAppend, columnAppend
  , takeRows, takeColumns, dropRows, dropColumns
  , transpose, matMul, flipud, fliprl
  ) where

import Prelude

import           Control.Exception.Assert.Sugar
import           Data.Maybe (fromJust)
import qualified Data.Vector.Generic as V
import           Numeric.LinearAlgebra (Matrix, Numeric, Vector)
import qualified Numeric.LinearAlgebra as HM

-- | A representation of a matrix as a product of a basic matrix
-- and an outer product of two vectors. Each component defaults to ones.
data MatrixOuter r = MatrixOuter (Maybe (Matrix r))
                                 (Maybe (Vector r)) (Maybe (Vector r))

emptyMatrixOuter :: MatrixOuter r
emptyMatrixOuter = MatrixOuter Nothing Nothing Nothing

nullMatrixOuter :: MatrixOuter r -> Bool
nullMatrixOuter (MatrixOuter Nothing Nothing Nothing) = True
nullMatrixOuter _ = False

rows :: Numeric r => MatrixOuter r -> Int
rows (MatrixOuter (Just m) _ _) = HM.rows m
rows (MatrixOuter Nothing (Just c) _) = V.length c
rows _ = error "rows: dimensions can't be determined"

cols :: Numeric r => MatrixOuter r -> Int
cols (MatrixOuter (Just m) _ _) = HM.cols m
cols (MatrixOuter Nothing _ (Just r)) = V.length r
cols _ = error "cols: dimensions can't be determined"

convertMatrixOuter :: (Numeric r, Num (Vector r)) => MatrixOuter r -> Matrix r
convertMatrixOuter (MatrixOuter (Just m) Nothing Nothing) = m
convertMatrixOuter (MatrixOuter (Just m) (Just c) Nothing) = m * HM.asColumn c
  -- This @asColumn@ is safe, because it's expanded via @m@.
  --
  -- Beware, depending on context, @m * outer c (konst 1 (cols m))@
  -- may allocate much less and perhaps @fromRows . toRowsMatrixOuter@
  -- may be even better; that's probably vector being picky about fusing.
  -- It doesn't matter if @m@ comes first.
convertMatrixOuter (MatrixOuter (Just m) Nothing (Just r)) = m * HM.asRow r
  -- This @asRow@ is safe, because it's expanded via @m@.
convertMatrixOuter (MatrixOuter (Just m) (Just c) (Just r)) = m * HM.outer c r
convertMatrixOuter (MatrixOuter Nothing (Just c) (Just r)) = HM.outer c r
convertMatrixOuter _ =
  error "convertMatrixOuter: dimensions can't be determined"

convertMatrixOuterOrNull :: (Numeric r, Num (Vector r))
                         => MatrixOuter r -> Matrix r
convertMatrixOuterOrNull (MatrixOuter Nothing Nothing Nothing) = HM.fromRows []
convertMatrixOuterOrNull m = convertMatrixOuter m

toRows :: (Numeric r, Num (Vector r)) => MatrixOuter r -> [Vector r]
toRows (MatrixOuter (Just m) Nothing Nothing) = HM.toRows m
toRows (MatrixOuter (Just m) mc Nothing) =
  maybe id (zipWith HM.scale . V.toList) mc $ HM.toRows m
toRows (MatrixOuter (Just m) mc (Just r)) =
  maybe (map (r *)) (zipWith (\s row -> r * HM.scale s row) . V.toList) mc
  $ HM.toRows m
toRows (MatrixOuter Nothing (Just c) (Just r)) =
  map (`HM.scale` r) $ V.toList c
toRows _ = error "toRows: dimensions can't be determined"

toColumns :: (Numeric r, Num (Vector r)) => MatrixOuter r -> [Vector r]
toColumns (MatrixOuter (Just m) Nothing Nothing) = HM.toColumns m
toColumns (MatrixOuter (Just m) Nothing mr) =
  maybe id (zipWith HM.scale . V.toList) mr $ HM.toColumns m
toColumns (MatrixOuter (Just m) (Just c) mr) =
-- traceShow (sizeMatrixOuter (MatrixOuter (Just m) (Just c) mr)) $
  maybe (map (c *)) (zipWith (\s col -> c * HM.scale s col) . V.toList) mr
  $ HM.toColumns m
toColumns (MatrixOuter Nothing (Just c) (Just r)) =
  map (`HM.scale` c) $ V.toList r
toColumns _ = error "toColumns: dimensions can't be determined"

size :: Numeric r => MatrixOuter r -> (Maybe (Int, Int), Maybe Int, Maybe Int)
size (MatrixOuter mm mc mr) = (HM.size <$> mm, HM.size <$> mc, HM.size <$> mr)

plus :: (Numeric r, Num (Vector r))
     => MatrixOuter r -> MatrixOuter r -> MatrixOuter r
plus o1 o2 =
  let !o = convertMatrixOuter o1 + convertMatrixOuter o2
  in MatrixOuter (Just o) Nothing Nothing
    -- TODO: Here we allocate up to 5 matrices, but we should allocate one
    -- and in-place add to it and multiply it, etc., ideally using raw FFI
    -- or at least vector streams/bundles; an easy option to have only
    -- one allocation is to do this pointwise, but then the speedup, if any,
    -- from blas/lapack does not apply

multiplyWithOuter :: (Numeric r, Num (Vector r))
                  => Matrix r -> MatrixOuter r -> MatrixOuter r
multiplyWithOuter k o@(MatrixOuter mm mc mr) =
  assert (HM.size k
          == maybe (HM.size $ fromJust mc, HM.size $ fromJust mr) HM.size mm
         `blame` (HM.size k, size o)) $
  let !m = maybe k (k *) mm
  in MatrixOuter (Just m) mc mr

konst :: Numeric r => r -> Int -> Int -> MatrixOuter r
konst x nrows ncols =
  let !c = HM.konst x nrows
      !r = HM.konst x ncols
  in MatrixOuter Nothing (Just c) (Just r)

asRow :: Numeric r => Vector r -> Int -> MatrixOuter r
asRow !r nrows =
  let !c = HM.konst 1 nrows
  in MatrixOuter Nothing (Just c) (Just r)

asColumn :: Numeric r => Vector r -> Int -> MatrixOuter r
asColumn !c ncols =
  let !r = HM.konst 1 ncols
  in MatrixOuter Nothing (Just c) (Just r)

rowAppend :: (Numeric r, Num (Vector r))
          => MatrixOuter r -> MatrixOuter r -> MatrixOuter r
rowAppend o1 o2 =
  let !o = convertMatrixOuter o1 HM.=== convertMatrixOuter o2
  in MatrixOuter (Just o) Nothing Nothing

columnAppend :: (Numeric r, Num (Vector r))
             => MatrixOuter r -> MatrixOuter r -> MatrixOuter r
columnAppend o1 o2 =
  let !o = convertMatrixOuter o1 HM.||| convertMatrixOuter o2
  in MatrixOuter (Just o) Nothing Nothing

takeRows :: Numeric r => Int -> MatrixOuter r -> MatrixOuter r
takeRows k (MatrixOuter mm mc mr) =
  MatrixOuter (HM.takeRows k <$> mm)
              (V.take k <$> mc)
              mr

takeColumns :: Numeric r => Int -> MatrixOuter r -> MatrixOuter r
takeColumns k (MatrixOuter mm mc mr) =
  MatrixOuter (HM.takeColumns k <$> mm)
              mc
              (V.take k <$> mr)

dropRows :: Numeric r => Int -> MatrixOuter r -> MatrixOuter r
dropRows k (MatrixOuter mm mc mr) =
  MatrixOuter (HM.dropRows k <$> mm)
              (V.drop k <$> mc)
              mr

dropColumns :: Numeric r => Int -> MatrixOuter r -> MatrixOuter r
dropColumns k (MatrixOuter mm mc mr) =
  MatrixOuter (HM.dropColumns k <$> mm)
              mc
              (V.drop k <$> mr)

transpose :: Numeric r => MatrixOuter r -> MatrixOuter r
transpose (MatrixOuter mm mc mr) = MatrixOuter (HM.tr' <$> mm) mr mc

matMul :: (Numeric r, Num (Vector r))
       => MatrixOuter r -> MatrixOuter r -> MatrixOuter r
matMul m1 m2 =
  MatrixOuter (Just $ convertMatrixOuter m1 HM.<> convertMatrixOuter m2)
              Nothing Nothing

flipud :: Numeric r => MatrixOuter r -> MatrixOuter r
flipud (MatrixOuter mm mc mr) =
  MatrixOuter (HM.flipud <$> mm) mr (V.reverse <$> mc)

fliprl :: Numeric r => MatrixOuter r -> MatrixOuter r
fliprl (MatrixOuter mm mc mr) =
  MatrixOuter (HM.fliprl <$> mm) (V.reverse <$> mr) mc


module Types where

import System.Random
import Numeric.Vector

data Camera = Camera { c_getPos :: Vec3d
                     , c_getForward :: Vec3d
                     , c_getUp :: Vec3d
                     , c_getRight :: Vec3d
                     , c_getFOV :: Double
                     }
            | OrthoCamera { c_getPos :: Vec3d
                          , c_getForward :: Vec3d
                          , c_getUp :: Vec3d
                          , c_getRight :: Vec3d
                          , c_getSize :: Double
                          } deriving Show

data Scene = Scene { s_getCamera :: Camera
                   , s_getObjects :: [Object]
                   , s_getLights :: [Light]
                   , s_getSkyColor :: Vec3d
                   , s_getWidth :: Int
                   , s_getHeight :: Int
                   , s_getSeed :: Int
                   , s_getSamples :: Int
                   , s_getBounces :: Int
                   } deriving Show

data SceneContext = SceneContext { ss_getScene :: Scene
                                 , ss_getGen :: StdGen
                                 , ss_getHit :: RayHit
                                 } deriving Show

data Ray = Ray Vec3d Vec3d deriving Show
r_getOrigin :: Ray -> Vec3d
r_getOrigin (Ray o _) = o
r_getDir :: Ray -> Vec3d
r_getDir (Ray _ d) = d

data RayHit = RayHit { rh_getPos :: Vec3d
                     , rh_getRay :: Ray
                     , rh_getDistance :: Double
                     , rh_getTri :: Tri
                     , rh_getObj :: Object
                     , rh_getShader :: BRDF
                     , rh_getTexCoord :: Vec2d
                     , rh_getNormal :: Vec3d
                     , rh_getTangent :: Vec3d
                     , rh_getBitangent :: Vec3d
                     , rh_getColor :: Vec3d
                     } deriving Show

data SVOp = VAdd | VSub | VDot | VCross | VMul | VDiv | VAbs | VMod deriving Show
data SOp = Add | Sub | Mul | Div | Abs | Mod | LessThan | GreaterThan deriving Show
data SVector = UV
             | Position
             | Normal
             | VecConst Double Double Double
             | VecMath SVOp SVector SVector
             | CombineXYZ SValue SValue SValue
             deriving Show
data SValue = ValConst Double
            | RayLength
            | ValMath SOp SValue SValue
            | SeparateX SVector | SeparateY SVector | SeparateZ SVector
            | SeparateR SColor | SeparateG SColor | SeparateB SColor
            deriving Show
data SColor = ColVertex
            | ColVector SVector
            | ColConst Double Double Double
            | CombineRGB SValue SValue SValue
            deriving Show
data BRDF = Emission SColor SValue
          | Diffuse SColor
          | Glossy SColor SValue
          | BRDFEmpty
          deriving Show

data Light = Directional { l_getAngle :: Double
                         , l_getStrength :: Double
                         , l_getDirection :: Vec3d
                         , l_getColor :: Vec3d
                         }
           | Point { l_getRadius :: Double
                   , l_getStrength :: Double
                   , l_getColor :: Vec3d
                   } deriving Show

data Object = Object { o_getTris :: [Tri]
                     , o_getShader :: BRDF
                     } deriving Show

data Tri = Tri { t_getP0 :: Vertex
               , t_getP1 :: Vertex
               , t_getP2 :: Vertex
               } deriving Show

data Vertex = Vertex { v_getPos :: Vec3d
                     , v_getColor :: Vec3d
                     , v_getTexCoord :: Vec2d
                     , v_getNormal :: Vec3d
                     , v_getTangent :: Vec3d
                     , v_getBitangent :: Vec3d
                     } deriving Show


{-# OPTIONS_GHC -Wall -Wno-partial-type-signatures #-}

{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE PartialTypeSignatures #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Sundials.CVode.ODE
-- Copyright   :  Dominic Steinitz 2018,
--                Novadiscovery 2018
-- License     :  BSD
-- Maintainer  :  Dominic Steinitz
-- Stability   :  provisional
--
-- Solution of ordinary differential equation (ODE) initial value problems.
--
-- <https://computation.llnl.gov/projects/sundials/sundials-software>
--
-- A simple example:
--
-- <<diagrams/brusselator.png#diagram=brusselator&height=400&width=500>>
--
-- @
-- import           Numeric.Sundials.CVode.ODE
-- import           Numeric.LinearAlgebra
--
-- import           Plots as P
-- import qualified Diagrams.Prelude as D
-- import           Diagrams.Backend.Rasterific
--
-- brusselator :: Double -> [Double] -> [Double]
-- brusselator _t x = [ a - (w + 1) * u + v * u * u
--                    , w * u - v * u * u
--                    , (b - w) / eps - w * u
--                    ]
--   where
--     a = 1.0
--     b = 3.5
--     eps = 5.0e-6
--     u = x !! 0
--     v = x !! 1
--     w = x !! 2
--
-- lSaxis :: [[Double]] -> P.Axis B D.V2 Double
-- lSaxis xs = P.r2Axis &~ do
--   let ts = xs!!0
--       us = xs!!1
--       vs = xs!!2
--       ws = xs!!3
--   P.linePlot' $ zip ts us
--   P.linePlot' $ zip ts vs
--   P.linePlot' $ zip ts ws
--
-- main = do
--   let res1 = odeSolve brusselator [1.2, 3.1, 3.0] (fromList [0.0, 0.1 .. 10.0])
--   renderRasterific "diagrams/brusselator.png"
--                    (D.dims2D 500.0 500.0)
--                    (renderAxis $ lSaxis $ [0.0, 0.1 .. 10.0]:(toLists $ tr res1))
-- @
--
-----------------------------------------------------------------------------
module Numeric.Sundials.CVode.ODE ( odeSolve
                                   , odeSolveV
                                   , odeSolveVWith
                                   , odeSolveVWith'
                                   , odeSolveRootVWith'
                                   , odeSolveWithEvents
                                   , ODEMethod(..)
                                   , StepControl(..)
                                   , SolverResult(..)
                                   , SundialsSolution(..)
                                   , EventSpec(..)
                                   , EventInfo(..)
                                   , CrossingDirection(..)
                                   ) where

import qualified Language.C.Inline as C
import qualified Language.C.Inline.Unsafe as CU

import           Data.Monoid ((<>))
import           Data.Maybe (isJust, fromJust)
import           Data.List (genericLength)

import           Foreign.C.Types (CDouble, CInt, CLong)
import           Foreign.Ptr (Ptr)
import           Foreign.Storable (peek, poke)

import qualified Data.Vector.Storable as V

import           Data.Coerce (coerce)
import           System.IO.Unsafe (unsafePerformIO)

import           Numeric.LinearAlgebra.Devel (createVector)

import           Numeric.LinearAlgebra.HMatrix (Vector, Matrix, toList, rows,
                                                cols, toLists, size, reshape,
                                                subVector, subMatrix, toColumns)

import           Numeric.Sundials.Arkode (cV_ADAMS, cV_BDF,
                                          vectorToC, cV_SUCCESS,
                                          SunVector(..))
import qualified Numeric.Sundials.Arkode as T
import           Numeric.Sundials.ODEOpts

C.context (C.baseCtx <> C.vecCtx <> C.funCtx <> T.sunCtx)

C.include "<stdlib.h>"
C.include "<stdio.h>"
C.include "<math.h>"
C.include "<cvode/cvode.h>"               -- prototypes for CVODE fcts., consts.
C.include "<nvector/nvector_serial.h>"    -- serial N_Vector types, fcts., macros
C.include "<sunmatrix/sunmatrix_dense.h>" -- access to dense SUNMatrix
C.include "<sunlinsol/sunlinsol_dense.h>" -- access to dense SUNLinearSolver
C.include "<cvode/cvode_direct.h>"        -- access to CVDls interface
C.include "<cvode/cvode_diag.h>"          -- CVDiag
C.include "<sundials/sundials_types.h>"   -- definition of type realtype
C.include "<sundials/sundials_math.h>"
C.include "../../../helpers.h"
C.include "Numeric/Sundials/Arkode_hsc.h"

-- | The direction in which a function should cross the x axis
data CrossingDirection = Upwards | Downwards | AnyDirection
  deriving (Eq, Show)

-- Contrary to the documentation, it appears that CVodeGetRootInfo
-- may use both 1 and -1 to indicate a root, depending on the
-- direction of the sign change. See near the end of cvRootfind.
intToDirection :: Integral d => d -> Maybe CrossingDirection
intToDirection d =
  case d of
    1  -> Just Upwards
    -1 -> Just Downwards
    _  -> Nothing

-- | Almost inverse of 'intToDirection'. Map 'Upwards' to 1, 'Downwards' to
-- -1, and 'AnyDirection' to 0.
directionToInt :: Integral d => CrossingDirection -> d
directionToInt d =
  case d of
    Upwards -> 1
    Downwards -> -1
    AnyDirection -> 0

data SundialsSolution =
  SundialsSolution
  { actualTimeGrid :: V.Vector Double                 -- ^ actual time grid returned by the solver (with duplicated event times)
  , solutionMatrix :: Matrix Double                -- ^ matrix of solutions: each column is an unknwown (add also the time vector?)
  , eventInfo      :: [EventInfo]         -- ^ event infos, as many items as triggered events during the simulation
  , diagnostics    :: SundialsDiagnostics -- ^ usual Sundials diagnostics
  }

data EventInfo =
  EventInfo
  { eventTime     :: !Double            -- ^ time at which event was triggered
  , eventIndex    :: !Int               -- ^ which index was triggered
  , rootDirection :: !CrossingDirection -- ^ in which direction ((+)->(-) or (-)->(+)) the root is crossed
  }
  deriving Show

getMethod :: ODEMethod -> Int
getMethod (ADAMS) = cV_ADAMS
getMethod (BDF)   = cV_BDF

getJacobian :: ODEMethod -> Maybe Jacobian
getJacobian _ = Nothing

-- | A version of 'odeSolveVWith' with reasonable default step control.
odeSolveV
    :: ODEMethod
    -> Maybe Double      -- ^ initial step size - by default, CVode
                         -- estimates the initial step size to be the
                         -- solution \(h\) of the equation
                         -- \(\|\frac{h^2\ddot{y}}{2}\| = 1\), where
                         -- \(\ddot{y}\) is an estimated value of the
                         -- second derivative of the solution at \(t_0\)
    -> Double            -- ^ absolute tolerance for the state vector
    -> Double            -- ^ relative tolerance for the state vector
    -> (Double -> Vector Double -> Vector Double) -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
    -> Vector Double     -- ^ initial conditions
    -> Vector Double     -- ^ desired solution times
    -> Matrix Double     -- ^ solution
odeSolveV meth hi epsAbs epsRel f y0 ts =
  odeSolveVWith meth (X epsAbs epsRel) hi g y0 ts
  where
    g t x0 = coerce $ f t x0

-- | A version of 'odeSolveV' with reasonable default parameters and
-- system of equations defined using lists. FIXME: we should say
-- something about the fact we could use the Jacobian but don't for
-- compatibility with hmatrix-gsl.
odeSolve :: (Double -> [Double] -> [Double]) -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
         -> [Double]                         -- ^ initial conditions
         -> Vector Double                    -- ^ desired solution times
         -> Matrix Double                    -- ^ solution
odeSolve f y0 ts =
  -- FIXME: These tolerances are different from the ones in GSL
  odeSolveVWith BDF (XX' 1.0e-6 1.0e-10 1 1)  Nothing g (V.fromList y0) (V.fromList $ toList ts)
  where
    g t x0 = V.fromList $ f t (V.toList x0)

-- | A version of 'odeSolveVWith'' with reasonable default solver
-- options.
odeSolveVWith ::
  ODEMethod
  -> StepControl
  -> Maybe Double -- ^ initial step size - by default, CVode
                  -- estimates the initial step size to be the
                  -- solution \(h\) of the equation
                  -- \(\|\frac{h^2\ddot{y}}{2}\| = 1\), where
                  -- \(\ddot{y}\) is an estimated value of the second
                  -- derivative of the solution at \(t_0\)
  -> (Double -> V.Vector Double -> V.Vector Double) -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
  -> V.Vector Double                     -- ^ Initial conditions
  -> V.Vector Double                     -- ^ Desired solution times
  -> Matrix Double                       -- ^ Error code or solution
odeSolveVWith method control initStepSize f y0 tt =
  case odeSolveVWith' opts f y0 tt of
    Left  (c, _v) -> error $ show c -- FIXME
    Right (v, _d) -> v
  where
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod   = method
                   , stepControl = control
                   , initStep    = initStepSize
                   }

odeSolveVWith' ::
  ODEOpts
  -> (Double -> V.Vector Double -> V.Vector Double) -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
  -> V.Vector Double                     -- ^ Initial conditions
  -> V.Vector Double                     -- ^ Desired solution times
  -> Either (Matrix Double, Int) (Matrix Double, SundialsDiagnostics) -- ^ Error code or solution
odeSolveVWith' opts f y0 tt =
  case solveOdeC (fromIntegral $ maxFail opts)
                  (fromIntegral $ maxNumSteps opts) (coerce $ minStep opts)
                  (fromIntegral . getMethod . odeMethod $ opts) (coerce $ initStep opts) jacH (scise $ stepControl opts)
                  (coerce f) (coerce y0)
                  0 (\_ x -> x) [] 0 (\_ _ y -> y) (coerce tt) of
    -- Remove the time column for backwards compatibility
    SolverError m c         -> Left
                               ( subMatrix (0, 1) (V.length tt, l) m
                               , fromIntegral c
                               )
    SolverSuccess _ m d     -> Right
                               ( subMatrix (0, 1) (V.length tt, l) m
                               , d
                               )
  where
    l = size y0
    scise (X aTol rTol)                          = coerce (V.replicate l aTol, rTol)
    scise (X' aTol rTol)                         = coerce (V.replicate l aTol, rTol)
    scise (XX' aTol rTol yScale _yDotScale)      = coerce (V.replicate l aTol, yScale * rTol)
    -- FIXME; Should we check that the length of ss is correct?
    scise (ScXX' aTol rTol yScale _yDotScale ss) = coerce (V.map (* aTol) ss, yScale * rTol)
    jacH = fmap (\g t v -> matrixToSunMatrix $ g (coerce t) (coerce v)) $
           getJacobian $ odeMethod opts

matrixToSunMatrix :: Matrix Double -> T.SunMatrix
matrixToSunMatrix m = T.SunMatrix { T.rows = nr, T.cols = nc, T.vals = vs }
  where
    nr = fromIntegral $ rows m
    nc = fromIntegral $ cols m
    -- FIXME: efficiency
    vs = V.fromList $ map coerce $ concat $ toLists m

solveOdeC ::
  CInt ->
  CLong ->
  CDouble ->
  CInt ->
  Maybe CDouble ->
  (Maybe (CDouble -> V.Vector CDouble -> T.SunMatrix)) ->
  (V.Vector CDouble, CDouble) ->
  (CDouble -> V.Vector CDouble -> V.Vector CDouble) -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
  -> V.Vector CDouble -- ^ Initial conditions
  -> CInt -- ^ Number of event equations
  -> (CDouble -> V.Vector CDouble -> V.Vector CDouble) -- ^ The event equations themselves
  -> [CrossingDirection] -- ^ The required crossing direction for each event
  -> CInt -- ^ Maximum number of events
  -> (Int -> CDouble -> V.Vector CDouble -> V.Vector CDouble)
      -- ^ Function to reset/update the state when an event occurs. The
      -- 'Int' argument is the 0-based number of the event that has
      -- occurred. If multiple events have occurred at the same time, they
      -- are handled in the increasing order of the event index. The other
      -- arguments are the time and the point in the state space. Return
      -- the updated point in the state space.
  -> V.Vector CDouble -- ^ Desired solution times
  -> SolverResult
solveOdeC maxErrTestFails maxNumSteps_ minStep_ method initStepSize
          jacH (aTols, rTol) fun f0 nr event_fn directions max_events apply_event ts =
  unsafePerformIO $ do

  let isInitStepSize :: CInt
      isInitStepSize = fromIntegral $ fromEnum $ isJust initStepSize
      ss :: CDouble
      ss = case initStepSize of
             -- It would be better to put an error message here but
             -- inline-c seems to evaluate this even if it is never
             -- used :(
             Nothing -> 0.0
             Just x  -> x

  let dim = V.length f0
      nEq :: CLong
      nEq = fromIntegral dim
      nTs :: CInt
      nTs = fromIntegral $ V.length ts
  output_mat_mut :: V.MVector _ CDouble <- V.thaw =<< createVector ((1 + fromIntegral dim) * (fromIntegral (2 * max_events) + fromIntegral nTs))
  diagMut :: V.MVector _ CLong <- V.thaw =<< createVector 10 -- FIXME
  -- We need the types that sundials expects.
  -- FIXME: The Haskell type is currently empty!
  let funIO :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr () -> IO CInt
      funIO t y f _ptr = do
        sv <- peek y
        poke f $ SunVector { sunVecN = sunVecN sv
                           , sunVecVals = fun t (sunVecVals sv)
                           }
        return 0

  let nrPre = fromIntegral nr
  gResults :: V.Vector CInt <- createVector nrPre
  -- FIXME: Do we need to do this here? Maybe as it will get GC'd and
  -- we'd have to do a malloc in C otherwise :(
  gResMut <- V.thaw gResults
  event_index_mut :: V.MVector _ CInt <- V.thaw =<< createVector (fromIntegral max_events)
  event_time_mut :: V.MVector _ CDouble <- V.thaw =<< createVector (fromIntegral max_events)
  -- Total number of events. This is *not* directly re
  n_events_mut :: V.MVector _ CInt <- V.thaw =<< createVector 1
  -- Total number of rows in the output_mat_mut matrix. It *cannot* be
  -- inferred from n_events_mut because when an event occurs k times, it
  -- contributes k to n_events_mut but only 2 to n_rows_mut.
  n_rows_mut :: V.MVector _ CInt <- V.thaw =<< createVector 1
  actual_event_direction_mut :: V.MVector _ CInt <- V.thaw =<< createVector (fromIntegral max_events)

  let event_fn_c :: CDouble -> Ptr T.SunVector -> Ptr CDouble -> Ptr () -> IO CInt
      event_fn_c x y f _ptr = do
        vals <- event_fn x <$> (sunVecVals <$> peek y)
        -- FIXME: We should be able to use poke somehow
        vectorToC vals nrPre f
        return 0

      apply_event_c :: CInt -> CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> IO CInt
      apply_event_c event_index t y y' = do
        sv <- peek y
        poke y' $ SunVector
          { sunVecN = sunVecN sv
          , sunVecVals = apply_event (fromIntegral event_index) t (sunVecVals sv)
          }
        return 0

      requested_event_directions :: V.Vector CInt
      requested_event_directions = V.fromList $ map directionToInt directions

  let isJac :: CInt
      isJac = fromIntegral $ fromEnum $ isJust jacH
      jacIO :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr T.SunMatrix ->
               Ptr () -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr T.SunVector ->
               IO CInt
      jacIO t y _fy jacS _ptr _tmp1 _tmp2 _tmp3 = do
        case jacH of
          Nothing   -> error "Numeric.Sundials.CVode.ODE: Jacobian not defined"
          Just jacI -> do j <- jacI t <$> (sunVecVals <$> peek y)
                          poke jacS j
                          -- FIXME: I don't understand what this comment means
                          -- Unsafe since the function will be called many times.
                          [CU.exp| int{ 0 } |]

  res :: Int <- fromIntegral <$> [C.block| int {
                         /* general problem variables */

                         int flag;                  /* reusable error-checking flag                 */

                         int i, j;                  /* reusable loop indices                        */
                         N_Vector y = NULL;         /* empty vector for storing solution            */
                         N_Vector tv = NULL;        /* empty vector for storing absolute tolerances */

                         SUNMatrix A = NULL;        /* empty matrix for linear solver               */
                         SUNLinearSolver LS = NULL; /* empty linear solver object                   */
                         void *cvode_mem = NULL;    /* empty CVODE memory structure                 */
                         realtype t;
                         long int nst, nfe, nsetups, nje, nfeLS, nni, ncfn, netf, nge;

                         realtype tout;

                         /* input_ind tracks the current index into the ts array */
                         int input_ind = 1;
                         /* output_ind tracks the current row into the output_mat_mut matrix.
                            If differs from input_ind because of the extra rows corresponding to events. */
                         int output_ind = 1;
                         /* We need to update n_rows_mut every time we update output_ind because
                            of the possibility of early return (in which case we still need to assemble
                            the partial results matrix). We could even work with n_rows_mut only and ditch
                            output_ind, but the inline-c expression is quite verbose, and output_ind is
                            more convenient to use in index calculations.
                         */
			 ($vec-ptr:(int *n_rows_mut))[0] = output_ind;
                         /* event_ind tracks the current event number */
                         int event_ind = 0;

                         /* general problem parameters */

                         realtype T0 = RCONST(($vec-ptr:(double *ts))[0]); /* initial time              */
                         sunindextype NEQ = $(sunindextype nEq);           /* number of dependent vars. */

                         /* Initialize data structures */

                         y = N_VNew_Serial(NEQ); /* Create serial vector for solution */
                         if (check_flag((void *)y, "N_VNew_Serial", 0)) return 1;
                         /* Specify initial condition */
                         for (i = 0; i < NEQ; i++) {
                           NV_Ith_S(y,i) = ($vec-ptr:(double *f0))[i];
                         };

                         cvode_mem = CVodeCreate($(int method), CV_NEWTON);
                         if (check_flag((void *)cvode_mem, "CVodeCreate", 0)) return(1);

                         /* Call CVodeInit to initialize the integrator memory and specify the
                          * user's right hand side function in y'=f(t,y), the inital time T0, and
                          * the initial dependent variable vector y. */
                         flag = CVodeInit(cvode_mem, $fun:(int (* funIO) (double t, SunVector y[], SunVector dydt[], void * params)), T0, y);
                         if (check_flag(&flag, "CVodeInit", 1)) return(1);

                         tv = N_VNew_Serial(NEQ); /* Create serial vector for absolute tolerances */
                         if (check_flag((void *)tv, "N_VNew_Serial", 0)) return 1;
                         /* Specify tolerances */
                         for (i = 0; i < NEQ; i++) {
                           NV_Ith_S(tv,i) = ($vec-ptr:(double *aTols))[i];
                         };

                         flag = CVodeSetMinStep(cvode_mem, $(double minStep_));
                         if (check_flag(&flag, "CVodeSetMinStep", 1)) return 1;
                         flag = CVodeSetMaxNumSteps(cvode_mem, $(long int maxNumSteps_));
                         if (check_flag(&flag, "CVodeSetMaxNumSteps", 1)) return 1;
                         flag = CVodeSetMaxErrTestFails(cvode_mem, $(int maxErrTestFails));
                         if (check_flag(&flag, "CVodeSetMaxErrTestFails", 1)) return 1;

                         /* Call CVodeSVtolerances to specify the scalar relative tolerance
                          * and vector absolute tolerances */
                         flag = CVodeSVtolerances(cvode_mem, $(double rTol), tv);
                         if (check_flag(&flag, "CVodeSVtolerances", 1)) return(1);

                         /* Call CVodeRootInit to specify the root function event_fn_c with nr components */
                         flag = CVodeRootInit(cvode_mem, $(int nr), $fun:(int (* event_fn_c) (double t, SunVector y[], double gout[], void * params)));

                         if (check_flag(&flag, "CVodeRootInit", 1)) return(1);

                         /* Set the initial step size if there is one */
                         if ($(int isInitStepSize)) {
                           /* FIXME: We could check if the initial step size is 0 */
                           /* or even NaN and then throw an error                 */
                           flag = CVodeSetInitStep(cvode_mem, $(double ss));
                           if (check_flag(&flag, "CVodeSetInitStep", 1)) return 1;
                         }

                         /* Set the Jacobian if there is one */
                         if ($(int isJac)) {
                           /* Initialize dense matrix data structure and solver */
                           A = SUNDenseMatrix(NEQ, NEQ);
                           if (check_flag((void *)A, "SUNDenseMatrix", 0)) return 1;
                           LS = SUNDenseLinearSolver(y, A);
                           if (check_flag((void *)LS, "SUNDenseLinearSolver", 0)) return 1;

                           /* Attach matrix and linear solver */
                           flag = CVDlsSetLinearSolver(cvode_mem, LS, A);
                           if (check_flag(&flag, "CVDlsSetLinearSolver", 1)) return 1;

                           /* Set the Jacobian */
                           flag = CVDlsSetJacFn(cvode_mem, $fun:(int (* jacIO) (double t, SunVector y[], SunVector fy[], SunMatrix Jac[], void * params, SunVector tmp1[], SunVector tmp2[], SunVector tmp3[])));
                           if (check_flag(&flag, "CVDlsSetJacFn", 1)) return 1;
                         } else /* Else we use CVDiag directly */ {
                           flag = CVDiag(cvode_mem);
                           if (check_flag(&flag, "CVDiag", 1)) return 1;
                         }

                         /* Store initial conditions */
			 ($vec-ptr:(double *output_mat_mut))[0 * (NEQ + 1) + 0] = ($vec-ptr:(double *ts))[0];
                         for (j = 0; j < NEQ; j++) {
                           ($vec-ptr:(double *output_mat_mut))[0 * (NEQ + 1) + (j + 1)] = NV_Ith_S(y,j);
                         }

                         while (1) {
                           flag = CVode(cvode_mem, ($vec-ptr:(double *ts))[input_ind], y, &t, CV_NORMAL); /* call integrator */
                           if (check_flag(&flag, "CVode solver failure, stopping integration", 1)) return 1;

                           /* Store the results for Haskell */
			   ($vec-ptr:(double *output_mat_mut))[output_ind * (NEQ + 1) + 0] = t;
                           for (j = 0; j < NEQ; j++) {
                             ($vec-ptr:(double *output_mat_mut))[output_ind * (NEQ + 1) + (j + 1)] = NV_Ith_S(y,j);
                           }
                           output_ind++;
                           ($vec-ptr:(int *n_rows_mut))[0] = output_ind;

                           if (flag == CV_ROOT_RETURN) {
                             if (event_ind >= $(int max_events)) {
                               /* We don't have any more space for events. Return an error. */
                               return 1;
                             }

                             /* Are we interested in this event?
                                If not, continue without any observable side-effects.
                             */
                             int good_event = 0;
			     flag = CVodeGetRootInfo(cvode_mem, ($vec-ptr:(int *gResMut)));
			     if (check_flag(&flag, "CVodeGetRootInfo", 1)) return 1;
                             for (i = 0; i < $(int nr); i++) {
                               int ev = ($vec-ptr:(int *gResMut))[i];
                               int req_dir = ($vec-ptr:(const int *requested_event_directions))[i];
                               if (ev != 0 && ev * req_dir >= 0) {
                                 good_event = 1;

                                 ($vec-ptr:(int *actual_event_direction_mut))[event_ind] = ev;
                                 ($vec-ptr:(int *event_index_mut))[event_ind] = i;
                                 ($vec-ptr:(double *event_time_mut))[event_ind] = t;
                                 event_ind++;

                                 /* Update the state with the supplied function */
                                 $fun:(int (* apply_event_c) (int, double, SunVector y[], SunVector z[]))(i, t, y, y);
                               }
                             }

                             if (good_event) {
                               ($vec-ptr:(double *output_mat_mut))[output_ind * (NEQ + 1) + 0] = t;
                               for (j = 0; j < NEQ; j++) {
                                 ($vec-ptr:(double *output_mat_mut))[output_ind * (NEQ + 1) + (j + 1)] = NV_Ith_S(y,j);
                               }
                               output_ind++;
                               ($vec-ptr:(int *n_rows_mut))[0] = output_ind;

                               flag = CVodeReInit(cvode_mem, t, y);
                               if (check_flag(&flag, "CVodeReInit", 1)) return(1);
                             } else {
                               /* Since this is not a wanted event, it shouldn't get a row */
                               output_ind--;
                               ($vec-ptr:(int *n_rows_mut))[0] = output_ind;
                             }
                           }
			   else {
			     if (++input_ind >= $(int nTs))
                               break;
			   }
                         }

			 /* The number of actual roots we found */
			 ($vec-ptr:(int *n_events_mut))[0] = event_ind;

                         /* Get some final statistics on how the solve progressed */
                         flag = CVodeGetNumSteps(cvode_mem, &nst);
                         check_flag(&flag, "CVodeGetNumSteps", 1);
                         ($vec-ptr:(long int *diagMut))[0] = nst;

                         /* FIXME */
                         ($vec-ptr:(long int *diagMut))[1] = 0;

                         flag = CVodeGetNumRhsEvals(cvode_mem, &nfe);
                         check_flag(&flag, "CVodeGetNumRhsEvals", 1);
                         ($vec-ptr:(long int *diagMut))[2] = nfe;
                         /* FIXME */
                         ($vec-ptr:(long int *diagMut))[3] = 0;

                         flag = CVodeGetNumLinSolvSetups(cvode_mem, &nsetups);
                         check_flag(&flag, "CVodeGetNumLinSolvSetups", 1);
                         ($vec-ptr:(long int *diagMut))[4] = nsetups;

                         flag = CVodeGetNumErrTestFails(cvode_mem, &netf);
                         check_flag(&flag, "CVodeGetNumErrTestFails", 1);
                         ($vec-ptr:(long int *diagMut))[5] = netf;

                         flag = CVodeGetNumNonlinSolvIters(cvode_mem, &nni);
                         check_flag(&flag, "CVodeGetNumNonlinSolvIters", 1);
                         ($vec-ptr:(long int *diagMut))[6] = nni;

                         flag = CVodeGetNumNonlinSolvConvFails(cvode_mem, &ncfn);
                         check_flag(&flag, "CVodeGetNumNonlinSolvConvFails", 1);
                         ($vec-ptr:(long int *diagMut))[7] = ncfn;

                         flag = CVDlsGetNumJacEvals(cvode_mem, &nje);
                         check_flag(&flag, "CVDlsGetNumJacEvals", 1);
                         ($vec-ptr:(long int *diagMut))[8] = ncfn;

                         flag = CVDlsGetNumRhsEvals(cvode_mem, &nfeLS);
                         check_flag(&flag, "CVDlsGetNumRhsEvals", 1);
                         ($vec-ptr:(long int *diagMut))[9] = ncfn;

                         /* Clean up and return */

                         N_VDestroy(y);          /* Free y vector          */
                         N_VDestroy(tv);         /* Free tv vector         */
                         CVodeFree(&cvode_mem);  /* Free integrator memory */
                         SUNLinSolFree(LS);      /* Free linear solver     */
                         SUNMatDestroy(A);       /* Free A matrix          */

                         return CV_SUCCESS;
                       } |]
  preD <- V.freeze diagMut
  let d = SundialsDiagnostics (fromIntegral $ preD V.!0)
                              (fromIntegral $ preD V.!1)
                              (fromIntegral $ preD V.!2)
                              (fromIntegral $ preD V.!3)
                              (fromIntegral $ preD V.!4)
                              (fromIntegral $ preD V.!5)
                              (fromIntegral $ preD V.!6)
                              (fromIntegral $ preD V.!7)
                              (fromIntegral $ preD V.!8)
                              (fromIntegral $ preD V.!9)
  n_rows <- fromIntegral . V.head <$> V.freeze n_rows_mut
  output_mat <- coerce . reshape (dim + 1) . subVector 0 ((dim + 1) * n_rows) <$>
    V.freeze output_mat_mut
  n_events <- fromIntegral . V.head <$> V.freeze n_events_mut
  event_time             :: V.Vector Double
    <- coerce . V.take n_events <$> V.freeze event_time_mut
  event_index            :: V.Vector Int
    <- V.map fromIntegral . V.take n_events <$> V.freeze event_index_mut
  actual_event_direction :: V.Vector CInt
    <- V.take n_events <$> V.freeze actual_event_direction_mut
  let
    events :: [EventInfo]
    events = zipWith3 EventInfo
      (V.toList event_time)
      (V.toList event_index)
      (map (fromJust . intToDirection) $ V.toList actual_event_direction)
  return $
    if res == cV_SUCCESS
      then
        SolverSuccess events output_mat d
      else
        SolverError output_mat res

data SolverResult
  = SolverError !(Matrix Double) !Int
      -- ^ Partial results and error code
  | SolverSuccess
      [EventInfo]
      !(Matrix Double)
      !SundialsDiagnostics
      -- ^ Times at which the event was triggered, information about which root and the
                                                   -- results and diagnostics.
    deriving Show

odeSolveRootVWith' ::
  ODEOpts
  -> (Double -> V.Vector Double -> V.Vector Double)
      -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
  -> Maybe (Double -> Vector Double -> Matrix Double)
      -- ^ The Jacobian (optional)
  -> V.Vector Double                      -- ^ Initial conditions
  -> [EventSpec]                          -- ^ Event specifications
  -> Int                                  -- ^ Maximum number of events
  -> V.Vector Double                      -- ^ Desired solution times
  -> SolverResult
odeSolveRootVWith' opts f mb_jacobian y0 event_specs nRootEvs tt =
  solveOdeC (fromIntegral $ maxFail opts)
                 (fromIntegral $ maxNumSteps opts) (coerce $ minStep opts)
                 (fromIntegral . getMethod . odeMethod $ opts) (coerce $ initStep opts) jacH (scise $ stepControl opts)
                 (coerce f) (coerce y0)
                 (genericLength event_specs) event_equations event_directions
                 (fromIntegral nRootEvs) reset_state
                 (coerce tt)
  where
    l = size y0
    scise (X aTol rTol)                          = coerce (V.replicate l aTol, rTol)
    scise (X' aTol rTol)                         = coerce (V.replicate l aTol, rTol)
    scise (XX' aTol rTol yScale _yDotScale)      = coerce (V.replicate l aTol, yScale * rTol)
    -- FIXME; Should we check that the length of ss is correct?
    scise (ScXX' aTol rTol yScale _yDotScale ss) = coerce (V.map (* aTol) ss, yScale * rTol)
    jacH = fmap (\g t v -> matrixToSunMatrix $ g (coerce t) (coerce v)) $ mb_jacobian
    event_equations :: CDouble -> Vector CDouble -> Vector CDouble
    event_equations t y = V.fromList $
      map (\ev -> coerce (eventCondition ev) t y) event_specs
    event_directions :: [CrossingDirection]
    event_directions = map eventDirection event_specs
    reset_state :: Int -> CDouble -> Vector CDouble -> Vector CDouble
    reset_state n_event = coerce $ eventUpdate (event_specs !! n_event)

data EventSpec = EventSpec
  { eventCondition :: Double -> V.Vector Double -> Double
  , eventDirection :: CrossingDirection
  , eventUpdate :: Double -> V.Vector Double -> V.Vector Double
  }

odeSolveWithEvents
  :: ODEOpts
  -> [EventSpec]
    -- ^ Event specifications
  -> Int
    -- ^ Maximum number of events
  -> (Double -> V.Vector Double -> V.Vector Double)
    -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
  -> Maybe (Double -> Vector Double -> Matrix Double)
    -- ^ The Jacobian (optional)
  -> V.Vector Double
    -- ^ Initial conditions
  -> V.Vector Double
    -- ^ Desired solution times
  -> Either Int SundialsSolution
    -- ^ Either an error code or a solution
odeSolveWithEvents opts event_specs max_events rhs mb_jacobian initial sol_times =
  let
    result :: SolverResult
    result =
      odeSolveRootVWith' opts rhs mb_jacobian initial event_specs
        max_events sol_times
  in
    case result of
      SolverError _ code -> Left code
      SolverSuccess events mx diagn ->
        Right $ SundialsSolution
            { actualTimeGrid = extractTimeGrid mx
            , solutionMatrix = mx
            , eventInfo = events
            , diagnostics = diagn
            }
  where
    -- The time grid is the first column of the result matrix
    extractTimeGrid :: Matrix Double -> Vector Double
    extractTimeGrid = head . toColumns