-- tiers.hs -- prints tiers of values up to a certain point
--
-- Copyright (c) 2015-2018 Rudy Matela.
-- Distributed under the 3-Clause BSD licence (see the file LICENSE).
import Test.LeanCheck
import Test.LeanCheck.Utils.Types
import Test.LeanCheck.Function ()
import Test.LeanCheck.Function.Eq ()
import Test.LeanCheck.Tiers (showTiers, finite)
import System.Environment
import Data.List (intercalate, nub)
import Data.Ratio ((%))
import Data.Int
import Data.Word
import Data.Complex

dropEmptyTiersTail :: [[a]] -> [[a]]
dropEmptyTiersTail ([]:[]:[]: []:[]:[]: _) = []
dropEmptyTiersTail (xs:xss) = xs:dropEmptyTiersTail xss
dropEmptyTiersTail []     = []

lengthT :: [[a]] -> Maybe Int
lengthT xss | finite xss' = Just . length $ concat xss'
            | otherwise   = Nothing
  where xss' = dropEmptyTiersTail xss

allUnique :: Eq a => [a] -> Bool
allUnique [] = True
allUnique (x:xs) = x `notElem` xs
                && allUnique (filter (/= x) xs)

countRepetitions :: Eq a => [a] -> Int
countRepetitions xs = length xs - length (nub xs)

ratioRepetitions :: Eq a => [a] -> Rational
ratioRepetitions [] = 0
ratioRepetitions xs = fromIntegral (countRepetitions xs) % fromIntegral (length xs)

showLengthT :: [[a]] -> String
showLengthT xss = case lengthT xss of
                    Nothing -> "Infinity"
                    Just x  -> show x

showDotsLongerThan :: Show a => Int -> [a] -> String
showDotsLongerThan n xs = "["
                       ++ intercalate "," (dotsLongerThan n $ map show xs)
                       ++ "]"
  where
  dotsLongerThan n xs = take n xs ++ ["..." | not . null $ drop n xs]

printTiers :: Show a => Int -> [[a]] -> IO ()
printTiers n = putStrLn . init . unlines . map ("  " ++) . lines . showTiers n

put :: (Show a, Eq a, Listable a) => String -> Int -> a -> IO ()
put t n a = do
  putStrLn $ "map length (tiers :: [[ " ++ t ++ " ]])  =  "
          ++ showDotsLongerThan n (map length $ tiers `asTypeOf` [[a]])
  putStrLn $ ""
  putStrLn $ "length (list :: [ " ++ t ++ " ])  =  "
          ++ showLengthT (tiers `asTypeOf` [[a]])
  putStrLn $ ""
  putStrLn $ "allUnique (list :: [ " ++ t ++ " ])  =  "
          ++ show (allUnique . concat . take n $ tiers `asTypeOf` [[a]])
  putStrLn $ ""
  putStrLn $ "ratioRepetitions (list :: [ " ++ t ++ " ])  =  "
          ++ show (ratioRepetitions . concat . take n $ tiers `asTypeOf` [[a]])
  putStrLn $ ""
  putStrLn $ "tiers :: [" ++ t ++ "]  ="
  printTiers n $ tiers `asTypeOf` [[a]]

u :: a
u = undefined

main :: IO ()
main = do
  as <- getArgs
  let (t,n) = case as of
              []      -> ("Int", 12)
              [t]     -> (t,     12)
              (t:n:_) -> (t, read n)
  case t of
    -- simple types
    "()"               -> put t n (u :: ()                   )
    "Int"              -> put t n (u :: Int                  )
    "Nat"              -> put t n (u :: Nat                  )
    "Integer"          -> put t n (u :: Integer              )
    "Bool"             -> put t n (u :: Bool                 )
    "Char"             -> put t n (u :: Char                 )
    "Float"            -> put t n (u :: Float                )
    "Double"           -> put t n (u :: Double               )
    "Rational"         -> put t n (u :: Rational             )
    -- fixed width integer types
    "Nat2"             -> put t n (u :: Nat2                 )
    "Nat3"             -> put t n (u :: Nat3                 )
    "Nat4"             -> put t n (u :: Nat4                 )
    "Word2"            -> put t n (u :: Word2                )
    "Word3"            -> put t n (u :: Word3                )
    "Word4"            -> put t n (u :: Word4                )
    "Word8"            -> put t n (u :: Word8                )
    "Word16"           -> put t n (u :: Word16               )
    "Word32"           -> put t n (u :: Word32               )
    "Word64"           -> put t n (u :: Word64               )
    "Int2"             -> put t n (u :: Int2                 )
    "Int3"             -> put t n (u :: Int3                 )
    "Int4"             -> put t n (u :: Int4                 )
    "Int8"             -> put t n (u :: Int8                 )
    "Int16"            -> put t n (u :: Int16                )
    "Int32"            -> put t n (u :: Int32                )
    "Int64"            -> put t n (u :: Int64                )
    -- complex numbers
    "Complex Double"   -> put t n (u :: Complex Double       )
    -- lists
    "[()]"             -> put t n (u :: [()]                 )
    "[Int]"            -> put t n (u :: [Int]                )
    "[Nat]"            -> put t n (u :: [Nat]                )
    "[Integer]"        -> put t n (u :: [Integer]            )
    "[Bool]"           -> put t n (u :: [Bool]               )
    "[Char]"           -> put t n (u :: [Char]               )
    "String"           -> put t n (u :: String               )
    -- pairs
    "((),())"          -> put t n (u :: ((),())              )
    "(Int,Int)"        -> put t n (u :: (Int,Int)            )
    "(Nat,Nat)"        -> put t n (u :: (Nat,Nat)            )
    "(Bool,Bool)"      -> put t n (u :: (Bool,Bool)          )
    "(Bool,Int)"       -> put t n (u :: (Bool,Int)           )
    "(Int,Bool)"       -> put t n (u :: (Int,Bool)           )
    "(Int,Int,Int)"    -> put t n (u :: (Int,Int,Int)        )
    "(Nat,Nat,Nat)"    -> put t n (u :: (Nat,Nat,Nat)        )
    -- lists & pairs
    "[((),())]"        -> put t n (u :: [((),())]            )
    "([()],[()])"      -> put t n (u :: ([()],[()])          )
    "([Bool],[Bool])"  -> put t n (u :: ([Bool],[Bool])      )
    "([Int],[Int])"    -> put t n (u :: ([Int],[Int])        )
    -- lists of lists
    "[[Int]]"          -> put t n (u :: [[Int]]              )
    -- functions
    "()->()"           -> put t n (u :: () -> ()             )
    "()->Bool"         -> put t n (u :: () -> Bool           )
    "Bool->()"         -> put t n (u :: Bool -> ()           )
    "Bool->Bool"       -> put t n (u :: Bool -> Bool         )
    "Bool->Bool->Bool" -> put t n (u :: Bool -> Bool -> Bool )
    "Int->Int"         -> put t n (u :: Int -> Int           )
    "Int->Int->Int"    -> put t n (u :: Int -> Int -> Int    )
    "()->Nat"          -> put t n (u :: () -> Nat            )
    "Nat->()"          -> put t n (u :: Nat -> ()            )
    "Nat->Nat"         -> put t n (u :: Nat -> Nat           )
    "Nat->Nat->Nat"    -> put t n (u :: Nat -> Nat -> Nat    )
    "(Nat,Nat)->Nat"   -> put t n (u :: (Nat,Nat) -> Nat     )
    "Bool->Maybe Bool" -> put t n (u :: Bool -> Maybe Bool   )
    "Maybe Bool->Bool" -> put t n (u :: Maybe Bool -> Bool   )
    "Maybe Bool->Maybe Bool" -> put t n (u :: Maybe Bool -> Maybe Bool)
    -- functions with lists
    "[()]->[()]"       -> put t n (u :: [()] -> [()]         )
    "[Bool]->[Bool]"   -> put t n (u :: [Bool] -> [Bool]     )
    "[Int]->[Int]"     -> put t n (u :: [Int] -> [Int]       )
    "[Nat]->[Nat]"     -> put t n (u :: [Nat] -> [Nat]       )
    -- more functions
    "Nat2->()"         -> put t n (u :: Nat2 -> ()           )
    "()->Nat2"         -> put t n (u :: () -> Nat2           )
    "Nat2->Nat2"       -> put t n (u :: Nat2 -> Nat2         )
    "Nat2->Nat3"       -> put t n (u :: Nat2 -> Nat3         )
    "Nat3->Nat2"       -> put t n (u :: Nat3 -> Nat2         )
    "Nat3->Nat3"       -> put t n (u :: Nat3 -> Nat3         )
    -- functions with mixed arguments
    "Bool->Int->Bool"  -> put t n (u :: Bool -> Int -> Bool  )
    "Int->Bool->Bool"  -> put t n (u :: Int -> Bool -> Bool  )
    -- functions with 3 arguments
    "Int->Int->Int->Int"     -> put t n (u :: Int -> Int -> Int -> Int)
    "Bool->Bool->Bool->Bool" -> put t n (u :: Bool -> Bool -> Bool -> Bool)
    -- special lists
    "Set Bool"         -> put t n (u :: Set Bool             )
    "Set ()"           -> put t n (u :: Set ()               )
    "Set Nat"          -> put t n (u :: Set Nat              )
    "Set Nat2"         -> put t n (u :: Set Nat2             )
    "Set Nat3"         -> put t n (u :: Set Nat3             )
    "Bag Bool"         -> put t n (u :: Bag Bool             )
    "Bag ()"           -> put t n (u :: Bag ()               )
    "Bag Nat"          -> put t n (u :: Bag Nat              )
    "Bag Nat2"         -> put t n (u :: Bag Nat2             )
    "Bag Nat3"         -> put t n (u :: Bag Nat3             )
    "NoDup Bool"       -> put t n (u :: NoDup Bool           )
    "NoDup ()"         -> put t n (u :: NoDup ()             )
    "NoDup Nat"        -> put t n (u :: NoDup Nat            )
    "NoDup Nat2"       -> put t n (u :: NoDup Nat2           )
    "NoDup Nat3"       -> put t n (u :: NoDup Nat3           )
    "Map Bool Bool"    -> put t n (u :: Map Bool Bool        )
    "Map () ()"        -> put t n (u :: Map () ()            )
    "Map Nat Nat"      -> put t n (u :: Map Nat Nat          )
    "Map Nat2 Nat2"    -> put t n (u :: Map Nat2 Nat2        )
    "Map Nat3 Nat3"    -> put t n (u :: Map Nat3 Nat3        )
    -- extreme integers
--  "X Int"            -> put t n (u :: X Int   ) -- device dependent
--  "X Word"           -> put t n (u :: X Word  ) -- device dependent
    "X Int4"           -> put t n (u :: X Int4  )
    "X Word4"          -> put t n (u :: X Word4 )
    "X Nat7"           -> put t n (u :: X Nat7  )
--  "Xs Int"           -> put t n (u :: Xs Int  ) -- device dependent
--  "Xs Word"          -> put t n (u :: Xs Word ) -- device dependent
    "Xs Int4"          -> put t n (u :: Xs Int4 )
    "Xs Word4"         -> put t n (u :: Xs Word4)
    "Xs Nat7"          -> put t n (u :: Xs Nat7 )
    -- unhandled
    _                  -> putStrLn $ "unknown/unhandled type `" ++ t ++ "'"


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
module TensorFlow.Operators
  ( (^+)
  , (^-)
  , (^*)
  , (^/)
  ) where

import Data.Int (Int8, Int16, Int32, Int64)
import Data.Word (Word8, Word16)
import Data.Complex (Complex)
import qualified TensorFlow.Core as TF (Tensor, Build, OneOf)
import qualified TensorFlow.GenOps.Core as TF (add, sub, mul, realDiv)

infixl 6 ^+, ^-
infixl 7 ^*, ^/

(^+) :: TF.OneOf '[Complex Double, Complex Float, Int16, Int32, Int64, Int8, Word16, Word8, Double, Float] t
     => TF.Tensor v'1 t
     -> TF.Tensor v'2 t
     -> TF.Tensor TF.Build t
a ^+ b = a `TF.add` b

(^-) :: TF.OneOf '[Complex Double, Complex Float, Int32, Int64, Word16, Double, Float] t
     => TF.Tensor v'1 t
     -> TF.Tensor v'2 t
     -> TF.Tensor TF.Build t
a ^- b = a `TF.sub` b

(^*) :: TF.OneOf '[Complex Double, Complex Float, Int16, Int32, Int64, Int8, Word16, Word8, Double, Float] t
     => TF.Tensor v'1 t
     -> TF.Tensor v'2 t
     -> TF.Tensor TF.Build t
a ^* b = a `TF.mul` b

(^/) :: TF.OneOf '[Complex Double, Complex Float, Int16, Int32, Int64, Int8, Word16, Word8, Double, Float] t
     => TF.Tensor v'1 t
     -> TF.Tensor v'2 t
     -> TF.Tensor TF.Build t
a ^/ b = a `TF.realDiv` b


{-# LANGUAGE BangPatterns, ScopedTypeVariables, CPP #-}
{-# OPTIONS_GHC -fwarn-unused-imports #-}

-- The goal of this benchmark is to demonstrate kernels that have private
-- state which has to be read from memory on each invocation.  Further,
-- the state forms the majority of the working set -- it's larger than
-- the streaming input.

-- In theory this means that keeping a kernel on one core and
-- exploiting pipeline parallelism is better than following the data
-- through the stream graph in a depth first traversal.

--module Main(main) where
module Main where

import Control.Monad as C
import Control.Monad.Par
import Control.Monad.Par.Stream as S
import Control.Monad.Par.OpenList
import Control.DeepSeq
import Control.Exception
import Control.Parallel.Strategies as Strat

-- import Data.Array.Unboxed as U
import Data.Complex
import Data.Int
import Data.Word
import Data.List  (intersperse)
import Data.List.Split (chunk)

import Prelude as P
import System.Environment
import System.Exit
import System.CPUTime
import System.CPUTime.Rdtsc
import GHC.Conc as Conc
import GHC.IO (unsafePerformIO, unsafeDupablePerformIO, unsafeInterleaveIO)

import Debug.Trace
import Control.Monad.Par.Logging

import qualified Data.Vector.Unboxed as UV
import           Data.Vector.Unboxed hiding ((++))

-- Performs some (presently meaningless) computation on a state & a
-- window (stream element) to compute a new state and new window.
--
-- Assumes statesize is a multiple of bufsize:
statefulKern :: Vector Double -> Vector Double -> (Vector Double, Vector Double)
statefulKern state buf = (newstate, newelem)
 where
  -- We could probably test the memory behavior we're interested in
  -- better with inplace update here... but for now, this:
  newstate = UV.map (\d -> d/sum + 2) state
  newelem  = UV.map (+sum) buf

  sum = P.sum partialSums
  partialSums = [ sumslice (cutslice n) | n <- [0..factor-1] ]

  cutslice n    = UV.slice (n*blen) blen state
  sumslice slce = UV.sum (UV.zipWith (+) buf slce)

  factor       = slen `quot` blen
  slen = UV.length state
  blen = UV.length buf

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

monadpar_version (_,numfilters, bufsize, statecoef, numwins) = do
  putStrLn$ "Running monad-par version."

  let statesize = bufsize * statecoef
  results <- evaluate $ runPar$ do
       strm1 :: Stream (UV.Vector Double) <- S.generate numwins (\n -> UV.replicate bufsize 0)
       -- Make a pipeline of numfilters stages:
       let initstate = UV.generate statesize fromIntegral
       pipe_end <- C.foldM (\s _ -> streamScan statefulKern initstate s) strm1 [1..numfilters]

       sums  <- streamMap UV.sum pipe_end
#if 0
       return sums

  -- This is tricky, but two different consumers shouldn't prevent
  -- garbage collection.
  ls <- toListSpin results
--  Just (Cons h _) <- pollIVar results
  putStrLn$ "Sum of first window: " ++ show (P.head ls)
  forkIO$ measureRateList ls
  putStrLn$ "Final sum = "++ show (P.sum ls)
#else

       streamFold (+) 0 sums

  putStrLn$ "Final sum = "++ show results
#endif

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

sparks_version (_,numfilters, bufsize, statecoef, numwins) = do
  putStrLn$ "Running sparks version."

  -- Here we represent the stream as a plain list.
  let
      statesize = bufsize * statecoef
      strm1 :: [UV.Vector Double] = P.replicate numwins $ UV.replicate bufsize 0
      initstate = UV.generate statesize fromIntegral
      applyKern = scan statefulKern initstate

-- This one has the problem that it fully evaluates the stream for the
-- first kernel before moving on to the second:
--      strm_last = (parRepeatFun numfilters applyKern) strm1

      pipe_end = applyNKernels statefulKern numfilters initstate strm1

      sums = P.map UV.sum pipe_end
-- #define SERIAL
#ifndef SERIAL
	     `using` (Strat.parBuffer numCapabilities rwhnf)
#endif

  putStrLn$ "Sum of first window: "++ show (P.head sums)
  measureRateList (sums)
--  measureRateList (strm_last)
--  measureRateList (forceList strm_last)
  putStrLn$ "Final Sum = " ++ show (P.sum sums)

-- Make sure the cars of a list are evaluated before following each cdr:
forceList [] = []
forceList (h:t) = rnf h `seq` forceList t

-- A slightly different version of Data.List.scanl
scan :: (a -> b -> (a,c)) -> a -> [b] -> [c]
scan f q [] = []
scan f q (h:t) = h' : scan f q t
 where
  (q',h') = f q h

type StatefulKernel s a b = s -> a -> (s,b)

-- applyNKernels _ _ _ [] = []
-- applyNKernels :: NFData a => StatefulKernel s a a -> Int -> s -> [a] -> [a]
applyNKernels :: (NFData a, NFData s) => StatefulKernel s a a -> Int -> s -> [a] -> [a]
applyNKernels _    0 _    ls = ls
applyNKernels kern n init ls =
  applyNKernels kern (n-1) init (loop init ls)
 where
  tasklog = unsafeNewTaskSeries (nameFromValue (n,kern))

  loop st [] = []
  loop st (h:t) =
    let (st', x) =
#if 0
	           timePure tasklog$
#endif
		   kern st h in
#ifndef SERIAL
    rnf x `par`
#endif
     x : loop st' t

-- Compose two stateful kernels in parallel.
composeStatefulKernels :: (NFData b, NFData s1) =>
			  StatefulKernel s1 a b -> StatefulKernel s2 b c
		       -> StatefulKernel (s1,s2) a c
-- composeStatefulKernels (f1,f2) (s1,s2) x =
composeStatefulKernels f1 f2 (s1,s2) x =
    rnf pr1 `par` (newstate, snd pr2)
 where
  pr1 = f1 s1 x
  pr2 = f2 s2 (snd pr1)
  newstate = (fst pr1, fst pr2)

parRepeatFun n f =
--  P.foldr (.) id (P.replicate n f)
  P.foldr (.|| rdeepseq) id (P.replicate n f)

--------------------------------------------------------------------------------
-- Main script

default_version = "monad"
default_numfilters = 4
default_bufsize    = 256
default_statecoef  = 10   -- in MULTIPLES of bufsize
default_numwins    = 10 * 1000

main = do
  args <- getArgs
  arg_tup@(version,_,_,_,_) <-
       case args of
	 []          -> return (default_version, default_numfilters, default_bufsize, default_statecoef, default_numwins)
	 [a,b,c,d,e] -> return (a, read b, read c, read d, read e)
	 _         -> do
	               putStrLn$ "ERROR: Invalid arguments, must take 0 or 5 args."
		       putStrLn$ "  Expected args: (version='monad'|'sparks' #filters, bufsize, stateSizeMultiplier, #bufsToProcess)"
		       putStrLn$ "  Received args: "++ show args
		       exitFailure

  putStrLn$ "numCapabilities: "++ show numCapabilities
  putStrLn$ "  Frequency in measurable ticks:  "++ commaint oneSecond ++ "\n"

  case version of
    "monad"  -> monadpar_version arg_tup
    "sparks" -> sparks_version  arg_tup
    _        -> error$ "unknown version: "++version

  putStrLn$ "Finally, dumping all logs:"
  printAllLogs

-- It is not necessary to evaluate every element in the case of an unboxed vector.
instance NFData a => NFData (UV.Vector a) where
 rnf !vec = ()

print_ msg = trace msg $ return ()

-- work pop 1 peek N push 1
-- float->float filter
-- firFilter n coefs =
-- {

--     float sum = 0;
--     for (int i = 0; i < N; i++)
--       sum += peek(i) * COEFF[N-1-i];
--     pop();
--     push(sum);
--   }
-- }

{-

Here's what cachegrind says (on 4 core nehalem):

  $ valgrind --tool=cachegrind ./stream/disjoint_working_sets_pipeline monad 4 768 10 1000 +RTS -N4
   .....
      [measureRate] current rate: 58  Total elems&time 916  181,988,055,721
      [measureRate] Hit end of stream after 1000 elements.
     Final sum = 1.560518243231086e22
     ==21202==
     ==21202== I   refs:      7,111,462,273
     ==21202== I1  misses:          374,190
     ==21202== L2i misses:          298,364
     ==21202== I1  miss rate:          0.00%
     ==21202== L2i miss rate:          0.00%
     ==21202==
     ==21202== D   refs:      3,882,935,974  (3,542,949,529 rd   + 339,986,445 wr)
     ==21202== D1  misses:       14,606,684  (    9,824,455 rd   +   4,782,229 wr)
     ==21202== L2d misses:        6,774,479  (    2,088,565 rd   +   4,685,914 wr)
     ==21202== D1  miss rate:           0.3% (          0.2%     +         1.4%  )
     ==21202== L2d miss rate:           0.1% (          0.0%     +         1.3%  )
     ==21202==
     ==21202== L2 refs:          14,980,874  (   10,198,645 rd   +   4,782,229 wr)
     ==21202== L2 misses:         7,072,843  (    2,386,929 rd   +   4,685,914 wr)
     ==21202== L2 miss rate:            0.0% (          0.0%     +         1.3%  )

Sparks version:
     Final Sum = 1.560518243231086e22
     ==21226==
     ==21226== I   refs:      5,898,314,238
     ==21226== I1  misses:          291,271
     ==21226== L2i misses:          246,518
     ==21226== I1  miss rate:          0.00%
     ==21226== L2i miss rate:          0.00%
     ==21226==
     ==21226== D   refs:      3,264,359,909  (3,206,394,437 rd   + 57,965,472 wr)
     ==21226== D1  misses:       16,003,068  (   10,905,138 rd   +  5,097,930 wr)
     ==21226== L2d misses:        9,177,043  (    4,207,106 rd   +  4,969,937 wr)
     ==21226== D1  miss rate:           0.4% (          0.3%     +        8.7%  )
     ==21226== L2d miss rate:           0.2% (          0.1%     +        8.5%  )
     ==21226==
     ==21226== L2 refs:          16,294,339  (   11,196,409 rd   +  5,097,930 wr)
     ==21226== L2 misses:         9,423,561  (    4,453,624 rd   +  4,969,937 wr)
     ==21226== L2 miss rate:            0.1% (          0.0%     +        8.5%  )

 -}


{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE TupleSections #-}
module Physics.Mechanics where

import           Numeric.LinearAlgebra (Matrix, Vector, scale, (#>), cross, cmap,dot, ident, tr, (<>), norm_2)
import qualified Number.Quaternion as Q
import           Control.Lens hiding ((<.>))
import qualified Data.Vector.Storable as VS
import           Control.Arrow hiding ((<+>))
import           Debug.Trace
import           Data.Maybe (fromMaybe)
import           Data.Function (on)

import Data.PhysicsData
import Data.GameObject
import Data.Mesh
import Physics.Collision
import qualified Data.Vector as V
import Data.List.Lens

import Utils(step,norm)
import Utils.Quaternions

infixr 1 <+>
(<+>) :: Vector Float -> Vector Float -> Vector Float
(<+>) = VS.zipWith (+)

-- (<*>) :: Vector Float -> Vector Float -> Vector Float
-- (<*>) = VS.zipWith (*)

(<.>) :: Vector Float -> Vector Float -> Float
(<.>) = (VS.sum .) . VS.zipWith (*)

addMaybeForce :: Maybe (Vector Float) -> Maybe PhysicsData -> Maybe PhysicsData
addMaybeForce mf mph = over acc . (<+>) <$> mf <*> mph

putForceLoc :: Vector Float -> Vector Float -> GameObject -> GameObject
putForceLoc point force obj =
          let m    = preview (physics._Just.massInv) obj
              lin  = flip scale (scale (force <.> point) force) <$> m
              ang  = cross force point
          in over physics (addMaybeForce lin)
           . over (physics._Just.torque) (<+>ang)
           $ obj

putForce :: Vector Float -> Vector Float -> GameObject -> GameObject
putForce point force obj =
          let p    = point - obj^.location
          in putForceLoc p force obj

putForceOnMid = over (physics._Just.acc) . (+)

putForceOnTip vec obj =
  let quat = obj^.mesh.quaternion
      tip  = rotateWithQuat quat [0,0,-0.1] + obj^.location
  in putForce tip vec obj

putForceOnTipLoc vec obj =
  let quat = obj^.mesh.quaternion
      f    = rotateWithQuat quat vec
  in putForceOnTip f obj

putForceOnSide vec obj =
  let quat = obj^.mesh.quaternion
      tip  = rotateWithQuat quat [1,0,0] <+> obj^.location
  in putForce tip vec obj

putForceOnSideLoc vec obj =
  let quat = obj^.mesh.quaternion
      f    = rotateWithQuat quat vec
  in putForceOnSide f obj

addScaled
  :: Setter GameObject GameObject (VS.Vector Float) (VS.Vector Float)
  -> Lens PhysicsData PhysicsData (VS.Vector Float) (VS.Vector Float)
  -> GameObject -> GameObject
addScaled lens1 lens2 = applyWithPhysics helper
  where helper ph = over lens1 (<+> VS.map (*step) (ph^.lens2))

applySpeed = addScaled location speed

gForce :: Vector Float
gForce = [0,9,0]

gravityF = over (physics._Just.acc) (subtract gForce)

gravCoeff = scale step gForce

getVec :: Lens PhysicsData PhysicsData (VS.Vector Float) (VS.Vector Float) -> GameObject -> VS.Vector Float
getVec l = fromMaybe [0,0,0] . preview (physics._Just.l)

getMat :: Lens PhysicsData PhysicsData (Matrix Float) (Matrix Float) -> GameObject -> Matrix Float
getMat l = fromMaybe (ident 3) . preview (physics._Just.l)

floatOnGround obj =
  let floorNormal = [0,1,0]
      loc       = obj^.location
      d         = dot floorNormal loc + 2
      f         = negate (2-d)^(2 :: Int)
  in if d < 1
  then putForceOnMid (scale f gForce) obj
  else obj

calculateRelPos contact obj = contact^.contactPoint - obj^.location

linearInertia obj =
  fromMaybe 0 $ obj^?physics._Just.massInv

adjustMove relPos p@(l, an) =
  let limit = norm relPos * 0.1
      newAn = signum an * limit
  in if abs an > limit
     then (l + an - newAn,newAn)
     else p

collisionMoves :: Contact -> (GameObject, Vector Float) -> (GameObject, Vector Float) -> ((Float, Vector Float), (Float, Vector Float))
collisionMoves contact p1 p2 =
  let ang1 = angularInertia contact p1
      lin1 = linearInertia $ fst p1
      ang2 = angularInertia contact p2
      lin2 = linearInertia $ fst p2
      coef = contact^.penetration / (ang1 + ang2 + lin1 + lin2)
      (fli1, fan1) = adjustMove (snd p1) (coef*lin1, coef*ang1)
      (fli2, fan2) = adjustMove (snd p2) (coef*lin2, coef*ang2)
      finang1 = getRotationPM contact p1 fan1 ang1
      finang2 = getRotationPM contact p1 fan2 ang2
  in ((fli1, finang1),(fli2, finang2))

getRotationPM :: Contact -> (GameObject, Vector Float) -> Float -> Float -> Vector Float
getRotationPM contact (obj, relPos) ang angIn =
  if angIn == 0
  then [0,0,0]
  else
    scale (ang) .
    scale (1/angIn) $
    getInertiaInv obj #> cross relPos (contact^.contactNormal)

applyMoves :: Contact -> (Float, Vector Float) -> GameObject -> GameObject
applyMoves contact (lin, ang) =
  over (mesh.quaternion) (qvMul ang) .
  over location (+ scale lin (contact^.contactNormal))

calculateDeltaVel :: Contact -> (GameObject, Vector Float) -> Float
calculateDeltaVel contact p@(obj, _) =
  case obj^.physics of
    Just pd ->
      (pd^.massInv +) $
      angularInertia contact p
      -- dot (contact^.contactNormal) .
      -- flip cross relPos .
      -- ((pd^.inertiaInv) #>) .
      -- flip cross (contact^.contactNormal) $ relPos
    Nothing -> 0

angCoeff = 1

angularInertia :: Contact ->  (GameObject, Vector Float) -> Float
angularInertia contact (obj, relPos) =
  case obj^.physics of
    Just pd ->
--      (*angCoeff) .
      dot (contact^.contactNormal) .        -- get component of speed pointing along the normal
      (\angS -> cross angS relPos) .        -- get linear speed of pt
      ((pd^.inertiaInv) #>) $               -- Get ang speed
      cross relPos (contact^.contactNormal) -- for unit impulse
    Nothing -> 0

calculateClosingVel :: Contact -> (GameObject, Vector Float) -> Float
calculateClosingVel contact (obj, relPos) =
  case obj^.physics of
    Just pd ->
      -- (makeContactBasis contact #>) .
      dot (contact^.contactNormal) .
      (pd^.speed +) . cross (pd^.angularS) $ relPos
    Nothing -> 0

calculateImpulse :: Contact -> (GameObject, Vector Float) -> (GameObject, Vector Float) -> Vector Float
calculateImpulse contact o1 o2 =
  let deltaV = (("Delta: "++) . show >>= trace) $
               ((+) `on` calculateDeltaVel contact) o1 o2
      closiV = (("Closing: "++) . show >>= trace) $
               ((+) `on` calculateClosingVel contact) o1 o2
      contV  = negate $ closiV
      restitution = if abs contV < 2
                    then 0.0
                    else 0.0
      impulse = -- if signum contV /= signum (contact^.sign)
                -- then
                  scale (contV * (1.0 + restitution)/deltaV) $ contact^.contactNormal
                -- else [0,0,0]
  in impulse

applyImpulse :: Vector Float -> (GameObject, Vector Float) -> GameObject
applyImpulse impulse (obj, relPos) =
  let vChange = -- impulse
        (("Linear impulse: "++) . show >>= trace) .
        fromMaybe [0,0,0] $
        flip scale impulse <$> (obj^?physics._Just.massInv)
      rotChange = (("Angular impulse: "++) . show >>= trace) .
--        scale angCoeff .
        fromMaybe [0,0,0] $
        (#> cross impulse relPos) <$> (obj^?physics._Just.inertiaInv)
  in over (physics._Just.speed) (+ vChange) .
     over (physics._Just.angularS) (+ rotChange) $ obj

resolveCollision :: (Int, Int) -> [Contact] -> V.Vector GameObject -> V.Vector GameObject
resolveCollision (ix1, ix2) (contact:cs) objs =
  let obj1 = objs V.! ix1
      obj2 = objs V.! ix2
      movement = scale (contact^.penetration*contact^.sign) $
                 contact^.contactNormal
      move = case obj1^.physics of
               Just _ -> over (ix 0._2.location) (+movement)
               Nothing -> over (ix 1._2.location) (subtract movement)
      relPos1 = (("RelPos1: "++) . show >>= trace) $ calculateRelPos contact obj1
      relPos2 = (("RelPos2: "++) . show >>= trace) $ calculateRelPos contact obj2
      impulse = (("Impulse: "++) . show >>= trace) .
                scale (contact^.sign) $
                calculateImpulse contact (obj1, relPos1) (obj2, relPos2)
      (m1, m2) = collisionMoves contact (obj1, relPos1) (obj2, relPos2)
      update i p =
        -- applyMoves contact p .
        applyImpulse i
      updates = move $
                [(ix1, update (impulse) m1 (obj1, relPos1))
                ,(ix2, update (-impulse) m2 (obj2, relPos2))]
      newObjs = V.unsafeUpd objs updates
      check   = V.unsafeIndex newObjs
      -- newColls =
      --   take (length cs) $
      --   detectCollision (prepareCPrim $ check ix1) (prepareCPrim $ check ix2)^.contacts
  in
    resolveCollision (ix1, ix2) cs newObjs
resolveCollision _ [] objs = objs

floorReaction obj =
  let floorHeight = [0,-2,0]
      floorNormal = [0,1,0] :: Vector Float
      vert      = obj^.mesh.vertices
      loc       = obj^.location
      r         = quatToMat3 $ obj^.mesh.quaternion
      m         = (1/) $ fromMaybe 1 $ obj^?physics._Just.massInv
      translate = (+loc) . subtract floorHeight
      collides = filter ((<0) . snd)
               . map (id &&& (dot floorNormal . translate) . (tr r #>))
               $ vert
      minD :: Float
      minD = negate . foldr min 0 . map snd $ collides
      response (point, depth) =
        let
          vrel = dot floorNormal $ getVec speed obj + cross (getVec angularS obj) point
          aobj = dot floorNormal $ cross (i #> cross point floorNormal) point
          i    = getMat inertiaInv obj
          j    = negate (0.8) * vrel/(m+aobj)
          force = scale j floorNormal
          torque = getMat inertiaInv obj #> cross point force
            -- getMat inertiaInv obj #> cross point force
        in (force, torque)
      (f,t) = foldr (\(a,b) (c,d) -> (a+c,b+d)) ([0,0,0],[0,0,0])
            . map response
            $ collides
      n = fromIntegral $ length collides
  in if n /= 0
  then over (physics._Just.angularS) (+ scale (1/n) t)
     . over (physics._Just.speed) (+ scale (1/n) f)
     . over (location) (+ scale minD floorNormal)
     $ obj
  else obj

applyAcc :: GameObject -> GameObject
applyAcc = addScaled (physics._Just.speed) acc

applyForce :: GameObject -> GameObject
applyForce = applyWithPhysics helper
  where helper ph obj =
         let tq    = ph^.torque
             om    = ph^.angularS
             inInv = ph^.inertiaInv
             newOm = om <+> (inInv #> scale step tq)
             quatD = qvMul (scale step newOm) $ obj^.mesh.quaternion
             newQ  = Q.normalize . addQuat quatD $ obj^.mesh.quaternion
             inInD = quatToLInv (ph^.inertiaD) newQ
         in set (physics._Just.torque) [0,0,0]
          . set (physics._Just.acc) [0,0,0]
          . applySpeed
--          . floorReaction
          . applyAcc
--          . floatOnGround
          . set (physics._Just.angularS) newOm
          . set (physics._Just.inertiaInv) inInD
          . set (mesh.quaternion) newQ
          $ obj

applyLinear :: GameObject -> GameObject
applyLinear = set (physics._Just.torque) [0,0,0]
              . set (physics._Just.acc) [0,0,0]
              . applySpeed
              -- . floorReaction
              . applyAcc
              -- . gravityF

applyWithPhysics :: (PhysicsData -> GameObject -> GameObject)
                   -> GameObject -> GameObject
applyWithPhysics f obj = case obj^.physics of
  Just phy -> f phy obj
  Nothing  -> obj

maxspeed = 50

yolo :: Float -> Vector Float -> Vector Float
yolo maxVal vec =
  let vecNorm = norm vec
      coeff   = (vecNorm/maxVal)**(1/2)
  in scale (0 - max 1 coeff) vec

friction = applyWithPhysics helper
  where helper phy =
          let maxspeed  = 1000000
              maxtorque = 100
              coeff1    = scale step . yolo maxspeed $ phy^.speed
              coeff2    = scale step . yolo maxtorque $ phy^.angularS
          in over (physics._Just.angularS) (<+> coeff2)
            . over (physics._Just.speed) (<+> coeff1)

  -- applyWithPhysics helper
  -- where helper phy obj =
  --         let quat   = obj^.mesh.quaternion
  --             rot    = rotateWithQuat quat [0,0,-1]
  --               -- (\(a,b,c) -> [a,b,c]) . imag . flip quatConcat (qInverse quat) $ quatConcat quat (0 +:: (0,0,-1))
  --             force  = 40 -- exp (maxspeed - sNorm) - 1
  --         in over (physics._Just.acc) (<+>cmap (*force) rot) obj


{-# LANGUAGE Haskell2010, TemplateHaskell #-}

module Numeric.Matrix.Sugar (
        iMatrix, dMatrix, cMatrix
    ) where

import Numeric.Matrix
import Language.Haskell.TH
import Language.Haskell.TH.Quote

import Data.Data
import Data.Complex
import Data.Ratio

type ReadM a = String -> Matrix a

iMatrix = QuasiQuoter
            (quoter (read :: ReadM Integer))
            undefined undefined undefined

dMatrix = QuasiQuoter
            (quoter (read :: ReadM Double))
            undefined undefined undefined

cMatrix = QuasiQuoter
            (quoter (read :: ReadM (Complex Double)))
            undefined undefined undefined

quoter :: (Data e, MatrixElement e)
       => (String -> Matrix e) -> String -> Q Exp

quoter read str = do
    let qExp = dataToExpQ (const Nothing) (toList (read str))
    [e| fromList $qExp |]


{-# LANGUAGE FlexibleContexts #-}

module Learn
  ( newManager
  , defaultManagerSettings
  , initNN
  , hotone
  , convertTrains
  , convertTests
  , trainingSimple
  , chunksOf
  ) where

import           Data.Bifunctor
import           Data.List
import           Data.List.Split
import           Debug.Trace
import           Layers
import           Mnist
import           Network.HTTP.Client
import           Numeric.LinearAlgebra

normalAffine :: Int -> Int -> IO (ForwardLayer R)
normalAffine nIn nOut = do
  weights <- rand nIn nOut
  bias <- flatten <$> rand nOut 1
  return $ AffineForward weights bias

initNN :: ForwardLayer R -> [Int] -> IO (ForwardNN R)
initNN eoa ns = do
  (lastAffine:affines) <- mapM (uncurry normalAffine) $ spans [] ns
  let layers = foldl' (\b a -> a ~> eoa ~> b) lastAffine affines
  return $ ForwardNN layers SoftmaxWithCrossForward
  where
    spans rs [a, b]   = (a, b) : rs
    spans rs (a:b:xs) = spans ((a, b) : rs) (b : xs)

convertTrains :: Int -> MnistData -> [TrainBatch R]
convertTrains batchSize (MnistData src) =
  map (mkTrainer . unzip) $ chunksOf batchSize vectors
  where
    vectors = map (bimap (hotone 10) flatten) src
    mkTrainer (a, b) = TrainBatch (fromRows a, fromZ (fromRows b) / 255)

hotone :: (Integral v, NElement a) => Int -> v -> Vector a
hotone n' v = fromList $ map fromIntegral list
  where
    n = if i < n' then n' else error ("Too large value: " ++ show i)
    list = replicate i 0 ++ 1 : replicate (n - i - 1) 0
    i = fromIntegral v

convertTests :: MnistData -> [(Int, Vector R)]
convertTests (MnistData src) =
  map (bimap fromIntegral $ (/ 255) . flatten . fromZ) src

trainingSimple ::
     Double
  -> Int
  -> Int
  -> ForwardNN R
  -> MnistData
  -> MnistData
  -> ([Double], Double)
trainingSimple rate batchSize nReplicate nn trainData testData =
  (losses, evaluate layers $ convertTests testData)
  where
    trainBatches = convertTrains batchSize trainData
    batches = concat $ replicate nReplicate trainBatches
    (ForwardNN layers _, losses) = learnAll rate nn batches


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Layers.Logit
Description : Exponential linear unit layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Elu (
    Elu (..)
  ) where

import           Control.DeepSeq              (NFData)
import           Data.Serialize
import           GHC.Generics                 (Generic)

import           GHC.TypeLits
import           Grenade.Core

import qualified Numeric.LinearAlgebra.Static as LAS

-- | An exponential linear unit.
--   A layer which can act between any shape of the same dimension, acting as a
--   diode on every neuron individually.
data Elu = Elu
  deriving (Generic, NFData, Show)

instance UpdateLayer Elu where
  type Gradient Elu = ()
  runUpdate _ _ _ = Elu

instance RandomLayer Elu where
  createRandomWith _ _ = return Elu

instance Serialize Elu where
  put _ = return ()
  get = return Elu

instance ( KnownNat i) => Layer Elu ('D1 i) ('D1 i) where
  type Tape Elu ('D1 i) ('D1 i) = LAS.R i

  runForwards _ (S1D y) = (y, S1D (elu y))
    where
      elu = LAS.dvmap (\a -> if a <= 0 then exp a - 1 else a)
  runBackwards _ y (S1D dEdy) = ((), S1D (elu' y * dEdy))
    where
      elu' = LAS.dvmap (\a -> if a <= 0 then exp a else 1)

instance (KnownNat i, KnownNat j) => Layer Elu ('D2 i j) ('D2 i j) where
  type Tape Elu ('D2 i j) ('D2 i j) = S ('D2 i j)

  runForwards _ (S2D y) = (S2D y, S2D (elu y))
    where
      elu = LAS.dmmap (\a -> if a <= 0 then exp a - 1 else a)
  runBackwards _ (S2D y) (S2D dEdy) = ((), S2D (elu' y * dEdy))
    where
      elu' = LAS.dmmap (\a -> if a <= 0 then exp a else 1)

instance (KnownNat i, KnownNat j, KnownNat k) => Layer Elu ('D3 i j k) ('D3 i j k) where

  type Tape Elu ('D3 i j k) ('D3 i j k) = S ('D3 i j k)

  runForwards _ (S3D y) = (S3D y, S3D (elu y))
    where
      elu = LAS.dmmap (\a -> if a <= 0 then exp a - 1 else a)
  runBackwards _ (S3D y) (S3D dEdy) = ((), S3D (elu' y * dEdy))
    where
      elu' = LAS.dmmap (\a -> if a <= 0 then exp a else 1)

-------------------- GNum instances --------------------

instance GNum Elu where
  _ |* Elu = Elu
  _ |+ Elu = Elu
  gFromRational _ = Elu


{-# LANGUAGE DataKinds #-}

import GHC.TypeLits
import Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra.HMatrix as LA

a = row (vec4 1 2 3 4)
u = vec4 10 20 30 40
v = vec2 5 0 & 0 & 3 & 7


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE OverloadedLabels #-}
{-# OPTIONS_GHC -Wall #-}

-- | Quantile chart patterns.
module Chart.Quantile
  (
    -- * quantile chart patterns
    quantileChart,
    histChart,

    -- * testing
    qua,
    qLines,
  )
where

import Chart
import Optics.Core
import Data.Foldable
import Data.Maybe
import Data.Text (Text)
import Prelude hiding (abs)
import Statistics.Distribution.Normal
import Statistics.Distribution
import Data.Bool

-- * charts

-- | Chart template for quantile data.
quantileChart ::
  [Text] ->
  [LineStyle] ->
  [AxisOptions] ->
  [[Double]] ->
  ChartSvg
quantileChart names ls as xs =
  mempty & #charts .~ chart0 & #hudOptions .~ hudOptions'
  where
    hudOptions' :: HudOptions
    hudOptions' =
      defaultHudOptions
        & ( #legends
              .~
                [(12, defaultLegendOptions
                    & #textStyle % #size .~ 0.1
                    & #vgap .~ 0.05
                    & #innerPad .~ 0.2
                    & #place .~ PlaceRight
                    & #content .~ zip names (fmap (\l -> LineChart l [[Point 0 0, Point 1 1]]) ls)
                )]
          )
        & set #axes ((5,) <$> as)

    chart0 = unnamed $
      zipWith (\s d -> LineChart s [d])
        ls
        (zipWith Point [0 ..] <$> xs)

-- | histogram chart
histChart ::
  Range Double ->
  Int ->
  [Double] ->
  ChartSvg
histChart r g xs =
  barChart defaultBarOptions barData'
  where
    barData' = BarData [freqs] xs'' []
    hcuts = grid OuterPos r g
    h = fill hcuts xs
    counts =
      (\(Rect _ _ _ w) -> w)
        <$> makeRects (IncludeOvers (width r / fromIntegral g)) h
    freqs = (/sum counts) <$> counts
    xs' =
      (\(Rect x x' _ _) -> (x + x') / 2)
        <$> makeRects (IncludeOvers (width r / fromIntegral g)) h
    xs'' = ["unders"] <> take (length xs' - 2) (drop 1 (comma (Just 2) <$> xs')) <> ["overs"]

qua :: Double -> Double -> Double -> Double -> Double
qua u s t p
  | t <= 0 = u*t
  | otherwise = Statistics.Distribution.quantile (normalDistr (u*t) (s*sqrt t)) p

qLines :: Double -> Int -> Colour -> [LineStyle]
qLines s n c = (\x -> defaultLineStyle & #color .~ x & #size .~ s) <$> cqs n c

cqs :: Int -> Colour -> [Colour]
cqs n c = fmap (\x -> mix x c (greyed c)) xs
  where
    xs = (\t -> fmap (\x -> fromIntegral (abs (x-(t `div` 2)-bool 0 1 (x>(t `div`2) && 1==t `mod` 2))) / fromIntegral (t `div` 2)) [0..t]) (n - 1)


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
module Numeric.QpOases where

import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Trans.Except
import qualified Data.Vector.Storable as VS
import Data.Matrix.CSC
import Foreign hiding (void)
import Foreign.C.Types
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel

data OasesError = SizeMismatch | OasesError Int deriving (Show)

data HessianType = HST_ZERO | HST_IDENTITY | HST_POSDEF | HST_POSDEF_NULLSPACE | HST_SEMIDEF | HST_INDEF | HST_UNKNOWN deriving (Enum, Show)

newtype SQProblem = SQProblem (ForeignPtr SQProblem) deriving (Eq, Show)

newtype SQProblemSchur = SQProblemSchur (ForeignPtr SQProblemSchur) deriving (Eq, Show)

setupSQP :: Int -> Int -> HessianType -> IO SQProblem
setupSQP nV nC h = do
    ptr2 <- alloca $ \ptr -> do
        !_ <- sqproblem_setup (fromIntegral nV) (fromIntegral nC) (fromIntegral $ fromEnum h) ptr
        peek ptr
    SQProblem <$> newForeignPtr sqproblem_cleanup ptr2

setupSQPSchur :: Int -> Int -> HessianType -> IO SQProblemSchur
setupSQPSchur nV nC h = do
    ptr2 <- alloca $ \ptr -> do
        !_ <- sqproblem_schur_setup (fromIntegral nV) (fromIntegral nC) (fromIntegral $ fromEnum h) ptr
        peek ptr
    SQProblemSchur <$> newForeignPtr sqproblem_schur_cleanup ptr2

withSQProblem :: SQProblem -> (Ptr SQProblem -> IO a) -> IO a
withSQProblem (SQProblem fPtr) f = withForeignPtr fPtr f

withSQProblemSchur :: SQProblemSchur -> (Ptr SQProblemSchur -> IO a) -> IO a
withSQProblemSchur (SQProblemSchur fPtr) f = withForeignPtr fPtr f

initSQP :: SQProblem -> Matrix Double -> Vector Double -> Matrix Double
        -> Vector Double -> Vector Double -> Maybe Double
        -> ExceptT OasesError IO (Vector Double, Double)
initSQP sqp h g a lbA ubA m'cpuTime =
  mat' h $ \hRow hCol hPtr ->
  vec' g $ \gSize gPtr ->
  mat' a $ \aRow aCol aPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = aRow
    unless (nVar == hRow) $ throwE SizeMismatch
    unless (nVar == hCol) $ throwE SizeMismatch
    unless (nVar == aCol) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr ->
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblem sqp $ \ptr ->
               sqproblem_init ptr hPtr gPtr aPtr nullPtr nullPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

initSparseSQP :: SQProblem -> CSC Double -> Vector Double -> Matrix Double
              -> Vector Double -> Vector Double -> Maybe Double
              -> ExceptT OasesError IO (Vector Double, Double)
initSparseSQP sqp h g a lbA ubA m'cpuTime =
  vec' (cscVals h) $ \_hValsSize hValsPtr ->
  vec' (cscCols h) $ \_hColsSize hColsPtr ->
  vec' (cscRows h) $ \_hRowsSize hRowsPtr ->
  vec' g $ \gSize gPtr ->
  mat' a $ \aRow aCol aPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = aRow
    unless (nVar == aCol) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr ->
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblem sqp $ \ptr ->
               sqproblem_sparse_init ptr hRowsPtr hColsPtr hValsPtr gPtr aPtr
               nullPtr nullPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

initSparseSQPSchur :: SQProblemSchur -> CSC Double -> Vector Double -> Matrix Double
                   -> Vector Double -> Vector Double -> Maybe Double
                   -> ExceptT OasesError IO (Vector Double, Double)
initSparseSQPSchur sqp h g a lbA ubA m'cpuTime =
  vec' (cscVals h) $ \_hValsSize hValsPtr ->
  vec' (cscCols h) $ \_hColsSize hColsPtr ->
  vec' (cscRows h) $ \_hRowsSize hRowsPtr ->
  vec' g $ \gSize gPtr ->
  mat' a $ \aRow aCol aPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = aRow
    unless (nVar == aCol) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr ->
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblemSchur sqp $ \ptr ->
               sqproblem_sparse_schur_init ptr hRowsPtr hColsPtr hValsPtr gPtr aPtr
               nullPtr nullPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

initSparseSparseSQPSchur :: SQProblemSchur -> CSC Double -> Vector Double -> CSC Double
                   -> Maybe (Vector Double) -> Maybe (Vector Double)
                   -> Vector Double -> Vector Double -> Maybe Double
                   -> ExceptT OasesError IO (Vector Double, Double)
initSparseSparseSQPSchur sqp h g a m'lb m'ub lbA ubA m'cpuTime =
  vec' (cscVals h) $ \_hValsSize hValsPtr ->
  vec' (cscCols h) $ \_hColsSize hColsPtr ->
  vec' (cscRows h) $ \_hRowsSize hRowsPtr ->
  vec' g $ \gSize gPtr ->
  vec' (cscVals a) $ \_aValsSize aValsPtr ->
  vec' (cscCols a) $ \_aColsSize aColsPtr ->
  vec' (cscRows a) $ \_aRowsSize aRowsPtr ->
  maybe (\f -> f undefined nullPtr) vec' m'lb $ \_lbSize lbPtr ->
  maybe (\f -> f undefined nullPtr) vec' m'ub $ \_ubSize ubPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = cscNRows a
    unless (nVar == cscNCols a) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr ->
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblemSchur sqp $ \ptr ->
               sqproblem_sparse_sparse_schur_init ptr hRowsPtr hColsPtr hValsPtr gPtr
               aRowsPtr aColsPtr aValsPtr
               lbPtr ubPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

hotstartSQP :: SQProblem -> Matrix Double -> Vector Double -> Matrix Double
            -> Vector Double -> Vector Double -> Maybe Double
            -> ExceptT OasesError IO (Vector Double, Double)
hotstartSQP sqp h g a lbA ubA m'cpuTime =
  mat' h $ \hRow hCol hPtr ->
  vec' g $ \gSize gPtr ->
  mat' a $ \aRow aCol aPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = aRow
    unless (nVar == hRow) $ throwE SizeMismatch
    unless (nVar == hRow) $ throwE SizeMismatch
    unless (nVar == hCol) $ throwE SizeMismatch
    unless (nVar == aCol) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr ->
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblem sqp $ \ptr ->
               sqproblem_hotstart ptr hPtr gPtr aPtr nullPtr nullPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

hotstartSparseSQP :: SQProblem -> CSC Double -> Vector Double -> Matrix Double
                  -> Vector Double -> Vector Double -> Maybe Double
                  -> ExceptT OasesError IO (Vector Double, Double)
hotstartSparseSQP sqp h g a lbA ubA m'cpuTime =
  vec' (cscVals h) $ \_hValsSize hValsPtr ->
  vec' (cscCols h) $ \_hColsSize hColsPtr ->
  vec' (cscRows h) $ \_hRowsSize hRowsPtr ->
  vec' g $ \gSize gPtr ->
  mat' a $ \aRow aCol aPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = aRow
    unless (nVar == aCol) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr -> do
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblem sqp $ \ptr ->
               sqproblem_sparse_hotstart ptr hRowsPtr hColsPtr hValsPtr gPtr
               aPtr nullPtr nullPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

hotstartSparseSQPSchur :: SQProblemSchur -> CSC Double -> Vector Double -> Matrix Double
                  -> Vector Double -> Vector Double -> Maybe Double
                  -> ExceptT OasesError IO (Vector Double, Double)
hotstartSparseSQPSchur sqp h g a lbA ubA m'cpuTime =
  vec' (cscVals h) $ \_hValsSize hValsPtr ->
  vec' (cscCols h) $ \_hColsSize hColsPtr ->
  vec' (cscRows h) $ \_hRowsSize hRowsPtr ->
  vec' g $ \gSize gPtr ->
  mat' a $ \aRow aCol aPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = aRow
    unless (nVar == aCol) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr ->
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblemSchur sqp $ \ptr ->
               sqproblem_sparse_schur_hotstart ptr hRowsPtr hColsPtr hValsPtr gPtr
               aPtr nullPtr nullPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

hotstartSparseSparseSQPSchur :: SQProblemSchur -> CSC Double -> Vector Double -> CSC Double
                  -> Maybe (Vector Double) -> Maybe (Vector Double) -> Vector Double -> Vector Double -> Maybe Double
                  -> ExceptT OasesError IO (Vector Double, Double)
hotstartSparseSparseSQPSchur sqp h g a m'lb m'ub lbA ubA m'cpuTime =
  vec' (cscVals h) $ \_hValsSize hValsPtr ->
  vec' (cscCols h) $ \_hColsSize hColsPtr ->
  vec' (cscRows h) $ \_hRowsSize hRowsPtr ->
  vec' g $ \gSize gPtr ->
  vec' (cscVals a) $ \_aValsSize aValsPtr ->
  vec' (cscCols a) $ \_aColsSize aColsPtr ->
  vec' (cscRows a) $ \_aRowsSize aRowsPtr ->
  maybe (\f -> f undefined nullPtr) vec' m'lb $ \_lbSize lbPtr ->
  maybe (\f -> f undefined nullPtr) vec' m'ub $ \_ubSize ubPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = cscNRows a
    unless (nVar == cscNCols a) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr ->
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblemSchur sqp $ \ptr ->
               sqproblem_sparse_sparse_schur_hotstart ptr hRowsPtr hColsPtr hValsPtr gPtr
               aRowsPtr aColsPtr aValsPtr lbPtr ubPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

data Options

foreign import ccall "sqproblem_init"
    sqproblem_init
        :: Ptr SQProblem -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_sparse_init"
    sqproblem_sparse_init
        :: Ptr SQProblem -> Ptr CInt -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_sparse_schur_init"
    sqproblem_sparse_schur_init
        :: Ptr SQProblemSchur -> Ptr CInt -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_sparse_sparse_schur_init"
    sqproblem_sparse_sparse_schur_init
        :: Ptr SQProblemSchur -> Ptr CInt -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr CInt -> Ptr CInt -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_hotstart"
    sqproblem_hotstart
        :: Ptr SQProblem -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_sparse_hotstart"
    sqproblem_sparse_hotstart
        :: Ptr SQProblem -> Ptr CInt -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_sparse_schur_hotstart"
    sqproblem_sparse_schur_hotstart
        :: Ptr SQProblemSchur -> Ptr CInt -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_sparse_sparse_schur_hotstart"
    sqproblem_sparse_sparse_schur_hotstart
        :: Ptr SQProblemSchur -> Ptr CInt -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr CInt -> Ptr CInt -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_setup"
    sqproblem_setup
        :: CInt -> CInt -> CInt -> Ptr (Ptr SQProblem)-> IO CInt

foreign import ccall "sqproblem_schur_setup"
    sqproblem_schur_setup
        :: CInt -> CInt -> CInt -> Ptr (Ptr SQProblemSchur)-> IO CInt

foreign import ccall "&sqproblem_cleanup"
    sqproblem_cleanup
        :: FunPtr (Ptr SQProblem -> IO ())

foreign import ccall "&sqproblem_schur_cleanup"
    sqproblem_schur_cleanup
        :: FunPtr (Ptr SQProblemSchur -> IO ())


module HashedExpression.Interp
  ( Evaluable (..),
    Approximable (..),
  )
where

import Data.Array
import Data.Complex
import qualified Data.IntMap.Strict as IM
import Data.List (intercalate)
import Data.Map (Map, fromList)
import qualified Data.Map as Map
import Debug.Trace (traceId, traceShowId)
import GHC.TypeLits (KnownNat)
import HashedExpression.Internal.Expression
  ( C,
    ET (..),
    Expression (..),
    ExpressionMap,
    Node (..),
    NodeID,
    NodeID,
    R,
    Scalar,
  )
import HashedExpression.Internal.Node
import HashedExpression.Internal.Utils
import HashedExpression.Prettify (prettify, showExp)
import HashedExpression.Value
import Text.Printf

-- | This operation emulates the mathematical operation
-- | Turn expression to the right type
expZeroR :: ExpressionMap -> NodeID -> Expression Scalar R
expZeroR = flip Expression

expZeroC :: ExpressionMap -> NodeID -> Expression Scalar C
expZeroC = flip Expression

-- | Choose branch base on condition value
chooseBranch :: [Double] -> Double -> [a] -> a
chooseBranch marks val branches
  | val < head marks = head branches
  | otherwise =
    snd . last . filter ((val >=) . fst) $ zip marks (tail branches)

-- | Approximable class
class
  Show a =>
  Approximable a where
  (~=) :: a -> a -> Bool
  prettifyShow :: a -> String

infix 4 ~=

-- |
relativeError :: Double -> Double -> Double
relativeError a b = abs (a - b) / max (abs a) (abs b)

instance Approximable Double where
  (~=) :: Double -> Double -> Bool
  a ~= b
    | abs (a - b) < 1.0e-5 = True
    | a == b = True
    | otherwise = relativeError a b < 0.01
  prettifyShow a
    | abs a < 1e-10 = "0"
    | otherwise = printf "%.2f" a

instance Approximable (Complex Double) where
  (~=) :: Complex Double -> Complex Double -> Bool
  a ~= b = (realPart a ~= realPart b) && (imagPart a ~= imagPart b)
  prettifyShow a =
    prettifyShow (realPart a) ++ " + " ++ prettifyShow (imagPart a) ++ "i"

instance Approximable (Array Int Double) where
  (~=) :: Array Int Double -> Array Int Double -> Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]"

instance Approximable (Array Int (Complex Double)) where
  (~=) :: Array Int (Complex Double) -> Array Int (Complex Double) -> Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]"

instance Approximable (Array (Int, Int) Double) where
  (~=) :: Array (Int, Int) Double -> Array (Int, Int) Double -> Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]"

instance Approximable (Array (Int, Int) (Complex Double)) where
  (~=) ::
    Array (Int, Int) (Complex Double) ->
    Array (Int, Int) (Complex Double) ->
    Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]"

instance Approximable (Array (Int, Int, Int) Double) where
  (~=) :: Array (Int, Int, Int) Double -> Array (Int, Int, Int) Double -> Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]"

instance Approximable (Array (Int, Int, Int) (Complex Double)) where
  (~=) ::
    Array (Int, Int, Int) (Complex Double) ->
    Array (Int, Int, Int) (Complex Double) ->
    Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]"

-- | These should be commented properly.
class Evaluable d rc output | d rc -> output where
  eval :: ValMaps -> Expression d rc -> output

-- |
instance Evaluable Scalar R Double where
  eval :: ValMaps -> Expression Scalar R -> Double
  eval valMap e@(Expression n mp)
    | [] <- retrieveShape n mp =
      case retrieveNode n mp of
        Var name ->
          case Map.lookup name valMap of
            Just (VScalar val) -> val
            _ -> error "no value associated with the variable"
        Const val -> val
        Sum R args -> sum . map (eval valMap . expZeroR mp) $ args
        Mul R args -> product . map (eval valMap . expZeroR mp) $ args
        Neg R arg -> - (eval valMap $ expZeroR mp arg)
        Scale R arg1 arg2 ->
          eval valMap (expZeroR mp arg1)
            * eval valMap (expZeroR mp arg2)
        Power x arg -> eval valMap (expZeroR mp arg) ^ x
        Div arg1 arg2 ->
          eval valMap (expZeroR mp arg1)
            / eval valMap (expZeroR mp arg2)
        Sqrt arg -> sqrt (eval valMap (expZeroR mp arg))
        Sin arg -> sin (eval valMap (expZeroR mp arg))
        Cos arg -> cos (eval valMap (expZeroR mp arg))
        Tan arg -> tan (eval valMap (expZeroR mp arg))
        Exp arg -> exp (eval valMap (expZeroR mp arg))
        Log arg -> log (eval valMap (expZeroR mp arg))
        Sinh arg -> sinh (eval valMap (expZeroR mp arg))
        Cosh arg -> cosh (eval valMap (expZeroR mp arg))
        Tanh arg -> tanh (eval valMap (expZeroR mp arg))
        Asin arg -> asin (eval valMap (expZeroR mp arg))
        Acos arg -> acos (eval valMap (expZeroR mp arg))
        Atan arg -> atan (eval valMap (expZeroR mp arg))
        Asinh arg -> asinh (eval valMap (expZeroR mp arg))
        Acosh arg -> acosh (eval valMap (expZeroR mp arg))
        Atanh arg -> atanh (eval valMap (expZeroR mp arg))
        RealPart arg -> realPart (eval valMap (expZeroC mp arg))
        ImagPart arg -> imagPart (eval valMap (expZeroC mp arg))
        InnerProd R arg1 arg2 ->
          case retrieveShape arg1 mp of
            [] ->
              eval valMap (expZeroR mp arg1)
                * eval valMap (expZeroR mp arg2)
            [size] ->
              let res1 = evaluate1DReal valMap $ (mp, arg1)
                  res2 = evaluate1DReal valMap $ (mp, arg2)
               in sum
                    [ x * y
                      | i <- [0 .. size - 1],
                        let x = res1 ! i,
                        let y = res2 ! i
                    ]
            [size1, size2] ->
              let res1 = evaluate2DReal valMap $ (mp, arg1)
                  res2 = evaluate2DReal valMap $ (mp, arg2)
               in sum
                    [ x * y
                      | i <- [0 .. size1 - 1],
                        j <- [0 .. size2 - 1],
                        let x = res1 ! (i, j),
                        let y = res2 ! (i, j)
                    ]
            [size1, size2, size3] ->
              let res1 = evaluate3DReal valMap $ (mp, arg1)
                  res2 = evaluate3DReal valMap $ (mp, arg2)
               in sum
                    [ x * y
                      | i <- [0 .. size1 - 1],
                        j <- [0 .. size2 - 1],
                        k <- [0 .. size3 - 1],
                        let x = res1 ! (i, j, k),
                        let y = res2 ! (i, j, k)
                    ]
            _ -> error "4D shape?"
        Piecewise marks conditionArg branchArgs ->
          let cdt = eval valMap $ expZeroR mp conditionArg
              branches = map (eval valMap . expZeroR mp) branchArgs
           in chooseBranch marks cdt branches
        _ ->
          error
            ("expression structure Scalar R is wrong " ++ prettify e)
    | otherwise = error "one r but shape is not [] ??"

instance Evaluable Scalar C (Complex Double) where
  eval :: ValMaps -> Expression Scalar C -> Complex Double
  eval valMap e@(Expression n mp)
    | [] <- retrieveShape n mp =
      case retrieveNode n mp of
        Sum C args -> sum . map (eval valMap . expZeroC mp) $ args
        Mul C args -> product . map (eval valMap . expZeroC mp) $ args
        Power x arg -> eval valMap (expZeroC mp arg) ^ x
        Neg C arg -> - (eval valMap $ expZeroC mp arg)
        Scale C arg1 arg2 ->
          case retrieveElementType arg1 mp of
            R ->
              fromR (eval valMap (expZeroR mp arg1))
                * eval valMap (expZeroC mp arg2)
            C ->
              eval valMap (expZeroC mp arg1)
                * eval valMap (expZeroC mp arg2)
        RealImag arg1 arg2 ->
          eval valMap (expZeroR mp arg1)
            :+ eval valMap (expZeroR mp arg2)
        InnerProd C arg1 arg2 ->
          case retrieveShape arg1 mp of
            [] ->
              eval valMap (expZeroC mp arg1)
                * conjugate (eval valMap (expZeroC mp arg2))
            [size] ->
              let res1 = evaluate1DComplex valMap $ (mp, arg1)
                  res2 = evaluate1DComplex valMap $ (mp, arg2)
               in sum
                    [ x * conjugate y
                      | i <- [0 .. size - 1],
                        let x = res1 ! i,
                        let y = res2 ! i
                    ]
            [size1, size2] ->
              let res1 = evaluate2DComplex valMap $ (mp, arg1)
                  res2 = evaluate2DComplex valMap $ (mp, arg2)
               in sum
                    [ x * conjugate y
                      | i <- [0 .. size1 - 1],
                        j <- [0 .. size2 - 1],
                        let x = res1 ! (i, j),
                        let y = res2 ! (i, j)
                    ]
            [size1, size2, size3] ->
              let res1 = evaluate3DComplex valMap $ (mp, arg1)
                  res2 = evaluate3DComplex valMap $ (mp, arg2)
               in sum
                    [ x * conjugate y
                      | i <- [0 .. size1 - 1],
                        j <- [0 .. size2 - 1],
                        k <- [0 .. size3 - 1],
                        let x = res1 ! (i, j, k),
                        let y = res2 ! (i, j, k)
                    ]
            _ -> error "4D shape?"
        Piecewise marks conditionArg branchArgs ->
          let cdt = eval valMap $ expZeroR mp conditionArg
              branches = map (eval valMap . expZeroC mp) branchArgs
           in chooseBranch marks cdt branches
        _ ->
          error
            ("expression structure Scalar C is wrong " ++ prettify e)
    | otherwise = error "One C but shape is not [] ??"

-- |
zipWithA :: Ix x => (a -> b -> c) -> Array x a -> Array x b -> Array x c
zipWithA f xs ys = listArray (bounds xs) $ zipWith f (elems xs) (elems ys)

foldrElementwise :: Ix ix => (a -> a -> a) -> [Array ix a] -> Array ix a
foldrElementwise f [x] = x
foldrElementwise f (x : xs) = zipWithA f x (foldrElementwise f xs)

-- |
evaluate1DReal :: ValMaps -> (ExpressionMap, NodeID) -> Array Int Double
evaluate1DReal valMap (mp, n)
  | [size] <- retrieveShape n mp =
    case retrieveNode n mp of
      Var name ->
        case Map.lookup name valMap of
          Just (V1D val) -> val
          _ -> error "no value associated with the variable"
      Const val -> listArray (0, size - 1) $ replicate size val
      Sum R args ->
        foldrElementwise (+) . map (evaluate1DReal valMap . (mp,)) $
          args
      Mul R args ->
        foldrElementwise (*) . map (evaluate1DReal valMap . (mp,)) $
          args
      Power x arg -> fmap (^ x) (evaluate1DReal valMap $ (mp, arg))
      Neg R arg -> fmap negate . evaluate1DReal valMap $ (mp, arg)
      Scale R arg1 arg2 ->
        let scalar = eval valMap $ expZeroR mp arg1
         in fmap (scalar *) . evaluate1DReal valMap $ (mp, arg2)
      Div arg1 arg2 ->
        zipWithA
          (/)
          (evaluate1DReal valMap $ (mp, arg2))
          (evaluate1DReal valMap $ (mp, arg2))
      Sqrt arg -> fmap sqrt . evaluate1DReal valMap $ (mp, arg)
      Sin arg -> fmap sin . evaluate1DReal valMap $ (mp, arg)
      Cos arg -> fmap cos . evaluate1DReal valMap $ (mp, arg)
      Tan arg -> fmap tan . evaluate1DReal valMap $ (mp, arg)
      Exp arg -> fmap exp . evaluate1DReal valMap $ (mp, arg)
      Log arg -> fmap log . evaluate1DReal valMap $ (mp, arg)
      Sinh arg -> fmap sinh . evaluate1DReal valMap $ (mp, arg)
      Cosh arg -> fmap cosh . evaluate1DReal valMap $ (mp, arg)
      Tanh arg -> fmap tanh . evaluate1DReal valMap $ (mp, arg)
      Asin arg -> fmap asin . evaluate1DReal valMap $ (mp, arg)
      Acos arg -> fmap acos . evaluate1DReal valMap $ (mp, arg)
      Atan arg -> fmap atan . evaluate1DReal valMap $ (mp, arg)
      Asinh arg -> fmap asinh . evaluate1DReal valMap $ (mp, arg)
      Acosh arg -> fmap acosh . evaluate1DReal valMap $ (mp, arg)
      Atanh arg -> fmap atanh . evaluate1DReal valMap $ (mp, arg)
      RealPart arg -> fmap realPart . evaluate1DComplex valMap $ (mp, arg)
      ImagPart arg -> fmap imagPart . evaluate1DComplex valMap $ (mp, arg)
      -- Rotate rA arg ->
      Piecewise marks conditionArg branchArgs ->
        let cdt = evaluate1DReal valMap $ (mp, conditionArg)
            branches = map (evaluate1DReal valMap . (mp,)) branchArgs
         in listArray
              (0, size - 1)
              [ chosen ! i
                | i <- [0 .. size - 1],
                  let chosen = chooseBranch marks (cdt ! i) branches
              ]
      Rotate [amount] arg ->
        rotate1D size amount (evaluate1DReal valMap $ (mp, arg))
      TwiceReFT arg ->
        let innerRes = evaluate1DReal valMap $ (mp, arg)
            scaleFactor = fromIntegral size / 2
         in listArray
              (0, size - 1)
              [ scaleFactor
                  * (innerRes ! i + innerRes ! ((size - i) `mod` size))
                | i <- [0 .. size - 1]
              ]
      TwiceImFT arg ->
        let innerRes = evaluate1DReal valMap $ (mp, arg)
            scaleFactor = fromIntegral size / 2
         in listArray
              (0, size - 1)
              [ scaleFactor
                  * (innerRes ! i - innerRes ! ((size - i) `mod` size))
                | i <- [0 .. size - 1]
              ]
      ReFT arg ->
        case retrieveElementType arg mp of
          R ->
            let inner =
                  fmap (:+ 0) . evaluate1DReal valMap $ (mp, arg)
                ftResult = fourierTransform1D size inner
             in fmap realPart ftResult
          C ->
            let inner = evaluate1DComplex valMap $ (mp, arg)
                ftResult = fourierTransform1D size inner
             in fmap realPart ftResult
      ImFT arg ->
        case retrieveElementType arg mp of
          R ->
            let inner =
                  fmap (:+ 0) . evaluate1DReal valMap $ (mp, arg)
                ftResult = fourierTransform1D size inner
             in fmap imagPart ftResult
          C ->
            let inner = evaluate1DComplex valMap $ (mp, arg)
                ftResult = fourierTransform1D size inner
             in fmap imagPart ftResult
      _ -> error "expression structure One R is wrong"
  | otherwise = error "one r but shape is not [size] ??"

instance (KnownNat n) => Evaluable n R (Array Int Double) where
  eval :: ValMaps -> Expression n R -> Array Int Double
  eval valMap (Expression n mp) = evaluate1DReal valMap (mp, n)

-- |
evaluate1DComplex ::
  ValMaps -> (ExpressionMap, NodeID) -> Array Int (Complex Double)
evaluate1DComplex valMap (mp, n)
  | [size] <- retrieveShape n mp =
    case retrieveNode n mp of
      Sum C args ->
        foldrElementwise (+) . map (evaluate1DComplex valMap . (mp,)) $
          args
      Mul C args ->
        foldrElementwise (*) . map (evaluate1DComplex valMap . (mp,)) $
          args
      Power x arg -> fmap (^ x) (evaluate1DComplex valMap $ (mp, arg))
      Neg C arg -> fmap negate . evaluate1DComplex valMap $ (mp, arg)
      Scale C arg1 arg2 ->
        case retrieveElementType arg1 mp of
          R ->
            let scalar = fromR . eval valMap $ expZeroR mp arg1
             in fmap (scalar *) . evaluate1DComplex valMap $
                  (mp, arg2)
          C ->
            let scalar = eval valMap $ expZeroC mp arg1
             in fmap (scalar *) . evaluate1DComplex valMap $
                  (mp, arg2)
      RealImag arg1 arg2 ->
        zipWithA
          (:+)
          (evaluate1DReal valMap $ (mp, arg1))
          (evaluate1DReal valMap $ (mp, arg2))
      Piecewise marks conditionArg branchArgs ->
        let cdt = evaluate1DReal valMap $ (mp, conditionArg)
            branches =
              map (evaluate1DComplex valMap . (mp,)) branchArgs
         in listArray
              (0, size - 1)
              [ chosen ! i
                | i <- [0 .. size - 1],
                  let chosen = chooseBranch marks (cdt ! i) branches
              ]
      Rotate [amount] arg ->
        rotate1D size amount (evaluate1DComplex valMap $ (mp, arg))
      _ -> error "expression structure One C is wrong"
  | otherwise = error "one C but shape is not [size] ??"

--                         in chooseBranch marks cdt branches
instance (KnownNat n) => Evaluable n C (Array Int (Complex Double)) where
  eval :: ValMaps -> Expression n C -> Array Int (Complex Double)
  eval valMap (Expression n mp) = evaluate1DComplex valMap (mp, n)

-- |
evaluate2DReal :: ValMaps -> (ExpressionMap, NodeID) -> Array (Int, Int) Double
evaluate2DReal valMap (mp, n)
  | [size1, size2] <- retrieveShape n mp =
    case retrieveNode n mp of
      Var name ->
        case Map.lookup name valMap of
          Just (V2D val) -> val
          _ -> error $ "no value associated with the variable" ++ name
      Const val ->
        listArray ((0, 0), (size1 - 1, size2 - 1)) $
          replicate (size1 * size2) val
      Sum R args ->
        foldrElementwise (+) . map (evaluate2DReal valMap . (mp,)) $
          args
      Mul R args ->
        foldrElementwise (*) . map (evaluate2DReal valMap . (mp,)) $
          args
      Power x arg -> fmap (^ x) (evaluate2DReal valMap $ (mp, arg))
      Neg R arg -> fmap negate . evaluate2DReal valMap $ (mp, arg)
      Scale R arg1 arg2 ->
        let scalar = eval valMap $ expZeroR mp arg1
         in fmap (scalar *) . evaluate2DReal valMap $ (mp, arg2)
      Div arg1 arg2 ->
        zipWithA
          (/)
          (evaluate2DReal valMap $ (mp, arg2))
          (evaluate2DReal valMap $ (mp, arg2))
      Sqrt arg -> fmap sqrt . evaluate2DReal valMap $ (mp, arg)
      Sin arg -> fmap sin . evaluate2DReal valMap $ (mp, arg)
      Cos arg -> fmap cos . evaluate2DReal valMap $ (mp, arg)
      Tan arg -> fmap tan . evaluate2DReal valMap $ (mp, arg)
      Exp arg -> fmap exp . evaluate2DReal valMap $ (mp, arg)
      Log arg -> fmap log . evaluate2DReal valMap $ (mp, arg)
      Sinh arg -> fmap sinh . evaluate2DReal valMap $ (mp, arg)
      Cosh arg -> fmap cosh . evaluate2DReal valMap $ (mp, arg)
      Tanh arg -> fmap tanh . evaluate2DReal valMap $ (mp, arg)
      Asin arg -> fmap asin . evaluate2DReal valMap $ (mp, arg)
      Acos arg -> fmap acos . evaluate2DReal valMap $ (mp, arg)
      Atan arg -> fmap atan . evaluate2DReal valMap $ (mp, arg)
      Asinh arg -> fmap asinh . evaluate2DReal valMap $ (mp, arg)
      Acosh arg -> fmap acosh . evaluate2DReal valMap $ (mp, arg)
      Atanh arg -> fmap atanh . evaluate2DReal valMap $ (mp, arg)
      RealPart arg -> fmap realPart . evaluate2DComplex valMap $ (mp, arg)
      ImagPart arg -> fmap imagPart . evaluate2DComplex valMap $ (mp, arg)
      Piecewise marks conditionArg branchArgs ->
        let cdt = evaluate2DReal valMap $ (mp, conditionArg)
            branches = map (evaluate2DReal valMap . (mp,)) branchArgs
         in listArray
              ((0, 0), (size1 - 1, size2 - 1))
              [ chosen ! (i, j)
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1],
                  let chosen =
                        chooseBranch marks (cdt ! (i, j)) branches
              ]
      Rotate [amount1, amount2] arg ->
        rotate2D
          (size1, size2)
          (amount1, amount2)
          (evaluate2DReal valMap $ (mp, arg))
      TwiceReFT arg ->
        let innerRes = evaluate2DReal valMap $ (mp, arg)
            scaleFactor = fromIntegral size1 * fromIntegral size2 / 2
         in listArray
              ((0, 0), (size1 - 1, size2 - 1))
              [ scaleFactor
                  * ( innerRes ! (i, j)
                        + innerRes
                        ! ((size1 - i) `mod` size1, (size2 - j) `mod` size2)
                    )
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1]
              ]
      TwiceImFT arg ->
        let innerRes = evaluate2DReal valMap $ (mp, arg)
            scaleFactor = fromIntegral size1 * fromIntegral size2 / 2
         in listArray
              ((0, 0), (size1 - 1, size2 - 1))
              [ scaleFactor
                  * ( innerRes ! (i, j)
                        - innerRes
                        ! ((size1 - i) `mod` size1, (size2 - j) `mod` size2)
                    )
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1]
              ]
      ReFT arg ->
        case retrieveElementType arg mp of
          R ->
            let inner =
                  fmap (:+ 0) . evaluate2DReal valMap $ (mp, arg)
                ftResult = fourierTransform2D (size1, size2) inner
             in fmap realPart ftResult
          C ->
            let inner = evaluate2DComplex valMap $ (mp, arg)
                ftResult = fourierTransform2D (size1, size2) inner
             in fmap realPart ftResult
      ImFT arg ->
        case retrieveElementType arg mp of
          R ->
            let inner =
                  fmap (:+ 0) . evaluate2DReal valMap $ (mp, arg)
                ftResult = fourierTransform2D (size1, size2) inner
             in fmap imagPart ftResult
          C ->
            let inner = evaluate2DComplex valMap $ (mp, arg)
                ftResult = fourierTransform2D (size1, size2) inner
             in fmap imagPart ftResult
      _ -> error "expression structure Two R is wrong"
  | otherwise = error "Two r but shape is not [size1, size2] ??"

instance
  (KnownNat m, KnownNat n) =>
  Evaluable '(m, n) R (Array (Int, Int) Double)
  where
  eval :: ValMaps -> Expression '(m, n) R -> Array (Int, Int) Double
  eval valMap (Expression n mp) = evaluate2DReal valMap (mp, n)

-- |
evaluate2DComplex ::
  ValMaps -> (ExpressionMap, NodeID) -> Array (Int, Int) (Complex Double)
evaluate2DComplex valMap (mp, n)
  | [size1, size2] <- retrieveShape n mp =
    case retrieveNode n mp of
      Sum C args ->
        foldrElementwise (+) . map (evaluate2DComplex valMap . (mp,)) $
          args
      Mul C args ->
        foldrElementwise (*) . map (evaluate2DComplex valMap . (mp,)) $
          args
      Power x arg -> fmap (^ x) (evaluate2DComplex valMap $ (mp, arg))
      Neg C arg -> fmap negate . evaluate2DComplex valMap $ (mp, arg)
      Scale C arg1 arg2 ->
        case retrieveElementType arg1 mp of
          R ->
            let scalar = fromR . eval valMap $ expZeroR mp arg1
             in fmap (scalar *) . evaluate2DComplex valMap $
                  (mp, arg2)
          C ->
            let scalar = eval valMap $ expZeroC mp arg1
             in fmap (scalar *) . evaluate2DComplex valMap $
                  (mp, arg2)
      RealImag arg1 arg2 ->
        zipWithA
          (:+)
          (evaluate2DReal valMap $ (mp, arg1))
          (evaluate2DReal valMap $ (mp, arg2))
      Piecewise marks conditionArg branchArgs ->
        let cdt = evaluate2DReal valMap $ (mp, conditionArg)
            branches =
              map (evaluate2DComplex valMap . (mp,)) branchArgs
         in listArray
              ((0, 0), (size1 - 1, size2 - 1))
              [ chosen ! (i, j)
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1],
                  let chosen =
                        chooseBranch marks (cdt ! (i, j)) branches
              ]
      Rotate [amount1, amount2] arg ->
        rotate2D
          (size1, size2)
          (amount1, amount2)
          (evaluate2DComplex valMap $ (mp, arg))
      _ -> error "expression structure Two C is wrong"
  | otherwise = error "Two C but shape is not [size1, size2] ??"

instance
  (KnownNat m, KnownNat n) =>
  Evaluable '(m, n) C (Array (Int, Int) (Complex Double))
  where
  eval ::
    ValMaps -> Expression '(m, n) C -> Array (Int, Int) (Complex Double)
  eval valMap (Expression n mp) = evaluate2DComplex valMap (mp, n)

evaluate3DReal ::
  ValMaps -> (ExpressionMap, NodeID) -> Array (Int, Int, Int) Double
evaluate3DReal valMap (mp, n)
  | [size1, size2, size3] <- retrieveShape n mp =
    case retrieveNode n mp of
      Var name ->
        case Map.lookup name valMap of
          Just (V3D val) -> val
          _ -> error "no value associated with the variable"
      Const val ->
        listArray ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1)) $
          replicate (size1 * size2 * size3) val
      Sum R args ->
        foldrElementwise (+) . map (evaluate3DReal valMap . (mp,)) $
          args
      Mul R args ->
        foldrElementwise (*) . map (evaluate3DReal valMap . (mp,)) $
          args
      Power x arg -> fmap (^ x) (evaluate3DReal valMap $ (mp, arg))
      Neg R arg -> fmap negate . evaluate3DReal valMap $ (mp, arg)
      Scale R arg1 arg2 ->
        let scalar = eval valMap $ expZeroR mp arg1
         in fmap (scalar *) . evaluate3DReal valMap $ (mp, arg2)
      Div arg1 arg2 ->
        zipWithA
          (/)
          (evaluate3DReal valMap $ (mp, arg2))
          (evaluate3DReal valMap $ (mp, arg2))
      Sqrt arg -> fmap sqrt . evaluate3DReal valMap $ (mp, arg)
      Sin arg -> fmap sin . evaluate3DReal valMap $ (mp, arg)
      Cos arg -> fmap cos . evaluate3DReal valMap $ (mp, arg)
      Tan arg -> fmap tan . evaluate3DReal valMap $ (mp, arg)
      Exp arg -> fmap exp . evaluate3DReal valMap $ (mp, arg)
      Log arg -> fmap log . evaluate3DReal valMap $ (mp, arg)
      Sinh arg -> fmap sinh . evaluate3DReal valMap $ (mp, arg)
      Cosh arg -> fmap cosh . evaluate3DReal valMap $ (mp, arg)
      Tanh arg -> fmap tanh . evaluate3DReal valMap $ (mp, arg)
      Asin arg -> fmap asin . evaluate3DReal valMap $ (mp, arg)
      Acos arg -> fmap acos . evaluate3DReal valMap $ (mp, arg)
      Atan arg -> fmap atan . evaluate3DReal valMap $ (mp, arg)
      Asinh arg -> fmap asinh . evaluate3DReal valMap $ (mp, arg)
      Acosh arg -> fmap acosh . evaluate3DReal valMap $ (mp, arg)
      Atanh arg -> fmap atanh . evaluate3DReal valMap $ (mp, arg)
      RealPart arg -> fmap realPart . evaluate3DComplex valMap $ (mp, arg)
      ImagPart arg -> fmap imagPart . evaluate3DComplex valMap $ (mp, arg)
      Piecewise marks conditionArg branchArgs ->
        let cdt = evaluate3DReal valMap $ (mp, conditionArg)
            branches = map (evaluate3DReal valMap . (mp,)) branchArgs
         in listArray
              ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
              [ chosen ! (i, j, k)
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1],
                  k <- [0 .. size3 - 1],
                  let chosen =
                        chooseBranch marks (cdt ! (i, j, k)) branches
              ]
      Rotate [amount1, amount2, amount3] arg ->
        rotate3D
          (size1, size2, size3)
          (amount1, amount2, amount3)
          (evaluate3DReal valMap $ (mp, arg))
      TwiceReFT arg ->
        let innerRes = evaluate3DReal valMap $ (mp, arg)
            scaleFactor =
              fromIntegral size1 * fromIntegral size2
                * fromIntegral size3
                / 2
         in listArray
              ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
              [ scaleFactor
                  * ( innerRes ! (i, j, k)
                        + innerRes
                        ! ( (size1 - i) `mod` size1,
                            (size2 - j) `mod` size2,
                            (size3 - k) `mod` size3
                          )
                    )
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1],
                  k <- [0 .. size3 - 1]
              ]
      TwiceImFT arg ->
        let innerRes = evaluate3DReal valMap $ (mp, arg)
            scaleFactor =
              fromIntegral size1 * fromIntegral size2
                * fromIntegral size3
                / 2
         in listArray
              ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
              [ scaleFactor
                  * ( innerRes ! (i, j, k)
                        - innerRes
                        ! ( (size1 - i) `mod` size1,
                            (size2 - j) `mod` size2,
                            (size3 - k) `mod` size3
                          )
                    )
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1],
                  k <- [0 .. size3 - 1]
              ]
      ReFT arg ->
        case retrieveElementType arg mp of
          R ->
            let inner =
                  fmap (:+ 0) . evaluate3DReal valMap $ (mp, arg)
                ftResult =
                  fourierTransform3D (size1, size2, size3) inner
             in fmap realPart ftResult
          C ->
            let inner = evaluate3DComplex valMap $ (mp, arg)
                ftResult =
                  fourierTransform3D (size1, size2, size3) inner
             in fmap realPart ftResult
      ImFT arg ->
        case retrieveElementType arg mp of
          R ->
            let inner =
                  fmap (:+ 0) . evaluate3DReal valMap $ (mp, arg)
                ftResult =
                  fourierTransform3D (size1, size2, size3) inner
             in fmap imagPart ftResult
          C ->
            let inner = evaluate3DComplex valMap $ (mp, arg)
                ftResult =
                  fourierTransform3D (size1, size2, size3) inner
             in fmap imagPart ftResult
      _ -> error "expression structure Three R is wrong"
  | otherwise = error "Three r but shape is not [size1, size2, size3] ??"

instance
  (KnownNat m, KnownNat n, KnownNat p) =>
  Evaluable '(m, n, p) R (Array (Int, Int, Int) Double)
  where
  eval :: ValMaps -> Expression '(m, n, p) R -> Array (Int, Int, Int) Double
  eval valMap (Expression n mp) = evaluate3DReal valMap (mp, n)

evaluate3DComplex ::
  ValMaps -> (ExpressionMap, NodeID) -> Array (Int, Int, Int) (Complex Double)
evaluate3DComplex valMap (mp, n)
  | [size1, size2, size3] <- retrieveShape n mp =
    case retrieveNode n mp of
      Sum C args ->
        foldrElementwise (+) . map (evaluate3DComplex valMap . (mp,)) $
          args
      Mul C args ->
        foldrElementwise (*) . map (evaluate3DComplex valMap . (mp,)) $
          args
      Power x arg -> fmap (^ x) (evaluate3DComplex valMap $ (mp, arg))
      Neg C arg -> fmap negate . evaluate3DComplex valMap $ (mp, arg)
      Scale C arg1 arg2 ->
        case retrieveElementType arg1 mp of
          R ->
            let scalar = fromR . eval valMap $ expZeroR mp arg1
             in fmap (scalar *) . evaluate3DComplex valMap $
                  (mp, arg2)
          C ->
            let scalar = eval valMap $ expZeroC mp arg1
             in fmap (scalar *) . evaluate3DComplex valMap $
                  (mp, arg2)
      RealImag arg1 arg2 ->
        zipWithA
          (:+)
          (evaluate3DReal valMap $ (mp, arg1))
          (evaluate3DReal valMap $ (mp, arg2))
      Piecewise marks conditionArg branchArgs ->
        let cdt = evaluate3DReal valMap $ (mp, conditionArg)
            branches =
              map (evaluate3DComplex valMap . (mp,)) branchArgs
         in listArray
              ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
              [ chosen ! (i, j, k)
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1],
                  k <- [0 .. size3 - 1],
                  let chosen =
                        chooseBranch marks (cdt ! (i, j, k)) branches
              ]
      Rotate [amount1, amount2, amount3] arg ->
        rotate3D
          (size1, size2, size3)
          (amount1, amount2, amount3)
          (evaluate3DComplex valMap $ (mp, arg))
      _ -> error "expression structure Three C is wrong"
  | otherwise = error "Three C but shape is not [size1, size2, size3] ??"

instance
  (KnownNat m, KnownNat n, KnownNat p) =>
  Evaluable '(m, n, p) C (Array (Int, Int, Int) (Complex Double))
  where
  eval ::
    ValMaps ->
    Expression '(m, n, p) C ->
    Array (Int, Int, Int) (Complex Double)
  eval valMap (Expression n mp) = evaluate3DComplex valMap (mp, n)

-- NOTE: `mod` in Haskell works with negative number, e.g, (-5) `mod` 3 = 1

-- | One dimension rotation
rotate1D ::
  -- | Size of the input array
  Int ->
  -- | amount of rotation
  Int ->
  -- | Input array
  Array Int a ->
  -- | Rotated array
  Array Int a
rotate1D size amount arr =
  listArray
    (0, size - 1)
    [arr ! ((i - amount) `mod` size) | i <- [0 .. size - 1]]

-- | Two dimension rotation
rotate2D ::
  -- | Size of the 2d input array
  (Int, Int) ->
  -- | amount of rotation for 2d array
  (Int, Int) ->
  -- | Input 2d array
  Array (Int, Int) a ->
  -- | Rotated 2d array
  Array (Int, Int) a
rotate2D (size1, size2) (amount1, amount2) arr =
  listArray
    ((0, 0), (size1 - 1, size2 - 1))
    [ arr ! ((i - amount1) `mod` size1, (j - amount2) `mod` size2)
      | i <- [0 .. size1 - 1],
        j <- [0 .. size2 - 1]
    ]

-- | Three dimension rotation
rotate3D ::
  -- | Size of 3d input array
  (Int, Int, Int) ->
  -- | Amount of Rotation for 3d array
  (Int, Int, Int) ->
  -- | Input 3d array
  Array (Int, Int, Int) a ->
  -- | Rotated 3d array
  Array (Int, Int, Int) a
rotate3D (size1, size2, size3) (amount1, amount2, amount3) arr =
  listArray
    ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
    [ arr
        ! ( (i - amount1) `mod` size1,
            (j - amount2) `mod` size2,
            (k - amount3) `mod` size3
          )
      | i <- [0 .. size1 - 1],
        j <- [0 .. size2 - 1],
        k <- [0 .. size3 - 1]
    ]

-- |
fourierTransform1D ::
  Int -> Array Int (Complex Double) -> Array Int (Complex Double)
fourierTransform1D size arr =
  listArray (0, size - 1) [computeX i | i <- [0 .. size - 1]]
  where
    computeX i = sum $ zipWithA (*) arr (fourierBasis i)
    fourierBasis i =
      let frequency n = 2 * pi * fromIntegral (i * n) / fromIntegral size
       in listArray
            (0, size - 1)
            [ cos (frequency n) :+ (- sin (frequency n))
              | n <- [0 .. size - 1]
            ]

-- |
fourierTransform2D ::
  (Int, Int) ->
  Array (Int, Int) (Complex Double) ->
  Array (Int, Int) (Complex Double)
fourierTransform2D (size1, size2) arr =
  listArray
    ((0, 0), (size1 - 1, size2 - 1))
    [computeX i j | i <- [0 .. size1 - 1], j <- [0 .. size2 - 1]]
  where
    computeX i j = sum $ zipWithA (*) arr (fourierBasis i j)
    fourierBasis i j =
      let frequency m n =
            2 * pi * fromIntegral (i * m) / fromIntegral size1
              + 2 * pi * fromIntegral (j * n) / fromIntegral size2
       in listArray
            ((0, 0), (size1 - 1, size2 - 1))
            [ cos (frequency m n) :+ (- sin (frequency m n))
              | m <- [0 .. size1 - 1],
                n <- [0 .. size2 - 1]
            ]

-- |
fourierTransform3D ::
  (Int, Int, Int) ->
  Array (Int, Int, Int) (Complex Double) ->
  Array (Int, Int, Int) (Complex Double)
fourierTransform3D (size1, size2, size3) arr =
  listArray
    ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
    [ computeX i j k
      | i <- [0 .. size1 - 1],
        j <- [0 .. size2 - 1],
        k <- [0 .. size3 - 1]
    ]
  where
    computeX i j k = sum $ zipWithA (*) arr (fourierBasis i j k)
    fourierBasis i j k =
      let frequency m n p =
            2 * pi * fromIntegral (i * m) / fromIntegral size1
              + 2 * pi * fromIntegral (j * n) / fromIntegral size2
              + 2 * pi * fromIntegral (k * p) / fromIntegral size3
       in listArray
            ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
            [ cos (frequency m n p) :+ (- sin (frequency m n p))
              | m <- [0 .. size1 - 1],
                n <- [0 .. size2 - 1],
                p <- [0 .. size3 - 1]
            ]


module SemiGr.SemiGrSpec (main,spec) where

import Test.Hspec
import Test.QuickCheck()
import Data.Complex(Complex((:+)))
import SemiGr.Semi(fromList, toList, LinkedList(..))

main :: IO ()
main = hspec spec

spec :: Spec
spec = do

  describe "Basic" $ do

    it "create linked list from list" $ do
      let a = [1, 2, 3, 4]
        in (fromList a) `shouldBe` LinkedList 1 (LinkedList 2 (LinkedList 3 (LinkedList 4 Empt)))

    it "join two linked list with semigroup opr" $ do
      let { a = fromList [1,2,3]
          ; b = fromList [4,5,6]}
          in (toList $ a <> b) `shouldBe` [1,4,2,5,3,6]