{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts      #-}
{-|
Module      : Grenade.Layers.Reshape
Description : Multipurpose reshaping layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Reshape (
    Reshape (..)
  ) where

import           Data.Serialize

import           Data.Singletons.TypeLits
import           GHC.TypeLits
import           Control.DeepSeq              (NFData (..))
import           GHC.Generics                 (Generic)

import           Numeric.LinearAlgebra.Static
import           Numeric.LinearAlgebra.Data as LA ( flatten )

import           Grenade.Core

-- | Reshape Layer
--
-- The Reshape layer can flatten any 2D or 3D image to 1D vector with the
-- same number of activations, as well as cast up from 1D to a 2D or 3D
-- shape.
--
-- Can also be used to turn a 3D image with only one channel into a 2D image
-- or vice versa.
data Reshape = Reshape
  deriving (Show,Generic,NFData)

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

instance RandomLayer Reshape where
  createRandomWith _ _ = return Reshape

instance (KnownNat a, KnownNat x, KnownNat y, a ~ (x * y)) => Layer Reshape ('D2 x y) ('D1 a) where
  type Tape Reshape ('D2 x y) ('D1 a) = ()
  runForwards _ (S2D y)   =  ((), fromJust' . fromStorable . flatten . extract $ y)
  runBackwards _ _ (S1D y) = ((), fromJust' . fromStorable . extract $ y)

instance (KnownNat a, KnownNat x, KnownNat y, KnownNat (x * z), KnownNat z, a ~ (x * y * z)) => Layer Reshape ('D3 x y z) ('D1 a) where
  type Tape Reshape ('D3 x y z) ('D1 a) = ()
  runForwards _ (S3D y)     = ((), fromJust' . fromStorable . flatten . extract $ y)
  runBackwards _ _ (S1D y) = ((), fromJust' . fromStorable . extract $ y)

instance (KnownNat y, KnownNat x, KnownNat z, z ~ 1) => Layer Reshape ('D3 x y z) ('D2 x y) where
  type Tape Reshape ('D3 x y z) ('D2 x y) = ()
  runForwards _ (S3D y)    = ((), S2D y)
  runBackwards _ _ (S2D y) = ((), S3D y)

instance (KnownNat y, KnownNat x, KnownNat z, z ~ 1) => Layer Reshape ('D2 x y) ('D3 x y z) where
  type Tape Reshape ('D2 x y) ('D3 x y z) = ()
  runForwards _ (S2D y)    = ((), S3D y)
  runBackwards _ _ (S3D y) = ((), S2D y)

instance (KnownNat a, KnownNat x, KnownNat y, a ~ (x * y)) => Layer Reshape ('D1 a) ('D2 x y) where
  type Tape Reshape ('D1 a) ('D2 x y) = ()
  runForwards _ (S1D y)   =  ((), fromJust' . fromStorable . extract $ y)
  runBackwards _ _ (S2D y) = ((), fromJust' . fromStorable . flatten . extract $ y)

instance (KnownNat a, KnownNat x, KnownNat y, KnownNat (x * z), KnownNat z, a ~ (x * y * z)) => Layer Reshape ('D1 a) ('D3 x y z) where
  type Tape Reshape ('D1 a) ('D3 x y z) = ()
  runForwards _ (S1D y)     = ((), fromJust' . fromStorable . extract $ y)
  runBackwards _ _ (S3D y)  = ((), fromJust' . fromStorable . flatten . extract $ y)

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

fromJust' :: Maybe x -> x
fromJust' (Just x) = x
fromJust' Nothing  = error $ "Reshape error: data shape couldn't be converted."

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

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


{-# LANGUAGE ScopedTypeVariables #-}
module Main where

import Util

import Text.ParserCombinators.Parsec
import Text.ParserCombinators.Parsec.Number

import Data.String.Utils

import Data.String
import Data.CSV
import Data.List

import Data.Ord

import Data.Int

import Statistics.Sample as S hiding(stdDev)
import Data.Vector(fromList)

import Numeric

import Data.Maybe
import Data.Either
import qualified Data.Map.Strict as M

import Debug.Trace

import System.Environment

data Speedup = Speedup {
    num :: String,
    name :: String,
    time :: Double,
    nCores :: String,
    speedup :: Double,
    stdDev :: Double,
    factor :: Double,
    factorStdDev :: Double,
    overhead :: Double,
    stdDevForOverhead :: Double,
    runtimeX :: Double,
    runtimeY :: Double
} deriving (Show)

textBFLaTeX :: String -> String
textBFLaTeX str = "\"\\textbf{" ++ str ++ "}\""

textITLaTeX :: String -> String
textITLaTeX str = "\"\\textit{" ++ str ++ "}\""

rndSignificant :: Int -> Double -> Double
rndSignificant cnt val
      | val == 0.0 = 0.0
      | otherwise = (fromIntegral shifted) / magnitude
        where d::Double = fromIntegral (ceiling (logBase 10 (if val < 0 then -val else val)))
              power::Int64 = (fromIntegral cnt) - (floor d)

              magnitude::Double = fromIntegral (10 ^ power)
              shifted::Int64 = (round (val * magnitude))

formatOverheadForLaTeX :: Double -> String
formatOverheadForLaTeX overhead
   | overhead < 0 = textBFLaTeX (formatPercent overhead)
   | overhead > 0 = textITLaTeX (formatPercent overhead)

formatPercent :: Double -> String
formatPercent x = ((showFFloat Nothing (rndPercent x) "") ++ "\\%")

rndPercent :: Double -> Double
rndPercent x = (rndSignificant 2 (x * 100))

main :: IO ()
main = do
    args <- getArgs
    if (length args < 3)
    then
        do
            putStrLn $ "usage: <program> file1 file2 output [calculateWorst]"
    else
        do
            let (file1:file2:output:rest) = args

                lineToSpeedup :: [String] -> Maybe Speedup
                lineToSpeedup (num:name:numStrings) | (length numStrings) == 4 =
                                                        let [timeString, nCores, speedupString, stdDevString] = numStrings
                                                            maybeDoubles = convDoubles [timeString, speedupString, stdDevString]
                                                        in
                                                            if (isNothing maybeDoubles)
                                                            then
                                                                Nothing
                                                            else
                                                                let (Just [time, speedup, stdDev]) = maybeDoubles
                                                                in
                                                                    Just $ Speedup {
                                                                        num = num,
                                                                        name = name,
                                                                        time = time,
                                                                        nCores = nCores,
                                                                        speedup = speedup,
                                                                        stdDev = stdDev,
                                                                        factor = 1,
                                                                        factorStdDev = 0,
                                                                        overhead = 0,
                                                                        stdDevForOverhead = 0,
                                                                        runtimeX = 0,
                                                                        runtimeY = 0
                                                                    }
                                                   | otherwise = Nothing
                lineToSpeedup _ = Nothing

                diff :: Speedup -> Speedup -> Maybe Speedup
                diff x y
                    | nCores x == nCores y && num x == num y = Just $ Speedup {
                            num = num x,
                            name = (name x) ++ "-" ++ (name y),
                            time = (time x) - (time y),
                            nCores = nCores x,
                            speedup = (speedup x) - (speedup y),
                            stdDev = max (stdDev x) (stdDev y),
                            -- care: different order than time. cause reasons
                            factor = (time y) / (time x),
                            factorStdDev = (max (stdDev x) (stdDev y)) / (time x),
                            overhead = (time y - time x) / (time y),
                            stdDevForOverhead = (max (stdDev x) (stdDev y)) / (time y),
                            runtimeX = time x,
                            runtimeY = time y
                        }
                    | otherwise = Nothing

                legend :: String
                legend = "\"\",\"name\",\"time\",\"nCores\",\"speedup\",\"max stddev\",\"factor\",\"factorStdDev\",\"overhead\",\"stdDevForOverhead\",\"runtimeX\",\"runtimeY\"" ++ "\n"

                str :: String -> String
                str st = "\"" ++ st ++ "\""

                toStringLn :: Speedup -> String
                toStringLn x = (toString x) ++ "\n"

                toString :: Speedup -> String
                toString x = (str (num x)) ++ "," ++ (str $ name x) ++ "," ++  (show $ time x) ++ "," ++
                                    (nCores x) ++ "," ++ (show $ speedup x) ++ "," ++ (show $ stdDev x) ++ "," ++
                                    (show $ factor x) ++ "," ++ (show $ factorStdDev x) ++ "," ++ (show $ overhead x) ++
                                    "," ++ (show $ stdDevForOverhead x) ++ "," ++ (show $ runtimeX x) ++  "," ++ (show $ runtimeY x)

                handleParse :: Either ParseError [[String]] -> [Speedup]
                handleParse (Right lines) = catMaybes $ map lineToSpeedup lines
                handleParse _  = []

            linesOrError1 <- parseFromFile csvFile file1
            linesOrError2 <- parseFromFile csvFile file2

            let speedups1 = handleParse $ linesOrError1
                speedups2 = handleParse $ linesOrError2
            if (length speedups1 == 0 || length speedups2 == 0 || length speedups1 /= length speedups2)
            then
                putStrLn "parse Error!"
            else
                do
                    let maybeDiffs = zipWith diff speedups1 speedups2
                        diffs = {-traceShowId $-} catMaybes maybeDiffs
                    if (length diffs /= length maybeDiffs)
                    then
                        putStrLn $ "parse Error!"
                    else
                        if (length args >= 4)
                        then
                            let
                                -- hacky, but nvm
                                (_:restOfRest) = rest

                                readBool :: String -> Bool
                                readBool = read

                                calcOutput :: Bool -> Speedup
                                calcOutput True = maximumBy (comparing factor) diffs
                                calcOutput False = minimumBy (comparing factor) diffs

                                rnd :: Speedup -> Speedup
                                rnd x = Speedup {
                                                    num = num x,
                                                    name = name x,
                                                    time = time x,
                                                    nCores = nCores x,
                                                    speedup = speedup x,
                                                    stdDev = stdDev x,
                                                    factor = rndVal (factor x),
                                                    factorStdDev = factorStdDev x,
                                                    overhead = rndVal $ overhead x,
                                                    stdDevForOverhead = stdDevForOverhead x,
                                                    runtimeX = runtimeX x,
                                                    runtimeY = runtimeY x
                                                }

                                rndVal = rndSignificant 3

                                meanMeanOverhead = S.mean . fromList

                                roundedGeometricMean = rndVal . S.geometricMean . fromList

                                roundedMean = rndVal . S.mean . fromList
                            in
                                do appendFile output $ (toString $ rnd $ calcOutput True)
                                    ++ ","
                                    ++ (toString $ rnd $ calcOutput False)
                                    ++ ","
                                    -- geometric mean of factor
                                    ++ (showFFloat Nothing $ roundedMean $ map (factor) diffs) ""
                                    ++ ","
                                    --
                                    ++ (showFFloat Nothing $ rndVal $ factorStdDev $ maximumBy (comparing factorStdDev) diffs) ""
                                    ++ ","
                                    ++ (formatOverheadForLaTeX $ meanMeanOverhead $ map (overhead) diffs)
                                    ++ ","
                                    ++ (formatPercent $ stdDevForOverhead $ maximumBy (comparing stdDevForOverhead) diffs)
                                    ++ ","
                                    ++ (showFFloat Nothing $ rndVal $ runtimeX $ maximumBy (comparing nCores) diffs) ""
                                    ++ "\n"

                        else
                            do writeFile output $ legend ++ (concat $ map toStringLn diffs)


{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeInType          #-}
{-# LANGUAGE TypeOperators       #-}

module TensorOps.Learn.NeuralNet.Recurrent
  ( Network
  , buildNet
  , netParams
  , runNetwork
  , runNetworkSt
  , genNet
  , fullyConnected
  , ffLayer
  , stateless
  , (~*~)
  , (*~)
  , (~*)
  , nmap
  , trainNetwork
  , trainNetwork'
  , networkGradient
  ) where

import           Control.Category
import           Control.DeepSeq
import           Control.Monad.Primitive
import           Control.Monad.State
import           Data.Kind
import           Data.Singletons
import           Data.Singletons.Prelude               (Sing(..))
import           Data.Type.Combinator
import           Data.Type.Conjunction
import           Data.Type.Length                      as TCL
import           Data.Type.Length.Util                 as TCL
import           Data.Type.Nat
import           Data.Type.Product                     as TCP
import           Data.Type.Product.Util                as TCP
import           Data.Type.Sing
import           Data.Type.Uniform
import           Data.Type.Vector                      as TCV
import           Data.Type.Vector.Util                 as TCV
import           Prelude hiding                        ((.), id)
import           Statistics.Distribution.Normal
import           System.Random.MWC
import           TensorOps.Learn.NeuralNet
import           TensorOps.NatKind
import           TensorOps.TOp                         as TO
import           TensorOps.Tensor                      as TT
import           TensorOps.Types
import           Type.Class.Higher
import           Type.Class.Higher.Util
import           Type.Class.Known
import           Type.Class.Witness
import           Type.Family.List
import           Type.Family.List.Util
import           Type.Family.Nat
import           Unsafe.Coerce
import qualified TensorOps.Learn.NeuralNet.FeedForward as FF

data Network :: ([k] -> Type) -> k -> k -> Type where
    N :: { _nsSs    :: !(Sing ss)
         , _nsPs    :: !(Sing ps)
         , _nOp     :: !(TOp ('[i] ': ss ++ ps) ('[o] ': ss))
         , _nState  :: !(Prod t ss)
         , _nParams :: !(Prod t ps)
         } -> Network t i o

instance NFData1 t => NFData (Network t i o) where
    rnf = \case
      N _ _ o p s -> o `seq` p `deepseq1` s `deepseq1` ()
    {-# INLINE rnf #-}

netParams
    :: Network t i o
    -> (forall ss ps. (SingI ss, SingI ps) => Prod t ss -> Prod t ps -> r)
    -> r
netParams = \case
    N sS sP _ s p -> \f -> f s p \\ sS \\ sP

buildNet
    :: forall ss ps i o t. (SingI ss, SingI ps)
    => TOp ('[i] ': ss ++ ps) ('[o] ': ss)  -- ^ Tensor operation for network
    -> Prod t ss        -- ^ Initial states
    -> Prod t ps        -- ^ Network parameters
    -> Network t i o
buildNet = N sing sing

fullyConnected
    :: forall k (i :: k) (o :: k) (m :: Type -> Type) (t :: [k] -> Type).
     ( SingI i
     , SingI o
     , PrimMonad m
     , Tensor t
     )
    => Activation k         -- ^ Activation function for internal state
    -> Gen (PrimState m)
    -> m (Network t i o)
fullyConnected a g = (\s w w' b -> buildNet @'[ '[o] ] @'[ '[o,o], '[o,i], '[o] ]
                                     fc (s :< Ø) (w' :< w :< b :< Ø)
                     )
          <$> genRand (normalDistr 0 0.5) g
          <*> genRand (normalDistr 0 0.5) g
          <*> genRand (normalDistr 0 0.5) g
          <*> genRand (normalDistr 0 0.5) g
  where
    fc  :: TOp '[ '[i], '[o], [o,o], '[o,i], '[o]] '[ '[o], '[o] ]
    fc = secondOp @'[ '[i] ] (
           firstOp @'[ '[o,i], '[o] ] (TO.swap >>> TO.matVec)
       >>> firstOp @'[ '[o] ]         TO.swap
         )
     >>> firstOp @'[ '[o], '[o] ] (TO.swap >>> TO.matVec)
     >>> TO.add3
     >>> TO.duplicate
     >>> secondOp @'[ '[o] ] (getAct a)
    {-# INLINE fc #-}
{-# INLINE fullyConnected #-}

-- | Convert a neural network from "TensorOps.Learn.NeuralNet.FeedForward"
-- to a stateless 'Network'.
--
-- Can be thought a functor from the category of stateless neural networks
-- to the category of stateful (recurrent) ones.
stateless
    :: FF.Network t i o
    -> Network t i o
stateless = \case
    FF.N sP o p -> N SNil sP o Ø p
{-# INLINE stateless #-}

ffLayer
    :: forall i o m t. (SingI i, SingI o, PrimMonad m, Tensor t)
    => Gen (PrimState m)
    -> m (Network t i o)
ffLayer g = stateless <$> FF.ffLayer g
{-# INLINE ffLayer #-}

genNet
    :: forall k o i m (t :: [k] -> Type). (SingI o, SingI i, PrimMonad m, Tensor t)
    => [(Integer, (Activation k, Maybe (Activation k)))]
    -> Activation k
    -> Maybe (Activation k)
    -> Gen (PrimState m)
    -> m (Network t i o)
genNet xs0 f fS g = go sing xs0
  where
    go  :: forall (j :: k). ()
        => Sing j
        -> [(Integer, (Activation k, Maybe (Activation k)))]
        -> m (Network t j o)
    go sj = (\\ sj) $ \case
      []        -> fmap (*~ getAct f) $ case fS of
        Just fS' -> fullyConnected fS' g
        Nothing  -> ffLayer g
      (x,(f', fS')):xs -> withNatKind x $ \sl -> (\\ sl) $ do
        n <- go sl xs
        l <- case fS' of
          Nothing   -> ffLayer g
          Just fS'' -> fullyConnected fS'' g
        return $ l *~ getAct f' ~*~ n
    {-# INLINE go #-}
{-# INLINE genNet #-}

instance Category (Network t) where
    id  = N SNil SNil idOp Ø Ø
    (.) = flip (~*~)

(~*~)
    :: forall k (t :: [k] -> Type) a b c. ()
    => Network t a b
    -> Network t b c
    -> Network t a c
(~*~) = \case
  N (sSs1 :: Sing ss1)
    (sPs1 :: Sing ps1)
    (o1   :: TOp ('[a] ': ss1 ++ ps1) ('[b] ': ss1))
    (s1   :: Prod t ss1)
    (p1   :: Prod t ps1) -> \case
    N (sSs2 :: Sing ss2)
      (sPs2 :: Sing ps2)
      (o2   :: TOp ('[b] ': ss2 ++ ps2) ('[c] ': ss2))
      (s2   :: Prod t ss2)
      (p2   :: Prod t ps2) ->
      let lSs1 :: Length ss1
          lSs1 = singLength sSs1
          lSs2 :: Length ss2
          lSs2 = singLength sSs2
          lPs1 :: Length ps1
          lPs1 = singLength sPs1
          lPs2 :: Length ps2
          lPs2 = singLength sPs2
          o :: TOp ('[a] ': (ss2 ++ ss1) ++ (ps1 ++ ps2)) ('[c] ': ss2 ++ ss1)
                    -- all these proofs lol
          o     = (\\ (unsafeCoerce Refl :: ((ss2 ++ ss1) ++ ps1) ++ ps2 :~: (ss2 ++ ss1) ++ (ps1 ++ ps2))) $
                  (\\ (unsafeCoerce Refl :: ((ss1 ++ ps1) ++ ss2) ++ ps2 :~: (ss1 ++ ps1) ++ (ss2 ++ ps2))) $
                  (\\ appendAssoc lSs1 lSs2 lPs2                  ) $
                  (\\ appendAssoc lSs2 lSs1 lPs1                  ) $
                  (\\ (lSs2 `TCL.append'` lSs1) `TCL.append'` lPs1) $
                  (\\ (lSs1 `TCL.append'` lPs1) `TCL.append'` lSs2) $
                  (\\ lSs1                                        ) $
                  (\\ lSs2                                        ) $
                  (\\ lSs1 `TCL.append'` lPs1                     ) $
                  (\\ lSs2 `TCL.append'` lPs2                     ) $
                    secondOp @'[ '[a] ]
                      (firstOp @ps2 (TO.swap' lSs2 (lSs1 `TCL.append'` lPs1))
                      )
                >>> firstOp @(ss2 ++ ps2) o1
                >>> secondOp @'[ '[b] ] (TO.swap' lSs1 (lSs2 `TCL.append'` lPs2))
                >>> firstOp @ss1 o2
      in  N (sSs2 %:++ sSs1)
            (sPs1 %:++ sPs2)
            o
            (s2 `TCP.append'` s1)
            (p1 `TCP.append'` p2)
infixr 4 ~*~
{-# INLINE (~*~) #-}

runNetwork
    :: (RealFloat (ElemT t), Tensor t)
    => Network t i o
    -> t '[i]
    -> (t '[o], Network t i o)
runNetwork (N sS sO o s p) x =
        (\case y :< s' -> (y, N sS sO o s' p))
      . runTOp o
      $ x :< (s `TCP.append'` p)
{-# INLINE runNetwork #-}

runNetworkSt
    :: (RealFloat (ElemT t), Tensor t, MonadState (Network t i o) m)
    => t '[i]
    -> m (t '[o])
runNetworkSt x = state $ flip runNetwork x

(~*) :: TOp '[ '[a] ] '[ '[b] ]
     -> Network t b c
     -> Network t a c
f ~* N sS sO o s p = N sS sO (f *>> o) s p
infixr 4 ~*
{-# INLINE (~*) #-}

(*~) :: Network t a b
     -> TOp '[ '[b] ] '[ '[c] ]
     -> Network t a c
N sS sO o s p *~ f = N sS sO (o >>> firstOp f) s p
infixl 5 *~
{-# INLINE (*~) #-}

nmap
     :: SingI o
     => (forall a. RealFloat a => a -> a)
     -> Network t i o
     -> Network t i o
nmap f n = n *~ TO.map f
{-# INLINE nmap #-}

netGrad
    :: forall k n (t :: [k] -> Type) (i :: k) (o :: k) ss ps.
     ( Tensor t
     , RealFloat (ElemT t)
     , SingI i
     , SingI o
     )
    => TOp '[ '[o], '[o] ] '[ '[] ]
    -> Vec n (t '[i])
    -> Vec n (t '[o])
    -> Sing ss
    -> Sing ps
    -> TOp ('[i] ': ss ++ ps) ('[o] ': ss)
    -> Prod t ss
    -> Prod t ps
    -> (Vec n (t '[i]), (Prod t ss, Prod t ps))
netGrad loss xs ys sS sP o s p =
      (prodToVec' I n grI, splitProd @ps lS grSP)
  where
    n :: Nat n
    n = known   \\ xs
    lS :: Length ss
    lS = singLength sS
    lP :: Length ps
    lP = singLength sP
    sO :: Sing (Replicate n '[o])
    sO = replicateSing (sing :: Sing '[o]) n
    lI :: Length (Replicate n '[i])
    lI = replicateLength @'[i] n
    lO :: Length (Replicate n '[o])
    lO = replicateLength @'[o] n
    unrolled
        :: TOp (Replicate n '[i] ++ ss ++ ps) (Replicate n '[o])
    unrolled = (\\ sS %:++ sO) $
               (\\ sS %:++ sP) $
                 unroll sS sP o n
             >>> TO.drop @(Replicate n '[o]) lS
    o'  :: TOp (Replicate n '[i] ++ ss ++ ps ++ Replicate n '[o]) '[ '[] ]
    o' = (\\ appendAssoc lI (lS `TCL.append'` lP) lO) $
         (\\ appendAssoc lS lP lO) $
         (\\ lI `TCL.append'` lS `TCL.append'` lP) $
         (\\ lO) $
            firstOp @(Replicate n '[o]) unrolled
        >>> rollup @k lS lP loss n
    xs' :: Prod t (Replicate n '[i] ++ ss ++ ps ++ Replicate n '[o])
    xs' = vecToProd' getI (TCV.reverse' xs) `TCP.append'` (s
            `TCP.append'` (
              p `TCP.append'` vecToProd' getI ys
            )
          )
    grad :: Prod t (Replicate n '[i] ++ ss ++ ps)
    grad = (\\ (unsafeCoerce Refl :: Replicate n '[i] ++ ss ++ ps ++ Replicate n '[o]
                                 :~: (Replicate n '[i] ++ ss ++ ps) ++ Replicate n '[o]
               )
           ) $
      takeProd @(Replicate n '[o]) (lI `TCL.append'` lS `TCL.append'` lP) $ gradTOp o' xs'
    grI   :: Prod t (Replicate n '[i])
    grSP  :: Prod t (ss ++ ps)
    (grI, grSP) = splitProd @(ss ++ ps) lI grad
{-# INLINE netGrad #-}

trainNetwork'
    :: forall k n (t :: [k] -> Type) (i :: k) (o :: k).
     ( Tensor t
     , RealFloat (ElemT t)
     , SingI i
     , SingI o
     )
    => TOp '[ '[o], '[o] ] '[ '[] ]     -- ^ loss function (input, target)
    -> ElemT t          -- ^ train rate for initial state
    -> ElemT t          -- ^ train rate for network parameters
    -> Vec n (t '[i])   -- ^ inputs
    -> Vec n (t '[o])   -- ^ targets
    -> Network t i o
    -> Network t i o
trainNetwork' loss rS rP xs ys = \case
    N sS sP o s p ->
      let (gS, gP) = snd $ netGrad loss xs ys sS sP o s p
          s' = map1 (f rS) $ zipProd3 (singProd sS) s gS
          p' = map1 (f rP) $ zipProd3 (singProd sP) p gP
      in  N sS sP o s' p'
  where
    f   :: forall ns. ()
        => ElemT t
        -> (Sing :&: t :&: t) ns
        -> t ns
    f r !(s1 :&: p1 :&: g1) =
      TT.zip (\(!p2) (!g2) -> p2 - r * g2) p1 g1 \\ s1
    {-# INLINE f #-}
{-# INLINE trainNetwork #-}

trainNetwork
    :: forall k (t :: [k] -> Type) (i :: k) (o :: k).
     ( Tensor t
     , RealFloat (ElemT t)
     , SingI i
     , SingI o
     )
    => TOp '[ '[o], '[o] ] '[ '[] ]     -- ^ loss function (input, target)
    -> ElemT t          -- ^ train rate for initial state
    -> ElemT t          -- ^ train rate for network parameters
    -> [(t '[i], t '[o])]   -- ^ inputs and targets
    -> Network t i o
    -> Network t i o
trainNetwork loss rS rP xsys n = withV xsys $ \v ->
    let (xs, ys) = TCV.unzip' v
    in  trainNetwork' loss rS rP xs ys n

networkGradient
    :: forall k n (t :: [k] -> Type) (i :: k) (o :: k) r.
     ( Tensor t
     , RealFloat (ElemT t)
     , SingI i
     , SingI o
     )
    => TOp '[ '[o], '[o] ] '[ '[] ]     -- ^ loss function (output, target)
    -> Vec n (t '[i])       -- ^ inputs
    -> Vec n (t '[o])       -- ^ targets
    -> Network t i o
    -> (forall ss ps. (SingI ss, SingI ps) => Vec n (t '[i]) -> Prod t ss -> Prod t ps -> r)
    -> r
networkGradient loss xs ys = \case
    N sS sP o s p -> \f -> case netGrad loss xs ys sS sP o s p of
      (gI, (gS, gP)) -> f gI gS gP \\ sS \\ sP
{-# INLINE networkGradient #-}

unroll
    :: forall ss ps i o n. (SingI (ss ++ ps), SingI i)
    => Sing ss
    -> Sing ps
    -> TOp ('[i] ': ss ++ ps) ('[o] ': ss)
    -> Nat n
    -> TOp (Replicate n '[i] ++ ss ++ ps) (ss ++ Replicate n '[o])
unroll sS sP o = \case
    Z_              -> TO.take lS lP    \\ appendNil lS
    S_ (m :: Nat m) -> (\\ (unsafeCoerce Refl :: Replicate m '[i] ++ '[i] ': ss ++ ps
                                             :~: '[i] ': Replicate m '[i] ++ ss ++ ps
                           )
                       ) $
                       (\\ (unsafeCoerce Refl :: (Replicate m '[i] ++ ss ++ ps) ++ '[ '[o] ]
                                             :~: Replicate m '[i] ++ ss ++ (ps >: '[o])
                           )
                       ) $
                       (\\ (unsafeCoerce Refl :: (ss ++ Replicate m '[o]) ++ '[ '[o] ]
                                             :~: ss ++ '[o] ': Replicate m '[o]
                           )
                       ) $
                       (\\ appendAssoc lS lP (LS LZ :: Length '[ '[o] ])) $
                       (\\ appendSnoc lP (Proxy @'[o])) $
                       (\\ replicateLength @'[i] m `TCL.append'` lS `TCL.append'` lP) $
                       (\\ lS `TCL.append'` replicateLength @'[o] m) $
                       (\\ replicateLength @'[i] m) $
                       (\\ lS) $
          secondOp @(Replicate m '[i]) @('[i] ': ss ++ ps) @(ss ++ ps >: '[o]) (
                (o &&& TO.drop @ps @('[i] ': ss) (LS lS))
            >>> TO.swap' (LS LZ) (lS `TCL.append'` lP)
          )
      >>> firstOp @'[ '[o] ] @(Replicate m '[i] ++ ss ++ ps) @(ss ++ Replicate m '[o]) (
            unroll sS sP o m
          )
  where
    lS :: Length ss
    lS = singLength sS
    lP :: Length ps
    lP = singLength sP
{-# INLINE unroll #-}

rollup
    :: forall k (ss :: [[k]]) (ps :: [[k]]) (o :: k) (n :: N). ()
    => Length ss
    -> Length ps
    -> TOp '[ '[o], '[o] ] '[ '[] ]
    -> Nat n
    -> TOp (Replicate n '[o] ++ Replicate n '[o]) '[ '[] ]
rollup lS lP loss = \case
    Z_                     -> TO.konst (US UØ) 0
    S_ Z_                  -> loss
    S_ (m@(S_ _) :: Nat m) ->
      let lO :: Length (Replicate m '[o])
          lO = replicateLength @'[o] m
      in  (\\ (unsafeCoerce Refl :: Replicate m '[o] ++ '[o] ': '[o] ': Replicate m '[o]
                                :~: '[o] ': Replicate m '[o] ++ '[o] ': Replicate m '[o]
              )
            ) $
            (\\ appendSnoc lO (Proxy @'[])) $
            (\\ lO) $
            (\\ lO `TCL.append'` lO) $
            (\\ appendAssoc lO lO (LS LZ :: Length '[ '[] ])) $
              secondOp @(Replicate m '[o]) @('[o] ': '[o] ': Replicate m '[o]) @(Replicate m '[o] >: '[]) (
                  firstOp @(Replicate m '[o]) @('[ '[o], '[o] ]) @('[ '[] ]) loss
                >>> TO.swap' @'[ '[] ] @(Replicate m '[o]) (LS LZ) lO
              )
            >>> firstOp @'[ '[] ] @(Replicate m '[o] ++ Replicate m '[o]) @'[ '[] ] (
                  rollup lS lP loss m
                )
            >>> TO.add
{-# INLINE rollup #-}


{-# LANGUAGE FlexibleContexts #-}
{-|
Module      : $Header$
CopyRight   : (c) 8c6794b6, 2011-2013
License     : BSD3
Maintainer  : 8c6794b6@gmail.com
Stability   : experimental
Portability : portable

Synthesis methods

-}
module Spectrofy.Synth
  ( sinsyn
  , fftsyn
  ) where

import Data.Complex (Complex(..), realPart)
import Data.Word (Word8)
import Data.List (foldl1')

import Data.Array.Repa
    ((:.)(..), Array, All(..), D, DIM2, DIM3, Source, U, Z(..))
import qualified Data.Array.Repa as R
import qualified Data.Array.Repa.FFTW as R

-- --------------------------------------------------------------------------
--
-- Manual sinusoids
--

-- | Sums up sinusoids manually.
--
-- Frequency is taken from y axis, time from x axis.
-- Amplitude for each frequency is result of luminated value in each pixel.
--
sinsyn ::
    Source s (Word8, Word8, Word8)
    => Int -> Int -> Array s DIM2 (Word8, Word8, Word8) -> Array D DIM2 Double
sinsyn rate delta = squash . sumSins . toSins rate . to3D delta . rgbamp
{-# INLINE sinsyn #-}

rgbamp ::
    Source s (Word8, Word8, Word8)
    => Array s DIM2 (Word8, Word8, Word8) -> Array D DIM2 Double
rgbamp arr = {-# SCC "rgbamp" #-} R.traverse arr f g where
  f = id
  g h ix = 10 ** ((red*0.21 + green*0.71 + blue*0.07) / (255*3)) - 1 where
    (red', green', blue') = h ix
    red = fromIntegral red'
    green = fromIntegral green'
    blue = fromIntegral blue'
{-# INLINE rgbamp #-}

to3D :: Source s Double => Int -> Array s DIM2 Double -> Array D DIM3 Double
to3D n = R.extend (Z:.All:.All:.n)
{-# INLINE to3D #-}

toSins :: Source s Double => Int -> Array s DIM3 Double -> Array D DIM3 Double
toSins rate arr = {-# SCC "toSins" #-} R.traverse arr id f where
  _:.y:._:._ = R.extent arr
  f g ix@(_:.i:._:.k)
    | amp <= 0  = 0
    | otherwise = amp * sin (frq * k' * 2 * pi / rate') / y'
    where
      amp = g ix
      frq = (rate'/2) - (((y'-i'+1)/y') * (rate'/2))
      rate' = fromIntegral rate
      i' = fromIntegral i
      y' = fromIntegral y
      k' = fromIntegral k
{-# INLINE toSins #-}

sumSins :: (Source s Double) => Array s DIM3 Double -> Array U DIM2 Double
sumSins arr = {-# SCC "sumSins" #-} R.sumS (R.backpermute sh' f arr) where
  f (_:.i:.j:.k) = Z:.k:.j:.i
  sh' = Z:.z:.y:.x
  (_:.x:.y:.z) = R.extent arr
{-# INLINE sumSins #-}

squash :: Source s Double => Array s DIM2 Double -> Array D DIM2 Double
squash arr = {-# SCC squash #-} R.backpermute sh' f arr where
  f (_:._:.j) = Z:.j `mod` x:.j `div` x
  (_:.x:.y) = R.extent arr
  sh' = Z:.1:.(x*y)
{-# INLINE squash #-}

-- --------------------------------------------------------------------------
--
-- Inversed FFT
--

-- | Performs 1-dimensional inverse fft on each column, then concatenate results.
--
-- RGB values are used for initial phases and magnitude.
--
fftsyn :: Int -> Array D DIM2 (Word8, Word8, Word8) -> Array D DIM2 Double
fftsyn fsz img =
  {-# SCC "fftsyn" #-}
  foldl1' R.append $ map (fftslice fsz img') [0..n-1] where
    img' = rgbamp img
    _:._:.n = R.extent img
{-# INLINE fftsyn #-}

fftslice :: Int -> Array D DIM2 Double -> Int -> Array D DIM2 Double
fftslice wsize arr n =
  {-# SCC "fs_reshape" #-}
  R.reshape (Z :. 1 :. wsize :: DIM2) $
  {-# SCC "fs_map" #-}
  R.map realPart $
  {-# SCC "fs_ifft" #-}
  R.ifft $
  {-# SCC "fs_computeS" #-}
  R.computeS $
  {-# SCC "fs_slice" #-}
  R.slice (R.map (\x -> x :+ x) $ grow2d wsize (zpad arr)) (Z:.All:.n)
{-# INLINE fftslice #-}

grow2d :: Int -> Array D DIM2 Double -> Array D DIM2 Double
grow2d n arr = {-# SCC "grow2d" #-} R.traverse arr f g where
  f _ = Z :. n :. y
  _:.x:.y = R.extent arr
  g h (_:.i:.j)
    | i' >= 4   = k/4
    | otherwise = h (Z:.i':.j)
    where
      k  = h (Z:.i'-3:.j) + h (Z:.i'-2:.j) + h (Z:.i'-1:.j) + h (Z:.i':.j)
      i' = fst $ properFraction $ (x' * fromIntegral i / n')
      n' = fromIntegral n :: Double
      x' = fromIntegral x
{-# INLINE grow2d #-}

zpad :: Array D DIM2 Double -> Array D DIM2 Double
zpad arr = {-# SCC "zpad" #-} R.reshape sh' arr' where
  _:.x:.y = R.extent arr
  len = x * y
  zeros = R.fromFunction zsh (const 0)
  zsh = Z :. len
  sh' = Z :. (2*x) :. y
  arr' = R.append (R.reshape zsh arr) zeros
{-# INLINE zpad #-}


{-# OPTIONS_GHC -fno-warn-unused-do-bind #-}

module Scheme48.Parsers where

import Scheme48.Types

import Control.Monad
import Numeric
import Data.Ratio
import Data.Complex
import Data.Array
import Text.ParserCombinators.Parsec hiding (spaces)

symbol :: Parser Char
symbol = oneOf "!#$%&|*+-/:<=>?@^_~"

spaces :: Parser ()
spaces = skipMany1 space

-- Strings --

escapedChars :: Parser Char
escapedChars = do
                _ <- char '\\'
                x <- oneOf "\\\"nrt"
                return $ case x of
                            '\\' -> x
                            '"'  -> x
                            'n'  -> '\n'
                            'r'  -> '\r'
                            't'  -> '\t'
                            _    -> error "impossible"

parseString :: Parser LispVal
parseString = try $ do
                _ <- char '"'
                x <- many (noneOf "\"")
                _ <- char '"'
                return $ String x

parseCharacter :: Parser LispVal
parseCharacter = do
                    try $ string "#\\"
                    x <- try (string "newline" <|> string "space")
                        <|> do { v <- anyChar; notFollowedBy alphaNum; return [v] }
                    return $ Character $ case x of
                            "newline" -> '\n'
                            "space"   -> ' '
                            _ -> head x

-- Atoms --

parseAtom :: Parser LispVal
parseAtom = try $ do
                first <- letter <|> symbol
                rest <- many (letter <|> digit <|> symbol)
                let atom = first:rest
                return $ Atom atom

--- Numbers ---

parseNumber :: Parser LispVal
parseNumber = try parseComplex
           <|> try parseRatio
           <|> try parseFloat
           <|> try parseDigit1
           <|> try parseDigit2
           <|> try parseHex
           <|> try parseOct
           <|> try parseBin

parseDigit1 :: Parser LispVal
parseDigit1 = try $ liftM (Number . read) (many1 digit)

parseDigit2 :: Parser LispVal
parseDigit2 = try $ string "#d" >> parseDigit1

parseHex :: Parser LispVal
parseHex =  try $ liftM (Number . hex2dig) (string "#x" >> many1 hexDigit)

hex2dig :: (Num a, Eq a) => String -> a
hex2dig = fst . head . readHex

parseOct :: Parser LispVal
parseOct = try $ liftM (Number . oct2dig) (string "#o" >> many1 octDigit)

oct2dig :: (Num a, Eq a) => String -> a
oct2dig = fst . head . readOct

parseBin :: Parser LispVal
parseBin = try $ liftM (Number . bin2dig) (string "#b" >> many1 (oneOf "01"))

bin2dig :: String -> Integer
bin2dig = bin2dig' 0

bin2dig' :: Num a => a -> String -> a
bin2dig' d ""     = d
bin2dig' d (x:xs) = let
                        n = case x of
                                '0' -> 0
                                '1' -> 1
                                _   -> error "Impossible"
                        o = 2 * d + n
                    in bin2dig' o xs

-- Floats --

parseFloat :: Parser LispVal
parseFloat = try $ do
  i <- many1 digit
  _ <- char '.'
  d <- many1 digit
  return $ Float $ fst . head $ readFloat $ i ++ "." ++ d

-- Ratios --

parseRatio :: Parser LispVal
parseRatio = try $ do
  n <- many1 digit
  _ <- char '/'
  d <- many1 digit
  return $ Ratio $ (read n) % (read d)

-- Complex --

parseComplex :: Parser LispVal
parseComplex = try $ do
  x <- (try parseFloat <|> parseDigit1)
  _ <- char '+'
  y <- (try parseFloat <|> parseDigit1)
  _ <- char 'i'
  return $ Complex $ toDouble x :+ toDouble y

toDouble :: LispVal -> Double
toDouble (Float f) = f
toDouble (Number n) = fromIntegral n
toDouble e = error $ "wrong type: " ++ show e

-- Booleans --

parseBool :: Parser LispVal
parseBool = try $ char '#' >>
  ((char 't' >> return (Bool True)) <|> (char 'f' >> return (Bool False)))

-- Lists --

parseList :: Parser LispVal
parseList = liftM List $ sepBy parseExpr spaces

parseDottedList :: Parser LispVal
parseDottedList = do
  h <- endBy parseExpr spaces
  t <- char '.' >> spaces >> parseExpr
  return $ DottedList h t

parseQuoted :: Parser LispVal
parseQuoted = try $ do
  _ <- char '\''
  x <- parseExpr
  return $ List [Atom "quote", x]

parseLists :: Parser LispVal
parseLists = try $ do
  _ <- char '('
  x <- try parseList <|> parseDottedList
  _ <- char ')'
  return x

-- quasiquote --

parseQuasi :: Parser LispVal
parseQuasi = try $ liftM (\x -> List [Atom "quasiquote", x]) (char '`' >> parseExpr)

parseUnQuote :: Parser LispVal
parseUnQuote = try $ liftM (\x -> List [Atom "unquote", x]) (char ',' >> parseExpr)

-- Vector --

parseVector :: Parser LispVal
parseVector = try $ do
  _ <- string "#("
  v <- sepBy parseExpr spaces
  _ <- char ')'
  return $ Vector $ listArray (0, (length v - 1)) v

-- Core Parser --

parseExpr :: Parser LispVal
parseExpr = parseNumber
        <|> parseBool
        <|> parseCharacter
        <|> parseLists
        <|> parseQuasi
        <|> parseUnQuote
        <|> parseVector
        <|> parseString
        <|> parseAtom
        <|> parseQuoted

parseExprs :: String -> Either ParseError LispVal
parseExprs = parse parseExpr "lisp"


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}

module Main (main) where

import qualified Data.Vector as V
import           Data.Vector.Unboxed.Deriving (derivingUnbox)
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as VGM
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Unboxed as VU
import qualified Numeric.LinearAlgebra as HMat

newtype Vec = Vec { unV :: HMat.Vector Double }
  deriving (Eq, Ord, Show)

-- Does crash with GHC 8.6.5

dims :: Int
dims = 3

data instance VU.MVector s Vec = MVector_Vec (VU.MVector s Double)
data instance VU.Vector    Vec = Vector_Vec  (VU.Vector    Double)

instance VGM.MVector VU.MVector Vec where
  basicLength (MVector_Vec v) = VGM.basicLength v `div` dims
  basicUnsafeSlice start size (MVector_Vec v) = MVector_Vec (VGM.basicUnsafeSlice (start * dims) (size * dims) v)
  basicOverlaps (MVector_Vec v1) (MVector_Vec v2) = VGM.basicOverlaps v1 v2
  basicUnsafeNew size = MVector_Vec <$> VGM.basicUnsafeNew (size * dims)
  basicInitialize (MVector_Vec vec) = VGM.basicInitialize vec
  basicUnsafeRead vec idx = fmap (Vec . VG.convert) $ VU.freeze sliceVals
    where MVector_Vec sliceVals = VGM.unsafeSlice idx 1 vec
  basicUnsafeWrite vec idx (Vec val) = VG.unsafeCopy sliceVals (VG.convert val)
    where MVector_Vec sliceVals = VGM.unsafeSlice idx 1 vec
  basicUnsafeCopy (MVector_Vec target) (MVector_Vec source) = VGM.basicUnsafeCopy target source
  basicUnsafeMove (MVector_Vec target) (MVector_Vec source) = VGM.basicUnsafeMove target source

instance VG.Vector VU.Vector Vec where
  basicUnsafeFreeze (MVector_Vec mvec) = Vector_Vec <$> VG.basicUnsafeFreeze mvec
  basicUnsafeThaw (Vector_Vec vec) = MVector_Vec <$> VG.basicUnsafeThaw vec
  basicLength (Vector_Vec vec) = VG.basicLength vec `div` dims
  basicUnsafeSlice idx size (Vector_Vec vec) = Vector_Vec $ VG.basicUnsafeSlice (idx * dims) (size * dims) vec
  basicUnsafeIndexM (Vector_Vec vec) idx = return $! Vec $ VS.convert $! VG.unsafeSlice (idx * dims) dims vec
  basicUnsafeCopy (MVector_Vec mvec) (Vector_Vec vec) = VG.unsafeCopy mvec vec
  elemseq _ x y = x `seq` y

instance VU.Unbox Vec

-- Does not crash:

-- $(derivingUnbox "Vec"
--   [t| Vec -> (Double, Double, Double) |]
--   [| \(Vec hm) -> let [a, b, c] = HMat.toList hm in (a, b, c) |]
--   [| \(a, b, c) -> Vec (HMat.fromList [a, b, c]) |]
--  )

newtype NormalVec = NormalVec { getNormalVec :: Vec }
  deriving (Eq, Ord, Show)

$(derivingUnbox "NormalVec"
  [t| NormalVec -> Vec |]
  [| getNormalVec |]
  [| NormalVec |]
 )

unboxedVectorVecRoundtrip :: V.Vector Vec -> Bool
unboxedVectorVecRoundtrip v =
  let
    !indexVector = VU.fromList [0,1,2]
    !v' = VG.map (v VG.!) indexVector
  in
    VG.convert v' == V.map (v VG.!) (VG.convert indexVector)

unboxedVectorNormalVecRoundtrip :: V.Vector NormalVec -> Bool
unboxedVectorNormalVecRoundtrip v =
  let
    !indexVector = VU.fromList [0,1,2]
    !v' = VG.map (v VG.!) indexVector
  in
    VG.convert v' == V.map (v VG.!) (VG.convert indexVector)

main :: IO ()
main = do
  putStrLn "Unboxed.Vector (Vec World) roundtrip (does not crash in repro)"

  let v = Vec (HMat.fromList [1,2,3])

  print $ unboxedVectorVecRoundtrip $
    V.fromList [ v, v, v ]

  putStrLn "Unboxed.Vector (NormalVec World) roundtrip (CRASHES in repro)"

  print $ unboxedVectorNormalVecRoundtrip $
    V.fromList [ NormalVec v, NormalVec v, NormalVec v ]


module Numeric.Loss where

import Dependent.Size
import Numeric.Vector.Sized

class Loss (size :: Size) fn where
  lossDerivative :: fn -> SizedArray size -> SizedArray size -> SizedArray size


{-# LANGUAGE TypeOperators #-}
module Types where

import           Data.Array.Repa
import           Data.Complex

type DIM6 = DIM5 :. Int

type R2S1Array   = Array U DIM3 (Complex Double) -- (Z :. numOrientations :. xLen :. yLen)
type R2S1RPArray = Array U DIM4 Double            -- (Z :. numOrientations :. numScales :. xLen :. yLen)
type R2S1T0Array = Array U DIM4 (Complex Double)  -- (Z :. numOrientations :. (L.length theta0freqs)  :. xLen :. yLen )
type R2Z1T0Array = Array U DIM4 (Complex Double) -- (Z :. (L.length thetafreqs) :. (L.length theta0Freqs) :. xLen :. yLen )

type R2Z1T0RadialArray = Array U DIM3 Double -- (Z :. (L.length thetafreqs) :. (L.length theta0Freqs) :. rLen)

type R2S1RPT0S0Array = Array U DIM6 (Complex Double)
-- (Z :. numOrientations :. numScales :. (L.length theta0freqs) :. (L.length scale0freqs)  :. xLen :. yLen )
type R2Z2T0S0Array = Array U DIM6 (Complex Double)
-- (Z :. (L.length thetafreqs)  :. (L.length scalefreqs) :. (L.length theta0Freqs) :. (L.length scale0Freqs) :. xLen :. yLen )

type R2Array   = Array U DIM2 (Complex Double) -- (Z :. xLen :. yLen)
type R2T0Array   = Array U DIM3 (Complex Double) -- (Z :. (L.length theta0freqs) :. xLen :. yLen)
type R2Z1Array   = Array U DIM3 (Complex Double) -- (Z :. (L.length thetafreqs) :. xLen :. yLen)
type R2Z2Array   = Array U DIM4 (Complex Double) -- (Z :. (L.length thetafreqs) :. (L.length scalefreqs) :. xLen :. yLen)

type R2T0S0Array   = Array U DIM4 (Complex Double) -- (Z :. (L.length theta0freqs) :. (L.length scale0freqs)  :. xLen :. yLen)

data R2S1RPPoint =
  R2S1RPPoint (Int, Int, Double, Double)
  deriving (Show, Read)

instance Eq R2S1RPPoint where
  (==) (R2S1RPPoint (x1, y1, _, _)) (R2S1RPPoint (x2, y2, _, _)) =
    x1 == x2 && y1 == y2


{-# LANGUAGE CPP #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE DeriveFunctor #-}

import Control.Monad ((>=>),forM)
import Control.Applicative (liftA2)
import Data.Complex
import Data.Either
import Data.Fixed
import Data.Functor.Const
import Data.Functor.Identity
import Data.HashMap.Strict
import Data.HashSet
import Data.Hashable
import Data.Int
import Data.IntMap.Strict
import Data.IntSet
import Data.Map
import Data.Monoid
import Data.Ord (Down(..))
import Data.Orphans ()
import Data.Proxy (Proxy(..))
import Data.Ratio
import Data.Semigroup
import Data.Semiring
import Data.Sequence
import Data.Set
import Data.Word
import GHC.Natural
import qualified Data.Foldable as F
import Prelude hiding (Num(..),(^))
import System.IO.Unsafe
import Test.QuickCheck hiding (Fixed)
import Test.QuickCheck.Gen (suchThat)
import Test.QuickCheck.Instances ()
import Test.Tasty (defaultMain, testGroup, TestTree)
import Text.Show.Functions ()
import qualified Data.Map as Map
import qualified Data.Sequence as Seq
import qualified GHC.Num as Num
import qualified Test.QuickCheck.Classes as QCC
import qualified Control.Exception as E

main :: IO ()
main = do
  QCC.lawsCheckMany namedTests
  F.fold pow_prop

forceError :: a -> IO (Either E.ErrorCall a)
forceError = E.try . E.evaluate

pow_prop :: [IO ()]
pow_prop = [quickCheck pow_prop1, quickCheck pow_prop2, quickCheck pow_prop3]

pow_prop1 :: Int -> Negative Int -> Property
pow_prop1 x y = ioProperty $ do
  p <- forceError (x ^ getNegative y)
  pure $ isLeft p

pow_prop2 :: Int -> Positive Int -> Property
pow_prop2 x y = ioProperty $ do
  p <- forceError (x ^ getPositive y)
  pure $ isRight p

pow_prop3 :: Int -> Property
pow_prop3 x = ioProperty $ do
  p <- forceError (x ^ 0)
  case p of
    Left e -> pure False
    Right x -> pure $ x == 1

type Laws = QCC.Laws

semiringLaws :: (Arbitrary a, Show a, Eq a, Semiring a) => Proxy a -> [Laws]
semiringLaws p = [QCC.semiringLaws p]

ringLaws :: (Arbitrary a, Show a, Eq a, Ring a) => Proxy a -> [Laws]
ringLaws p = [QCC.semiringLaws p, QCC.ringLaws p]

namedTests :: [(String, [Laws])]
namedTests =
  [ ("Bool", semiringLaws pBool)
--  , ("Double", ringLaws pDouble) -- needs to be within some epsilon
--  , ("Float", ringLaws pFloat)   -- needs to be within some epsilon
--  , ("Complex", ringLaws pComplex) -- needs to be within some epsilon
  , ("Int", ringLaws pInt)
  , ("Int8", ringLaws pInt8)
  , ("Int16", ringLaws pInt16)
  , ("Int32", ringLaws pInt32)
  , ("Int64", ringLaws pInt64)
  , ("Word", ringLaws pWord)
  , ("Word8", ringLaws pWord8)
  , ("Word16", ringLaws pWord16)
  , ("Word32", ringLaws pWord32)
  , ("Word64", ringLaws pWord64)
  , ("()", ringLaws pUnit)
  , ("Maybe", semiringLaws pMaybe)
  , ("PosRatio", semiringLaws pPosRatio)
  , ("IO", ringLaws pIO)
  , ("Fixed", ringLaws pFixed)
  , ("Identity", ringLaws pIdentity)
  , ("Dual", ringLaws pDual)
  , ("(->)", ringLaws pFunction)
  , ("Down", ringLaws pDown)
  , ("Const", ringLaws pConst)
--  , ("IntMap", semiringLaws pIntMap) -- needs newtypes
  , ("Set", semiringLaws pSet)
--  , ("IntSet", semiringLaws pIntSet) -- needs newtypes
  , ("HashSet", semiringLaws pHashSet)
  , ("HashMap", semiringLaws pHashMap)
  , ("Map", semiringLaws pMap)
  , ("Predicate", semiringLaws pPredicate)
  , ("Equivalence", semiringLaws pEquivalence)
  , ("Op", ringLaws pOp)
  , ("Ap", ringLaws pAp)

  , ("IntSet Sum", semiringLaws pIntSetSum)
  , ("IntSet Product", semiringLaws pIntSetProduct)
  , ("IntSet Min", semiringLaws pIntSetMin)
  , ("IntSet Max", semiringLaws pIntSetMax)

  , ("IntMap Sum", semiringLaws pIntMapSum)
  , ("IntMap Product", semiringLaws pIntMapProduct)
  , ("IntMap Min", semiringLaws pIntMapMin)
  , ("IntMap Max", semiringLaws pIntMapMax)
  ]

#if !(MIN_VERSION_base(4,12,0))
newtype Ap f a = Ap { getAp :: f a }
  deriving (Eq, Functor, Applicative, Show)
instance (Semiring a, Applicative f) => Semiring (Ap f a) where
  zero = pure zero
  one = pure one
  plus = liftA2 plus
  times = liftA2 times
#endif
deriving instance (Arbitrary (f a)) => Arbitrary (Ap f a)

newtype Predicate a = Predicate (a -> Bool)
  deriving (Eq, Show)
deriving instance Semiring (Predicate a)
deriving instance (Arbitrary a, CoArbitrary a) => Arbitrary (Predicate a)

newtype Equivalence a = Equivalence (a -> a -> Bool)
  deriving (Eq, Show)
deriving instance Semiring a => Semiring (Equivalence a)
deriving instance (Arbitrary a, CoArbitrary a) => Arbitrary (Equivalence a)

newtype Op a b = Op (b -> a)
  deriving (Eq, Show)
deriving instance Semiring a => Semiring (Op a b)
deriving instance Ring a => Ring (Op a b)
deriving instance (Arbitrary a, CoArbitrary a, CoArbitrary b) => Arbitrary (Op a b)

pAp = p @(Ap Identity Int)
pPredicate = p @(Predicate Int)
pEquivalence = p @(Equivalence Int)
pOp = p @(Op Int Int)

deriving instance Arbitrary a => Arbitrary (Down a)

-- | Rationals constrained to being positive
newtype PosRatio a = PosRatio { getPosRatio :: Ratio a }
  deriving (Eq, Show, Semiring)

instance (Arbitrary a, Integral a, Ord a) => Arbitrary (PosRatio a) where
  arbitrary =
    let nm = suchThat (arbitrary :: Gen a) (> 0)
        dm = suchThat (arbitrary :: Gen a) (> 0)
    in PosRatio <$> liftA2 (%) nm dm

p :: forall a. Proxy a
p = Proxy

instance (Arbitrary a, Eq b) => Eq (a -> b) where
  (==) f g =
    let x = generate arbitrary :: IO a
        x' = unsafePerformIO x
    in f x' == g x'

instance Show (Endo a) where
  show _ = "<<<Endo>>>"

deriving instance (Arbitrary a, Eq a) => Eq (Endo a)

-- no
instance Show (IO a) where
  show _ = "<<<IO action>>>"

-- no no
instance Eq a => Eq (IO a) where
  x == y = unsafePerformIO $ liftA2 (==) x y

-- this is fine
instance Arbitrary a => Arbitrary (IO a) where
  arbitrary = fmap pure (arbitrary :: Gen a)

deriving instance Hashable a => Hashable (Sum a)

deriving instance Arbitrary (IntSetOf a)
deriving instance Arbitrary v => Arbitrary (IntMapOf k v)

pBool = p @Bool
pDouble = p @Double
pFloat = p @Float
pInt = p @Int
pInt8 = p @Int8
pInt16 = p @Int16
pInt32 = p @Int32
pInt64 = p @Int64
pInteger = p @Integer
pNatural = p @Natural
pWord = p @Word
pWord8 = p @Word8
pWord16 = p @Word16
pWord32 = p @Word32
pWord64 = p @Word64
pUnit = p @()
pMaybe = p @(Maybe Int)
pPosRatio = p @(PosRatio Int)
pIO = p @(IO Int)
pComplex = p @(Complex Double)
pFixed = p @(Fixed E0)
pMin = p @(Min Int)
pMax = p @(Max Int)
pIdentity = p @(Identity Int)
pDual = p @(Dual Int)
pEndo = p @(Endo (Sum Int))
pSum  = p @(Sum Int)
pProduct = p @(Product Int)
pDown = p @(Down Int)
pSeq = p @(Seq Int)
pSet = p @(Set (Sum Int))
pHashSet = p @(HashSet (Sum Int))
pFunction = p @(Int -> Int)
pMap = p @(Map (Sum Int) Int)
pHashMap = p @(HashMap (Sum Int) Int)
pConst = p @(Const Int Int)
pAlt = p @(Alt Maybe Int)

pIntSetSum = p @(IntSetOf (Sum Int))
pIntSetProduct = p @(IntSetOf (Product Int))
pIntSetMin = p @(IntSetOf (Min Int))
pIntSetMax = p @(IntSetOf (Max Int))

pIntMapSum = p @(IntMapOf (Sum Int) Int)
pIntMapProduct = p @(IntMapOf (Product Int) Int)
pIntMapMin = p @(IntMapOf (Min Int) Int)
pIntMapMax = p @(IntMapOf (Max Int) Int)


{-# LANGUAGE ApplicativeDo #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE EmptyDataDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -Wall #-}

-- import FMeme.Funks
-- import Moo.GeneticAlgorithm

-- import Numeric.LinearAlgebra hiding (Range, append)

import Chart
-- import qualified Data.HashMap.Strict as HashMap

import Control.Category (id)
import Control.Lens
import qualified Data.Attoparsec.Text as A
import Data.Biapplicative
import Data.Scientific
import Data.Sv
import qualified Data.Sv.Decode as D
import Lucid hiding (b_)
import Network.JavaScript
import Network.Wai.Middleware.Static ((>->), addBase, noDots, staticPolicy)
import Numeric.Backprop
import Protolude hiding ((<<*>>), Meta, Rep, replace)
import System.IO.Unsafe (unsafePerformIO)
import Web.Page
import Web.Scotty
import qualified Data.Text as Text

data AB a = AB {alpha :: a, beta :: a} deriving (Eq, Show)

fromAB :: AB a -> [a]
fromAB (AB a b) = [a, b]

toAB :: [a] -> Maybe (AB a)
toAB (a : b : _) = Just (AB a b)
toAB _ = Nothing

repAB :: (Monad m) => Int -> AB Double -> AB (Range Double) -> SharedRep m (AB Double)
repAB steps (AB a0 b0) (AB (Range amin amax) (Range bmin bmax)) = do
  a' <- slider (Just "alpha") amin amax ((amax - amin) / fromIntegral steps) a0
  b' <- slider (Just "beta") bmin bmax ((amax - amin) / fromIntegral steps) b0
  pure $ AB a' b'

data Meta a = Meta {maxn :: Int, tol :: a, rate :: a} deriving (Eq, Show, Generic)

repMeta :: (Monad m) => Meta Double -> SharedRep m (Meta Double)
repMeta m = do
  n' <- sliderI (Just "max iterations") 1 10000 100 (m ^. #maxn)
  t' <- slider (Just "tolerance (t)") 0 0.1 0.001 (m ^. #tol)
  r' <- slider (Just "convergence rate (r)") 0 0.1 0.001 (m ^. #rate)
  pure $ Meta n' t' r'

ex1File :: FilePath
ex1File = "app/ex1data1.txt"

opts :: ParseOptions
opts = ParseOptions (fromIntegral $ fromEnum ' ') Unheaded

data Ex1Data = Ex1Data {v1 :: Double, v2 :: Double} deriving (Show, Eq)

toLR :: Ex1Data -> ([Double], Double)
toLR (Ex1Data x y) = ([x], y)

ex1Decoder :: Decode' ByteString Ex1Data
ex1Decoder = Ex1Data <$> D.double <*> D.double

ex1 :: IO [Ex1Data]
ex1 = do
  r <- parseDecodeFromFile ex1Decoder opts ex1File
  pure $ case r of
    Success xs -> xs
    Failure _ -> []

xy :: [([Double], Double)]
{-# NOINLINE xy #-}
xy = unsafePerformIO $ do
  ex1' <- ex1
  pure $ (\(Ex1Data x y) -> ([x], y)) <$> ex1'

repXY :: (Monad m) => [(Double, Double)] -> SharedRep m [Spot Double]
repXY d = do
  a <-
    dropdown
      A.takeText
      show
      (Just "type")
      [ "exercise 1"
      ]
      "exercise 1"
  pure
    ( case a of
        "exercise 1" -> uncurry SP <$> d
        _ -> uncurry SP <$> d
    )

abLine :: Range Double -> AB Double -> [Spot Double]
abLine (Range x0 x1) (AB a b) = [SP x0 (a + b * x0), SP x1 (a + b * x1)]

modelText :: AB Double -> Text
modelText (AB a b) =
  Text.pack $ "y = "
    <> formatScientific Exponent (Just 3) (fromFloatDigits a)
    <> " + "
    <> formatScientific Exponent (Just 3) (fromFloatDigits b)
    <> " * x"

midChart ::
  SharedRep IO Text ->
  Application ->
  Application
midChart sr = midShared sr initChartRender updateChart

initChartRender ::
  Engine ->
  Rep Text ->
  StateT (HashMap Text Text) IO ()
initChartRender e r =
  void $
    oneRep
      r
      ( \(Rep h fa) m -> do
          append e "input" (toText h)
          replace e "output" (either id id . snd $ fa m)
      )

updateChart :: Engine -> Either Text (HashMap Text Text, Either Text Text) -> IO ()
updateChart e (Left err) = append e "log" ("map error: " <> err)
updateChart e (Right (_, Left err)) = append e "log" ("parse error: " <> err)
updateChart e (Right (_, Right c)) = replace e "output" c

repMain ::
  (Monad m) =>
  ChartSvgStyle ->
  GlyphStyle ->
  HudConfig ->
  LineStyle ->
  LineStyle ->
  Meta Double ->
  [Ex1Data] ->
  SharedRep m Text
repMain cscfg gcfg hcfg lcfg gdcfg mcfg xs =
  bimap hmap mmap cs <<*>> gs <<*>> d <<*>> h <<*>> ls <<*>> ls2 <<*>> ab <<*>> meta
  where
    p0 = [0, 0]
    ab = repAB 100 (AB 0 0) (AB (Range -5 5) (Range -2 2))
    h = repHudConfigDefault hcfg
    cs = repChartSvgStyle cscfg
    gs = repGlyphStyle gcfg
    ls = repLineStyle lcfg
    ls2 = repLineStyle gdcfg
    meta = repMeta mcfg
    d = repXY ((\(Ex1Data x y) -> (x, y)) <$> xs)
    xrange = space1 $ v1 <$> xs :: Range Double
    mmap cs' gs' d' h' ls' ls2' ab' meta' =
      renderHudChartWith
        cs'
        h''
        [ Chart (GlyphA gs') d',
          Chart (LineA ls') (abLine xrange ab'),
          Chart (LineA ls2') (gdLine xrange gdp)
        ]
        <> renderChartSvg 400 400 (chartSvg (Rect -1.0 1.0 -1.0 1.0) (pixel'' (toLR <$> xs)))
      where
        gdp = linBatch meta' p0 (toLR <$> xs)
        h'' =
          h'
            & #hudTitles %~ (<> [modelTitle ab'] <> [gdTitle gdp])
    hmap cs' gs' d' h' ls' ls2' ab' meta' =
      accordion_
        "acca"
        (Just "Chart Configuration")
        [ ("Chart Svg Stylings", cs'),
          ( "Chart",
            accordion_
              "accb"
              Nothing
              [ ("Scatter Style", gs'),
                ("Line Style", ls'),
                ("GD Line Style", ls2'),
                ("Data", d')
              ]
          ),
          ("Chart Hud", h'),
          ("Meta", meta'),
          ("Model", ab')
        ]

svgStyle :: ChartSvgStyle
svgStyle =
  defaultChartSvgStyle
    & #sizex .~ 900
    & #sizey .~ 600
    & #chartAspect .~ 1.5
    & #outerPad .~ Nothing
    & #chartFrame .~ Nothing
    & #orig .~ Nothing

hudConfig :: HudConfig
hudConfig =
  defaultHudConfig
    & #hudCanvas .~ Nothing
    & #hudTitles
      .~ [ defaultTitle "Training Data",
           (defaultTitle "Profit in $10,0000s" :: Title Double)
             & #place .~ PlaceRight
             & #style . #size .~ 0.05,
           (defaultTitle "Population of City in 10,000s" :: Title Double)
             & #place .~ PlaceBottom
             & #style . #size .~ 0.05
         ]
    & #hudAxes
      .~ [ defaultAxisConfig
             & #abar .~ Just b,
           defaultAxisConfig
             & #abar .~ Just b'
             & #place .~ PlaceLeft
         ]
  where
    b =
      (defaultBar :: Bar Double)
        & #wid .~ (0.01 :: Double)
        & #buff .~ 0.01
        & #rstyle .~ rs
    b' =
      (defaultBar :: Bar Double)
        & #wid .~ (0.01 :: Double)
        & #buff .~ (0.01 :: Double)
        & #rstyle .~ rs
    rs = RectStyle 0 grey 0 (PixelRGB8 105 48 157) 0.5

glyphStyle :: GlyphStyle
glyphStyle =
  defaultGlyphStyle
    & #color .~ PixelRGB8 217 124 32
    & #borderOpacity .~ 1
    & #shape .~ SquareGlyph
    & #borderSize .~ 0.002

modelStyle :: LineStyle
modelStyle =
  defaultLineStyle
    & #color .~ PixelRGB8 173 33 23
    & #width .~ 0.005

gdStyle :: LineStyle
gdStyle =
  defaultLineStyle
    & #color .~ PixelRGB8 50 33 173
    & #width .~ 0.005

modelTitle :: AB Double -> Title Double
modelTitle ab =
  defaultTitle (modelText ab)
    & (#place .~ PlaceBottom :: Title Double -> Title Double)
    & #anchor .~ AnchorStart
    & #style . #color .~ PixelRGB8 173 33 23
    & #style . #size .~ 0.05

gdTitle :: Converge [Double] -> Title Double
gdTitle ab =
  defaultTitle (gdText ab)
    & #style . #color .~ PixelRGB8 50 33 173
    & #style . #size .~ 0.05
    & (#place .~ PlaceBottom :: Title Double -> Title Double)
    & #anchor .~ AnchorStart

gdLine :: Range Double -> Converge [Double] -> [Spot Double]
gdLine _ Nada = []
gdLine r (Unstable (x : y : _)) = abLine r (AB x y)
gdLine _ (Unstable _) = []
gdLine r (Convergent (x : y : _)) = abLine r (AB x y)
gdLine _ (Convergent _) = []
gdLine _ Divergent = []

gdText :: Converge [Double] -> Text
gdText Nada = "No Data"
gdText Divergent = "Divergent"
gdText (Unstable (x : y : _)) = modelText (AB x y) <> " (Unstable)"
gdText (Unstable _) = "Can't happen"
gdText (Convergent (x : y : _)) = modelText (AB x y) <> " (Converged)"
gdText (Convergent _) = "Can't happen"

testPage :: Text -> [(Text, Html ())] -> Page
testPage title' sections =
  bootstrapPage
    <> bridgePage
    & #htmlHeader .~ title_ "mlh charts"
    & #htmlBody
      .~ b_
        "container"
        ( mconcat
            [ b_ "row" (h1_ (toHtml title')),
              b_ "row" $ mconcat $ (\(t, h) -> b_ "col" (h2_ (toHtml t) <> with div_ [id_ t] h)) <$> sections
            ]
        )

main :: IO ()
main = do
  ex1' <- ex1
  scotty 3000 $ do
    middleware $ staticPolicy (noDots >-> addBase "other")
    middleware
      ( midChart
          ( repMain
              svgStyle
              glyphStyle
              hudConfig
              modelStyle
              gdStyle
              (Meta 1000 0.0001 0.01)
              ex1'
          )
      )
    servePageWith
      "/chart/ex1"
      defaultPageConfig
      ( testPage
          "exercise 1"
          [ ("input", mempty),
            ("output", mempty)
          ]
      )

-- backprop
type Model p a b = p -> a -> b

type ModelG p a b =
  forall z.
  Reifies z W =>
  BVar z p ->
  BVar z a ->
  BVar z b

linReg :: ModelG [Double] [Double] Double
linReg bs xs = sum (zipWith (*) (sequenceVar bs) (auto 1 : sequenceVar xs))

sqerr :: (Backprop p, Backprop b, Num b) => ModelG p a b -> a -> b -> p -> p
sqerr f x targ = gradBP $ \p -> (f p (auto x) - auto targ) ^ (2 :: Int)

sqerrb :: (Backprop p, Backprop b, Num b) => ModelG p a b -> a -> b -> p -> (b, p)
sqerrb f x targ = backprop $ \p -> (f p (auto x) - auto targ) ^ (2 :: Int)

sqerre :: (Num b) => ModelG p a b -> a -> b -> p -> b
sqerre f x targ = evalBP $ \p -> (f p (auto x) - auto targ) ^ (2 :: Int)

sqerrs :: (Foldable f, Functor f, Backprop p, Backprop b, Num b, Floating p) => ModelG [p] [a] b -> f ([a], b) -> [p] -> [p]
sqerrs f xs p = (/ fromIntegral (length xs)) <$> foldl' (zipWith (+)) (repeat 0) ((\(x, targ) -> sqerr f x targ p) <$> xs)

sgd :: (Backprop b, Backprop p, Num b, Num p) => [p] -> ModelG [p] [a] b -> [p] -> [([a], b)] -> [p]
sgd r f = foldl' (\p (x, y) -> zipWith (-) p (zipWith (*) r (sqerr f x y p)))

data Converge a = Nada | Unstable a | Convergent a | Divergent deriving (Eq, Show)

converge :: Int -> (a -> a -> Bool) -> (a -> Bool) -> [a] -> Converge a
converge _ _ _ [] = Nada
converge n p d (x : xs) = go 0 p d x xs
  where
    go :: Int -> (a -> a -> Bool) -> (a -> Bool) -> a -> [a] -> Converge a
    go _ _ d' x' [] = bool (Unstable x') Divergent (d' x')
    go k p' d' x0 (x1 : xs')
      | d' x1 = Divergent
      | k > n = Unstable x1
      | p' x0 x1 = Convergent x1
      | otherwise = go (k + 1) p' d' x1 xs'

batchgd ::
  (Backprop b, Backprop p, Num b, Floating p) =>
  Int ->
  ([p] -> [p] -> Bool) ->
  ([p] -> Bool) ->
  p ->
  ModelG [p] [a] b ->
  [p] ->
  [([a], b)] ->
  Converge [p]
batchgd n pr d r f p0 xs' =
  converge n pr d $
    iterate (\ps -> zipWith (-) ps ((r *) <$> sqerrs f xs' ps)) p0

linBatch :: Meta Double -> [Double] -> [([Double], Double)] -> Converge [Double]
linBatch m =
  batchgd
    (m ^. #maxn)
    (\p p' -> (< (m ^. #tol)) $ sum $ zipWith (\x y -> abs (x - y)) p p')
    (any isNaN)
    (m ^. #rate)
    linReg

pixel' :: (Point Double -> Double) -> AB (Range Double) -> [Chart Double]
pixel' f (AB (Range amin amax) (Range bmin bmax)) =
  (\(r, c) -> Chart (RectA (RectStyle 0 black 0 c 1)) [SpotRect r])
    <$> pixelate f (toRect $ SR amin amax bmin bmax) (Point 100 100) blue grey

fsqerr :: [([Double], Double)] -> Point Double -> Double
fsqerr xs (Point a b) = sqrt $ sum $ (\(x, y) -> sqerre linReg x y [a, b]) <$> xs

pixel'' :: [([Double], Double)] -> [Chart Double]
pixel'' xs = pixel' (fsqerr xs) (AB (Range -5 0) (Range 0 2))


{-# LANGUAGE DeriveGeneric     #-}
{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms   #-}
{-# LANGUAGE RecordWildCards   #-}

module Main where

import           BDSCOD.Aggregation         (AggregatedObservations (..),
                                             AggregationTimes,
                                             aggregateUnscheduledObservations,
                                             maybeAggregationTimes,
                                             pattern AggTimes)
import           BDSCOD.Llhd                (initLlhdState, llhdAndNB,
                                             logPdeStatistics, pdeGF,
                                             pdeStatistics)
import           BDSCOD.Types               (LogLikelihood (..),
                                             MCMCConfiguration (..), MWCSeed,
                                             NegativeBinomial (..), NumLineages,
                                             Observation (..),
                                             ObservedEvent (..),
                                             PDESolution (..), Parameters (..),
                                             packParameters, putLambda, putMu,
                                             putNus, putOmega, putPsi, putRhos,
                                             scheduledTimes, unpackParameters)
import           BDSCOD.Utility             (eventsAsObservations, invLogit)
import           Control.Monad              (liftM2, when)
import           Control.Monad.Except       (ExceptT, runExceptT, throwError)
import           Control.Monad.Reader       (ReaderT, asks, liftIO, runReaderT)
import qualified Data.Aeson                 as Json
import qualified Data.ByteString.Builder    as BBuilder
import qualified Data.ByteString.Lazy       as L
import           Data.ByteString.Lazy.Char8 (pack, singleton)
import qualified Data.Csv                   as Csv
import           Data.List                  (intercalate)
import           Data.Maybe                 (fromJust)
import qualified Data.Vector.Unboxed        as Unboxed
import qualified Epidemic.BDSCOD            as SimBDSCOD
import           Epidemic.Types.Events      (EpidemicEvent (..), asNewickString,
                                             maybeEpidemicTree,
                                             maybeReconstructedTree)
import           Epidemic.Types.Parameter   (AbsoluteTime (..), Probability,
                                             Rate, TimeDelta (..), Timed (..))
import           Epidemic.Types.Population  (Identifier (..), People (..),
                                             Person (..), numPeople)
import           Epidemic.Utility           ()
import qualified Epidemic.Utility           as SimUtil
import           GHC.Generics               (Generic)
import           Numeric.GSL.Minimization   (MinimizeMethod (NMSimplex2),
                                             minimizeV)
import           Numeric.LinearAlgebra      (dot)
import           Numeric.LinearAlgebra.Data (Vector (..), fromList, linspace,
                                             toList)
import           Numeric.MCMC.Metropolis    (Chain (..), asGenIO, chain')
import           System.Environment         (getArgs)
import           System.Random.MWC          (initialize)

-- | Record of the scheduled observation times.
data ScheduledTimes =
  ScheduledTimes
    { stRhoTimes :: [AbsoluteTime]
    , stNuTimes  :: [AbsoluteTime]
    }
  deriving (Show)

-- | These objects define the specifics of the evaluation of LLHD profiles. If a
-- point estimate is given, then that is the central point of the profiles,
-- otherwise the parameters are estimated first. In every case the natural death
-- rate is set to the true values of the simulation. The boolean value is true
-- when the unscheduled observations are to be aggregated, otherwise they are
-- left as they are.
data InferenceConfiguration =
  InferenceConfiguration
    { reconstructedTreeOutputFiles :: (FilePath, FilePath)
    , observationsOutputCsv        :: FilePath
    , llhdOutputCsv                :: FilePath
    , pointEstimatesCsv            :: FilePath
    , maybePointEstimate           :: Maybe Parameters
    , icMaybeTimesForAgg           :: Maybe ([AbsoluteTime], [AbsoluteTime])
    , icMaybeMCMCConfig            :: Maybe MCMCConfiguration
    }
  deriving (Show, Generic)

instance Json.FromJSON InferenceConfiguration

-- | This object configures the whole evaluation of this program and is to be
-- read in from a suitable JSON file.
--
--     * A CSV to write the whole simulation to
--     * The parameters to use for the simulation
--     * The duration of the simulation
--     * The bounds on the size of an acceptable simulation
--     * One inference configuration for the regular data (with and without
--     estimated parameters) and another for the aggregated data with estimated
--     parameters.
--     * A Boolean to toggle printing progress on and off with the
--     @isVerbosePutStrLn@ function.
data Configuration =
  Configuration
    { simulatedEventsOutputCsv :: FilePath
    , simulationParameters :: Parameters
    , simulationDuration :: TimeDelta
    , simulationSizeBounds :: (Int, Int)
    , inferenceConfigurations :: ( InferenceConfiguration
                                 , InferenceConfiguration
                                 , InferenceConfiguration)
    , isVerbose :: Bool
    , configSimulationSeed :: MWCSeed
    }
  deriving (Show, Generic)

instance Json.FromJSON Configuration

-- | Monad for running a simulation, it uses the ExceptT for error messages and
-- the ReaderT for holding program configuration.
type Simulation x = ReaderT Configuration (ExceptT String IO) x

-- | A convenience function for printing output if in verbose mode when in the
-- 'Simulation' monad.
ifVerbosePutStrLn :: String -> Simulation ()
ifVerbosePutStrLn msg =
  do
    beLoud <- asks isVerbose
    when beLoud $ liftIO (putStrLn msg)

-- | This type is used to indicate if parameters are the true ones used in the
-- simulation or estimates parameters and in the case of estimated parameters,
-- what sort of data was used to inform the estimate.
data AnnotatedParameter
  = TrueParameters Parameters
  | EstimatedParametersRegularData Parameters
  | EstimatedParametersAggregatedData Parameters
  deriving (Show, Eq)

-- | A BDSCOD simulation configuration based on the parameters in the
-- environment.
bdscodConfiguration = do
  simParams@(Parameters (pLambda, pMu, pPsi, Timed pRhos, pOmega, Timed pNus)) <-
    asks simulationParameters
  if pLambda > 0 && pMu > 0 && pPsi > 0 && pOmega > 0 && null pRhos && null pNus
    then do
      (TimeDelta simDur) <- asks simulationDuration
      let bdscodConfig =
            SimBDSCOD.configuration (AbsoluteTime simDur) (unpackParameters simParams)
      case bdscodConfig of
        Nothing       -> throwError "Could not construct BDSCOD configuration"
        (Just config) -> return config
    else throwError "Simulation parameters not acceptable for this program..."

-- | Check if the LTT ever returns to 1 after being larger than 1.
pMultipleOrigins :: [EpidemicEvent] -> Bool
pMultipleOrigins = go 1
  where
    numLineage :: People -> NumLineages
    numLineage = fromIntegral . numPeople
    go :: NumLineages -> [EpidemicEvent] -> Bool
    go _ []                         = False
    go n (Removal _ _:ees)          = (n < 3) || go (n-1) ees
    go n (Sampling _ _:ees)         = (n < 3) || go (n-1) ees
    go n (Occurrence _ _:ees)       = (n < 3) || go (n-1) ees
    go n (Catastrophe _ people:ees) = if m > 1 then go m ees else True where m = n - numLineage people
    go n (Disaster _ people:ees)    = if m > 1 then go m ees else True where m = n - numLineage people
    go n (Infection {}:ees)      = go (n+1) ees

-- | Simulate the actual epidemic making sure that the results are acceptable
-- before returning, otherwise try using a different seed.
simulateEpidemic seedInt bdscodConfig = do
  ifVerbosePutStrLn "Running simulateEpidemic..."
  genIO <- liftIO $ prngGen seedInt
  simEvents <-
    liftIO $ SimUtil.simulation' bdscodConfig SimBDSCOD.allEvents genIO
  (sizeLowerBound, sizeUpperBound) <- asks simulationSizeBounds
  if length simEvents > sizeLowerBound && length simEvents < sizeUpperBound && not (pMultipleOrigins simEvents)
    then do
      ifVerbosePutStrLn "simulated an acceptable epidemic..."
      simEventsCsv <- asks simulatedEventsOutputCsv
      ifVerbosePutStrLn $
        "\twriting events to: " ++
        simEventsCsv ++ "\n\tthere are " ++ show (length simEvents) ++ " events."
      liftIO $ L.writeFile simEventsCsv (Csv.encode simEvents)
      return simEvents
    else do
      ifVerbosePutStrLn $ "\tthere where multiple origins: " ++ show (pMultipleOrigins simEvents)
      ifVerbosePutStrLn $ "\trepeating the simulation with seed: " ++ show seedInt
      simulateEpidemic (succ seedInt) bdscodConfig

-- | Take a simulated epidemic and generate the observations, first with full
-- resolution of the event times and the true epidemic parameters, second with
-- the event times and the estimated parameters and third with the sampling
-- times aggregated as defined in the inference configuration. This is done all
-- at the same time because the raw epidemic events are needed to genereate the
-- observations and the aggregated observations.
observeEpidemicThrice ::
     [EpidemicEvent] -- ^ the raw simulated epidemic events
  -> Simulation ( (InferenceConfiguration, [Observation])
                , (InferenceConfiguration, [Observation])
                , (InferenceConfiguration, AggregatedObservations))
observeEpidemicThrice simEvents = do
  (regInfConfig, regInfConfig', aggInfConfig) <- asks inferenceConfigurations
  let maybeRegObs = eventsAsObservations <$> SimBDSCOD.observedEvents simEvents
      maybeAggTimes =
        icMaybeTimesForAgg aggInfConfig >>= uncurry maybeAggregationTimes
      maybeAggObs =
        liftM2 aggregateUnscheduledObservations maybeAggTimes maybeRegObs
      (reconNewickTxt, reconNewickCsv) =
        reconstructedTreeOutputFiles regInfConfig
      maybeNewickData =
        asNewickString (AbsoluteTime 0, Person (Identifier 1)) =<<
        maybeReconstructedTree =<< maybeEpidemicTree simEvents
  case maybeNewickData of
    Just (newickBuilder, newickMetaData) -> do
      ifVerbosePutStrLn $
        "Writing reconstructed tree Newick to " ++ reconNewickTxt
      liftIO $
        L.writeFile reconNewickTxt (BBuilder.toLazyByteString newickBuilder)
      ifVerbosePutStrLn $
        "Writing reconstructed tree tip labels to " ++ reconNewickCsv
      liftIO $ L.writeFile reconNewickCsv (Csv.encode newickMetaData)
    Nothing -> throwError "Could not reconstruct tree..."
  case (maybeRegObs, maybeAggObs) of
    (Just regObs, Just aggObs@(AggregatedObservations _ unboxedAggObs)) -> do
      ifVerbosePutStrLn "Writing the regular observations to CSV"
      liftIO $
        L.writeFile (observationsOutputCsv regInfConfig) (Csv.encode regObs)
      ifVerbosePutStrLn "Writing the aggregated observations to CSV"
      liftIO $
        L.writeFile
          (observationsOutputCsv aggInfConfig)
          (Csv.encode unboxedAggObs)
      return
        ( (regInfConfig, regObs)
        , (regInfConfig', regObs)
        , (aggInfConfig, aggObs))
    (Just _, Nothing) ->
      throwError "Could not evaluate aggregated observations..."
    (Nothing, Just _) -> throwError "Could not evaluate regular observations..."
    (Nothing, Nothing) ->
      throwError "Could not evaluate either set of observations..."

-- | Evaluate the NB posterior approximation of the prevalence for a single
-- point in parameter space.
recordPresentPrevalenceEstimate ::
     InferenceConfiguration
  -> [Observation]
  -> AnnotatedParameter
  -> Simulation ()
recordPresentPrevalenceEstimate InferenceConfiguration {..} obs centerParam =
  let (parametersUsed,singleParams) = case centerParam of
                                        (TrueParameters x) -> ("true_parameters_regular_data" :: L.ByteString,x)
                                        (EstimatedParametersRegularData x) -> ("estimated_parameters_regular_data",x)
                                        (EstimatedParametersAggregatedData x) -> ("estimated_parameters_aggregated_data",x)
      nBVal = pure (parametersUsed, snd $ llhdAndNB obs singleParams initLlhdState)
   in do ifVerbosePutStrLn $ "\twriting NB values to: " ++ pointEstimatesCsv
         liftIO $ L.writeFile pointEstimatesCsv (Csv.encode nBVal)

-- | Run the estimation of the prevalence (at present) on the given set of regular
-- (i.e., disaggregated) observations at the parameters used to simulate the
-- data set and write the result to file.
evaluateLLHD :: InferenceConfiguration
             -> [Observation] -- ^ the regular observations
             -> Simulation ()
evaluateLLHD infConfig obs = do
  ifVerbosePutStrLn "Running evaluateLLHD..."
  simParams <- asks simulationParameters
  ifVerbosePutStrLn "\tUsing non-aggregated data with the true parameters..."
  ifVerbosePutStrLn $ show simParams
  recordPresentPrevalenceEstimate infConfig obs (TrueParameters simParams)

-- | Using regular (i.e., disaggregated) observations, estimate the parameters
-- used in the simulation and run the estimation of the prevalence (at present)
-- using this estimate. Finally generate MCMC samples of the posterior
-- distribution to understand the likelihood surface.
--
-- __NOTE__ This uses the actual simulation parameters as a way to get the
-- scheduled observation times, /they are not used in the inference/. The
-- optimisation starts at a fixed initial condition.
--
estimateLLHD :: InferenceConfiguration
             -> [Observation] -- ^ the regular observations
             -> Simulation ()
estimateLLHD infConfig obs = do
  ifVerbosePutStrLn "Running estimateLLHD..."
  Parameters (_, deathRate, _, _, _, _) <- asks simulationParameters
  let mleParams@(Parameters (estBirthRate, _, _, _, _, _)) = estimateRegularParameters deathRate obs
      annotatedMLE = EstimatedParametersRegularData mleParams
  ifVerbosePutStrLn "\tUsing non-aggregated data the computed MLE is..."
  ifVerbosePutStrLn $ "\tThe estimated birth rates is" ++ show estBirthRate
  recordPresentPrevalenceEstimate infConfig obs annotatedMLE
  runUnscheduledObservationMCMC infConfig deathRate obs annotatedMLE

-- | Using __aggregated__ observations, estimate the parameters of the process
-- and run the estimation of the prevalence (at present) using this estimate.
-- Finally generate MCMC samples of the posterior distribution to understand the
-- likelihood surface.
estimateLLHDAggregated :: InferenceConfiguration
                       -> AggregatedObservations -- ^ the aggregated observations
                       -> Simulation ()
estimateLLHDAggregated infConfig (AggregatedObservations _ obs) = do
  ifVerbosePutStrLn "Running estimateLLHDAggregated..."
  Parameters (_, deathRate, _, _, _, _) <- asks simulationParameters
  let (Just (rhoTimes,nuTimes)) = icMaybeTimesForAgg infConfig
      schedTimes = ScheduledTimes {stRhoTimes = rhoTimes, stNuTimes = nuTimes}
      mleParams = estimateAggregatedParameters deathRate schedTimes obs
      annotatedMLE = EstimatedParametersAggregatedData mleParams
  ifVerbosePutStrLn "\tUsing aggregated data the computed MLE is..."
  ifVerbosePutStrLn $ "\t" ++ show mleParams
  recordPresentPrevalenceEstimate infConfig obs annotatedMLE
  runScheduledObservationMCMC infConfig deathRate schedTimes obs annotatedMLE

-- | Use GSL to estimate the MLE based on the observations given. This uses a
-- simplex method because it seems to be faster and more accurate than the
-- simulated annealing and does not require a derivative.
--
-- __NOTE__ This will fit the model assuming that we have __aggregated__
-- observations, i.e., it will crash if there are unscheduled samples in the
-- data.
--
-- __NOTE__ we fix the death rate to the true value, which is given as an
-- argument, because this is assumed to be known a priori also, since this is
-- being applied to aggregated data we can set the value of psi and omega to
-- zero.
--
-- __NOTE__ The @energyFunc@ uses a regularisation cost to prevent the
-- parameters from wandering off which can occur with smaller data sets.
--
estimateAggregatedParameters :: Rate -- ^ the known death rate
                             -> ScheduledTimes -- ^ the time of scheduled observations
                             -> [Observation] -- ^ the aggregated observations
                             -> Parameters
estimateAggregatedParameters deathRate ScheduledTimes{..} obs =
  let maxIters = 20
      desiredPrec = 1e-1
      initBox = fromList ([0.1, 0.1, 0.1] :: [Rate]) -- lambda, rho and nu (because mu, psi, omega are fixed)
      randInit = fromList $ exp <$> [-1,-2,-2]

      vecAsParams :: Vector Rate -> Maybe Parameters
      vecAsParams x =
        let [lnR1, lnP1, lnP2] = toList x
            p1 = invLogit lnP1
            p2 = invLogit lnP2
            r = exp lnR1
            rts = [(t, p1) | t <- stRhoTimes]
            nts = [(t, p2) | t <- stNuTimes]
            packedParams = packParameters (r, deathRate, 0, rts, 0, nts)
          in if r < 5
             then Just packedParams
             else Nothing

      negLlhd :: Parameters -> LogLikelihood
      negLlhd ps = negate . fst $ llhdAndNB obs ps initLlhdState

      energyFunc :: Vector Rate -> Double
      energyFunc x = maybe 1e6 negLlhd $ vecAsParams x

      (est, _) = minimizeV NMSimplex2 desiredPrec maxIters initBox energyFunc randInit
   in fromJust $ vecAsParams est

-- | This is the main entry point to the actual simulation study, it is its own
-- function to avoid stressing about configuration IO. Since this is within the
-- simulation monad it has access to all the configuration data and can perform
-- IO.
simulationStudy :: Simulation ()
simulationStudy = do
  bdscodConfig <- bdscodConfiguration
  prngSeed <- asks configSimulationSeed
  epiSim <- simulateEpidemic prngSeed bdscodConfig
  (regObs, regObs', aggObs) <- observeEpidemicThrice epiSim
  uncurry evaluateLLHD regObs
  uncurry estimateLLHD regObs'
  uncurry estimateLLHDAggregated aggObs

-- | We use an external file to configure the simulation so that it is easier to
-- try different parameter values and random seeds. The @simulationStudy@
-- function is real driver of this and runs in the @Simulation@ monad.
main :: IO ()
main = do
  configFilePath <- head <$> getArgs
  maybeConfig <- Json.decode <$> L.readFile configFilePath
  case maybeConfig of
    Nothing ->
      putStrLn $ "Could not get configuration from file: " ++ configFilePath
    Just config -> do
      result <- runExceptT (runReaderT simulationStudy config)
      case result of
        Left errMsg -> putStrLn errMsg
        Right _     -> return ()

-- | Use GSL to estimate the MLE based on the observations given. This uses a
-- simplex method because it seems to be faster and more accurate than the
-- simulated annealing.
--
-- __NOTE__ This will fit the model assuming that we have __disaggregated__
-- observations.
--
-- __NOTE__ we fix the death rate to the true value because
-- this is assumed to be known a priori.
--
estimateRegularParameters ::
     Rate -> [Observation] -> Parameters
estimateRegularParameters deathRate obs =
  let maxIters = 50
      desiredPrec = 1e-2
      -- the only parameters are lambda, psi and omega (because mu, rho and nu
      -- are fixed)
      initBox = fromList [0.1, 0.1, 0.1]

      randInit :: Vector Double
      randInit = fromList [-1.5,-2.5,-2.5]

      vecAsParams :: Vector Double -> Maybe Parameters
      vecAsParams x =
        let [lnR1, lnR2, lnR3] = toList x
            r1 = exp lnR1
            r2 = exp lnR2
            r3 = exp lnR3
            packedParams = packParameters (r1, deathRate, r2, [], r3, [])
            rNaught = r1 / (deathRate + r2 + r3)
        in if r1 < 5 && 1 < rNaught && rNaught < 6
              -- constrain the R0 to be less than 5 and a sensible birth rate.
           then Just packedParams
           else Nothing

      negLlhd :: Parameters -> LogLikelihood
      negLlhd ps = negate . fst $ llhdAndNB obs ps initLlhdState

      energyFunc :: Vector Double -> Double
      energyFunc x = maybe 1e6 negLlhd $ vecAsParams x

      est :: Vector Double
      (est, path) =
        minimizeV NMSimplex2 desiredPrec maxIters initBox energyFunc randInit

   in fromJust (vecAsParams est)

-- | Action to run an MCMC analysis based on regular unscheduled observations.
runUnscheduledObservationMCMC :: InferenceConfiguration
                              -> Rate
                              -> [Observation]
                              -> AnnotatedParameter
                              -> Simulation ()
runUnscheduledObservationMCMC InferenceConfiguration {..} deathRate obs (EstimatedParametersRegularData (Parameters (mleR1, _, mleR2, _, mleR3, _))) =
  case icMaybeMCMCConfig of
    (Just mcmcConfig) ->
      let numIters = mcmcNumIters mcmcConfig
          stepSd = mcmcStepSD mcmcConfig
          variableNames = ["lambda", "psi", "omega", "nbSize", "nbProb"]
          x0 = [mleR1, mleR2, mleR3]
          listAsParams [r1, r2, r3] =
            packParameters (r1, deathRate, r2, [], r3, [])
          logPost x = fst $ llhdAndNB obs (listAsParams x) initLlhdState
          prngSeed = mcmcSeed mcmcConfig
          maybeGenQuantityFunc = Just (\x -> snd $ llhdAndNB obs (listAsParams x) initLlhdState)
       in do ifVerbosePutStrLn "Running runUnscheduledObservationMCMC..."
             genIO <- liftIO $ prngGen prngSeed
             chainVals <-
               liftIO $ (asGenIO $ chain' numIters stepSd x0 logPost maybeGenQuantityFunc) genIO
             liftIO $
               L.writeFile
                 (mcmcOutputCSV mcmcConfig)
                 (chainAsByteString variableNames chainVals)
    Nothing -> ifVerbosePutStrLn "No MCMC configuration found!"

runScheduledObservationMCMC ::
     InferenceConfiguration
  -> Rate -- ^ the /known/ death rate
  -> ScheduledTimes
  -> [Observation]
  -> AnnotatedParameter
  -> Simulation ()
runScheduledObservationMCMC InferenceConfiguration {..} deathRate ScheduledTimes {..} obs (EstimatedParametersAggregatedData (Parameters (mleR1, _, _, _, _, _))) =
  case icMaybeMCMCConfig of
    (Just mcmcConfig) ->
      let numIters = mcmcNumIters mcmcConfig
          stepSd = mcmcStepSD mcmcConfig
          variableNames = ["lambda", "rho", "nu", "nbSize", "nbProb"]
          x0 = [mleR1, 0.5, 0.5]
          listAsParams [r1, p1, p2] =
            packParameters
              ( r1
              , deathRate
              , 0
              , [(rt, p1) | rt <- stRhoTimes]
              , 0
              , [(nt, p2) | nt <- stNuTimes])
          logPost :: [Double] -> Double
          logPost x@[r1, p1, p2] =
            if r1 > 0 && 0 < p1 && p1 < 1 && 0 < p2 && p2 < 1
              then fst $ llhdAndNB obs (listAsParams x) initLlhdState
              else (-1 / 0)
          prngSeed = mcmcSeed mcmcConfig
          maybeGenQuantityFunc =
            Just (\x -> snd $ llhdAndNB obs (listAsParams x) initLlhdState)
       in do ifVerbosePutStrLn "Running runScheduledObservationMCMC..."
             genIO <- liftIO $ prngGen prngSeed
             chainVals <-
               liftIO $
               (asGenIO $ chain' numIters stepSd x0 logPost maybeGenQuantityFunc)
                 genIO
             liftIO $
               L.writeFile
                 (mcmcOutputCSV mcmcConfig)
                 (chainAsByteString variableNames chainVals)
             return ()
    Nothing -> ifVerbosePutStrLn "No MCMC configuration found!"

-- | A bytestring representation of the MCMC samples.
chainAsByteString :: [String] -- ^ the names of the elements of the chain
                  -> [Chain [Double] NegativeBinomial] -- ^ the samples in the chain
                  -> L.ByteString
chainAsByteString varNames chainVals =
  let header = pack $ intercalate "," ("llhd" : varNames)
      nbParams cv =
        case chainTunables cv of
          Just (NegBinomSizeProb r p) -> [r, p]
          Just Zero                   -> [0, 0]
          Nothing                     -> [-1, -1]
      records =
        Csv.encode
          [chainScore cv : (chainPosition cv ++ nbParams cv) | cv <- chainVals]
      linebreak = singleton '\n'
   in mconcat [header, linebreak, records]

-- | A generator for random numbers from a seed.
prngGen seed = initialize (Unboxed.fromList [seed])


module SchemeParser.ParserSpec where

import Test.Hspec

import Data.Ratio ((%))
import Data.Complex (Complex(..))

import qualified Data.Map as M
import Text.Parsec (runParser)

import SchemeParser.Types
import SchemeParser.Parser
import SchemeParser.Printer

main :: IO ()
main = hspec spec

parse :: String -> Either () LispVal
parse input = case runParser parseExpr M.empty "scheme-test" input of
  Left err -> Left ()
  Right val -> Right val

spec :: Spec
spec =
  describe "parser" $ do
    it "parses atoms" $ do
      parse "atom" `shouldBe` Right (LAtom "atom")
      parse "atom123" `shouldBe` Right (LAtom "atom123")
      parse "atom!$%" `shouldBe` Right (LAtom "atom!$%")
      parse "!atom" `shouldBe` Right (LAtom "!atom")
      parse "!123" `shouldBe` Right (LAtom "!123")
      parse "!atom123" `shouldBe` Right (LAtom "!atom123")
      parse "!123atom" `shouldBe` Right (LAtom "!123atom")
      parse "!%123$atom%$" `shouldBe` Right (LAtom "!%123$atom%$")
    it "parses strings" $ do
      parse "\"string\"" `shouldBe` Right (LString "string")
      parse "\"string123\"" `shouldBe` Right (LString "string123")
      parse "\"string!$%\"" `shouldBe` Right (LString "string!$%")
      parse "\"!string\"" `shouldBe` Right (LString "!string")
      parse "\"!123\"" `shouldBe` Right (LString "!123")
      parse "\"!string123\"" `shouldBe` Right (LString "!string123")
      parse "\"!123string\"" `shouldBe` Right (LString "!123string")
      parse "\"!%123$string%$\"" `shouldBe` Right (LString "!%123$string%$")
      parse "\"string\\\"nested\\\"123\"" `shouldBe` Right (LString "string\"nested\"123")
      parse "\"string'nested singlequote'123\"" `shouldBe` Right (LString "string'nested singlequote'123")
    it "parses bools" $ do
      parse "#t" `shouldBe` Right (LBool True)
      parse "#f" `shouldBe` Right (LBool False)
    it "parses chars" $ do
      parse "#\\a" `shouldBe` Right (LChar 'a')
      parse "#\\ " `shouldBe` Right (LChar ' ')
      parse "#\\space" `shouldBe` Right (LChar ' ')
      parse "#\\newline" `shouldBe` Right (LChar '\n')
      parse "#\\n" `shouldBe` Right (LChar 'n')
    it "parses numbers" $ do
      parse "123" `shouldBe` Right (LNumber 123)
      parse "#d123" `shouldBe` Right (LNumber 123)
      parse "#b1010" `shouldBe` Right (LNumber 10)
      parse "#o123" `shouldBe` Right (LNumber 83)
      parse "#xff" `shouldBe` Right (LNumber 255)
    it "parses floats" $ do
      parse "123.123" `shouldBe` Right (LFloat 123.123)
      parse "0.123" `shouldBe` Right (LFloat 0.123)
      parse "123.0" `shouldBe` Right (LFloat 123.0)
    it "parses ratios" $ do
      parse "10/3" `shouldBe` Right (LRatio (10 % 3))
      parse "25/5" `shouldBe` Right (LRatio (25 % 5))
      parse "25/5" `shouldBe` Right (LRatio (5 % 1))
    it "parses complex numbers" $ do
      parse "1+1i" `shouldBe` Right (LComplex (1 :+ 1))
      parse "123+456i" `shouldBe` Right (LComplex (123 :+ 456))
      -- TODO: test floats
    it "parses lists" $ do
      parse "()" `shouldBe` Right (LList [])
      parse "(atom)" `shouldBe` Right (LList [LAtom "atom"])
      parse "(one two three)" `shouldBe` Right (LList [LAtom "one", LAtom "two", LAtom "three"])
      parse "(\"foo\")" `shouldBe` Right (LList [LString "foo"])
      parse "(\"one\" \"two\" \"three\")" `shouldBe` Right (LList [LString "one", LString "two", LString "three"])
      parse "(#\\a)" `shouldBe` Right (LList [LChar 'a'])
      parse "(#\\a #\\  #\\space #\\newline #\\n)" `shouldBe` Right (LList [LChar 'a', LChar ' ', LChar ' ', LChar '\n', LChar 'n'])
      parse "(#t)" `shouldBe` Right (LList [LBool True])
      parse "(1)" `shouldBe` Right (LList [LNumber 1])
      parse "(1 2 3)" `shouldBe` Right (LList [LNumber 1, LNumber 2, LNumber 3])
      parse "(#d1)" `shouldBe` Right (LList [LNumber 1])
      parse "(#d1 #d2 #d3)" `shouldBe` Right (LList [LNumber 1, LNumber 2, LNumber 3])
      parse "(#o1)" `shouldBe` Right (LList [LNumber 1])
      parse "(#o1 #o2 #o3)" `shouldBe` Right (LList [LNumber 1, LNumber 2, LNumber 3])
      parse "(#b1)" `shouldBe` Right (LList [LNumber 1])
      parse "(#b1 #b10 #b11)" `shouldBe` Right (LList [LNumber 1, LNumber 2, LNumber 3])
      parse "(#x1)" `shouldBe` Right (LList [LNumber 1])
      parse "(#x1 #x2 #x3)" `shouldBe` Right (LList [LNumber 1, LNumber 2, LNumber 3])
      parse "(12.34)" `shouldBe` Right (LList [LFloat 12.34])
      parse "(12.34 56.78 90.12)" `shouldBe` Right (LList [LFloat 12.34, LFloat 56.78, LFloat 90.12])
      parse "(12/34)" `shouldBe` Right (LList [LRatio (12 % 34)])
      parse "(12/34 56/78 90/12)" `shouldBe` Right (LList [LRatio (12 % 34), LRatio (56 % 78), LRatio (90 % 12)])
      parse "(12+34i)" `shouldBe` Right (LList [LComplex (12 :+ 34)])
      parse "(12+34i 56+78i 90+12i)" `shouldBe` Right (LList [LComplex (12 :+ 34), LComplex (56 :+ 78), LComplex (90 :+ 12)])
      -- TODO: test complex with floats
      -- TODO: test more stuff


module Generators
  (
  -- * Molecular Generators
    generateOneMolecule
  , generateManyMolecules
  ) where

import Control.Lens
import System.Directory
import DistanceGeometry
import ReadWrite
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data
import Types

generateOneMolecule ::
     FilePath -- ^ Input directory
  -> FilePath -- ^ Output dorectory
  -> FilePath -- ^ Molecule
  -> Double -- ^ Error
  -> IO ()
generateOneMolecule idir odir mol err = generateManyMolecules idir odir mol err 1

generateManyMolecules ::
     FilePath -- ^ Input directory
  -> FilePath -- ^ Output directory
  -> FilePath -- ^ Molecule
  -> Double -- ^ Error
  -> Int -- ^ The number of molecules generated
  -> IO ()
generateManyMolecules idir odir mol err n = do
  let odir' = odir ++ "/" ++ mol
  doesPathExist odir' >>= (`createDirectoryIfMissing` odir')
  step odir' 1
  where
    step odir s
      | s > n = return ()
      | otherwise = do
        let (atoms, bonds) = readMolV2000 (idir ++ "/" ++ mol ++ ".mol")
        let s0@(u, l) = (triangleSmooth . generateDistBoundsMatr atoms) bonds
        s1 <- randomDistMatr s0
        let s2 = (generateCoorFromEigValVec . largestEigValVec . distMatrToMetricMatr) s1
        let derr = distanceErrorFunction s2 u l bonds
            cerr = chiralErrorFunction s2 u l bonds
        if derr + cerr > err
          then step odir s
          else do
            -- Write result molecule in output files
            let newmole = updateCoord s2 atoms
            writeXYZ (odir ++ "/" ++ mol ++ "-" ++ show s ++ ".xyz") "Comment" newmole

            -- Write logs in output files
            let logs =
                 show newmole ++ "\n" ++
                 "Distance error = " ++ show derr ++ "\n" ++
                 "Chiral error = " ++ show cerr ++ "\n"
            writeFile (odir ++ "/" ++ mol ++ "-" ++ show s ++ ".log") logs

            -- Print to console
            putStrLn ("generateManyMolecules: " ++ mol ++ ": step " ++ show s ++ ": OK!")
            step odir (s+1)