{-# LANGUAGE GADTs #-}
module Math.IRT.Model.FourPLM
  ( FourPLM (..)
  ) where

import Numeric.AD (Mode, Scalar, auto)
import Numeric.AD.Mode.Forward.Double
import qualified Numeric.AD.Mode.Tower as T

import Statistics.Distribution

import Math.IRT.Internal.Distribution
import Math.IRT.Internal.LogLikelihood
import Math.IRT.Model.Generic

data FourPLM = FourPLM { discrimination :: !Double
                       , difficulty     :: !Double
                       , pseudoGuessing :: !Double
                       , asymptote      :: !Double
                       } deriving (Show)

instance Distribution FourPLM where
    cumulative = cumulative4PL

instance ContDistr FourPLM where
    density  x = diff (cumulative4PL x)
    quantile _ = error "This shouldn't be needed"

instance DensityDeriv FourPLM where
    densityDeriv x = (!! 2) . T.diffs (cumulative4PL x)

instance GenericModel FourPLM where
    fromRasch         b = FourPLM 1.0 b 0.0 1.0
    fromOnePLM        b = FourPLM 1.7 b 0.0 1.0
    fromTwoPLM      a b = FourPLM   a b 0.0 1.0
    fromThreePLM  a b c = FourPLM   a b   c 1.0
    fromFourPLM         = FourPLM

instance LogLikelihood FourPLM where
    logLikelihood = logLikeFunc cumulative4PL

cumulative4PL :: (Mode a, Floating a, Scalar a ~ Double) =>
       FourPLM -- ^The IRT parameters
    -> a       -- ^Theta
    -> a
cumulative4PL params theta =
    let (a, b, c, d) = makeParamAutos params in
    c + (              (d - c)
         -----------------------------------
         / (1 + exp ((-a) * (theta - b))))
    where makeParamAutos (FourPLM sa sb sc sd) =
              let a = auto sa
                  b = auto sb
                  c = auto sc
                  d = auto sd
              in (a, b, c, d)


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}

module Utils where
import Control.Monad.Fix
import Control.Monad
import System.IO
import Signal
import Circuit
import AmpOp
    ( SignalState(SignalTwoNodes),
      AmpOp(openLoopGain),
      Impedance,
      Output,
      Input,
      dc100,
      sen,
      ground,
      lm741,
      r1,
      r2,
      eps )
import Data.List
import Data.Complex
import Simulation

type Iteration = Int

makePlotIterations = do
    outh <- openFile "outputIterations.txt" WriteMode
    mapM_ (\(i, o) -> hPutStrLn outh (show i ++ " " ++ show o)) (ampOpInvertingList' lm741 12 r1 r2)
    hClose outh

ampOpInvertingList' :: AmpOp -> Input -> Impedance -> Impedance -> [(Iteration, Output)]
ampOpInvertingList' model input r1 r2 = fix calculate (100, 0, [], 1)
    where calculate f (vXOld, vOutOld, l, it) = if abs(realPart vOutOld - realPart vOut) <= 0.000001 && abs(realPart vXOld - realPart vX) <= 0.000001 then l ++ [(it, vOut)] else f (vOut, vX, l ++ [(it, vOut)], it + 1)
            where vX = (- vOutOld) / openLoopGain model `at` 0
                  vOut = (vXOld * (r1 + r2) `at` 0 - input * r2 `at` 0) / r1 `at` 0

checkDataIterations :: IO ()
checkDataIterations = print (ampOpInvertingList' lm741 12 r1 r2)

ampOpInvertingList :: AmpOp -> Impedance -> Impedance -> SignalState -> Circuit Input [(Iteration, SignalState)]
ampOpInvertingList model r1 r2 (SignalTwoNodes (vXInit, vOutInit)) =
    ($ (vXInit, vOutInit, [], 1))<$> mfix
    (\f -> Circuit $ \vIn -> Signal $ \time (vXOld, vOutOld, l, it) ->
        let vX = (- vOutOld) / openLoopGain model
            vOut = ((vXOld * (r1 + r2)) - (vIn * r2)) / r1
        in if abs((realPart <$> vOutOld) - (realPart <$> vOut)) <= eps && abs((realPart <$> vXOld) - (realPart <$> vX)) <= eps
            then l ++ [(it, SignalTwoNodes (vX, vOut))]
            else f (vX, vOut, l ++ [(it, SignalTwoNodes (vX, vOut))], it + 1))

makeSamplesInput :: Signal Input -> [Metric] -> [(Metric, Input)]
makeSamplesInput input samples = do time <- samples
                                    let result = input `at` time
                                    return (time, result)

makeSamplesOutput :: SignalState -> (SignalState -> Circuit Input [(Iteration, SignalState)]) -> Signal Input -> [Metric] -> [[(Metric, Iteration, Output)]]
makeSamplesOutput _ _ _ [] = []
makeSamplesOutput initialState circuit input (t:ts) = map (\(i, SignalTwoNodes (_, o)) -> (t, i, o `at` t)) results : rest
    where results = circuit initialState `simulate` input `at` t
          rest = makeSamplesOutput (snd $ last results) circuit input ts

ampOpInvertingList2 :: AmpOp -> Impedance -> Impedance -> Circuit Input [(Iteration, Output)]
ampOpInvertingList2 model r1 r2 =
    ($ (dc100, ground, [], 1))<$> mfix
    (\f -> Circuit $ \vIn -> Signal $ \time (vXOld, vOutOld, l, it) ->
        let vX = (- vOutOld) / openLoopGain model
            vOut = ((vXOld * (r1 + r2)) - (vIn * r2)) / r1
        in if abs((realPart <$> vOutOld) - (realPart <$> vOut)) <= eps && abs((realPart <$> vXOld) - (realPart <$> vX)) <= eps
            then l ++ [(it, vOut `at` time)]
            else f (vX, vOut, l ++ [(it, vOut `at` time)], it + 1))

makeSamplesOutput2 :: Signal Input -> [Metric] -> [[(Metric, Iteration, Output)]]
makeSamplesOutput2 input samples = do time <- samples
                                      let results = (ampOpInvertingList2 lm741 r1 r2 `simulate` input) `at` time
                                      mapM (\(i, o) -> [(time, i, o)]) results

testCircuit = ampOpInvertingList lm741 r1 r2

makePlotMetricOutput = do
    outh <- openFile "outputMetric.txt" WriteMode
    mapM_ (\(t, i, o) -> hPutStrLn outh (show t ++ " " ++ show i ++ " " ++ show o)) (concat $ makeSamplesOutput iSignalState testCircuit sen (map (:+ 0) [0.1,0.2..10.0]))
    hClose outh

makePlotMetricInput = do
    outh <- openFile "inputMetric.txt" WriteMode
    mapM_ (\(t, o) -> hPutStrLn outh (show t ++ " " ++ show o)) (makeSamplesInput sen (map (:+ 0) [0.1,0.2..10.0]))
    hClose outh


{-# LANGUAGE TemplateHaskell #-}

module Shapes where

import Control.Lens hiding (element)
import Graphics.Gloss.Data.Picture
import Data.Complex

data ShapeType = Kite | Dart deriving (Show)

data Shape = Shape {_shape :: ShapeType
                   ,_translate :: Complex Double
                   ,_rotate :: Complex Double
} deriving (Show)

newtype Vertices = Vertices {_nodes :: [Complex Double]} deriving (Show)

makeLenses ''Shape
makeLenses ''Vertices

standardKite :: Shape
standardKite = Shape {_shape = Kite, _translate = 0, _rotate = 1}

standardDart :: Shape
standardDart = Shape {_shape = Dart, _translate = 0, _rotate = 1}

z :: Complex Double
z = cis (pi/5)

zPower :: Int -> Complex Double
zPower n = cis((pi/5) * fromIntegral (n `mod` 10))

longEdge :: Vertices
longEdge = Vertices {_nodes = [0,(0.5 :+ (-0.1))*z,z]}

shortEdge :: Vertices
shortEdge = Vertices {_nodes = [0,(0.5 :+ 0.1)*(1-z),1-z]}

rotranslate :: Complex Double -> Complex Double -> Vertices -> Vertices
rotranslate z w = over nodes (map (\u -> u*z+w))

reverseVertices :: Vertices -> Vertices
reverseVertices = over nodes reverse

concatVertices :: [Vertices] -> Vertices
concatVertices p = Vertices {_nodes = concatMap (tail._nodes) p}

shapeToVertices :: Shape -> Vertices
shapeToVertices (Shape Kite z x) = rotranslate x z kiteVertices
shapeToVertices (Shape Dart z x) = rotranslate x z dartVertices

complexToPoint :: Complex Double -> Point
complexToPoint z = (realToFrac x, realToFrac y)
  where
    x = realPart z
    y = imagPart z

verticesToPicture :: Vertices -> Picture
verticesToPicture = lineLoop . map complexToPoint . view nodes

shapeToPicture :: Shape -> Picture
shapeToPicture = verticesToPicture . shapeToVertices

refine :: Shape -> [Shape]
refine s@(Shape t w z) =
  let phi = (sqrt 5 - 1)/2 :+ 0
      tr t' u v = Shape t' (z*phi*v+w) (z*u*phi)
  in case t of
  Kite -> [tr Dart 1 0
          ,tr Kite (zPower 3) (1 - zPower 3)
          ,tr Kite (zPower 7) (1 + zPower 2)
          ]
  Dart -> [tr Kite (zPower 9) 0
          ,tr Dart (zPower 6) (1+phi)
          ]

tilings :: [[Shape]]
tilings = [standardKite] : map (concatMap refine) tilings

tilingPics :: [Picture]
tilingPics = map (pictures . map shapeToPicture) tilings

kiteVertices :: Vertices
kiteVertices = concatVertices
  [longEdge
  ,rotranslate 1 z shortEdge
  ,reverseVertices $ rotranslate (zPower 4) (zPower 9) shortEdge
  ,reverseVertices $ rotranslate (zPower 8) 0 longEdge
  ]

dartVertices :: Vertices
dartVertices = concatVertices
  [reverseVertices $ rotranslate (zPower 4) 1 longEdge
  ,reverseVertices $ rotranslate (zPower 3) x shortEdge
  ,rotranslate (zPower 9) x shortEdge
  ,rotranslate (zPower 2) (zPower 8) longEdge
  ]
  where
    x = (1 + (zPower 4) - (zPower 3))

printVertices :: Vertices -> IO ()
printVertices = mapM_ print . _nodes


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE StandaloneDeriving #-}

module HLearn.Models.Distributions.Multivariate.MultiNormal
    ( MultiNormal (..)
    )
    where

-- import qualified Control.ConstraintKinds as CK
import qualified Data.Foldable as F
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import Data.Vector.Unboxed.Deriving

import Control.DeepSeq
import GHC.ST
import GHC.TypeLits

import Foreign.Storable
import Numeric.LinearAlgebra hiding ((<>))
import qualified Numeric.LinearAlgebra as LA

import HLearn.Algebra
import HLearn.Models.Distributions.Common

-------------------------------------------------------------------------------
-- data types

data MultiNormalVec (n::Nat) prob = MultiNormalVec
    { q0 :: !prob
    , q1 :: !(VU.Vector prob)
    , q2 :: !(V.Vector (VU.Vector prob))
    }
    deriving (Read,Show,Eq,Ord)

instance NFData (MultiNormalVec n prob) where
    rnf mn = seq mn ()

newtype MultiNormal prob (xs::[*]) = MultiNormal (MultiNormalVec (Length xs) prob)
    deriving (Read,Show,Eq,Ord,NFData)

deriving instance (Monoid  (MultiNormalVec (Length xs) prob)) => Monoid (MultiNormal prob xs)
deriving instance (Abelian (MultiNormalVec (Length xs) prob)) => Abelian (MultiNormal prob xs)
deriving instance (Group   (MultiNormalVec (Length xs) prob)) => Group (MultiNormal prob xs)
deriving instance (Module  (MultiNormalVec (Length xs) prob)) => Module (MultiNormal prob xs)

-------------------------------------------------------------------------------
-- algebra

instance (Num prob, VU.Unbox prob, SingI n) => Abelian (MultiNormalVec n prob)
instance (Num prob, VU.Unbox prob, SingI n) => Monoid (MultiNormalVec n prob) where
    mempty = MultiNormalVec
        { q0 = 0
        , q1 = VU.replicate n 0
        , q2 = V.replicate n (VU.replicate n 0)
        }
        where
            n = fromIntegral $ fromSing (sing :: Sing n)
    mn1 `mappend` mn2 = MultiNormalVec
        { q0 = (q0 mn1) + (q0 mn2)
        , q1 = VU.zipWith (+) (q1 mn1) (q1 mn2)
        , q2 = V.zipWith (VU.zipWith (+)) (q2 mn1) (q2 mn2)
        }

instance (Num prob, VU.Unbox prob, SingI n) => Group (MultiNormalVec n prob) where
    inverse mn = MultiNormalVec
        { q0 = negate $ q0 mn
        , q1 = VU.map negate (q1 mn)
        , q2 = V.map (VU.map negate) (q2 mn)
        }

instance (Num prob) => HasRing (MultiNormalVec n prob) where
    type Ring (MultiNormalVec n prob) = prob

instance (Num prob, VU.Unbox prob, SingI n) => Module (MultiNormalVec n prob) where
    r .* mn = MultiNormalVec
        { q0 = r * q0 mn
        , q1 = VU.map (r*) (q1 mn)
        , q2 = V.map (VU.map (r*)) (q2 mn)
        }

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

instance (Num prob) => HasRing (MultiNormal prob xs) where
    type Ring (MultiNormal prob xs) = prob

-------------------------------------------------------------------------------
-- training

instance (SingI n, Num prob, VU.Unbox prob) => HomTrainer (MultiNormalVec n prob) where
    type Datapoint (MultiNormalVec n prob) = (VU.Vector prob)
    train1dp dp = MultiNormalVec
        { q0 = 1
        , q1 = dp
        , q2 = V.generate n (\i -> VU.generate n (\j -> (dp VU.! i)*(dp VU.! j)))
        }
        where
            n = fromIntegral $ fromSing (sing :: Sing n)

instance
    ( SingI (Length xs)
    , Num prob
    , VU.Unbox prob
    , HList2List (Datapoint (MultiNormal prob xs)) prob
    ) => HomTrainer (MultiNormal prob xs)
        where
    type Datapoint (MultiNormal prob xs) = HList xs
    train1dp dp = MultiNormal $ train1dp $ VU.fromList $ hlist2list dp

instance (Num prob) => NumDP (MultiNormal prob xs) where
    numdp (MultiNormal mn) = q0 mn

-------------------------------------------------------------------------------
-- distribution

class (Probabilistic dist) => Covariance dist where
    covar :: dist -> Matrix (Probability dist)

instance
    ( VU.Unbox prob
    , SingI k
    , Num prob
    ) => Probabilistic (MultiNormalVec k prob)
        where
    type Probability (MultiNormalVec k prob) = prob

instance
    ( VU.Unbox prob
    , SingI k
    , Fractional prob
    , Enum prob
    , Storable prob
    ) => Covariance (MultiNormalVec k prob)
        where
    covar mn = (k><k) $
            [ (1/(n-1))*( mij - ((q1 mn VU.! j)*((q1 mn) VU.! i))/n )
            | ((i,j),mij) <- assocs
            ]
        where
            assocs = zip [(i,j) | i<-[0..k-1],j<-[0..k-1]] (concat $ V.toList $ fmap VU.toList $ q2 mn)
            mean i = ((q1 mn) VU.! i)/n
            n = q0 mn
            k = fromIntegral $ fromSing (sing :: Sing k)

instance
    ( HList2List (HList dpL) prob
    , VU.Unbox prob
    , Num prob
    , SingI (FromNat1 (Length1 dpL))
    ) => Probabilistic (MultiNormal prob dpL)
        where
    type Probability (MultiNormal prob dpL) = prob

instance
    ( HList2List (HList dpL) prob
    , VU.Unbox prob
    , Floating prob
    , Field prob
    , Enum prob
    , SingI (FromNat1 (Length1 dpL))
--     , Covariance (MultiNormal dpL prob)
    , Storable prob
    ) => PDF (MultiNormal prob dpL)
        where
    pdf (MultiNormal dist) dpL = 1/(sqrt $ (2*pi)^(k)*(det sigma))*(exp $ (-1/2)*(top) )
        where
            top=minElement $ ((trans $ x `sub` mu) LA.<> (inv sigma) LA.<> (x `sub` mu))

            k = fromIntegral $ fromSing (sing :: Sing (Length dpL)) :: Int
            x = k><1 $ hlist2list dpL
            n = q0 dist
            sigma = covar dist
            mu = k><1 $ map (/n) $ VU.toList (q1 $ dist)
--             covarM = (k><k) $ elems $ covar dist

-------------------------------------------------------------------------------
-- tests

-- ds = map (listArray (0,2))
--     [[1,2,4]
--     ,[2,5,6]
--     ,[3,1,1]
--     ]

-- test = train ds :: MultiNormalArray Double 3

ds2 = map VU.fromList
    [[1,2,4]
    ,[2,5,6]
    ,[3,1,1]
    ]

test2 = train ds2 :: MultiNormalVec 3 Double

ds =
    [ 1:::2:::3:::HNil
    , 2:::5:::6:::HNil
    , 3:::1:::1:::HNil
    , 3:::2:::1:::HNil
    ]
test = train ds :: MultiNormal Double '[Double,Double,Double]


{-# LANGUAGE DataKinds
           , GADTs
           , TypeOperators #-}
module Aqua
    (
      module Aqua
    ) where

import Control.Arrow
import Data.Complex
import GHC.TypeLits

import QBasis
import Vect

type QReg n = Vect C (QBasis n)
type QCirc n = Vect C (QBasis n) -> Vect C (QBasis n)

runCircuit :: KnownNat n => QCirc n -> QReg n
runCircuit = ($ unit)

tensor :: QReg m -> QReg n -> QReg (m + n)
tensor vec1 vec2 = vec1 `te` vec2

(⊗) :: KnownNat m => QCirc m -> QCirc n -> QCirc (m + n)
c1 ⊗ c2 = c1 `tf` c2

unit :: KnownNat n => QReg n
unit = V [(unit', 1)]

idGate :: QCirc n
idGate = id

notGate :: QCirc 1
notGate = linear n
    where n :: QBasis 1 -> QReg 1
          n (Cons False Nil) = V [(Cons True Nil, 1)]
          n (Cons True Nil) = V [(Cons False Nil, 1)]

hGate :: QCirc 1
hGate = linear h
    where h :: QBasis 1 -> QReg 1
          h (Cons False Nil) = V [(Cons False Nil, 1 / sqrt 2)
                                 ,(Cons True Nil, 1 / sqrt 2)]
          h (Cons True Nil) = V [(Cons False Nil, 1 / sqrt 2)
                                ,(Cons True Nil, -1 / sqrt 2)]

rotGate :: Double -> QCirc 1
rotGate t = linear r
    where r :: QBasis 1 -> QReg 1
          r (Cons False Nil) = V [(Cons False Nil, 1)]
          r (Cons True Nil) = V [(Cons True Nil, cis t)]

cnotGate :: QCirc 2
cnotGate = linear n
    where n :: QBasis 2 -> QReg 2
          n (Cons False (Cons False Nil)) = V [(Cons False (Cons False Nil), 1)]
          n (Cons False (Cons True Nil)) = V [(Cons False (Cons True Nil), 1)]
          n (Cons True (Cons False Nil)) = V [(Cons True (Cons True Nil), 1)]
          n (Cons True (Cons True Nil)) = V [(Cons True (Cons False Nil), 1)]

swapGate :: QCirc 2
swapGate = linear s
    where s :: QBasis 2 -> QReg 2
          s (Cons False (Cons False Nil)) = V [(Cons False (Cons False Nil), 1)]
          s (Cons False (Cons True Nil)) = V [(Cons True (Cons False Nil), 1)]
          s (Cons True (Cons False Nil)) = V [(Cons False (Cons True Nil), 1)]
          s (Cons True (Cons True Nil)) = V [(Cons True (Cons True Nil), 1)]


{-# LANGUAGE FlexibleContexts #-}

module Blash where

import qualified Data.Vector.Generic as V
import Data.Vector.Generic ((!))
import Data.Complex
import BlashImpl (fZERO, stride, isample, map_reduce, Size, Inc, Index, copyM, axpyM)

copy :: (V.Vector v a)
         => Size
         -> v a
         -> Inc
         -> v a
         -> Inc
         -> v a
copy n dx incx dy incy = V.modify modifier dy
  where
    modifier ys = copyM n dx incx ys incy

axpy :: (Ord a, Num a, V.Vector v a)
         => Size
         -> a
         -> v a
         -> Inc
         -> v a
         -> Inc
         -> v a
axpy n da dx incx dy incy = V.modify modifier dy
  where
    modifier ys = axpyM n da dx incx ys incy

-- /* dot product dx dot dy. */
dot :: (Num a, V.Vector v a)
        => Size
        -> v a
        -> Inc
        -> v a
        -> Inc
        -> a
dot n _ _ _ _ | n <= 0 = fZERO
dot n dx incx dy incy = sum prods
  where
    prods = flip map [0..(n-1)] f
    f i = (dx ! ix i) * (dy ! iy i)
    ix = stride n incx
    iy = stride n incy

-- /* dot product dx dot dy where dx & dy are complex. */
dotc :: (RealFloat a, V.Vector v (Complex a))
        => Size
        -> v (Complex a)
        -> Inc
        -> v (Complex a)
        -> Inc
        -> Complex a
dotc n _ _ _ _ | n <= 0 = fZERO :+ fZERO
dotc n dx incx dy incy = sum prods
  where
    prods = flip map [0..(n-1)] f
    f i = (conjugate (dx ! ix i)) * (dy ! iy i)
    ix = stride n incx
    iy = stride n incy

-- /* compute the L2 norm of array DX of length N, stride INCX */
-- TODO should really be (Num a) - or split between Complex and Real
nrm2 :: (Ord a, Floating a, V.Vector v a)
        => Size
        -> v a
        -> Inc
        -> a
nrm2 n _  incx | n <= 0 || incx <= 0 = fZERO
nrm2 n dx _ | n == 1 = abs $ dx ! 0
nrm2 n dx incx =
  let dx' = isample n dx incx
      xmax = V.maximum $ V.map abs dx'
  in
    case xmax == fZERO of
    True -> fZERO
    False ->
      let scale = 1.0 / xmax
          scaled_dx' = V.map (* scale) dx'
          n' = V.length scaled_dx'
          scaled_drnm2 = sqrt $ dot n' scaled_dx' 1 scaled_dx' 1
      in
        xmax * scaled_drnm2

-- DASUM takes the sum of the absolute values.
asum :: (Num a, V.Vector v a)
        => Size
        -> v a
        -> Inc
        -> a
asum = map_reduce abs V.sum

-- DAMAX takes the max of the absolute values.
iamax :: (Ord a, Num a, V.Vector v a)
        => Size
        -> v a
        -> Inc
        -> Index
iamax = map_reduce abs V.maxIndex


{-# LANGUAGE FlexibleContexts, BangPatterns #-}

-- |
-- Module    : Statistics.Sample.Histogram
-- Copyright : (c) 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Functions for computing histograms of sample data.

module Statistics.Sample.Histogram
    (
      histogram
    -- * Building blocks
    , histogram_
    , range
    ) where

import Numeric.MathFunctions.Constants (m_epsilon,m_tiny)
import Statistics.Function (minMax)
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM

-- | /O(n)/ Compute a histogram over a data set.
--
-- The result consists of a pair of vectors:
--
-- * The lower bound of each interval.
--
-- * The number of samples within the interval.
--
-- Interval (bin) sizes are uniform, and the upper and lower bounds
-- are chosen automatically using the 'range' function.  To specify
-- these parameters directly, use the 'histogram_' function.
histogram :: (G.Vector v0 Double, G.Vector v1 Double, Num b, G.Vector v1 b) =>
             Int                -- ^ Number of bins (must be positive).
          -> v0 Double          -- ^ Sample data (cannot be empty).
          -> (v1 Double, v1 b)
histogram numBins xs = (G.generate numBins step, histogram_ numBins lo hi xs)
    where (lo,hi)    = range numBins xs
          step i     = lo + d * fromIntegral i
          d          = (hi - lo) / fromIntegral numBins
{-# INLINE histogram #-}

-- | /O(n)/ Compute a histogram over a data set.
--
-- Interval (bin) sizes are uniform, based on the supplied upper
-- and lower bounds.
histogram_ :: (Num b, RealFrac a, G.Vector v0 a, G.Vector v1 b) =>
              Int
           -- ^ Number of bins.  This value must be positive.  A zero
           -- or negative value will cause an error.
           -> a
           -- ^ Lower bound on interval range.  Sample data less than
           -- this will cause an error.
           -> a
           -- ^ Upper bound on interval range.  This value must not be
           -- less than the lower bound.  Sample data that falls above
           -- the upper bound will cause an error.
           -> v0 a
           -- ^ Sample data.
           -> v1 b
histogram_ numBins lo hi xs0 = G.create (GM.replicate numBins 0 >>= bin xs0)
  where
    bin xs bins = go 0
     where
       go i | i >= len = return bins
            | otherwise = do
         let x = xs `G.unsafeIndex` i
             b = truncate $ (x - lo) / d
         write' bins b . (+1) =<< GM.read bins b
         go (i+1)
       write' bins' b !e = GM.write bins' b e
       len = G.length xs
       d = ((hi - lo) / fromIntegral numBins) * (1 + realToFrac m_epsilon)
{-# INLINE histogram_ #-}

-- | /O(n)/ Compute decent defaults for the lower and upper bounds of
-- a histogram, based on the desired number of bins and the range of
-- the sample data.
--
-- The upper and lower bounds used are @(lo-d, hi+d)@, where
--
-- @d = (maximum sample - minimum sample) / ((bins - 1) * 2)@
--
-- If all elements in the sample are the same and equal to @x@ range
-- is set to @(x - |x|/10, x + |x|/10)@. And if @x@ is equal to 0 range
-- is set to @(-1,1)@. This is needed to avoid creating histogram with
-- zero bin size.
range :: (G.Vector v Double) =>
         Int                    -- ^ Number of bins (must be positive).
      -> v Double               -- ^ Sample data (cannot be empty).
      -> (Double, Double)
range numBins xs
    | numBins < 1 = error "Statistics.Histogram.range: invalid bin count"
    | G.null xs   = error "Statistics.Histogram.range: empty sample"
    | lo == hi    = case abs lo / 10 of
                      a | a < m_tiny -> (-1,1)
                        | otherwise  -> (lo - a, lo + a)
    | otherwise   = (lo-d, hi+d)
  where
    d | numBins == 1 = 0
      | otherwise    = (hi - lo) / ((fromIntegral numBins - 1) * 2)
    (lo,hi)          = minMax xs
{-# INLINE range #-}


-----------------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2010-2014
-- License     :  BSD3
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- Portability :  GHC only
--
-- Root finding using Halley's rational method (the second in
-- the class of Householder methods). Assumes the function is three
-- times continuously differentiable and converges cubically when
-- progress can be made.
--
-----------------------------------------------------------------------------

module Numeric.AD.Rank1.Halley
  (
  -- * Halley's Method (Tower AD)
    findZero
  , inverse
  , fixedPoint
  , extremum
  ) where

import Prelude hiding (all)
import Numeric.AD.Internal.Forward (Forward)
import Numeric.AD.Internal.On
import Numeric.AD.Internal.Tower (Tower)
import Numeric.AD.Mode
import Numeric.AD.Rank1.Tower (diffs0)
import Numeric.AD.Rank1.Forward (diff)

-- $setup
-- >>> import Data.Complex

-- | The 'findZero' function finds a zero of a scalar function using
-- Halley's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes constant
-- ("it converges"), no further elements are returned.
--
-- Examples:
--
-- >>> take 10 $ findZero (\x->x^2-4) 1
-- [1.0,1.8571428571428572,1.9997967892704736,1.9999999999994755,2.0]
--
-- >>> last $ take 10 $ findZero ((+1).(^2)) (1 :+ 1)
-- 0.0 :+ 1.0
findZero :: (Fractional a, Eq a) => (Tower a -> Tower a) -> a -> [a]
findZero f = go where
  go x = x : if x == xn then [] else go xn where
    (y:y':y'':_) = diffs0 f x
    xn = x - 2*y*y'/(2*y'*y'-y*y'')
{-# INLINE findZero #-}

-- | The 'inverse' function inverts a scalar function using
-- Halley's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes constant
-- ("it converges"), no further elements are returned.
--
-- Note: the @take 10 $ inverse sqrt 1 (sqrt 10)@ example that works for Newton's method
-- fails with Halley's method because the preconditions do not hold!
inverse :: (Fractional a, Eq a) => (Tower a -> Tower a) -> a -> a -> [a]
inverse f x0 y = findZero (\x -> f x - auto y) x0
{-# INLINE inverse  #-}

-- | The 'fixedPoint' function find a fixedpoint of a scalar
-- function using Halley's method; its output is a stream of
-- increasingly accurate results.  (Modulo the usual caveats.)
--
-- If the stream becomes constant ("it converges"), no further
-- elements are returned.
--
-- >>> last $ take 10 $ fixedPoint cos 1
-- 0.7390851332151607
fixedPoint :: (Fractional a, Eq a) => (Tower a -> Tower a) -> a -> [a]
fixedPoint f = findZero (\x -> f x - x)
{-# INLINE fixedPoint #-}

-- | The 'extremum' function finds an extremum of a scalar
-- function using Halley's method; produces a stream of increasingly
-- accurate results.  (Modulo the usual caveats.) If the stream becomes
-- constant ("it converges"), no further elements are returned.
--
-- >>> take 10 $ extremum cos 1
-- [1.0,0.29616942658570555,4.59979519460002e-3,1.6220740159042513e-8,0.0]
extremum :: (Fractional a, Eq a) => (On (Forward (Tower a)) -> On (Forward (Tower a))) -> a -> [a]
extremum f = findZero (diff (off . f . On))
{-# INLINE extremum #-}


{-# LANGUAGE OverloadedLists #-}

module Main where

import Numeric.GSL.ODE
import Numeric.LinearAlgebra

-- Differential equation
f :: Double -> [Double] -> [Double]
f t [x, v] = [v, - x + mu * v * (1 - x ^ 2)]

-- Mu scalar, dampening strenth
mu :: Double
mu = 0.1

-- Boundary conditions
ts :: Vector Double
ts = linspace 1000 (0, 50)

-- Use default solver: Embedded Runge-Kutta-Fehlberg (4, 5) method.
vanderpol1 :: [Vector Double]
vanderpol1 = toColumns $ odeSolve f [1, 0] ts

-- Use Runge-Kutta (2,3) solver
vanderpol2 :: [Vector Double]
vanderpol2 = toColumns $ odeSolveV RK2 hi epsAbs epsRel (l2v f) [1, 0] ts
  where
    epsAbs = 1.49012e-08
    epsRel = epsAbs
    hi = (ts ! 1 - ts ! 0) / 100
    l2v f = \t -> fromList . f t . toList

main :: IO ()
main = do
  print vanderpol1
  print vanderpol2


{-# language FlexibleInstances, FlexibleContexts, TypeFamilies #-}
{-# language ConstrainedClassMethods #-}
{-# language CPP #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Core.Numeric.BLAS.Class
-- Copyright   :  (c) Marco Zocca 2017-2018
-- License     :  see the file LICENSE
--
-- Maintainer  :  ocramz fripost org
-- Stability   :  experimental
-- Portability :  portable
--
-- Typeclasses for linear algebra and related concepts
--
-----------------------------------------------------------------------------
module Core.Numeric.BLAS.Class where

import Data.Complex
import Foreign.C.Types (CFloat, CDouble)

-- * Matrix and vector elements (optionally Complex)
--
-- Uses 'FlexibleContexts' and 'TypeFamilies'
class (Eq e , Fractional e, Floating e, Num (EltMag e), Ord (EltMag e)) => Elt e where
  type EltMag e :: *
  -- | Complex conjugate, or identity function if its input is real-valued
  conj :: e -> e
  conj = id
  -- | Magnitude
  mag :: e -> EltMag e

instance Elt Double where {type EltMag Double = Double ; mag = id}
instance Elt Float where {type EltMag Float = Float; mag = id}
instance (RealFloat e) => Elt (Complex e) where
  type EltMag (Complex e) = e
  conj = conjugate
  mag = magnitude

infixl 6 ^+^, ^-^

-- * Additive group
class AdditiveGroup v where
  -- | The zero element: identity for '(^+^)'
  zeroV :: v
  -- | Add vectors
  (^+^) :: v -> v -> v
  -- | Additive inverse
  negateV :: v -> v
  -- | Group subtraction
  (^-^) :: v -> v -> v
  (^-^) x y = x ^+^ negateV y

infixr 7 .*

-- * Vector space @v@.
class (AdditiveGroup v, Num (Scalar v)) => VectorSpace v where
  type Scalar v :: *
  -- | Scale a vector
  (.*) :: Scalar v -> v -> v

-- | Adds inner (dot) products.
class VectorSpace v => InnerSpace v where
  -- | Inner/dot product
  (<.>) :: v -> v -> Scalar v

-- | Inner product
dot :: InnerSpace v => v -> v -> Scalar v
dot = (<.>)

-- | Scale a vector by the reciprocal of a number (e.g. for normalization)
(./) :: (VectorSpace v, s ~ Scalar v, Fractional s) => v -> s -> v
v ./ s = recip s .* v

-- | Vector multiplied by scalar
(*.) :: (VectorSpace v, s ~ Scalar v) => v -> s -> v
(*.) = flip (.*)

-- | Convex combination of two vectors (NB: 0 <= `a` <= 1).
cvx :: VectorSpace v => Scalar v -> v -> v -> v
cvx a u v = a .* u ^+^ ((1-a) .* v)

-- ** Hilbert-space distance function

-- |`hilbertDistSq x y = || x - y ||^2` computes the squared L2 distance between two vectors
hilbertDistSq :: InnerSpace v => v -> v -> Scalar v
hilbertDistSq x y = t <.> t where
  t = x ^-^ y

-- * Normed vector spaces
--
-- Uses 'ConstrainedClassMethods'
class (InnerSpace v, Num (RealScalar v), Eq (RealScalar v), Epsilon (Magnitude v), Show (Magnitude v), Ord (Magnitude v)) => Normed v where
  type Magnitude v :: *
  type RealScalar v :: *
  -- | L1 norm
  norm1 :: v -> Magnitude v
  -- | Euclidean (L2) norm squared
  norm2Sq :: v -> Magnitude v
  -- | Lp norm (p > 0)
  normP :: RealScalar v -> v -> Magnitude v
  -- | Normalize w.r.t. Lp norm
  normalize :: RealScalar v -> v -> v
  -- | Normalize w.r.t. L2 norm
  normalize2 :: v -> v
  -- | Normalize w.r.t. norm2' instead of norm2
  normalize2' :: Floating (Scalar v) => v -> v
  normalize2' x = x ./ norm2' x
  -- | Euclidean (L2) norm
  norm2 :: Floating (Magnitude v) => v -> Magnitude v
  norm2 x = sqrt (norm2Sq x)
  -- | Euclidean (L2) norm; returns a Complex (norm :+ 0) for Complex-valued vectors
  norm2' :: Floating (Scalar v) => v -> Scalar v
  norm2' x = sqrt $ x <.> x
  -- | Lp norm (p > 0)
  norm :: Floating (Magnitude v) => RealScalar v -> v -> Magnitude v
  norm p v
    | p == 1 = norm1 v
    | p == 2 = norm2 v
    | otherwise = normP p v

-- | Infinity-norm (Real)
normInftyR :: (Foldable t, Ord a) => t a -> a
normInftyR x = maximum x

-- | Infinity-norm (Complex)
normInftyC :: (Foldable t, RealFloat a, Functor t) => t (Complex a) -> a
normInftyC x = maximum (magnitude <$> x)

-- | Lp inner product (p > 0)
dotLp :: (Set t, Foldable t, Floating a) => a -> t a -> t a ->  a
dotLp p v1 v2 = sum u**(1/p) where
  f a b = (a*b)**p
  u = liftI2 f v1 v2

-- | Reciprocal
reciprocal :: (Functor f, Fractional b) => f b -> f b
reciprocal = fmap recip

-- |Scale
scale :: (Num b, Functor f) => b -> f b -> f b
scale n = fmap (* n)

-- * Matrix ring

-- | A matrix ring is any collection of matrices over some ring R that form a ring under matrix addition and matrix multiplication

class (AdditiveGroup m, Epsilon (MatrixNorm m)) => MatrixRing m where
  type MatrixNorm m :: *
  -- | Matrix-matrix product
  (##) :: m -> m -> m
  -- | Matrix times matrix transpose (A B^T)
  (##^) :: m -> m -> m
  -- | Matrix transpose times matrix (A^T B)
  (#^#) :: m -> m -> m
  a #^# b = transpose a ## b
  -- | Matrix transpose (Hermitian conjugate in the Complex case)
  transpose :: m -> m
  -- | Frobenius norm
  normFrobenius :: m -> MatrixNorm m

-- * Linear vector space

class (VectorSpace v {-, MatrixRing (MatrixType v)-}) => LinearVectorSpace v where
  type MatrixType v :: *
  -- | Matrix-vector action
  (#>) :: MatrixType v -> v -> v
  -- | Dual matrix-vector action
  (<#) :: v -> MatrixType v -> v

-- * Typeclasses for internal use (not to be exposed to the user)

-- ** FiniteDim : finite-dimensional objects
class FiniteDim f where
  type FDSize f
  -- | Dimension (i.e. Int for SpVector, (Int, Int) for SpMatrix)
  dim :: f -> FDSize f

-- ** HasData : accessing inner data (do not export)
class HasData f where
  type HDData f
  -- | Number of nonzeros
  nnz :: f -> Int
  dat :: f -> HDData f

-- ** Sparse : sparse datastructures
class (FiniteDim f, HasData f) => Sparse f where
  -- | Sparsity (fraction of nonzero elements)
  spy :: Fractional b => f -> b

-- ** Set : types that behave as sets
--
-- Example : sparse vector addition operates on the union of the nonzero indices, whereas the sparse vector inner product uses
class Functor f => Set f where
  -- | Union binary lift : apply function on _union_ of two "sets"
  liftU2 :: (a -> a -> a) -> f a -> f a -> f a
  -- | Intersection binary lift : apply function on _intersection_ of two "sets"
  liftI2 :: (a -> a -> b) -> f a -> f a -> f b

-- * instances

-- | Instances for builtin types
--
-- Uses 'CPP'
#define ScalarType(t) \
instance AdditiveGroup (t) where {zeroV = 0; (^+^) = (+); negateV = negate};\
instance VectorSpace (t) where {type Scalar (t) = t; (.*) = (*) };

-- ScalarType(Int)
-- ScalarType(Integer)
ScalarType(Float)
ScalarType(Double)
ScalarType(Complex Float)
ScalarType(Complex Double)

#undef ScalarType

instance InnerSpace Float  where {(<.>) = (*)}
instance InnerSpace Double where {(<.>) = (*)}
instance InnerSpace (Complex Float)  where {x <.> y = x * conjugate y}
instance InnerSpace (Complex Double) where {x <.> y = x * conjugate y}

#define SimpleNormedInstance(t) \
instance Normed (t) where {type Magnitude (t) = t; type RealScalar (t) = t;\
 norm1 = abs; norm2Sq = (**2); normP _ = abs; normalize _ = signum;\
 normalize2 = signum; normalize2' = signum; norm2 = abs; norm2' = abs; norm _ = abs};

SimpleNormedInstance(Float)
SimpleNormedInstance(Double)

#undef SimpleNormedInstance

#define ComplexNormedInstance(t) \
instance Normed (Complex (t)) where {\
 type Magnitude  (Complex (t)) = t;\
 type RealScalar (Complex (t)) = t;\
 norm1   (r :+ i) = abs r + abs i;\
 norm2Sq (r :+ i) = r*r + i*i;\
 normP p (r :+ i) = (r**p + i**p)**(1/p);\
 normalize p c = toC (1 / normP p c) * c;\
 normalize2  c = (1 / norm2' c) * c;\
 norm2  = magnitude;\
 norm2' = toC . magnitude;};

ComplexNormedInstance(Float)
ComplexNormedInstance(Double)

#undef ComplexNormedInstance

-- * Utilities

-- | Lift a real number onto the complex plane
toC :: Num a => a -> Complex a
toC r = r :+ 0

-- | Provides a test to see if a quantity is near zero.
--
-- Appears originally in ekmett/linear
--
-- TODO[ocramz] Uses 'FlexibleInstances', so might be a good idea to move it in a separate module
--
-- >>> nearZero (1e-11 :: Double)
-- False
--
-- >>> nearZero (1e-17 :: Double)
-- True
--
-- >>> nearZero (1e-5 :: Float)
-- False
--
-- >>> nearZero (1e-7 :: Float)
-- True
class (Floating a, Num a) => Epsilon a where
  -- | Determine if a quantity is near zero.
  nearZero :: a -> Bool

-- | @'abs' a '<=' 1e-6@
instance Epsilon Float where
  nearZero a = abs a <= 1e-6

-- | @'abs' a '<=' 1e-12@
instance Epsilon Double where
  nearZero a = abs a <= 1e-12

-- | @'abs' a '<=' 1e-6@
instance Epsilon CFloat where
  nearZero a = abs a <= 1e-6

-- | @'abs' a '<=' 1e-12@
instance Epsilon CDouble where
  nearZero a = abs a <= 1e-12

-- | @'magnitude' a '<=' 1e-6@
instance Epsilon (Complex Float) where
  nearZero a = magnitude a <= 1e-6

-- | @'magnitude' a '<=' 1e-12@
instance Epsilon (Complex Double) where
  nearZero a = magnitude a <= 1e-12

-- | @'magnitude' a '<=' 1e-6@
instance Epsilon (Complex CFloat) where
  nearZero a = magnitude a <= 1e-6

-- | @'magnitude' a '<=' 1e-12@
instance Epsilon (Complex CDouble) where
  nearZero a = magnitude a <= 1e-12

-- * Rounding operations

-- | Is this quantity close to 1 ?
nearOne :: Epsilon a => a -> Bool
nearOne x = nearZero (1 - x)

-- | Is this quantity distinguishable from 0 ?
isNz :: Epsilon a => a -> Bool
isNz x = not (nearZero x)

withDefault :: (t -> Bool) -> t -> t -> t
withDefault q d x | q x = d
                  | otherwise = x

roundZero, roundOne, roundZeroOne :: Epsilon a => a -> a
roundZero = withDefault nearZero (fromIntegral (0 :: Int))
roundOne = withDefault nearOne (fromIntegral (1 :: Int))

with2Defaults :: (t -> Bool) -> (t -> Bool) -> t -> t -> t -> t
with2Defaults q1 q2 d1 d2 x | q1 x = d1
                            | q2 x = d2
                            | otherwise = x

-- | Round to respectively 0 or 1
roundZeroOne = with2Defaults nearZero nearOne (fromIntegral (0 :: Int)) (fromIntegral (1 :: Int))


module Numeric.Sundials.ODEOpts where

import           Data.Int (Int32)
import qualified Data.Vector.Storable as VS

import           Numeric.LinearAlgebra.HMatrix (Vector, Matrix)

type Jacobian = Double -> Vector Double -> Matrix Double

-- | Adaptive step-size control
-- functions.
--
-- [GSL](https://www.gnu.org/software/gsl/doc/html/ode-initval.html#adaptive-step-size-control)
-- allows the user to control the step size adjustment using
-- \(D_i = \epsilon^{abs}s_i + \epsilon^{rel}(a_{y} |y_i| + a_{dy/dt} h |\dot{y}_i|)\) where
-- \(\epsilon^{abs}\) is the required absolute error, \(\epsilon^{rel}\)
-- is the required relative error, \(s_i\) is a vector of scaling
-- factors, \(a_{y}\) is a scaling factor for the solution \(y\) and
-- \(a_{dydt}\) is a scaling factor for the derivative of the solution \(dy/dt\).
--
-- [CVode](https://computation.llnl.gov/projects/sundials/cvode)
-- allows the user to control the step size adjustment using
-- \(\eta^{rel}|y_i| + \eta^{abs}_i\). For compatibility with
-- [hmatrix-gsl](https://hackage.haskell.org/package/hmatrix-gsl),
-- tolerances for \(y\) and \(\dot{y}\) can be specified but the latter have no
-- effect.
data StepControl = X     Double Double -- ^ absolute and relative tolerance for \(y\); in GSL terms, \(a_{y} = 1\) and \(a_{dy/dt} = 0\); in ARKode terms, the \(\eta^{abs}_i\) are identical
                 | X'    Double Double -- ^ absolute and relative tolerance for \(\dot{y}\); in GSL terms, \(a_{y} = 0\) and \(a_{dy/dt} = 1\); in ARKode terms, the latter is treated as the relative tolerance for \(y\) so this is the same as specifying 'X' which may be entirely incorrect for the given problem
                 | XX'   Double Double Double Double -- ^ include both via relative tolerance
                                                     -- scaling factors \(a_y\), \(a_{{dy}/{dt}}\); in ARKode terms, the latter is ignored and \(\eta^{rel} = a_{y}\epsilon^{rel}\)
                 | ScXX' Double Double Double Double (Vector Double) -- ^ scale absolute tolerance of \(y_i\); in ARKode terms, \(a_{{dy}/{dt}}\) is ignored, \(\eta^{abs}_i = s_i \epsilon^{abs}\) and \(\eta^{rel} = a_{y}\epsilon^{rel}\)
  deriving (Eq, Ord, Show, Read)

-- | Stepping functions
data ODEMethod = ADAMS
               | BDF
  deriving (Eq, Ord, Show, Read)

data ODEOpts = ODEOpts {
    maxNumSteps :: Int32
  , minStep     :: Double
  , maxFail     :: Int32
  , odeMethod   :: ODEMethod
  , stepControl :: StepControl
  , initStep    :: 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\)
  } deriving (Read, Show, Eq, Ord)

data SundialsDiagnostics = SundialsDiagnostics {
    odeGetNumSteps               :: Int
  , odeGetNumStepAttempts        :: Int
  , odeGetNumRhsEvals_fe         :: Int
  , odeGetNumRhsEvals_fi         :: Int
  , odeGetNumLinSolvSetups       :: Int
  , odeGetNumErrTestFails        :: Int
  , odeGetNumNonlinSolvIters     :: Int
  , odeGetNumNonlinSolvConvFails :: Int
  , dlsGetNumJacEvals            :: Int
  , dlsGetNumRhsEvals            :: Int
  } deriving Show


{-# LANGUAGE FlexibleContexts #-}

module Main where

import Wigner
import BraKet
import Operators
import QUtil
import LinAlg

import Data.Complex
import Data.List.Split

main :: IO ()
main = do
    let expd = computeExpD dim ps dZ (-bound, bound) states
    let n = round . sqrt . fromIntegral . length $ expd
    -- Wigner transform the expected D
    let cplane = getPlane dZ (-bound) bound
    writeComplex "wigner" $ wignerTransform cplane (dZ^2) cplane expd
    --writeComplex "wigner" . mconcat . wt2d . chunksOf n $ expd
    writeComplex "cplane" cplane

-- dimension of vector space to work in
dim = 40
-- number of steps in the complex plane to use
-- computation scales as O(nsteps^2)
nsteps = 50
-- bounds on the complex plane
bound = 5.0
-- step size in complex plane
dZ = 2.0 * bound / fromIntegral nsteps

-- list of states to use
-- really a list of functions, I use applicative to
-- convert to a list of sized state vectors
states :: [Int -> Ket (Complex Double)]
-- vacuum
states = [vacuum]
-- squeezed vacuum
--states = [\dim -> act (squeeze dim 0.75) $ vacuum dim]
-- one photon Fock
--states = [flip fockN 1]
-- Linear combination
--states = [\dim -> vacuum dim <+> fockN dim 1]
-- coherent state
--states = [flip coherent (2)]
-- coherent superposition (AKA schrodinger's cat)
--states = [\dim -> coherent dim 2 <+> coherent dim (-2.0)]
-- Mixed schrodinger cat
--states = [flip coherent ((-2.0) :+ 0.0), flip coherent (2.0 :+ 0.0)]
-- Absorbed
--states = [vacuum, flip fockN 1]
-- classical probabilities of mixed state
ps = [1.0]

-- Construct the complex plane from a range and step size
getPlane :: Double -> Double -> Double -> [Complex Double]
getPlane dz zMin zMax = (:+) <$> range <*> range
  where
    range = [zMin, zMin+dz..zMax]

-- Compute the expected value of the displacement operator
-- over the complex plane
computeExpD :: Int -> [Complex Double] -> Double -> (Double, Double)
                   -> [Int -> Ket (Complex Double)] -> [Complex Double]
computeExpD dim ps dz (zMin, zMax) states = (\z -> expectDisp dim z ps (states <*> [dim])) <$> plane
  where plane = getPlane dz zMin zMax

-- Write complex array to file
writeComplex :: FilePath -> [Complex Double] -> IO ()
writeComplex fpath = writeFile fpath . complexToPlane

-- Make list of complex numbers easy to parse by python
complexToPlane :: [Complex Double] -> String
complexToPlane [] = ""
complexToPlane ((x :+ y):zs) = show x ++ " " ++ show y ++ "\n" ++ complexToPlane zs


{-# OPTIONS_GHC -Wall #-}

{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Sundials.ARKode.ODE
-- Copyright   :  Dominic Steinitz 2018,
--                Novadiscovery 2018
-- License     :  BSD
-- Maintainer  :  Dominic Steinitz
-- Stability   :  provisional
--
-- Solution of ordinary differential equation (ODE) initial value problems.
-- See <https://computation.llnl.gov/projects/sundials/sundials-software> for more detail.
--
-- A simple example:
--
-- <<diagrams/brusselator.png#diagram=brusselator&height=400&width=500>>
--
-- @
-- import           Numeric.Sundials.ARKode.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))
-- @
--
-- With Sundials ARKode, it is possible to retrieve the Butcher tableau for the solver.
--
-- @
-- import           Numeric.Sundials.ARKode.ODE
-- import           Numeric.LinearAlgebra
--
-- import           Data.List (intercalate)
--
-- import           Text.PrettyPrint.HughesPJClass
--
--
-- butcherTableauTex :: ButcherTable -> String
-- butcherTableauTex (ButcherTable m c b b2) =
--   render $
--   vcat [ text ("\n\\begin{array}{c|" ++ (concat $ replicate n "c") ++ "}")
--        , us
--        , text "\\hline"
--        , text bs <+> text "\\\\"
--        , text b2s <+> text "\\\\"
--        , text "\\end{array}"
--        ]
--   where
--     n = rows m
--     rs = toLists m
--     ss = map (\r -> intercalate " & " $ map show r) rs
--     ts = zipWith (\i r -> show i ++ " & " ++ r) (toList c) ss
--     us = vcat $ map (\r -> text r <+> text "\\\\") ts
--     bs  = " & " ++ (intercalate " & " $ map show $ toList b)
--     b2s = " & " ++ (intercalate " & " $ map show $ toList b2)
--
-- main :: IO ()
-- main = do
--
--   let res = butcherTable (SDIRK_2_1_2 undefined)
--   putStrLn $ show res
--   putStrLn $ butcherTableauTex res
--
--   let resA = butcherTable (KVAERNO_4_2_3 undefined)
--   putStrLn $ show resA
--   putStrLn $ butcherTableauTex resA
--
--   let resB = butcherTable (SDIRK_5_3_4 undefined)
--   putStrLn $ show resB
--   putStrLn $ butcherTableauTex resB
-- @
--
--  Using the code above from the examples gives
--
-- KVAERNO_4_2_3
--
-- \[
-- \begin{array}{c|cccc}
-- 0.0 & 0.0 & 0.0 & 0.0 & 0.0 \\
-- 0.871733043 & 0.4358665215 & 0.4358665215 & 0.0 & 0.0 \\
-- 1.0 & 0.490563388419108 & 7.3570090080892e-2 & 0.4358665215 & 0.0 \\
-- 1.0 & 0.308809969973036 & 1.490563388254106 & -1.235239879727145 & 0.4358665215 \\
-- \hline
--  & 0.308809969973036 & 1.490563388254106 & -1.235239879727145 & 0.4358665215 \\
--  & 0.490563388419108 & 7.3570090080892e-2 & 0.4358665215 & 0.0 \\
-- \end{array}
-- \]
--
-- SDIRK_2_1_2
--
-- \[
-- \begin{array}{c|cc}
-- 1.0 & 1.0 & 0.0 \\
-- 0.0 & -1.0 & 1.0 \\
-- \hline
--  & 0.5 & 0.5 \\
--  & 1.0 & 0.0 \\
-- \end{array}
-- \]
--
-- SDIRK_5_3_4
--
-- \[
-- \begin{array}{c|ccccc}
-- 0.25 & 0.25 & 0.0 & 0.0 & 0.0 & 0.0 \\
-- 0.75 & 0.5 & 0.25 & 0.0 & 0.0 & 0.0 \\
-- 0.55 & 0.34 & -4.0e-2 & 0.25 & 0.0 & 0.0 \\
-- 0.5 & 0.2727941176470588 & -5.036764705882353e-2 & 2.7573529411764705e-2 & 0.25 & 0.0 \\
-- 1.0 & 1.0416666666666667 & -1.0208333333333333 & 7.8125 & -7.083333333333333 & 0.25 \\
-- \hline
--  & 1.0416666666666667 & -1.0208333333333333 & 7.8125 & -7.083333333333333 & 0.25 \\
--  & 1.2291666666666667 & -0.17708333333333334 & 7.03125 & -7.083333333333333 & 0.0 \\
-- \end{array}
-- \]
-----------------------------------------------------------------------------
module Numeric.Sundials.ARKode.ODE ( odeSolve
                                   , odeSolveV
                                   , odeSolveVWith
                                   , odeSolveVWith'
                                   , ButcherTable(..)
                                   , butcherTable
                                   , ODEMethod(..)
                                   , StepControl(..)
                                   ) where

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

import           Data.Monoid ((<>))
import           Data.Maybe (isJust)

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

import qualified Data.Vector.Storable as V

import           Data.Coerce (coerce)
import           System.IO.Unsafe (unsafePerformIO)
import           GHC.Generics (C1, Constructor, (:+:)(..), D1, Rep, Generic, M1(..),
                               from, conName)

import           Numeric.LinearAlgebra.Devel (createVector)

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

import           Numeric.Sundials.ODEOpts (ODEOpts(..), Jacobian, SundialsDiagnostics(..))
import qualified Numeric.Sundials.Arkode as T
import           Numeric.Sundials.Arkode (arkSMax,
                                          sDIRK_2_1_2,
                                          bILLINGTON_3_3_2,
                                          tRBDF2_3_3_2,
                                          kVAERNO_4_2_3,
                                          aRK324L2SA_DIRK_4_2_3,
                                          cASH_5_2_4,
                                          cASH_5_3_4,
                                          sDIRK_5_3_4,
                                          kVAERNO_5_3_4,
                                          aRK436L2SA_DIRK_6_3_4,
                                          kVAERNO_7_4_5,
                                          aRK548L2SA_DIRK_8_4_5,
                                          hEUN_EULER_2_1_2,
                                          bOGACKI_SHAMPINE_4_2_3,
                                          aRK324L2SA_ERK_4_2_3,
                                          zONNEVELD_5_3_4,
                                          aRK436L2SA_ERK_6_3_4,
                                          sAYFY_ABURUB_6_3_4,
                                          cASH_KARP_6_4_5,
                                          fEHLBERG_6_4_5,
                                          dORMAND_PRINCE_7_4_5,
                                          aRK548L2SA_ERK_8_4_5,
                                          vERNER_8_5_6,
                                          fEHLBERG_13_7_8)

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

C.include "<stdlib.h>"
C.include "<stdio.h>"
C.include "<math.h>"
C.include "<arkode/arkode.h>"                 -- prototypes for ARKODE 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 "<arkode/arkode_direct.h>"          -- access to ARKDls interface
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"

-- | Stepping functions
data ODEMethod = SDIRK_2_1_2            Jacobian
               | SDIRK_2_1_2'
               | BILLINGTON_3_3_2       Jacobian
               | BILLINGTON_3_3_2'
               | TRBDF2_3_3_2           Jacobian
               | TRBDF2_3_3_2'
               | KVAERNO_4_2_3          Jacobian
               | KVAERNO_4_2_3'
               | ARK324L2SA_DIRK_4_2_3  Jacobian
               | ARK324L2SA_DIRK_4_2_3'
               | CASH_5_2_4             Jacobian
               | CASH_5_2_4'
               | CASH_5_3_4             Jacobian
               | CASH_5_3_4'
               | SDIRK_5_3_4            Jacobian
               | SDIRK_5_3_4'
               | KVAERNO_5_3_4          Jacobian
               | KVAERNO_5_3_4'
               | ARK436L2SA_DIRK_6_3_4  Jacobian
               | ARK436L2SA_DIRK_6_3_4'
               | KVAERNO_7_4_5          Jacobian
               | KVAERNO_7_4_5'
               | ARK548L2SA_DIRK_8_4_5  Jacobian
               | ARK548L2SA_DIRK_8_4_5'
               | HEUN_EULER_2_1_2         Jacobian
               | HEUN_EULER_2_1_2'
               | BOGACKI_SHAMPINE_4_2_3   Jacobian
               | BOGACKI_SHAMPINE_4_2_3'
               | ARK324L2SA_ERK_4_2_3     Jacobian
               | ARK324L2SA_ERK_4_2_3'
               | ZONNEVELD_5_3_4          Jacobian
               | ZONNEVELD_5_3_4'
               | ARK436L2SA_ERK_6_3_4     Jacobian
               | ARK436L2SA_ERK_6_3_4'
               | SAYFY_ABURUB_6_3_4       Jacobian
               | SAYFY_ABURUB_6_3_4'
               | CASH_KARP_6_4_5          Jacobian
               | CASH_KARP_6_4_5'
               | FEHLBERG_6_4_5         Jacobian
               | FEHLBERG_6_4_5'
               | DORMAND_PRINCE_7_4_5     Jacobian
               | DORMAND_PRINCE_7_4_5'
               | ARK548L2SA_ERK_8_4_5     Jacobian
               | ARK548L2SA_ERK_8_4_5'
               | VERNER_8_5_6            Jacobian
               | VERNER_8_5_6'
               | FEHLBERG_13_7_8         Jacobian
               | FEHLBERG_13_7_8'
  deriving Generic

constrName :: (HasConstructor (Rep a), Generic a)=> a -> String
constrName = genericConstrName . from

class HasConstructor (f :: * -> *) where
  genericConstrName :: f x -> String

instance HasConstructor f => HasConstructor (D1 c f) where
  genericConstrName (M1 x) = genericConstrName x

instance (HasConstructor x, HasConstructor y) => HasConstructor (x :+: y) where
  genericConstrName (L1 l) = genericConstrName l
  genericConstrName (R1 r) = genericConstrName r

instance Constructor c => HasConstructor (C1 c f) where
  genericConstrName x = conName x

instance Show ODEMethod where
  show x = constrName x

-- FIXME: We can probably do better here with generics
getMethod :: ODEMethod -> Int
getMethod (SDIRK_2_1_2 _)            = sDIRK_2_1_2
getMethod (SDIRK_2_1_2')             = sDIRK_2_1_2
getMethod (BILLINGTON_3_3_2 _)       = bILLINGTON_3_3_2
getMethod (BILLINGTON_3_3_2')        = bILLINGTON_3_3_2
getMethod (TRBDF2_3_3_2 _)           = tRBDF2_3_3_2
getMethod (TRBDF2_3_3_2')            = tRBDF2_3_3_2
getMethod (KVAERNO_4_2_3  _)         = kVAERNO_4_2_3
getMethod (KVAERNO_4_2_3')           = kVAERNO_4_2_3
getMethod (ARK324L2SA_DIRK_4_2_3 _)  = aRK324L2SA_DIRK_4_2_3
getMethod (ARK324L2SA_DIRK_4_2_3')   = aRK324L2SA_DIRK_4_2_3
getMethod (CASH_5_2_4 _)             = cASH_5_2_4
getMethod (CASH_5_2_4')              = cASH_5_2_4
getMethod (CASH_5_3_4 _)             = cASH_5_3_4
getMethod (CASH_5_3_4')              = cASH_5_3_4
getMethod (SDIRK_5_3_4 _)            = sDIRK_5_3_4
getMethod (SDIRK_5_3_4')             = sDIRK_5_3_4
getMethod (KVAERNO_5_3_4 _)          = kVAERNO_5_3_4
getMethod (KVAERNO_5_3_4')           = kVAERNO_5_3_4
getMethod (ARK436L2SA_DIRK_6_3_4 _)  = aRK436L2SA_DIRK_6_3_4
getMethod (ARK436L2SA_DIRK_6_3_4')   = aRK436L2SA_DIRK_6_3_4
getMethod (KVAERNO_7_4_5 _)          = kVAERNO_7_4_5
getMethod (KVAERNO_7_4_5')           = kVAERNO_7_4_5
getMethod (ARK548L2SA_DIRK_8_4_5 _)  = aRK548L2SA_DIRK_8_4_5
getMethod (ARK548L2SA_DIRK_8_4_5')   = aRK548L2SA_DIRK_8_4_5
getMethod (HEUN_EULER_2_1_2 _)       = hEUN_EULER_2_1_2
getMethod (HEUN_EULER_2_1_2')        = hEUN_EULER_2_1_2
getMethod (BOGACKI_SHAMPINE_4_2_3 _) = bOGACKI_SHAMPINE_4_2_3
getMethod (BOGACKI_SHAMPINE_4_2_3')  = bOGACKI_SHAMPINE_4_2_3
getMethod (ARK324L2SA_ERK_4_2_3 _)   = aRK324L2SA_ERK_4_2_3
getMethod (ARK324L2SA_ERK_4_2_3')    = aRK324L2SA_ERK_4_2_3
getMethod (ZONNEVELD_5_3_4 _)        = zONNEVELD_5_3_4
getMethod (ZONNEVELD_5_3_4')         = zONNEVELD_5_3_4
getMethod (ARK436L2SA_ERK_6_3_4 _)   = aRK436L2SA_ERK_6_3_4
getMethod (ARK436L2SA_ERK_6_3_4')    = aRK436L2SA_ERK_6_3_4
getMethod (SAYFY_ABURUB_6_3_4 _)     = sAYFY_ABURUB_6_3_4
getMethod (SAYFY_ABURUB_6_3_4')      = sAYFY_ABURUB_6_3_4
getMethod (CASH_KARP_6_4_5 _)        = cASH_KARP_6_4_5
getMethod (CASH_KARP_6_4_5')         = cASH_KARP_6_4_5
getMethod (FEHLBERG_6_4_5 _)         = fEHLBERG_6_4_5
getMethod (FEHLBERG_6_4_5' )         = fEHLBERG_6_4_5
getMethod (DORMAND_PRINCE_7_4_5 _)   = dORMAND_PRINCE_7_4_5
getMethod (DORMAND_PRINCE_7_4_5')    = dORMAND_PRINCE_7_4_5
getMethod (ARK548L2SA_ERK_8_4_5 _)   = aRK548L2SA_ERK_8_4_5
getMethod (ARK548L2SA_ERK_8_4_5')    = aRK548L2SA_ERK_8_4_5
getMethod (VERNER_8_5_6 _)           = vERNER_8_5_6
getMethod (VERNER_8_5_6')            = vERNER_8_5_6
getMethod (FEHLBERG_13_7_8 _)        = fEHLBERG_13_7_8
getMethod (FEHLBERG_13_7_8')         = fEHLBERG_13_7_8

getJacobian :: ODEMethod -> Maybe Jacobian
getJacobian (SDIRK_2_1_2 j)            = Just j
getJacobian (BILLINGTON_3_3_2 j)       = Just j
getJacobian (TRBDF2_3_3_2 j)           = Just j
getJacobian (KVAERNO_4_2_3  j)         = Just j
getJacobian (ARK324L2SA_DIRK_4_2_3 j)  = Just j
getJacobian (CASH_5_2_4 j)             = Just j
getJacobian (CASH_5_3_4 j)             = Just j
getJacobian (SDIRK_5_3_4 j)            = Just j
getJacobian (KVAERNO_5_3_4 j)          = Just j
getJacobian (ARK436L2SA_DIRK_6_3_4 j)  = Just j
getJacobian (KVAERNO_7_4_5 j)          = Just j
getJacobian (ARK548L2SA_DIRK_8_4_5 j)  = Just j
getJacobian (HEUN_EULER_2_1_2 j)       = Just j
getJacobian (BOGACKI_SHAMPINE_4_2_3 j) = Just j
getJacobian (ARK324L2SA_ERK_4_2_3 j)   = Just j
getJacobian (ZONNEVELD_5_3_4 j)        = Just j
getJacobian (ARK436L2SA_ERK_6_3_4 j)   = Just j
getJacobian (SAYFY_ABURUB_6_3_4 j)     = Just j
getJacobian (CASH_KARP_6_4_5 j)        = Just j
getJacobian (FEHLBERG_6_4_5 j)         = Just j
getJacobian (DORMAND_PRINCE_7_4_5 j)   = Just j
getJacobian (ARK548L2SA_ERK_8_4_5 j)   = Just j
getJacobian (VERNER_8_5_6 j)           = Just j
getJacobian (FEHLBERG_13_7_8 j)        = Just j
getJacobian _                          = Nothing

-- | A version of 'odeSolveVWith' with reasonable default step control.
odeSolveV
    :: ODEMethod
    -> Maybe Double      -- ^ initial step size - by default, ARKode
                         -- 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 SDIRK_5_3_4' (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)

odeSolveVWith ::
  ODEMethod
  -> StepControl
  -> Maybe Double -- ^ initial step size - by default, ARKode
                  -- 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 method control initStepSize 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   = error "ARKode: unexpected use of ODEOpts.odeMethod"
                   , stepControl = error "ARKode: unexpected use of ODEOpts.stepControl"
                   , initStep    = error "ARKode: unexpected use of ODEOpts.initStep"
                   }

odeSolveVWith' ::
  ODEOpts
  -> ODEMethod
  -> StepControl
  -> Maybe Double -- ^ initial step size - by default, ARKode
                  -- 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
  -> Either (Matrix Double, Int) (Matrix Double, SundialsDiagnostics) -- ^ Error code or solution
odeSolveVWith' opts method control initStepSize f y0 tt =
  case solveOdeC (fromIntegral $ maxFail opts)
                 (fromIntegral $ maxNumSteps opts) (coerce $ minStep opts)
                 (fromIntegral $ getMethod method) (coerce initStepSize) jacH (scise control)
                 (coerce f) (coerce y0) (coerce tt) of
    Left  (v, c) -> Left  (reshape l (coerce v), fromIntegral c)
    Right (v, d) -> Right (reshape l (coerce v), 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 method
    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
  -> V.Vector CDouble -- ^ Desired solution times
  -> Either (V.Vector CDouble, CInt) (V.Vector CDouble, SundialsDiagnostics) -- ^ Partial solution and error code or
                                                                             -- solution and diagnostics
solveOdeC maxErrTestFails maxNumSteps_ minStep_ method initStepSize
          jacH (aTols, rTol) fun f0 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
  quasiMatrixRes <- createVector ((fromIntegral dim) * (fromIntegral nTs))
  qMatMut <- V.thaw quasiMatrixRes
  diagnostics :: V.Vector CLong <- createVector 10 -- FIXME
  diagMut <- V.thaw diagnostics
  -- We need the types that sundials expects. These are tied together
  -- in 'CLangToHaskellTypes'. 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 $ T.SunVector { T.sunVecN = T.sunVecN sv
                             , T.sunVecVals = fun t (T.sunVecVals sv)
                             }
        [CU.exp| int{ 0 } |]
  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.ARKode.ODE: Jacobian not defined"
          Just jacI -> do j <- jacI t <$> (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 <- [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 *arkode_mem = NULL;   /* empty ARKode memory structure                */
                         realtype t;
                         long int nst, nst_a, nfe, nfi, nsetups, nje, nfeLS, nni, ncfn, netf;

                         /* 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];
                         };

                         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];
                         };

                         arkode_mem = ARKodeCreate(); /* Create the solver memory */
                         if (check_flag((void *)arkode_mem, "ARKodeCreate", 0)) return 1;

                         /* Call ARKodeInit to initialize the integrator memory and specify the */
                         /* right-hand side function in y'=f(t,y), the inital time T0, and      */
                         /* the initial dependent variable vector y.  Note: we treat the        */
                         /* problem as fully implicit and set f_E to NULL and f_I to f.         */

                         /* Here we use the C types defined in helpers.h which tie up with */
                         /* the Haskell types defined in CLangToHaskellTypes                             */
                         if ($(int method) < MIN_DIRK_NUM) {
                           flag = ARKodeInit(arkode_mem, $fun:(int (* funIO) (double t, SunVector y[], SunVector dydt[], void * params)), NULL, T0, y);
                           if (check_flag(&flag, "ARKodeInit", 1)) return 1;
                         } else {
                           flag = ARKodeInit(arkode_mem, NULL, $fun:(int (* funIO) (double t, SunVector y[], SunVector dydt[], void * params)), T0, y);
                           if (check_flag(&flag, "ARKodeInit", 1)) return 1;
                         }

                         flag = ARKodeSetMinStep(arkode_mem, $(double minStep_));
                         if (check_flag(&flag, "ARKodeSetMinStep", 1)) return 1;
                         flag = ARKodeSetMaxNumSteps(arkode_mem, $(long int maxNumSteps_));
                         if (check_flag(&flag, "ARKodeSetMaxNumSteps", 1)) return 1;
                         flag = ARKodeSetMaxErrTestFails(arkode_mem, $(int maxErrTestFails));
                         if (check_flag(&flag, "ARKodeSetMaxErrTestFails", 1)) return 1;

                         /* Set routines */
                         flag = ARKodeSVtolerances(arkode_mem, $(double rTol), tv);
                         if (check_flag(&flag, "ARKodeSVtolerances", 1)) return 1;

                         /* 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 = ARKDlsSetLinearSolver(arkode_mem, LS, A);
                         if (check_flag(&flag, "ARKDlsSetLinearSolver", 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 = ARKodeSetInitStep(arkode_mem, $(double ss));
                           if (check_flag(&flag, "ARKodeSetInitStep", 1)) return 1;
                         }

                         /* Set the Jacobian if there is one */
                         if ($(int isJac)) {
                           flag = ARKDlsSetJacFn(arkode_mem, $fun:(int (* jacIO) (double t, SunVector y[], SunVector fy[], SunMatrix Jac[], void * params, SunVector tmp1[], SunVector tmp2[], SunVector tmp3[])));
                           if (check_flag(&flag, "ARKDlsSetJacFn", 1)) return 1;
                         }

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

                         /* Explicitly set the method */
                         if ($(int method) >= MIN_DIRK_NUM) {
                           flag = ARKodeSetIRKTableNum(arkode_mem, $(int method));
                           if (check_flag(&flag, "ARKodeSetIRKTableNum", 1)) return 1;
                         } else {
                           flag = ARKodeSetERKTableNum(arkode_mem, $(int method));
                           if (check_flag(&flag, "ARKodeSetERKTableNum", 1)) return 1;
                         }

                         /* Main time-stepping loop: calls ARKode to perform the integration */
                         /* Stops when the final time has been reached                       */
                         for (i = 1; i < $(int nTs); i++) {

                           flag = ARKode(arkode_mem, ($vec-ptr:(double *ts))[i], y, &t, ARK_NORMAL); /* call integrator */
                           if (check_flag(&flag, "ARKode solver failure, stopping integration", 1)) return 1;

                           /* Store the results for Haskell */
                           for (j = 0; j < NEQ; j++) {
                             ($vec-ptr:(double *qMatMut))[i * NEQ + j] = NV_Ith_S(y,j);
                           }
                         }

                         /* Get some final statistics on how the solve progressed */

                         flag = ARKodeGetNumSteps(arkode_mem, &nst);
                         check_flag(&flag, "ARKodeGetNumSteps", 1);
                         ($vec-ptr:(long int *diagMut))[0] = nst;

                         flag = ARKodeGetNumStepAttempts(arkode_mem, &nst_a);
                         check_flag(&flag, "ARKodeGetNumStepAttempts", 1);
                         ($vec-ptr:(long int *diagMut))[1] = nst_a;

                         flag = ARKodeGetNumRhsEvals(arkode_mem, &nfe, &nfi);
                         check_flag(&flag, "ARKodeGetNumRhsEvals", 1);
                         ($vec-ptr:(long int *diagMut))[2] = nfe;
                         ($vec-ptr:(long int *diagMut))[3] = nfi;

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

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

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

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

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

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

                         /* Clean up and return */
                         N_VDestroy(y);            /* Free y vector          */
                         N_VDestroy(tv);           /* Free tv vector         */
                         ARKodeFree(&arkode_mem);  /* Free integrator memory */
                         SUNLinSolFree(LS);        /* Free linear solver     */
                         SUNMatDestroy(A);         /* Free A matrix          */

                         return flag;
                       } |]
  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)
  m <- V.freeze qMatMut
  if res == 0
    then do
      return $ Right (m, d)
    else do
      return $ Left  (m, res)

data ButcherTable = ButcherTable { am  :: Matrix Double
                                 , cv  :: Vector Double
                                 , bv  :: Vector Double
                                 , b2v :: Vector Double
                                 }
  deriving Show

data ButcherTable' a = ButcherTable' { am'  :: V.Vector a
                                     , cv'  :: V.Vector a
                                     , bv'  :: V.Vector a
                                     , b2v' :: V.Vector a
                                     }
  deriving Show

butcherTable :: ODEMethod -> ButcherTable
butcherTable method =
  case getBT method of
    Left c -> error $ show c -- FIXME
    Right (ButcherTable' v w x y, sqp) ->
      ButcherTable { am = subMatrix (0, 0) (s, s) $ (arkSMax >< arkSMax) (V.toList v)
                   , cv = subVector 0 s w
                   , bv = subVector 0 s x
                   , b2v = subVector 0 s y
                   }
      where
        s = fromIntegral $ sqp V.! 0

getBT :: ODEMethod -> Either Int (ButcherTable' Double, V.Vector Int)
getBT method = case getButcherTable method of
                 Left c ->
                   Left $ fromIntegral c
                 Right (ButcherTable' a b c d, sqp) ->
                   Right $ ( ButcherTable' (coerce a) (coerce b) (coerce c) (coerce d)
                           , V.map fromIntegral sqp )

getButcherTable :: ODEMethod
                -> Either CInt (ButcherTable' CDouble, V.Vector CInt)
getButcherTable method = unsafePerformIO $ do
  -- ARKode seems to want an ODE in order to set and then get the
  -- Butcher tableau so here's one to keep it happy
  let funI :: CDouble -> V.Vector CDouble -> V.Vector CDouble
      funI _t ys = V.fromList [ ys V.! 0 ]
  let funE :: CDouble -> V.Vector CDouble -> V.Vector CDouble
      funE _t ys = V.fromList [ ys V.! 0 ]
      f0        = V.fromList [ 1.0 ]
      ts        = V.fromList [ 0.0 ]
      dim = V.length f0
      nEq :: CLong
      nEq = fromIntegral dim
      mN :: CInt
      mN = fromIntegral $ getMethod method

  btSQP :: V.Vector CInt <- createVector 3
  btSQPMut <- V.thaw btSQP
  btAs :: V.Vector CDouble <- createVector (arkSMax * arkSMax)
  btAsMut <- V.thaw btAs
  btCs  :: V.Vector CDouble <- createVector arkSMax
  btBs  :: V.Vector CDouble <- createVector arkSMax
  btB2s :: V.Vector CDouble <- createVector arkSMax
  btCsMut  <- V.thaw btCs
  btBsMut  <- V.thaw btBs
  btB2sMut <- V.thaw btB2s
  let funIOI :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr () -> IO CInt
      funIOI t y f _ptr = do
        sv <- peek y
        poke f $ T.SunVector { T.sunVecN = T.sunVecN sv
                             , T.sunVecVals = funI t (T.sunVecVals sv)
                             }
        -- FIXME: I don't understand what this comment means
        -- Unsafe since the function will be called many times.
        [CU.exp| int{ 0 } |]
  let funIOE :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr () -> IO CInt
      funIOE t y f _ptr = do
        sv <- peek y
        poke f $ T.SunVector { T.sunVecN = T.sunVecN sv
                             , T.sunVecVals = funE t (T.sunVecVals sv)
                             }
        -- FIXME: I don't understand what this comment means
        -- Unsafe since the function will be called many times.
        [CU.exp| int{ 0 } |]
  res <- [C.block| int {
                         /* general problem variables */

                         int flag;                /* reusable error-checking flag      */
                         N_Vector y = NULL;       /* empty vector for storing solution */
                         void *arkode_mem = NULL; /* empty ARKode memory structure     */
                         int i, j;                /* reusable loop indices             */

                         /* 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];
                         };
                         arkode_mem = ARKodeCreate(); /* Create the solver memory */
                         if (check_flag((void *)arkode_mem, "ARKodeCreate", 0)) return 1;

                         flag = ARKodeInit(arkode_mem, $fun:(int (* funIOE) (double t, SunVector y[], SunVector dydt[], void * params)), $fun:(int (* funIOI) (double t, SunVector y[], SunVector dydt[], void * params)), T0, y);
                         if (check_flag(&flag, "ARKodeInit", 1)) return 1;

                         if ($(int mN) >= MIN_DIRK_NUM) {
                         flag = ARKodeSetIRKTableNum(arkode_mem, $(int mN));
                         if (check_flag(&flag, "ARKodeSetIRKTableNum", 1)) return 1;
                         } else {
                         flag = ARKodeSetERKTableNum(arkode_mem, $(int mN));
                         if (check_flag(&flag, "ARKodeSetERKTableNum", 1)) return 1;
                         }

                         int s, q, p;
                         realtype *ai = (realtype *)malloc(ARK_S_MAX * ARK_S_MAX * sizeof(realtype));
                         realtype *ae = (realtype *)malloc(ARK_S_MAX * ARK_S_MAX * sizeof(realtype));
                         realtype *ci = (realtype *)malloc(ARK_S_MAX * sizeof(realtype));
                         realtype *ce = (realtype *)malloc(ARK_S_MAX * sizeof(realtype));
                         realtype *bi = (realtype *)malloc(ARK_S_MAX * sizeof(realtype));
                         realtype *be = (realtype *)malloc(ARK_S_MAX * sizeof(realtype));
                         realtype *b2i = (realtype *)malloc(ARK_S_MAX * sizeof(realtype));
                         realtype *b2e = (realtype *)malloc(ARK_S_MAX * sizeof(realtype));
                         flag = ARKodeGetCurrentButcherTables(arkode_mem, &s, &q, &p, ai, ae, ci, ce, bi, be, b2i, b2e);
                         if (check_flag(&flag, "ARKode", 1)) return 1;
                         $vec-ptr:(int *btSQPMut)[0] = s;
                         $vec-ptr:(int *btSQPMut)[1] = q;
                         $vec-ptr:(int *btSQPMut)[2] = p;
                         for (i = 0; i < s; i++) {
                           for (j = 0; j < s; j++) {
                             /* FIXME: double should be realtype */
                             ($vec-ptr:(double *btAsMut))[i * ARK_S_MAX + j] = ai[i * ARK_S_MAX + j];
                           }
                         }

                         for (i = 0; i < s; i++) {
                           ($vec-ptr:(double *btCsMut))[i]  = ci[i];
                           ($vec-ptr:(double *btBsMut))[i]  = bi[i];
                           ($vec-ptr:(double *btB2sMut))[i] = b2i[i];
                         }

                         /* Clean up and return */
                         N_VDestroy(y);            /* Free y vector */
                         ARKodeFree(&arkode_mem);  /* Free integrator memory */

                         return flag;
                       } |]
  if res == 0
    then do
      x <- V.freeze btAsMut
      y <- V.freeze btSQPMut
      z <- V.freeze btCsMut
      u <- V.freeze btBsMut
      v <- V.freeze btB2sMut
      return $ Right (ButcherTable' { am' = x, cv' = z, bv' = u, b2v' = v }, y)
    else do
      return $ Left res

-- | Adaptive step-size control
-- functions.
--
-- [GSL](https://www.gnu.org/software/gsl/doc/html/ode-initval.html#adaptive-step-size-control)
-- allows the user to control the step size adjustment using
-- \(D_i = \epsilon^{abs}s_i + \epsilon^{rel}(a_{y} |y_i| + a_{dy/dt} h |\dot{y}_i|)\) where
-- \(\epsilon^{abs}\) is the required absolute error, \(\epsilon^{rel}\)
-- is the required relative error, \(s_i\) is a vector of scaling
-- factors, \(a_{y}\) is a scaling factor for the solution \(y\) and
-- \(a_{dydt}\) is a scaling factor for the derivative of the solution \(dy/dt\).
--
-- [ARKode](https://computation.llnl.gov/projects/sundials/arkode)
-- allows the user to control the step size adjustment using
-- \(\eta^{rel}|y_i| + \eta^{abs}_i\). For compatibility with
-- [hmatrix-gsl](https://hackage.haskell.org/package/hmatrix-gsl),
-- tolerances for \(y\) and \(\dot{y}\) can be specified but the latter have no
-- effect.
data StepControl = X     Double Double -- ^ absolute and relative tolerance for \(y\); in GSL terms, \(a_{y} = 1\) and \(a_{dy/dt} = 0\); in ARKode terms, the \(\eta^{abs}_i\) are identical
                 | X'    Double Double -- ^ absolute and relative tolerance for \(\dot{y}\); in GSL terms, \(a_{y} = 0\) and \(a_{dy/dt} = 1\); in ARKode terms, the latter is treated as the relative tolerance for \(y\) so this is the same as specifying 'X' which may be entirely incorrect for the given problem
                 | XX'   Double Double Double Double -- ^ include both via relative tolerance
                                                     -- scaling factors \(a_y\), \(a_{{dy}/{dt}}\); in ARKode terms, the latter is ignored and \(\eta^{rel} = a_{y}\epsilon^{rel}\)
                 | ScXX' Double Double Double Double (Vector Double) -- ^ scale absolute tolerance of \(y_i\); in ARKode terms, \(a_{{dy}/{dt}}\) is ignored, \(\eta^{abs}_i = s_i \epsilon^{abs}\) and \(\eta^{rel} = a_{y}\epsilon^{rel}\)


import Control.Parallel

import Numeric.Container hiding (linspace)
import Numeric.LinearAlgebra.Util
import Numeric.LinearAlgebra.Data hiding (linspace)
import Numeric.LinearAlgebra.Algorithms
import Numeric.LinearAlgebra.Devel
import Data.List hiding (find)
import DSP.Basic

--gridsize
n = 75

--index function
index :: (Int, Int) -> Int -> Double
index (i,j) n
	| i == j       =  6
	| i == j + n*n = -1
	| i == j - n*n = -1
	| i == j + n   = -1
	| i == j - n   = -1
	| i == j + 1   = -1
	| i == j - 1   = -1
	| otherwise    =  0

-- Laplace Matrix in 3 dimensions
t = buildMatrix (n*n*n) (n*n*n) (\(i,j) -> index (i,j) n)
eigen = eigenvalues t

--Calculate corresponding Eigenfunction to n-th energy-eigenvalue
psi :: Int -> Vector Double
psi x = nullVector (a x)
	where a x = t - (mapMatrix (* (realPart (eigen @> x))) (ident (n*n*n)) )

--For further image processing
v = asColumn (psi 56)

main :: IO ()
main = do
-- write out matrix to visualize with python
saveMatrix "vmat.txt" "%f" v


{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE InstanceSigs         #-}
{-# LANGUAGE KindSignatures       #-}
{-# LANGUAGE LambdaCase           #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE UndecidableInstances #-}

module TensorOps.BLAS.HMat
  ( HMat
  , HMatD
  ) where

import           Control.DeepSeq
import           Data.Kind
import           Data.Singletons
import           Data.Singletons.TypeLits
import           Data.Type.Combinator
import           Data.Type.Vector            (Vec, VecT(..))
import           Data.Type.Vector.Util       (curryV2', curryV3')
import           Numeric.LinearAlgebra
import           Numeric.LinearAlgebra.Data  as LA
import           Numeric.LinearAlgebra.Devel
import           TensorOps.BLAS
import           Type.Class.Higher
import           Type.Class.Higher.Util
import qualified Data.Finite                 as DF
import qualified Data.Finite.Internal        as DF
import qualified Data.Vector.Storable        as VS

type HMatD = HMat Double

data HMat :: Type -> BShape Nat -> Type where
    HMV :: { unHMV :: !(Vector a) } -> HMat a ('BV n)
    HMM :: { unHMM :: !(Matrix a) } -> HMat a ('BM n m)

instance (VS.Storable a, Show a, Element a) => Show (HMat a s) where
    showsPrec p = \case
      HMV x -> showParen (p > 10) $ showString "HMV "
                                  . showsPrec 11 x
      HMM x -> showParen (p > 10) $ showString "HMM "
                                  . showsPrec 11 x

instance (VS.Storable a, Show a, Element a) => Show1 (HMat a)

instance (VS.Storable a, NFData a) => NFData (HMat a s) where
    rnf = \case
      HMV xs -> rnf xs
      HMM xs -> rnf xs
    {-# INLINE rnf #-}

instance (VS.Storable a, NFData a) => NFData1 (HMat a)

instance (SingI s, Container Vector a, Container Matrix a, Num a) => Num (HMat a s) where
    (+) = unsafeZipH add add
    (*) = unsafeZipH (VS.zipWith (*)) (liftMatrix2 (VS.zipWith (*)))
    (-) = unsafeZipH (VS.zipWith (-)) (liftMatrix2 (VS.zipWith (-)))
    negate = unsafeMapH (scale (-1)) (scale (-1))
    abs    = unsafeMapH (cmap abs) (cmap abs)
    signum = unsafeMapH (cmap signum) (cmap signum)
    fromInteger = case (sing :: Sing s) of
        SBV n   -> HMV . flip konst (fromIntegral (fromSing n)) . fromInteger
        SBM n m -> HMM . flip konst (fromIntegral (fromSing n)
                                    ,fromIntegral (fromSing m)
                                    ) . fromInteger

-- | WARNING!! Functions should assume equal sized inputs and return
-- outputs of the same size!  This is not checked!!!
unsafeZipH
    :: (Vector a -> Vector a -> Vector a)
    -> (Matrix a -> Matrix a -> Matrix a)
    -> HMat a s -> HMat a s -> HMat a s
unsafeZipH f g = \case
    HMV x -> \case
      HMV y -> HMV $ f x y
    HMM x -> \case
      HMM y -> HMM $ g x y

-- | WARNING!! Functions should return outputs of the same size!  This is
-- not checked!!!
unsafeMapH
    :: (Vector a -> Vector a)
    -> (Matrix a -> Matrix a)
    -> HMat a s -> HMat a s
unsafeMapH f g = \case
    HMV x -> HMV $ f x
    HMM x -> HMM $ g x

liftB'
    :: (Numeric a)
    => Sing s
    -> (Vec n a -> a)
    -> Vec n (HMat a s)
    -> HMat a s
liftB' s f xs = bgen s $ \i -> f (indexB i <$> xs)
{-# INLINE liftB' #-}

instance (Container Vector a, Numeric a) => BLAS (HMat a) where
    type ElemB (HMat a) = a

    -- TODO: rewrite rules
    -- write in parallel?
    liftB
        :: forall n s. ()
        => Sing s
        -> (Vec n a -> a)
        -> Vec n (HMat a s)
        -> HMat a s
    liftB s f = \case
        ØV -> case s of
          SBV sN    -> HMV $ konst (f ØV) ( fromIntegral (fromSing sN) )
          SBM sN sM -> HMM $ konst (f ØV) ( fromIntegral (fromSing sN)
                                          , fromIntegral (fromSing sM)
                                          )
        I x :* ØV -> case x of
          HMV x' -> HMV (cmap (f . (:* ØV) . I) x')
          HMM x' -> HMM (cmap (f . (:* ØV) . I) x')
        I x :* I y :* ØV -> case x of
          HMV x' -> case y of
            HMV y' -> HMV $ VS.zipWith (curryV2' f) x' y'
          HMM x' -> case y of
            HMM y' -> HMM $ liftMatrix2 (VS.zipWith (curryV2' f)) x' y'
        xs@(I x :* I y :* I z :* ØV) -> case x of
          HMV x' -> case y of
            HMV y' -> case z of
              HMV z' -> HMV $ VS.zipWith3 (curryV3' f) x' y' z'
          _ -> liftB' s f xs
        xs@(_ :* _ :* _ :* _ :* _) -> liftB' s f xs

    axpy α (HMV x) my
        = HMV
        . maybe id (add . unHMV) my
        . scale α
        $ x
    {-# INLINE axpy #-}
    dot (HMV x) (HMV y)
        = x <.> y
    {-# INLINE dot #-}
    ger (HMV x) (HMV y)
        = HMM $ x `outer` y
    {-# INLINE ger #-}
    gemv α (HMM a) (HMV x) mβy
        = HMV
        . maybe id (\(β, HMV y) -> add (scale β y)) mβy
        . (a #>)
        . scale α
        $ x
    {-# INLINE gemv #-}
    gemm α (HMM a) (HMM b) mβc
        = HMM
        . maybe id (\(β, HMM c) -> add (scale β c)) mβc
        . (a <>)
        . scale α
        $ b
    {-# INLINE gemm #-}
    scaleB α = unsafeMapH (scale α) (scale α)
    {-# INLINE scaleB #-}
    addB = unsafeZipH add add
    {-# INLINE addB #-}
    indexB = \case
        PBV i -> \case
          HMV x -> x `atIndex` fromInteger (DF.getFinite i)
        PBM i j -> \case
          HMM x -> x `atIndex` ( fromInteger (DF.getFinite i)
                               , fromInteger (DF.getFinite j)
                               )
    {-# INLINE indexB #-}
    indexRowB i (HMM x) = HMV (x ! fromInteger (DF.getFinite i))
    {-# INLINE indexRowB #-}
    transpB (HMM x) = HMM (tr x)
    {-# INLINE transpB #-}
    iRowsB f (HMM x) = fmap (HMM . fromRows)
                     . traverse (\(i,r) -> unHMV <$> f (DF.Finite i) (HMV r))
                     . zip [0..]
                     . toRows
                     $ x
    {-# INLINE iRowsB #-}
    iElemsB f = \case
        HMV x -> fmap (HMV . fromList)
               . traverse (\(i,e) -> f (PBV (DF.Finite i)) e)
               . zip [0..]
               . LA.toList
               $ x
        HMM x -> fmap (HMM . fromLists)
               . traverse (\(i,rs) ->
                     traverse (\(j, e) -> f (PBM (DF.Finite i) (DF.Finite j)) e)
                   . zip [0..]
                   $ rs
                 )
               . zip [0..]
               . toLists
               $ x
    {-# INLINE iElemsB #-}
    -- TODO: can be implemented in parallel maybe?
    bgenA = \case
      SBV sN -> \f -> fmap (HMV . fromList)
                    . traverse (\i -> f (PBV (DF.Finite i)))
                    $ [0 .. fromSing sN - 1]
      SBM sN sM -> \f -> fmap (HMM . fromLists)
                       . traverse (\(i, js) ->
                           traverse (\j -> f (PBM (DF.Finite i) (DF.Finite j))) js
                         )
                       . zip [0 .. fromSing sN - 1]
                       $ repeat [0 .. fromSing sM - 1]
    {-# INLINE bgenA #-}
    bgenRowsA
        :: forall f n m. (Applicative f, SingI n)
        => (DF.Finite n -> f (HMat a ('BV m)))
        -> f (HMat a ('BM n m))
    bgenRowsA f = fmap (HMM . fromRows)
                . traverse (fmap unHMV . f . DF.Finite)
                $ [0 .. fromSing (sing @Nat @n) - 1]
    {-# INLINE bgenRowsA #-}

    eye = HMM . ident . fromIntegral . fromSing
    {-# INLINE eye #-}
    diagB = HMM . diag . unHMV
    {-# INLINE diagB #-}
    getDiagB = HMV . takeDiag . unHMM
    {-# INLINE getDiagB #-}
    traceB = sumElements . takeDiag . unHMM
    {-# INLINE traceB #-}
    sumB = \case
      HMV xs -> sumElements xs
      HMM xs -> sumElements xs
    {-# INLINE sumB #-}