{-# 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}\)


{-# LANGUAGE FlexibleContexts #-}

import Data.List
import Data.Function
import Control.Lens
import Control.Monad
import Numeric.LinearAlgebra

grid :: Integer -> Matrix Double
grid n = build (300, 300) (builder `on` (+1) . round)
  where builder x y = let id = x + 10
                       in fromInteger . subtract 5
                        $ (id * y + n) * id `mod` 1000 `div` 100

main = getContents
   >>= putStrLn
     . (\(z, (x, y)) -> show (x+1) ++ "," ++ show (y+1) ++ "," ++ show z)
     . (_2 %~ snd)
     . maximumBy (compare `on` snd)
     . zip ([1..300] :: [Int])
     . map (liftM2 (,) maxElement maxIndex)
     . zipWith corr2 (map (konst 1 . join (,)) [1..300])
     . repeat . grid . read


-- |
-- Module    : GRN.Density
-- Copyright : (c) 2011 Jason Knight
-- License   : BSD3
--
-- Maintainer  : jason@jasonknight.us
-- Stability   : experimental
-- Portability : portable
--
-- Simulates many different random values in the unknown Kmap entries and
-- estimates the density of this distribution for each gene.
--
-- Then we will convolve this with the 'measurement kernel' to provide a non
-- dirac prior for every gene (even deterministic ones).
--

module GRN.Density where

import GRN.StateTransition
import GRN.Render
import GRN.Sparse
import GRN.Types

import qualified Data.Map as M
import Data.List
import Control.Monad
import Text.Printf
import System.Console.ParseArgs

import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Generic as G

import Control.Parallel.Strategies

{-import Numeric.FFT.Vector.Unitary-}

import Data.Graph.Analysis.Algorithms.Common (componentsOf)

import Statistics.Sample.KernelDensity
import Statistics.Sample
import Numeric.MathFunctions.Constants
import Statistics.Distribution hiding (stdDev)
import Statistics.Distribution.Normal
import Graphics.Gnuplot.Simple hiding (Points)

step = 0.01 -- Step size
nl = 101 :: Int -- Number of poitns
smoothingVar = 0.001 -- Variance of smoothing gaussian

points = U.enumFromStepN 0.00 step nl

normSmooth mu = G.map (density $ normalDistr mu smoothingVar) points
stdNorm = normSmooth 0.5

myKernelEst :: U.Vector Double -> U.Vector Double
myKernelEst sample = fst $ kde_ nl 0.0 1.0 sample

myEst :: U.Vector Double -> U.Vector Double
myEst xv  | stdDev xv < 0.001 = normalize.normSmooth $ (G.sum xv / (fromIntegral.G.length $ xv))
          | otherwise = normalize.clamp.convolve stdNorm.myKernelEst $ xv

normalize xv = let scale = 1.0 / (step * G.sum xv) in G.map (*scale) xv
clamp xv = G.map (\x->if x<0 then 0 else x) xv

-- Assume that both xs and ys are the same size
convolve xs ys = error "FFTW out of date." --G.slice start l uncentered
  {-where-}
    {-xsF = run dftR2C (xs G.++ zeros)-}
    {-ysF = run dftR2C (ys G.++ zeros)-}
    {-l = G.length xs-}
    {-start = floor $ (fromIntegral l)/2-}
    {-zeros = U.replicate l 0.0-}
    {-mult = G.zipWith (*) xsF ysF-}
    {-uncentered = run dftC2R mult-}

calcEsts :: GeneMC -> AttractorDensity
calcEsts = M.map myEst

simRuns :: Args String -> ParseData -> IO () --U.Vector SSA
simRuns args p = do
  let
    titles = map (\x->defaultStyle {lineSpec = CustomStyle [LineTitle x]} ) (M.keys p)
    titleAll = zip titles
    prepForPlot = titleAll.map (G.toList.G.zip points.snd).M.toList
    ks = buildKmaps p
    dm = getRequiredArg args "submode" -- Number of simruns

    vals = concatMap (\(Kmap _ _ x)->M.elems x) (M.elems ks)
    unknowns = [C x | C x <- vals] ++ [X | X <- vals]

  print unknowns

  when (dm == "m") $ do
    let
      allssas = calcSSAsDOK args p ks
      allests = calcEsts allssas
    --print allssas
    plotPathsStyle [XRange (0.0,1.0)] (prepForPlot allests)

  when (dm == "g") $ do
    let
      grComps = componentsOf $ simulate args $ kmapToStateGraph ks
      grNComps = map normalizeGraph grComps
      allests = map (\(d,g)->(d, g, calcEsts.calcSSAs args p ks $ g)) grNComps

    mapM_ (plotPathsStyle [XRange (0.0,1.0)] . prepForPlot . (\(_,_,x)->x)) allests
    --when gen $ drawStateGraph snd.G.head $ simVec) args


{-# LANGUAGE BangPatterns #-}
{-|
Provides a parser for text formatted MNIST data.
-}
module ML.MnistLoader (parseMnist) where

import           ML.Sample (Sample(..))

import           Data.Attoparsec.ByteString.Char8 (Parser, char, double, parseOnly, sepBy)
import qualified Data.ByteString.Char8 as BS
import qualified Data.Vector as V
import           Numeric.LinearAlgebra (R, assoc, vector)

makeSample :: ([R], [R]) -> Sample
makeSample (image, [expected]) = Sample (vector image) (assoc 10 0 [(floor expected, 1)])

-- | Break the list of samples into components.
--
-- These components are implicit in the file format.
toSampleSet :: V.Vector Sample -> (V.Vector Sample, V.Vector Sample, V.Vector Sample)
toSampleSet vec = (train, validation, test)
    where (train, rest) = V.splitAt 50000 vec
          (validation, test) = V.splitAt 10000 rest

-- | Parse an MNIST sample image.
sample :: Parser Sample
sample = do
    !x <- double `sepBy` (char ' ')
    let !s = makeSample $ splitAt 784 x
    pure s

-- | Parse the given bytestring into (training data, validation data, test data).
--
-- Data format:
--
--   * 70000 lines lines of 785 doubles.
--   * first 1     - 50000 lines are test data.
--   * lines 50001 - 60000 are validation data.
--   * lines 60001 - 70000 are test data.
--   * first 784 double of each line are the image.
--   * last double of each is the expected value.
parseMnist :: BS.ByteString -> (V.Vector Sample, V.Vector Sample, V.Vector Sample)
parseMnist = toSampleSet . V.unfoldr parseLines . BS.lines
    where parseLines :: [BS.ByteString] -> Maybe (Sample, [BS.ByteString])
          parseLines []     = Nothing
          parseLines (l:ls) = Just (fromRight $ parseOnly sample l, ls)

          fromRight (Right b) = b


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 #-}


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeApplications #-}

module NormalizeSpec where

import Commons
import Data.Array (listArray)
import Data.Complex (Complex (..))
import qualified Data.IntMap.Strict as IM
import qualified Data.Map as Map
import Data.Maybe (fromJust)
import HashedExpression.Internal.Expression
import HashedExpression.Internal.Normalize
import HashedExpression.Internal.Utils
import HashedExpression.Interp (eval, (~=))
import HashedExpression.Operation
import HashedExpression.Prettify
import HashedExpression.Value
import Test.Hspec
import Var
import Prelude hiding ((^))

reFT :: (DimensionType d) => Expression d R -> Expression d R
reFT = xRe . ft

imFT :: (DimensionType d) => Expression d R -> Expression d R
imFT = xIm . ft

spec :: Spec
spec = do
  describe "Normalize spec" $ do
    specify "normalize scalar one zero" $ do
      constant 0.0 *. constant 9.0 `shouldNormalizeTo` constant 0.0
      x * one `shouldNormalizeTo` x
      one * x `shouldNormalizeTo` x
      x * zero `shouldNormalizeTo` zero
      zero * x `shouldNormalizeTo` zero
      y * (x * zero) `shouldNormalizeTo` zero
      zero * (x * one) `shouldNormalizeTo` zero
      zero * x * one `shouldNormalizeTo` zero
      zero * (x * y) `shouldNormalizeTo` zero
      (x * y) * zero `shouldNormalizeTo` zero
      (x * zero) * one `shouldNormalizeTo` zero
      ((x * y) * one) `shouldNormalizeTo` (x * y)
      x * y * z * one `shouldNormalizeTo` x * y * z
      product [x, y, z, t, w, zero] `shouldNormalizeTo` zero
    specify "normalize log and exponential" $ do
      log (exp x) `shouldNormalizeTo` x
      exp (log x) `shouldNormalizeTo` x
    specify "complex related" $ do
      ((x +: y) * (z +: w)) `shouldNormalizeTo` ((x * z - y * w) +: (x * w + y * z))
      xRe (x +: y) `shouldNormalizeTo` x
      xIm (x +: y) `shouldNormalizeTo` y
      (x +: y) + (u +: v) `shouldNormalizeTo` (x + u) +: (y + v)
      s *. (x +: y) `shouldNormalizeTo` (s *. x) +: (s *. y)
      (x +: y) * (z +: w) `shouldNormalizeTo` (x * z - y * w) +: (x * w + y * z)
    specify "dot product" $ do
      x <.> zero `shouldNormalizeTo` zero
      zero <.> x `shouldNormalizeTo` zero
      ((s *. x) <.> y) `shouldNormalizeTo` (s *. (x <.> y))
      x <.> (s *. y) `shouldNormalizeTo` s *. (x <.> y)
    specify "distributivity" $ do
      x * (y + z) `shouldNormalizeTo` (x * y + x * z)
      (y + z) * x `shouldNormalizeTo` (x * y + x * z)
      (x *. (y + z)) `shouldNormalizeTo` (x *. y + x *. z)
      (x <.> (y + z)) `shouldNormalizeTo` ((x <.> y) + (x <.> z))
      (y + z) <.> x `shouldNormalizeTo` (y <.> x) + (z <.> x)
      x * sum [y, z, t, u, v] `shouldNormalizeTo` sum (map (x *) [y, z, t, u, v])
      sum [y, z, t, u, v] * x `shouldNormalizeTo` sum (map (x *) [y, z, t, u, v])
      x *. sum [y, z, t, u, v] `shouldNormalizeTo` sum (map (x *.) [y, z, t, u, v])
      x <.> sum [y, z, t, u, v] `shouldNormalizeTo` sum (map (x <.>) [y, z, t, u, v])
      sum [y, z, t, u, v] <.> x `shouldNormalizeTo` sum (map (<.> x) [y, z, t, u, v])
      product [a, b, c, sum [x, y, z]] `shouldNormalizeTo` sum (map (product . (: [a, b, c])) [x, y, z])
      (x + y) * (z + t) * a * b `shouldNormalizeTo` (a * b * x * z + a * b * x * t + a * b * y * z + a * b * y * t)
    specify "flatten sum and product" $ do
      product [x * y, product [z, t, w], one] `shouldNormalizeTo` product [x, y, z, t, w]
      sum [x + y, sum [z, t, w + s], zero] `shouldNormalizeTo` sum [x, y, z, t, w, s]
    specify "group constants together" $ do
      product [one, one, x, y, one, z] `shouldNormalizeTo` product [x, y, z]
      sum [one, one, x, y, one, z] `shouldNormalizeTo` sum [constant 3, x, y, z]
      product [constant 1, constant 2, x, y, constant 3, z] `shouldNormalizeTo` product [constant 6, x, y, z]
    specify "combine same terms" $ do
      sum [one *. x1, x1, x1, constant 3 *. y1, y1] `shouldNormalizeTo` sum [constant 3 *. x1, constant 4 *. y1]
      sum [constant (-1) *. x1, x1, constant 3 *. y1, y1, z1] `shouldNormalizeTo` sum [constant 4 *. y1, z1]
      x1 - x1 `shouldNormalizeTo` zero1
      sum [one *. x, x, x, constant 3 *. y, y] `shouldNormalizeTo` sum [constant 3 *. x, constant 4 *. y]
      sum [constant (-1) *. x, x, constant 3 *. y, y, z] `shouldNormalizeTo` sum [constant 4 *. y, z]
      x - x `shouldNormalizeTo` zero
    specify "scale rules" $ do
      x *. (y *. v) `shouldNormalizeTo` (x * y) *. v
    specify "negate rules" $ do
      negate (negate x) `shouldNormalizeTo` normalize x
      negate (negate (x + y)) `shouldNormalizeTo` (x + y)
      negate zero `shouldNormalizeTo` zero
  describe "Normalize spec higher dimension" $ do
    specify "normalize one d one zero" $ do
      x1 * one1 `shouldNormalizeTo` x1
      one1 * x1 `shouldNormalizeTo` x1
      x1 * zero1 `shouldNormalizeTo` zero1
      zero1 * x1 `shouldNormalizeTo` zero1
      y1 * (x1 * zero1) `shouldNormalizeTo` zero1
      zero1 * (x1 * one1) `shouldNormalizeTo` zero1
      zero1 * x1 * one1 `shouldNormalizeTo` zero1
      zero1 * (x1 * y1) `shouldNormalizeTo` zero1
      (x1 * y1) * zero1 `shouldNormalizeTo` zero1
      (x1 * zero1) * one1 `shouldNormalizeTo` zero1
      (x1 * y1) * one1 `shouldNormalizeTo` (x1 * y1)
      x1 * y1 * z1 * one1 `shouldNormalizeTo` x1 * y1 * z1
    specify "dot product higher dimension with scaling and point wise" $ do
      x1 <.> zero1 `shouldNormalizeTo` zero
      zero1 <.> x1 `shouldNormalizeTo` zero
      (s *. x1) <.> y1 `shouldNormalizeTo` s *. (x1 <.> y1)
      x1 <.> (s *. y1) `shouldNormalizeTo` s *. (x1 <.> y1)
      x1 * (y1 + z1) `shouldNormalizeTo` x1 * y1 + x1 * z1
      (y1 + z1) * x1 `shouldNormalizeTo` x1 * y1 + x1 * z1
      s *. (y1 + z1) `shouldNormalizeTo` s *. y1 + s *. z1
      x1 <.> (y1 + z1) `shouldNormalizeTo` (x1 <.> y1) + (x1 <.> z1)
      (y1 + z1) <.> x1 `shouldNormalizeTo` (y1 <.> x1) + (z1 <.> x1)
    specify "log and exp higher" $ do
      log (exp x1) `shouldNormalizeTo` x1
      exp (log x1) `shouldNormalizeTo` x1
      log (exp x2) `shouldNormalizeTo` x2
      exp (log x2) `shouldNormalizeTo` x2
    specify "rotate rules" $ do
      rotate 1 (rotate 2 x1) `shouldNormalizeTo` rotate 3 x1
      rotate 2 (rotate 3 x1) `shouldNormalizeTo` rotate 5 x1
      rotate (1, 1) (rotate (2, 3) x2)
        `shouldNormalizeTo` rotate (3, 4) x2
      rotate (0, 1) (rotate (2, -3) x2)
        `shouldNormalizeTo` rotate (2, -2) x2
      rotate (0, 0, 0) x3 `shouldNormalizeTo` x3
      rotate (2, 3, 4) (s *. y3) `shouldNormalizeTo` s
        *. rotate (2, 3, 4) y3
      rotate (2, 3) (sum [x2, y2, z2])
        `shouldNormalizeTo` sum (map (rotate (2, 3)) [x2, y2, z2])
      rotate (2, 3) (product [x2, y2, z2])
        `shouldNormalizeTo` product (map (rotate (2, 3)) [x2, y2, z2])
  describe "Normalize piecewise" $ specify "some piecewise rules" $ do
    piecewise [1] c [x, y] `shouldNormalizeTo` piecewise [1] c [x, zero] + piecewise [1] c [zero, y]
    piecewise [1] c [x + z, y] `shouldNormalizeTo` piecewise [1] c [x, zero] + piecewise [1] c [z, zero] + piecewise [1] c [zero, y]
    piecewise [1] c [x *. y, zero] `shouldNormalizeTo` x *. piecewise [1] c [y, zero]
    piecewise [1, 2] c [zero, x *. y, zero] `shouldNormalizeTo` x *. piecewise [1, 2] c [zero, y, zero]
    piecewise [4, 5] x1 [zero1, x1 * (s *. y1), zero1] `shouldNormalizeTo` s *. piecewise [4, 5] x1 [zero1, x1 * y1, zero1]
    piecewise [1] c1 [c1 +: y1, z1 +: t1] `shouldNormalizeTo` piecewise [1] c1 [c1, z1] +: piecewise [1] c1 [y1, t1]
  describe "Fourier transform" $ specify "some Ft rules" $ do
    reFT (reFT x1) + imFT (imFT x1) `shouldNormalizeTo` constant (fromIntegral defaultDim1D) *. x1
    reFT (reFT x2) + imFT (imFT x2) `shouldNormalizeTo` constant (fromIntegral (default1stDim2D * default2ndDim2D)) *. x2
    let x = variable1D @10 "x"
        y = variable1D @10 "y"
        z = variable1D @10 "z"
        t = variable1D @10 "t"
        valMap =
          Map.fromList
            [ ("x", V1D $ listArray (0, 9) [1 ..]),
              ("y", V1D $ listArray (0, 9) [2, 5 ..]),
              ("z", V1D $ listArray (0, 9) [3, 8 ..]),
              ("t", V1D $ listArray (0, 9) [0, -1 ..])
            ]
    let res1 = eval valMap $ reFT (reFT (x * y + z * t))
        res2 = eval valMap . normalize $ reFT (reFT (x * y + z * t))
    res1 `shouldApprox` res2
    let res1 = eval valMap $ imFT (imFT (x * y + z * t))
        res2 = eval valMap . normalize $ imFT (imFT (x * y + z * t))
    res1 `shouldApprox` res2
    let res1 = eval valMap $ reFT (reFT (x * y + z * t)) + imFT (imFT (x * y + z * t))
        res2 = eval valMap . normalize $ reFT (reFT (x * y + z * t)) + imFT (imFT (x * y + z * t))
    res1 `shouldApprox` res2
  describe "Unit tests" $ specify "normalize scalar one zero" $ do
    x `shouldNormalizeTo` x
    constant 1 / x `shouldNormalizeTo` x ^ (-1)
    x + x `shouldNormalizeTo` constant 2 *. x
    x - x `shouldNormalizeTo` constant 0
    x * x `shouldNormalizeTo` x ^ 2
    x / x `shouldNormalizeTo` constant 1
    x + y `shouldNormalizeTo` x + y
    x - y `shouldNormalizeTo` x - y
    x * y `shouldNormalizeTo` x * y
    x / y `shouldNormalizeTo` x * y ^ (-1)
    (x + y) * (x + y) `shouldNormalizeTo` (constant 2.0 *. (x * y)) + (x ^ 2) + (y ^ 2)
    (x + y) ^ 2 `shouldNormalizeTo` (constant 2.0 *. (x * y)) + (x ^ 2) + (y ^ 2)
    (x - y) ^ 2 `shouldNormalizeTo` (constant (-2.0) *. (x * y)) + (x ^ 2) + (y ^ 2)
    (x - y) * (x - y) `shouldNormalizeTo` (constant (-2.0) *. (x * y)) + (x ^ 2) + (y ^ 2)
    (x * y) ^ 2 `shouldNormalizeTo` (x ^ 2) * (y ^ 2)
    (x * y) * (x * y) `shouldNormalizeTo` (x ^ 2) * (y ^ 2)
    (x / y) * (x / y) `shouldNormalizeTo` (y ^ (-2)) * (x ^ 2)
    (x / y) ^ 2 `shouldNormalizeTo` (y ^ (-2)) * (x ^ 2)
    (constant 1 / x) * (constant 1 / x) `shouldNormalizeTo` x ^ (-2)
    (constant 1 / x) ^ 2 `shouldNormalizeTo` x ^ (-2)
    x * x `shouldNormalizeTo` x ^ 2
    (x * x) * x `shouldNormalizeTo` x ^ 3
    constant 1 / x `shouldNormalizeTo` x ^ (-1)
    x * x / x `shouldNormalizeTo` normalize x
    x / x `shouldNormalizeTo` constant 1
    x / x / x `shouldNormalizeTo` x ^ (-1)
    x / x / y `shouldNormalizeTo` y ^ (-1)
    (x * y) ^ 3 `shouldNormalizeTo` (x ^ 3) * (y ^ 3)
    ((x * y) ^ 3) * x / y `shouldNormalizeTo` (y ^ 2) * (x ^ 4)
    ((x * y) ^ 3) * x + y `shouldNormalizeTo` y + ((y ^ 3) * (x ^ 4))
    ((x * y) ^ 3) * (x + y) `shouldNormalizeTo` ((y ^ 3) * (x ^ 4)) + ((x ^ 3) * (y ^ 4))
    x / (x ^ 3) `shouldNormalizeTo` x ^ (-2)
    (x ^ 2) ^ 2 `shouldNormalizeTo` x ^ 4
    (x + y) * (x + y) `shouldNormalizeTo` (constant 2.0 *. (x * y)) + (x ^ 2) + (y ^ 2)
    ((x + y) * (x + y)) * (x + y) `shouldNormalizeTo` (constant 3.0 *. (y * (x ^ 2))) + (constant 3.0 *. (x * (y ^ 2))) + (x ^ 3) + (y ^ 3)
    ((x * y) ^ 3) * (x + y) ^ 2 `shouldNormalizeTo` ((y ^ 3) * (x ^ 5)) + ((x ^ 3) * (y ^ 5)) + (constant 2.0 *. ((x ^ 4) * (y ^ 4)))
    (x + y) ^ 2 `shouldNormalizeTo` (constant 2.0 *. (x * y)) + (x ^ 2) + (y ^ 2)
    x / (x ^ 2) `shouldNormalizeTo` x ^ (-1)
    x / (x * x) `shouldNormalizeTo` x ^ (-1)
    x / x * x `shouldNormalizeTo` normalize x
    x / (x ^ 3) `shouldNormalizeTo` x ^ (-2)
    (x * x) / y `shouldNormalizeTo` (y ^ (-1)) * (x ^ 2)
    (x ^ 2) ^ 3 `shouldNormalizeTo` x ^ 6
    ((x ^ 3) ^ 3) ^ 2 `shouldNormalizeTo` x ^ 18
    y * (y ^ 2) ^ 2 `shouldNormalizeTo` y ^ 5
    x / ((y ^ 2) ^ 3) `shouldNormalizeTo` x * y ^ (-6)
    ((x * y) ^ 2) / x `shouldNormalizeTo` x * y ^ 2
    ((x + y) ^ 2) * ((x - y) ^ 2) `shouldNormalizeTo` (constant (-2.0) *. ((x ^ 2) * (y ^ 2))) + (x ^ 4) + (y ^ 4)
    ((x + y) ^ 2) + ((x - y) ^ 2) `shouldNormalizeTo` (constant 2.0 *. (x ^ 2)) + (constant 2.0 *. (y ^ 2))
    ((x + y) ^ 2) - ((x - y) ^ 2) `shouldNormalizeTo` constant 4.0 *. (x * y)
    (x + y) * (x - y) `shouldNormalizeTo` (constant (-1.0) *. (y ^ 2)) + (x ^ 2)
    (x + y) + (x - y) `shouldNormalizeTo` constant 2.0 *. x
    (x + y) - (x - y) `shouldNormalizeTo` constant 2.0 *. y
    (x + y) / (x - y) `shouldNormalizeTo` (x * ((x + (constant (-1.0) *. y)) ^ (-1))) + (y * ((x + (constant (-1.0) *. y)) ^ (-1)))
    x <.> y `shouldNormalizeTo` x * y
    x <.> constant 1 `shouldNormalizeTo` normalize x
    x <.> constant (-1) `shouldNormalizeTo` constant (-1) * x
    x <.> constant 0 `shouldNormalizeTo` constant 0
    x <.> x `shouldNormalizeTo` normalize x ^ 2
    x <.> x * y `shouldNormalizeTo` y * (x ^ 2)
    x <.> x + x `shouldNormalizeTo` x + (x ^ 2)
    x <.> (x + x) `shouldNormalizeTo` constant 2.0 *. (x ^ 2)
    x <.> (constant 1 / x) `shouldNormalizeTo` constant 1
    x <.> (y / x) `shouldNormalizeTo` normalize y
    (x <.> (x + y)) ^ 2 `shouldNormalizeTo` ((x ^ 2) * (y ^ 2)) + (constant 2.0 *. (y * (x ^ 3))) + (x ^ 4)