{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE CPP                  #-}
{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}
{-|
Module      : Grenade.Core.Shape
Description : Dependently typed shapes of data which are passed between layers of a network
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental

-}
module Grenade.Core.Shape (
    S (..)
  , Shape (..)
#if MIN_VERSION_singletons(2,6,0)
  , SShape (..)
#else
  , Sing (..)
#endif

  , randomOfShape
  , fromStorable
  , fromStorableV
  , isVectorShape
  ) where

#if MIN_VERSION_singletons(2,6,0)
import           Data.Kind                    (Type)
import           Data.Singletons.TypeLits     (SNat (..))
#endif

import           Control.DeepSeq              (NFData (..))
import           Data.Maybe                   (fromMaybe)
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons
import           Data.Singletons.TypeLits
import           Data.Vector.Storable         (Vector)
import qualified Data.Vector.Storable         as V
import           GHC.TypeLits                 hiding (natVal)
import qualified Numeric.LinearAlgebra        as NLA
import           Numeric.LinearAlgebra.Static hiding (zipWithVector)
import qualified Numeric.LinearAlgebra.Static as H
import           System.Random.MWC
import           Unsafe.Coerce                (unsafeCoerce)

import           Debug.Trace

import           Grenade.Types
import           Grenade.Utils.Vector

-- | The current shapes we accept.
--   at the moment this is just one, two, and three dimensional
--   Vectors/Matricies.
--
--   These are only used with DataKinds, as Kind `Shape`, with Types 'D1, 'D2, 'D3.
data Shape
  = D1 Nat
  -- ^ One dimensional vector
  | D2 Nat Nat
  -- ^ Two dimensional matrix. Row, Column.
  | D3 Nat Nat Nat
  -- ^ Three dimensional matrix. Row, Column, Channels.

-- | Concrete data structures for a Shape.
--
--   All shapes are held in contiguous memory.
--   3D is held in a matrix (usually row oriented) which has height depth * rows.
data S (n :: Shape) where
  S1D :: ( KnownNat len )
      => R len
      -> S ('D1 len)

  S2D :: ( KnownNat rows, KnownNat columns )
      => L rows columns
      -> S ('D2 rows columns)

  S3D :: ( KnownNat rows
         , KnownNat columns
         , KnownNat depth
         , KnownNat (rows * depth))
      => L (rows * depth) columns
      -> S ('D3 rows columns depth)

  -- HMatrix instances
  S1DV :: (KnownNat len) -- ^ Always a row-vector
       => !(V.Vector RealNum)
       -> S ('D1 len)

  -- HMatrix instances
  S2DV :: (KnownNat rows, KnownNat columns)
       => !(V.Vector RealNum) -- ^ Vector in row or column major accoring to @order@ in @Grenade.Utils.Conversion@. Use Data.Matrix here?
       -> S ('D2 rows columns)

  -- HMatrix instances
  -- S3DH :: (KnownNat rows, KnownNat columns, KnownNat depth, KnownNat (rows * depth))
  --      => HBLAS.MDenseMatrix RealWorld 'HBLAS.Row Double
  --      -> S ('D3 rows columns depth)

isVectorShape :: S x -> Bool
isVectorShape S1DV{} = True
isVectorShape S2DV{} = True
isVectorShape _      = False

instance Show (S n) where
  show (S1D x) = "S1D " ++ show x
  show (S2D x) = "S2D " ++ show x
  show (S3D x) = "S3D " ++ show x
  show (S1DV x) = "S1DV " ++ show x
  show inp@(S2DV x) = "S2DV" ++ show (sz inp) ++ " " ++ show x
    where
      sz ::
           forall rows cols. (KnownNat rows, KnownNat cols)
        => (S ('D2 rows cols))
        -> (Int, Int)
      sz _ =
        let rows = fromIntegral $ natVal (Proxy :: Proxy rows)
            cols = fromIntegral $ natVal (Proxy :: Proxy cols)
         in (rows, cols)

-- Singleton instances.
--
-- These could probably be derived with template haskell, but this seems
-- clear and makes adding the KnownNat constraints simple.
-- We can also keep our code TH free, which is great.
#if MIN_VERSION_singletons(2,6,0)
-- In singletons 2.6 Sing switched from a data family to a type family.
type instance Sing = SShape

data SShape :: Shape -> Type where
  D1Sing :: Sing a -> SShape ('D1 a)
  D2Sing :: Sing a -> Sing b -> SShape ('D2 a b)
  D3Sing :: KnownNat (a * c) => Sing a -> Sing b -> Sing c -> SShape ('D3 a b c)
#else
data instance Sing (n :: Shape) where
  D1Sing :: Sing a -> Sing ('D1 a)
  D2Sing :: Sing a -> Sing b -> Sing ('D2 a b)
  D3Sing :: KnownNat (a * c) => Sing a -> Sing b -> Sing c -> Sing ('D3 a b c)
#endif

instance KnownNat a => SingI ('D1 a) where
  sing = D1Sing sing
instance (KnownNat a, KnownNat b) => SingI ('D2 a b) where
  sing = D2Sing sing sing
instance (KnownNat a, KnownNat b, KnownNat c, KnownNat (a * c)) => SingI ('D3 a b c) where
  sing = D3Sing sing sing sing

instance SingI x => Num (S x) where
  (+) = n2 (+)
  (-) = n2 (-)
  (*) = n2 (*)
  abs = n1 abs
  signum = n1 signum
  fromInteger x = nk (fromInteger x)

instance SingI x => Fractional (S x) where
  (/) = n2 (/)
  recip = n1 recip
  fromRational x = nk (fromRational x)

instance SingI x => Floating (S x) where
  pi = nk pi
  exp = n1 exp
  log = n1 log
  sqrt = n1 sqrt
  (**) = n2 (**)
  logBase = n2 logBase
  sin = n1 sin
  cos = n1 cos
  tan = n1 tan
  asin = n1 asin
  acos = n1 acos
  atan = n1 atan
  sinh = n1 sinh
  cosh = n1 cosh
  tanh = n1 tanh
  asinh = n1 asinh
  acosh = n1 acosh
  atanh = n1 atanh

--
-- I haven't made shapes strict, as sometimes they're not needed
-- (the last input gradient back for instance)
--
instance NFData (S x) where
  rnf (S1D x)   = rnf x
  rnf (S2D x)   = rnf x
  rnf (S3D x)   = rnf x
  rnf (S1DV !v) = rnf v
  rnf (S2DV !v) = rnf v
  -- rnf (S3DH !_) = ()

-- | Generate random data of the desired shape
randomOfShape :: forall x . (SingI x) => IO (S x)
randomOfShape = do
  seed :: Int <- withSystemRandom . asGenST $ \gen -> uniform gen
  return $ case (sing :: Sing x) of
    D1Sing SNat ->
        S1D (randomVector seed Uniform * 2 - 1)

    D2Sing SNat SNat ->
        S2D (uniformSample seed (-1) 1)

    D3Sing SNat SNat SNat ->
        S3D (uniformSample seed (-1) 1)

-- | Generate a shape from a Storable Vector.
--
--   Returns Nothing if the vector is of the wrong size.
fromStorable :: forall x. SingI x => Vector RealNum -> Maybe (S x)
fromStorable xs =
  case (sing :: Sing x) of
    D1Sing SNat           -> S1D <$> H.create xs
    D2Sing SNat SNat      -> S2D <$> mkL xs
    D3Sing SNat SNat SNat -> S3D <$> mkL xs
  where
    mkL ::
         forall rows columns. (KnownNat rows, KnownNat columns)
      => Vector RealNum
      -> Maybe (L rows columns)
    mkL v =
      let rows = fromIntegral $ natVal (Proxy :: Proxy rows)
          columns = fromIntegral $ natVal (Proxy :: Proxy columns)
       in if rows * columns == V.length v
            then H.create $ NLA.reshape columns v
            else Nothing

fromStorableV :: forall x. SingI x => Vector RealNum -> S x
fromStorableV v =
  case (sing :: Sing x) of
    D1Sing SNat           -> S1DV v
    D2Sing SNat SNat      -> S2DV v
    D3Sing SNat SNat SNat -> error "unexpected case in fromStorableS"

instance SingI x => Serialize (S x) where
  put i =
    case i of
      S1D x -> put (1 :: Int) >> (putListOf put . NLA.toList . H.extract $ x)
      S2D x -> put (1 :: Int) >> (putListOf put . NLA.toList . NLA.flatten . H.extract $ x)
      S3D x -> put (1 :: Int) >> (putListOf put . NLA.toList . NLA.flatten . H.extract $ x)
      S1DV x -> put (2 :: Int) >> putListOf put (V.toList x)
      S2DV x -> put (2 :: Int) >> putListOf put (V.toList x)
  get = do
    (nr :: Int) <- get
    case nr of
      1 -> do
        Just i <- fromStorable . V.fromList <$> getListOf get
        return i
      2 -> fromStorableV . V.fromList <$> getListOf get
      _ -> error "unexpected case in get in Serialize instance in Shape.hs"

-- Helper function for creating the number instances
n1 :: ( forall a. Floating a => a -> a ) -> S x -> S x
n1 f (S1D x)  = S1D (f x)
n1 f (S2D x)  = S2D (f x)
n1 f (S3D x)  = S3D (f x)
n1 f (S1DV x) = S1DV (mapVector f x)
n1 f (S2DV x) = S2DV (mapVector f x)

-- helper function for creating the number instances
n2 :: ( forall a. Floating a => a -> a -> a ) -> S x -> S x -> S x
n2 f (S1D x) (S1D y)    = S1D (f x y)
n2 f (S2D x) (S2D y)    = S2D (f x y)
n2 f (S3D x) (S3D y)    = S3D (f x y)
n2 f (S1DV x) (S1DV y)  = S1DV (zipWithVector f x y)
n2 f (S2DV x) (S2DV y)  = S2DV (zipWithVector f x y)
n2 f (S1D x) y@S1DV {}  = n2 f (S1DV $ extract x) y
n2 f x@S1DV {} (S1D y)  = n2 f x (S1DV $ unsafeCoerce y)
n2 f x@S2D{} y@(S2DV _) = n2 f (toS2DV x) y
n2 f x@(S2DV _) y@S2D{} = n2 f x (toS2DV y)

toS2DV :: S ('D2 i j) -> S ('D2 i j)
toS2DV (S2D x)  = S2DV $ V.concat $ map H.extract . H.toColumns $ x
toS2DV (S2DV x) = S2DV x

-- Helper function for creating the number instances
nk :: forall x. SingI x => RealNum -> S x
nk x = case (sing :: Sing x) of
  D1Sing SNat ->
    S1D (konst x)

  D2Sing SNat SNat ->
    S2D (konst x)

  D3Sing SNat SNat SNat ->
    S3D (konst x)


{-# LANGUAGE CPP              #-}
{-# LANGUAGE DataKinds        #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs            #-}
{-# LANGUAGE TemplateHaskell  #-}
{-# LANGUAGE TypeOperators    #-}

module Test.Grenade.Sys.Utils where

import           Numeric.LinearAlgebra        hiding (R, konst, randomVector,
                                               uniformSample, (===))
import qualified Numeric.LinearAlgebra        as LA
import qualified Numeric.LinearAlgebra.Data   as D
import           Numeric.LinearAlgebra.Static (L, R)
import qualified Numeric.LinearAlgebra.Static as H

import           Hedgehog
import qualified Hedgehog.Gen                 as Gen
import qualified Hedgehog.Range               as Range

import           Test.Hedgehog.Compat
import           Test.Hedgehog.Hmatrix

import           GHC.TypeLits
import           Data.List
import           Data.Maybe
import           Data.Function
import           Data.Either
import           Control.Monad                (guard)
import           System.FilePath
import qualified Data.ByteString.Lazy as BS
import           Data.Binary

import           Grenade
import           Grenade.Utils.PascalVoc
import           Grenade.Utils.ImageNet

imagesDir :: FilePath
imagesDir = (takeDirectory __FILE__) </> "Images"

loadSerializedImage :: (KnownNat d, KnownNat c, KnownNat (d * c), c ~ 3) => FilePath -> IO (Maybe (S ('D3 d d c)))
loadSerializedImage path = do
  bs <- BS.readFile path
  let mat = decode bs :: LA.Matrix Double
  return $ S3D <$> (H.create (D.cmap doubleToRealNum mat))

loadSerializedChannel :: (KnownNat d, KnownNat c, KnownNat (d * c), c ~ 1) => FilePath -> IO (Maybe (S ('D3 d d c)))
loadSerializedChannel path = do
  bs <- BS.readFile path
  let mat = decode bs :: LA.Matrix Double
  return $ S3D <$> (H.create (D.cmap doubleToRealNum mat))

-- Performs linear regression on data points, where
-- the x-values consist of the data set [1..n]
linearRegression :: Int -> [Double] -> Double
linearRegression n ys = m
  where
    n'    = fromIntegral n :: Double
    xs    = [1..n] :: [Int]
    x     = (n' + 1) / 2
    x2    = (fromIntegral $ sum (map (^ 2) xs)) / n'
    xyNum = (sum $ zipWith (\l i -> l * (fromIntegral i)) ys xs) :: Double
    xy    = xyNum / n' :: Double
    y     = (sum ys) / n'
    m     = (xy - (x * y)) / (x2 - (x * x)) :: Double


{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}

-- Import all models under maintenance.
-- Models not imported here will not be compiled
-- when invoking `stack bench`.
import qualified BetaBin
import Control.Applicative ((<$>))
import Control.Arrow (first, second)
import Control.Monad (unless, when)
import Control.Monad.Bayes.Class
import Control.Monad.Bayes.Enumerator
import Control.Monad.Bayes.Inference
import Control.Monad.Bayes.Population
import Control.Monad.Bayes.Sampler
import Control.Monad.Bayes.Simple
import Control.Monad.Bayes.Weighted
import Control.Monad.IO.Class (MonadIO, liftIO)
import qualified DPmixture
import qualified Data.Vector as Vector
import qualified Dice
import qualified Gamma
import Graphics.Rendering.Chart.Backend.Cairo
import Graphics.Rendering.Chart.Easy
import qualified HMM
import Numeric.LogDomain
import Options.Applicative
import Plotting
import Statistics.Sample
import System.Directory
import System.IO

opts :: ParserInfo Bool
opts =
  flip info fullDesc $
    switch
      ( long "trial"
          <> help "Run a quick version of benchmarks to check that all is working correctly"
      )

tryCache :: (MonadIO m, Read a, Show a) => FilePath -> m a -> m a
tryCache filepath fresh = do
  exists <- liftIO $ doesFileExist filepath
  if exists
    then liftIO $ fmap read (readFile filepath)
    else do
      value <- fresh
      liftIO $ writeFile filepath (show value)
      return value

main = do
  -- make sure `putStrLn` prints to console immediately
  hSetBuffering stdout LineBuffering
  trial <- execParser opts
  when trial $ putStrLn "Trial run"
  sampleIO hmmBenchmark

meanVar :: (MonadDist m, CustomReal m ~ Double) => Int -> m Double -> m (Double, Double)
meanVar n d = do
  xs <- Vector.replicateM n d
  return $ meanVariance xs

smcParams :: [Int]
smcParams = [10, 20, 50, 100, 200, 500, 1000]

ns :: [Int]
--ns = [10,20,50,100,200,500,1000]
ns = [10, 20 .. 1000]

smcParamsDouble :: [Double]
smcParamsDouble = map fromIntegral smcParams

smcResults :: (MonadDist m, CustomReal m ~ Double) => [m (Vector.Vector Double)]
smcResults = map (\p -> Vector.replicateM 10 $ fmap HMM.hmmKL $ explicitPopulation $ smcMultinomial (length HMM.values) p HMM.hmm) smcParams

hmmBenchmark :: SamplerIO ()
hmmBenchmark = do
  liftIO $ putStrLn "running HMM benchmark"
  isSamples <- fmap (drop 5000) $ explicitPopulation $ importance 10000 HMM.hmm
  let isRes = map (\n -> HMM.hmmKL $ take n isSamples) ns
  -- mhSamples <- fmap (drop 5000) $ traceMH 10000 HMM.hmm
  -- let mhRes = map (\n -> HMM.hmmKL $ take n $ map (,1) mhSamples) ns
  mhPriorSamples <- drop 5000 <$> mhPrior HMM.hmm 10000
  let mhPriorRes = map (\n -> HMM.hmmKL $ take n $ map (,1) mhPriorSamples) ns
  pimhSamples <- pimh (length HMM.values) 100 1000 HMM.hmm
  let pimhRes = map (\n -> HMM.hmmKL $ take n $ map (,1) pimhSamples) ns
  liftIO $ toFile (fo_format .~ PDF $ def) "anytime.pdf" $ do
    layout_title .= "HMM"
    anytimePlot
      "#samples"
      "KL"
      ns
      [ ("IS", isRes),
        -- ("MHtrace", mhRes),
        ("MHprior", mhPriorRes),
        ("PIMH", pimhRes)
      ]
  smcRes <- sequence smcResults
  liftIO $ toFile (fo_format .~ PDF $ def) "smc.pdf" $ do
    layout_title .= "HMM"
    oneShotPlot
      "#particles"
      "KL"
      [ ("SMC", zip smcParamsDouble (map (errBars 2) smcRes))
      ]


{-# LANGUAGE CPP #-}

module Main (main) where

import Data.Complex (Complex ((:+)))
import Data.Complex.Polar
#ifdef __GLASGOW_HASKELL__
import GHC.Stack (HasCallStack)
#endif
import System.Exit
  ( ExitCode (ExitFailure, ExitSuccess),
    exitWith,
  )
import Test.HUnit
  ( Assertion,
    Test (TestCase, TestList),
    assertBool,
    assertEqual,
    errors,
    failures,
    runTestTT,
  )

tests :: Test
tests =
  TestList
    [ TestCase $ rAndThetaAre 1 (0.1 - pi) $ mkPolar (- 1) 0.1,
      TestCase $ rAndThetaAre 1 (pi / 2) $ mkPolar 1 (-3 * pi / 2),
      TestCase $ rAndThetaAre 1 (0.5 - pi) $ mkPolar 1 (pi + 0.5),
#if __GLASGOW_HASKELL__ >= 736
      TestCase $ rAndThetaAre 1 (0.1 - pi) $ (- 1) :< 0.1,
      TestCase $ rAndThetaAre 1 (pi / 2) $ 1 :< (- 3 * pi / 2),
      TestCase $ rAndThetaAre 1 (0.5 - pi) $ 1 :< (pi + 0.5),
#endif
      TestCase $ rAndThetaAre 1 0.3 $ cis 0.3,
      TestCase $ rAndThetaAre 1 (0.3 - pi) $ cis (pi + 0.3),
      TestCase $ rAndThetaAre 0 0 $ mkPolar 0 pi,
      TestCase $ rAndThetaAre 1 pi $ signum (mkPolar 1 pi),
      TestCase $ rAndThetaAre 0 0 $ signum (mkPolar 0 pi),
      TestCase $ rAndThetaAre 1 (-3 * pi / 4) $ negate (mkPolar 1 (pi / 4)),
      TestCase $ rAndThetaAre 1 pi $ fromRational (- 1),
      TestCase $ rAndThetaAre 1 0 $ exp $ mkPolar (10 * pi) (pi / 2),
      TestCase $ rAndThetaAre (sqrt 2) (pi / 4) $ fromComplex $ 1 :+ 1,
      TestCase $ isApproximately "realPart" 1 $ realPart $ mkPolar 1 0,
      TestCase $ isApproximately "imagPart" 1 $ imagPart $ mkPolar 1 (pi / 2),
      TestCase $ rAndThetaAre 1 (- pi / 4) $ conjugate $ mkPolar 1 (pi / 4),
      TestCase $
        rAndThetaAre (sqrt 2) (pi / 4) $
          mkPolar 1 0 + mkPolar 1 (pi / 2),
      TestCase $
        rAndThetaAre (sqrt 2) (- pi / 4) $
          mkPolar 1 0 - mkPolar 1 (pi / 2),
      TestCase $ rAndThetaAre 1 0.7 $ mkPolar 1 0.3 * mkPolar 1 0.4,
      TestCase $ rAndThetaAre (sqrt 2) 0 $ abs $ mkPolar (sqrt 2) 0.3,
      TestCase $ rAndThetaAre 1 0 1, -- Test Num fromInteger
      TestCase $ rAndThetaAre 1 0 $ mkPolar 1 0 * mkPolar 2 0.2 / mkPolar 2 0.2,
      TestCase $ rAndThetaAre pi 0 pi, -- Test Floating pi
      TestCase $ rAndThetaAre 54.5984 2.00000 $ exp z,
      TestCase $ rAndThetaAre 1.56798 0.30018 $ log z,
      TestCase $ rAndThetaAre 2.11474 0.23182 $ sqrt z,
      TestCase $ rAndThetaAre 3.70497 (- 2.4472) $ sin z,
      TestCase $ rAndThetaAre 3.68529 2.30135 $ cos z,
      TestCase $ rAndThetaAre 1.00534 1.53455 $ tan z,
      TestCase $ rAndThetaAre 27.3050 1.99974 $ sinh z,
      TestCase $ rAndThetaAre 27.2930 2.00025 $ cosh z,
      TestCase $ rAndThetaAre 1.00043 (- 0.0005) $ tanh z,
      TestCase $ rAndThetaAre 2.44362 1.10527 $ asin z,
      TestCase $ rAndThetaAre 2.23441 (- 1.3570) $ acos z,
      TestCase $ rAndThetaAre 1.37491 0.07018 $ atan z,
      TestCase $ rAndThetaAre 2.24493 0.20357 $ asinh z,
      TestCase $ rAndThetaAre 2.23441 0.21370 $ acosh z,
      TestCase $ rAndThetaAre 1.48069 1.43491 $ atanh z,
      TestCase $ assertBool "Eq instance works" $ mkPolar 1 pi == mkPolar 1 pi,
      TestCase $
        assertEqual "Show output good" "mkPolar 1.0 1.0" $
          show $ mkPolar 1 1,
      TestCase $ readShowInverse z,
      TestCase $ readShowInverse [z, z],
      TestCase $ readShowInverse $ Just z
    ]
  where
    z = mkPolar 4.47214 0.46364

approximatelyEqualFloat :: Float -> Float -> Bool
approximatelyEqualFloat x1 x2 = abs (x1 - x2) < 0.001

isApproximately ::
#ifdef __GLASGOW_HASKELL__
  HasCallStack =>
#endif
  String -> Float -> Float -> Assertion
isApproximately name x1 x2 =
  assertBool (name ++ " = " ++ show x1 ++ " is approximately " ++ show x2) $
    approximatelyEqualFloat x1 x2

rAndThetaAre ::
#ifdef __GLASGOW_HASKELL__
  HasCallStack =>
#endif
  Float -> Float -> Polar Float -> Assertion
rAndThetaAre r theta z = do
  isApproximately "magnitude" r (magnitude z)
  isApproximately "phase" theta (phase z)
  let (r', theta') = polar z
  isApproximately "magnitude from polar" r r'
  isApproximately "phase from polar" theta theta'
#if __GLASGOW_HASKELL >= 786
  let (r'' :< theta') = polar z
  isApproximately "magnitude from :<" r r''
  isApproximately "phase from :<" theta theta''
#endif

readShowInverse ::
  (Show a, Read a,
#ifdef __GLASGOW_HASKELL__
  HasCallStack,
#endif
  Eq a) =>
  a -> Assertion
readShowInverse a = assertEqual ("read . show = id on " ++ show a)
                                (read $ show a)
                                a

main :: IO ()
main = do
  counts <- runTestTT tests
  if errors counts + failures counts == 0
    then exitWith ExitSuccess
    else exitWith (ExitFailure 1)


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

import FrequencyResponse

import CLaSH.Prelude
import CLaSH.Signal.MultiSignal
import Data.Complex as C
import qualified Prelude as P
import Data.Maybe
import Control.Applicative
import Graphics.EasyPlot
import Data.Map.Strict as Map

{-
Ciruits and stuff
-}

registerP = prepend

firP coeffs x = dotp coeffs (windowP x)
  where
    dotp as bs = sum (zipWith (*) as bs)

iirP cA cB x = r where
    oB =  firP cB x
    oA =  firP cA (registerP def r)
    r = oB - oA

-- fir coeficients
firCoef :: Fractional a => Vec 201 a
firCoef = -6.86056317e-19 :> 1.08643703e-04 :> -6.04051550e-18 :> -2.03379801e-04 :> -2.01638331e-04 :> 1.96946883e-18 :> 2.90791068e-18 :> -4.62180925e-04 :> -1.00665733e-03 :> -9.14037108e-04 :> 3.86774364e-18 :> 1.01950810e-03 :> 1.25131737e-03 :> 6.39186084e-04 :> -8.00554425e-18 :> 2.62712857e-18 :> 3.77239653e-04 :> 4.17356596e-04 :> -1.21745475e-17 :> -2.64065618e-04 :> -5.23918059e-18 :> 2.99576343e-04 :> -2.73859684e-18 :> -6.08801988e-04 :> -6.23370976e-04 :> 1.63484509e-17 :> -7.03525331e-18 :> -1.52438502e-03 :> -3.36082345e-03 :> -3.07661933e-03 :> -1.30788363e-17 :> 3.45259832e-03 :> 4.23241069e-03 :> 2.15427395e-03 :> 0.00000000e+00 :> 2.70383448e-17 :> 1.24466846e-03 :> 1.36380726e-03 :> -2.92823518e-17 :> -8.44354542e-04 :> -4.24392593e-17 :> 9.35427798e-04 :> 2.64177628e-17 :> -1.85473668e-03 :> -1.87576073e-03 :> 2.24314248e-18 :> -5.19238250e-17 :> -4.42290170e-03 :> -9.63775624e-03 :> -8.72271164e-03 :> 3.66725201e-17 :> 9.57794343e-03 :> 1.16209519e-02 :> 5.85688426e-03 :> -6.95248052e-17 :> 1.32593176e-17 :> 3.29422893e-03 :> 3.58089966e-03 :> 9.08965047e-18 :> -2.18539647e-03 :> -2.12919221e-17 :> 2.39194534e-03 :> -8.31327663e-17 :> -4.69672484e-03 :> -4.73138591e-03 :> 5.79766610e-17 :> -5.28939995e-18 :> -1.10712423e-02 :> -2.40989091e-02 :> -2.18047819e-02 :> -6.11496326e-17 :> 2.39909500e-02 :> 2.91782046e-02 :> 1.47558239e-02 :> 3.11706634e-17 :> 1.33673427e-18 :> 8.44218666e-03 :> 9.25308362e-03 :> -9.62845645e-17 :> -5.76785543e-03 :> -2.84709264e-17 :> 6.49485159e-03 :> -4.19102125e-17 :> -1.32387105e-02 :> -1.36437840e-02 :> 1.16146829e-16 :> -4.98075231e-17 :> -3.49290893e-02 :> -7.90710907e-02 :> -7.48583001e-02 :> -3.01478553e-33 :> 9.23333495e-02 :> 1.20794568e-01 :> 6.66720749e-02 :> -9.64385396e-17 :> 4.65544367e-17 :> 5.76578486e-02 :> 7.99902785e-02 :> -8.32650050e-17 :> -1.34035769e-01 :> 8.00710266e-01 :> -1.34035769e-01 :> -8.32650050e-17 :> 7.99902785e-02 :> 5.76578486e-02 :> 4.65544367e-17 :> -9.64385396e-17 :> 6.66720749e-02 :> 1.20794568e-01 :> 9.23333495e-02 :> -3.01478553e-33 :> -7.48583001e-02 :> -7.90710907e-02 :> -3.49290893e-02 :> -4.98075231e-17 :> 1.16146829e-16 :> -1.36437840e-02 :> -1.32387105e-02 :> -4.19102125e-17 :> 6.49485159e-03 :> -2.84709264e-17 :> -5.76785543e-03 :> -9.62845645e-17 :> 9.25308362e-03 :> 8.44218666e-03 :> 1.33673427e-18 :> 3.11706634e-17 :> 1.47558239e-02 :> 2.91782046e-02 :> 2.39909500e-02 :> -6.11496326e-17 :> -2.18047819e-02 :> -2.40989091e-02 :> -1.10712423e-02 :> -5.28939995e-18 :> 5.79766610e-17 :> -4.73138591e-03 :> -4.69672484e-03 :> -8.31327663e-17 :> 2.39194534e-03 :> -2.12919221e-17 :> -2.18539647e-03 :> 9.08965047e-18 :> 3.58089966e-03 :> 3.29422893e-03 :> 1.32593176e-17 :> -6.95248052e-17 :> 5.85688426e-03 :> 1.16209519e-02 :> 9.57794343e-03 :> 3.66725201e-17 :> -8.72271164e-03 :> -9.63775624e-03 :> -4.42290170e-03 :> -5.19238250e-17 :> 2.24314248e-18 :> -1.87576073e-03 :> -1.85473668e-03 :> 2.64177628e-17 :> 9.35427798e-04 :> -4.24392593e-17 :> -8.44354542e-04 :> -2.92823518e-17 :> 1.36380726e-03 :> 1.24466846e-03 :> 2.70383448e-17 :> 0.00000000e+00 :> 2.15427395e-03 :> 4.23241069e-03 :> 3.45259832e-03 :> -1.30788363e-17 :> -3.07661933e-03 :> -3.36082345e-03 :> -1.52438502e-03 :> -7.03525331e-18 :> 1.63484509e-17 :> -6.23370976e-04 :> -6.08801988e-04 :> -2.73859684e-18 :> 2.99576343e-04 :> -5.23918059e-18 :> -2.64065618e-04 :> -1.21745475e-17 :> 4.17356596e-04 :> 3.77239653e-04 :> 2.62712857e-18 :> -8.00554425e-18 :> 6.39186084e-04 :> 1.25131737e-03 :> 1.01950810e-03 :> 3.86774364e-18 :> -9.14037108e-04 :> -1.00665733e-03 :> -4.62180925e-04 :> 2.90791068e-18 :> 1.96946883e-18 :> -2.01638331e-04 :> -2.03379801e-04 :> -6.04051550e-18 :> 1.08643703e-04 :> -6.86056317e-19 :> Nil

-- iir coeficients
iirCoefA :: Fractional a => Vec 8 a
iirCoefB :: Fractional a => Vec 9 a
iirCoefA = -1.8933239813032006 :> 3.422153467072297 :> -3.907588921469733 :> 3.6401868030176012 :> -2.5428541496659465 :> 1.3470500696022936 :> -0.49424202583584054 :> 0.10443313376764539 :> Nil
iirCoefB = 0.0025798085212212327 :> 0.02063846816976986 :> 0.07223463859419452 :> 0.14446927718838903 :> 0.18058659648548628 :> 0.14446927718838903 :> 0.07223463859419452 :> 0.02063846816976986 :> 0.0025798085212212327 :> Nil

{-
Drawing helpers
-}

simFuncLin f skip ph =  snd $ getSpecLin <$> P.last $ P.take skip $ getResponseLin f ph
simFunc f skip ph =  (Map.! 1) $ getSpectrum <$> P.last $ P.take skip $ getResponse f ph
simFuncAcyclic f ph =  (Map.! 1) $ getSpectrum $ getResponseAcyclic f ph

genData f sk pn = (x,y) where
    x = fmap (\a -> (1.0 * pi * fromInteger a / fromInteger pn)) [0, 1 .. pn]
    y = fmap (simFunc f sk) x

-- how is mapping over result like Complex Double -> Double
-- f is ciruit
-- sk skip first sk values from result, for acyclic circuit this should be more than circuitry delay
-- pn - number of linear points from 0 to 1 (nyquist freq)
plt how f sk pn = plot' [] X11 $ Data2D [Style Lines,Title "plot"] [] (P.zip ((/pi) <$> x) (how <$> y))
  where
    (x,y) = genData f sk pn

genDataLin f sk pn = (x,y) where
    x = fmap (\a -> (1.0 * pi * fromInteger a / fromInteger pn)) [0, 1 .. pn]
    y = fmap (simFuncLin f sk) x

pltLin how f sk pn = plot' [] X11 $ Data2D [Style Lines,Title "plot"] [] (P.zip ((/pi) <$> x) (how <$> y))
  where
    (x,y) = genDataLin f sk pn

genDataAcyclic f pn = (x,y) where
    x = fmap (\a -> (1.0 * pi * fromInteger a / fromInteger pn)) [0, 1 .. pn]
    y = fmap (simFuncAcyclic f) x

pltDataAcyclic how f pn = plot' [] X11 $ Data2D [Style Lines,Title "plot"] [] (P.zip ((/pi) <$> x) (how <$> y))
  where
    (x,y) = genDataAcyclic f pn

magLog x = logBase 10 (magnitude x) * 20
magLogMin m x = max m (magLog x)

integrator i = r where
    r = (registerP 0 r) + i

{-
Examples
-}

-- plots working with linear circuits
plotFirLin = pltLin magnitude f 210 100 where
 f = firP firCoef

-- slower and generic version of samplePlotLin
plotFir = plt magnitude f 210 100 where
 f = firP firCoef

-- plot acyclic ciruit
plotFirAcyclic = pltDataAcyclic magnitude f 1000 where
 f = firP firCoef

plotIir = plt magnitude (iirP iirCoefA iirCoefB) 200 100
plotIirLog = pltLin magLog (iirP iirCoefA iirCoefB) 200 100

plotInteg = plt magnitude integrator 100 100

main = plotIir


-- | Main module for loading all parts of Hoqus package and dependencies from
-- 'hmatrix' package.
module Hoqus where

import Numeric.LinearAlgebra.Data
import Numeric.LinearAlgebra

import Hoqus.Dirac
import Hoqus.Gates
import Hoqus.MtxFun
import Hoqus.Fidelity


module Main where

import           Control.Lens
import           Data.H3.Colour                 (ordinalColours, toCSSColour)
import           Data.H3.Extent                 (Extent, extent, fromTuple,
                                                 resize, toTuple)
import           Data.H3.Scalable               (Scalable (..))
import           Data.H3.Scales                 (IncludeZeroPolicy (..),
                                                 Product (..), continuous)
import           Data.H3.Svg                    (ViewBoxMode (..), renderSvg)
import           Data.H3.Visuals                (ChartVisuals (..),
                                                 FontSize (..),
                                                 LabelOffset (..), Pixel (..),
                                                 Shape (..), Vis (..),
                                                 backgroundShapes, cartesian,
                                                 foregroundShapes, mapColour,
                                                 noGrid)
import           Data.List.NonEmpty             (NonEmpty (..))
import           Data.Semigroup.Foldable        (Foldable1 (..))
import qualified Data.Text                      as Text
import qualified Data.Text.Lazy.IO              as TextIO
import qualified Graphics.Svg                   as SVG
import           Statistics.Distribution        (ContDistr (density),
                                                 Mean (mean), Variance (stdDev))
import qualified Statistics.Distribution.Normal as Distribution.Normal

newtype PlotData = PlotData { getPlots :: [(Double, Double)] }

_PlotData :: Iso' PlotData [(Double, Double)]
_PlotData = iso getPlots PlotData

main :: IO ()
main = do
  let
    plts  = [
      PlotData $ sample 1.0 0.25,
      PlotData $ sample 2.0 0.44]
    dims = (fromTuple (0, 900), fromTuple (500, 0))
    r    = renderSvg (AddViewBox 1.25) (mapColour Text.pack . plotShape plts) dims
  TextIO.writeFile "out.svg" $ SVG.prettyText r

-- | Sample from a normal distribution
sample :: Double -> Double -> [(Double, Double)]
sample theMean theStdDev = vls where
  d = Distribution.Normal.normalDistr theMean theStdDev
  (f, mn, mx) = let m = mean d
                    std = stdDev d in
                    (density d, m - 3 * std, m + 3 * std)
  samples = fmap (\i -> mn + i * delta) [0 .. n] where
    delta = (mx - mn) / n
    n = 100
  vls = fmap (over _2 f) (zip samples samples)

plotShape :: [PlotData] -> (Extent Double, Extent Double) -> Shape String (Pixel Double, Pixel Double)
plotShape ps tgt = result where
  allPoints = foldMap getPlots ps
  result = case allPoints of
    [] -> EmptyShape
    (x:xs) -> over both Pixel <$> AGroup [backgroundShapes s, plotLines, foregroundShapes s] where
      (xExtent, yExtent) = foldMap1 (over both extent) (x :| xs)
      xOpts = continuous xExtent DontIncludeZero
      yOpts = continuous yExtent DontIncludeZero
      tickLength = 5
      ((_, _), (yLw, _)) = over both toTuple tgt
      xyOpts = cartesian
        (fromTuple (yLw, yLw + tickLength))
        (fromTuple (negate tickLength, 0))
        (FontSize "10pt")
        (LabelOffset (-20, 15))
        xOpts
        yOpts
      s = visuals (Vis xyOpts) tgt
      colors = toCSSColour . runIdentity <$> scale (ordinalColours [1 .. length xs]) ()
      plotLines =
        fmap (over both runIdentity . getProduct . combinedScale)
        $ AGroup
        $ mkPlot <$> zip ps [1..]
      mkPlot (PlotData pts, col) = AColouredShape (colors col) $ ALine pts
      combinedScale = scale xyOpts tgt


-----------------------------------------------------------------------------
-- |
-- Module      :  FlashADC
-- Copyright   :  José Edil Guimarães de Medeiros
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  j.edil@ene.unb.br
-- Stability   :  experimental
-- Portability :  portable
--
-- This module shows the model of a flash architecture ADC.
--
-- =Running the demo
--
-- To run the demo you need the ForSyDe installed in your
-- environment. For plotting, you need the
-- <http://gnuplot.sourceforge.net Gnuplot package>.
--
-- To calculate 10000 samples for R=1e3 Omhs, C=1e-6 F, T=1e-6run in @ghci@:
--
-- >>> simulate 4
--
-- To plot the response for a sine wave input, run in @ghci@:
--
-- >>> plotOutput 4 0.01 1000
-----------------------------------------------------------------------------

module ForSyDe.Shallow.Example.Synchronous.FlashADC where

import ForSyDe.Shallow
import Data.Complex

type Resistance = Double
type Voltage = Double
type Code = Integer
type Bits = Integer
type Time = Double
type Frequency = Double

-- | 'flashADC' is the top level module.
flashADC :: [Resistance]                -- ^ Resistance values
          -> Signal Voltage             -- ^ Input signal
          -> Signal Code                -- ^ Output signal
flashADC resistors input = decoder $ compNetwork input $ resNetwork resistors

-- | 'decoder' takes the thermometer code and outputs integers.
decoder :: [Signal Bits]                -- ^ Bit inputs
        -> Signal Code                  -- ^ Output signal
decoder = foldl1 (zipWithSY (+))

-- | 'compNetwork' implements the comparator array.
compNetwork :: Signal Voltage           -- ^ ADC input signal
            -> [Voltage]                -- ^ Voltage thresholds
            -> [Signal Bits]            -- ^ Bit outputs
compNetwork input = zipWith (\i v -> mapSY (comparator v) i) (repeat input)

-- | 'comparator' is the one bit quantizer.
comparator :: Voltage                   -- ^ (+) input
           -> Voltage                   -- ^ (-) input
           -> Bits                      -- ^ Output
comparator i v
  | v <= i = 0
  | otherwise = 1

-- | 'resNetwork' implements the resistor voltage scaling network.
resNetwork :: [Resistance]              -- ^ Resistor values
           -> [Voltage]                 -- ^ Threshold voltages
resNetwork resistors = init $ tail $ scanl (\v r -> v + vdd * r / (sumR)) 0 resistors
  where vdd = 1
        sumR = sum resistors

-- | 'simulate' takes the system parameters and runs the simulation
-- with a sine wave input.
simulate :: Int                         -- ^ ADC resolution
         -> Int                         -- ^ Number of samples
         -> Signal Code                 -- ^ ADC output
simulate res n = flashADC resistors input
  where resistors = replicate (2^res) 1
        input = sineWave' 0.5 0.5 50 n

-- | 'plotFFT' runs the simulation with a sine wave input
plotFFT :: Int                  -- ^ ADC resolution
        -> Int                  -- ^ FFT depth (power of 2)
        -> IO String            -- ^ plot
plotFFT res n = plotCT' (toRational 1.0) [(g_out, "g")]
  where g_out = d2aConverter DAhold (toRational 1.0) $ signal g
        a = fromSignal $ simulate res n
        b = (:+) <$> (map ((\x -> x/2^res - 0.5).fromIntegral) a)
        c = zipWith ($) b $ repeat 0.0
        d = vector c
        e = fromVector $ fft n d
        f = map ((*20).(logBase 10).(\x -> 2*x/(fromIntegral n)) . magnitude) e
        g = (take (n `div` 2) f)

-- | 'plotOutput' uses the CTLib plot capabilities to plot the
-- output. In a later version, a plotter to Synchornous signals will be
-- developed.
plotOutput :: Int                       -- ^ Resolution
           -> Double                    -- ^ Discretization timestep
           -> Int                       -- ^ Number of samples
           -> IO String                 -- ^ plot
plotOutput res t n = plotCT' (toRational t) [(output, "output")]
  where output = d2aConverter DAhold (toRational t) adcSignal
        adcSignal = signal $ map (toRational) $ fromSignal $
                    simulate res n

-- | 'sineWave' is an auxiliary function that creates a sine wave signal for
-- simulation purposes.
sineWave' :: Voltage                     -- ^ Amplitude
          -> Voltage                     -- ^ Offset
          -> Frequency                   -- ^ Frequency
          -> Int                         -- ^ Number of samples
          -> Signal Time                 -- ^ Output signal
sineWave' amp offset freq n = signal sineW
  where sineW = map (\t -> (amp/2) * sin(2*pi*freq*t) + offset) grid
        grid = linspace 0 1 n

-- | 'linspace' is an auxiliary function that creates an uniform
-- spaced sampling grid for simulation purposes.
linspace :: Time                        -- ^ Start time
         -> Time                        -- ^ Stop time
         -> Int                         -- ^ Number of samples
         -> [Time]                      -- ^ Sampling grid
linspace start stop n = map (scale) nodes
  where n' = fromIntegral n
        nodes = map (fromIntegral) [0..(n-1)]
        scale t = t * length / (n' - 1) + start
        length = stop - start


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies     #-}
-- |
-- Module      : Main
-- Description : Test runner
-- Copyright   : (c) Tom Westerhout, 2017
-- License     : BSD3
-- Maintainer  : t.westerhout@student.ru.nl
-- Stability   : experimental

module Main where

import qualified System.Random.MWC as MWC
import           Control.Lens
import           Control.Monad.Reader
import           Control.Monad.State
import qualified Data.List            as L
import Data.Complex
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA
import Data.Vector.Storable((!))

import PSO.Random
import PSO.Energy
import PSO.Swarm

class TestFunction f where
  function :: f -> V.Vector Float -> Float
  initBounds :: f -> (Float, Float)
  dim :: f -> Int

data RosenbrockFn = RosenbrockFn

instance TestFunction RosenbrockFn where
  function _ x =
    let p1 = (100.0 *) . V.sum . V.map (^^2)
                $ V.zipWith (-) (V.tail x) (V.map (^^2) $ V.init x)
        p2 = V.sum . V.map (^^2) . V.map (1.0 -) $ V.init x
    in p1 + p2
  initBounds _ = (15.0, 30.0)
  dim _ = 30

unpack :: (Float, Float) -> Int -> (V.Vector Float, V.Vector Float)
unpack (x, y) d = (f x, f y)
  where f = V.fromList . replicate d

wpg :: (Float, Float, Float)
wpg = (0.7298 :: Float, 1.49618, 1.49618)

cmUpdater :: (Num χ, RandomScalable m χ, VectorSpace χ Float)
        => PhaseUpdater m (SwarmGuide χ r) (BeeGuide χ r) (CMState χ) r
cmUpdater = standardUpdater wpg

qmUpdater :: (RandomScalable m χ, Randomisable m Float, χ ~ V.Vector Float)
        => PhaseUpdater m (SwarmGuide χ r) (BeeGuide χ r) (QMState χ) r
qmUpdater = (PhaseUpdater $ deltaWellUpdater (2 * log 2 * 0.7 :: Float))

cmRunND :: IO ()
cmRunND = do
  let xs = optimiseND
            (mkCMState (unpack (initBounds RosenbrockFn) (dim RosenbrockFn)))
            cmUpdater
            (function RosenbrockFn)
            20
            (\s -> (s^.guide.val) < 1.0
                   || (s^.guide.iteration == 10000))
  gen <- mkMWCGen (Just 123)
  swarms <- runReaderT xs gen
  writeEnergies2TSV "Function.dat" (view val) swarms
  let swarm = last swarms
  -- mapM_ (print . (!!0) . (view bees)) $ swarms
  putStrLn ""
  putStrLn $ "[+] Best[f] = " ++ show (swarm ^. guide . val)

qmRunND :: IO ()
qmRunND = do
  let xs = optimiseND
            (mkQMState (unpack (initBounds RosenbrockFn) (dim RosenbrockFn)))
            qmUpdater
            (function RosenbrockFn)
            20
            (\s -> (s^.guide.val) < 1.0
                   || (s^.guide.iteration == 5000))
  gen <- mkMWCGen (Just 123)
  swarms <- runReaderT xs gen
  writeEnergies2TSV "Function.dat" (view val) swarms
  let swarm = last swarms
  -- mapM_ (print . (!!0) . (view bees)) $ swarms
  putStrLn ""
  putStrLn $ "[+] Best[f] = " ++ show (swarm ^. guide . val)

main :: IO ()
main = qmRunND


import Codec.Picture
import Codec.Picture.Types
import Data.Word
import Data.Vector as V
import Data.Complex

adjust :: Double -> Int -> Int -> Double
adjust s xr x = (fromIntegral x - (fromIntegral xr)/2)*s

cmpmake :: Double -> Int -> Int -> Int -> Int -> Complex Double
cmpmake s xr yr x y = adjust s xr x :+ adjust s yr y

polartolch :: (Double, Double) -> Double -> (Double, Double, Double)
polartolch (mag, ph) high = (linearInter high 60 mag, 90, ph)

lchtolab :: (Double, Double, Double) -> (Double, Double, Double)
lchtolab (l, c, h) = (l, c * (cos h), c * (sin h))

labtoxyz :: (Double, Double, Double) -> (Double, Double, Double)
labtoxyz (l, a, b) = let  eps = 216/24389
                          kap = 24389/27
                          fy  = (l + 16)/116
                          fz  = fy - (b/200)
                          fx  = (a/500) + fy
                          x   = if fx**3 > eps then fx**3 else (116*fx - 16)/kap
                          y   = if l > kap*eps then ((l + 16)/116)**3 else l/kap
                          z   = if fz**3 > eps then fz**3 else (116*fz - 16)/kap
                            in (x, y, z)

--https://www.mathworks.com/help/images/ref/lab2xyz.html

a =   (1.0985,  1,          0.3558)
c =   (0.9807,  1,          1.1822)
e =   (1,       1,          1)
d50 = (0.9642,  1.0,        0.8251)
d55 = (0.9568,  1.0,        0.9214)
d65 = (0.95047, 1.00,       1.08883)
icc = (0.9642,  1,          0.8249)

data RGBgamut = RGBgamut {row1 :: (Double, Double, Double),
                          row2 :: (Double, Double, Double),
                          row3 :: (Double, Double, Double),
                          white :: (Double, Double, Double),
                          gamma :: Double}

sRGBgamut = RGBgamut (3.2404542, -1.5371385, -0.4985314)
                      (-0.9692660, 1.8760108, 0.0415560)
                      (0.0556434, -0.2040259, 1.0572252)
                      d65
                      2.4
adobeRGBgamut = RGBgamut (2.0413690, -0.5649464, -0.3446944)
                          (-0.9692660, 1.8760108, 0.0415560)
                          (0.0134474, -0.1183897, 1.0154096)
                          d65
                          (563/256)

adobeWideRGBgamut = RGBgamut (1.4628067, -0.1840623, -0.2743606)
                              (-0.5217933, 1.4472381, 0.0677227)
                              (0.0349342, -0.0968930, 1.2884099)
                              d50
                              (563/256)

proPhotoRGBgamut = RGBgamut (1.3459433, -0.2556075, -0.0511118)
                              (-0.5445989, 1.5081673, 0.0205351)
                              (0,       0,            1.2118128)
                              d50
                              1.8

pointwise :: (Double, Double, Double) -> (Double, Double, Double) -> (Double, Double, Double)
pointwise (a, b, c) (x, y, z) = (a*x, b*y, c*z)

matrixMult :: (Double, Double, Double) -> (Double, Double, Double) -> (Double, Double, Double) -> (Double, Double, Double) -> (Double, Double, Double)
matrixMult (a, b, c) (d, e, f) (g, h, i) (x, y, z) = (a*x + b*y + c*z, d*z + e*y + f*z, g*x + h*y + i*z)

xyztorgb :: RGBgamut -> (Double, Double, Double) -> (Double, Double, Double)
xyztorgb gam (x, y, z) =  let (xl, yl, zl)  = matrixMult (row1 gam) (row2 gam) (row3 gam) $ pointwise (x, y, z) (white gam)
                              compand c     = c**(1/(gamma gam))
                                in (compand xl, compand yl, compand zl)

linearInter :: Double -> Double -> Double -> Double
linearInter highin highout x = highout/highin * x

expInter :: Double -> Double -> Double -> Double
expInter low high x = 1 - base ** (x) where
  base = 0.5

makeSampleSpace :: Int -> Int -> Double -> Vector (Vector (Double, Double))
makeSampleSpace w h s = generate (2*h + 1) (\y -> generate (2*w + 1) (\x -> (s*(fromIntegral x - fromIntegral w), -s*(fromIntegral y - fromIntegral h))))

inmap :: (a -> b) -> Vector (Vector a) -> Vector (Vector b)
inmap f l = V.map (V.map f) l

findextrema :: Vector (Vector (Double, Double)) -> Double
findextrema lst =
  let lst' = V.concat $ toList lst
      lst'' = V.map (\(x, y) -> x) lst'
          in V.maximum lst''

highest :: Vector Double -> Double
highest vec = if V.length vec == 1 then V.head vec else
            (if V.head vec > (highest $ slice 1 (V.length vec - 1) vec) then V.head vec else (highest $ slice 1 (V.length vec - 1) vec))

urx = 1000
ury = 1000
iscale = 0.003

f :: Complex Double -> Complex Double
f z =  cexp $ z

cexp :: Complex Double -> Complex Double
cexp (a :+ b) = (exp a :+ 0) * (cis b)

newimage :: Vector (Vector PixelRGB16)
newimage = let  sampling = makeSampleSpace urx ury iscale
                polars = inmap (\(x,y) -> polar $ f $ x :+ y) sampling
                high = findextrema polars
                lch = inmap (\(x,y) -> polartolch (x,y) high) polars
                lab = inmap lchtolab lch
                xyz = inmap labtoxyz lab
                rgb = inmap (xyztorgb proPhotoRGBgamut) xyz
                rounded = inmap (\(r, g, b) -> PixelRGB16 (round $ r*65530) (round $ g*65530) (round $ b*65530)) rgb
                  in rounded

main :: IO ()
main = writePng "./blah.png" $ generateImage (\x y -> newimage ! y ! x) (2*urx + 1) (2*ury + 1)


-- Fingerprint
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE StrictData #-}

module Fingerprint (
  Fingerprint
, ColorVector
, fp2colvec
, makeMap
, findSimilarImage
, colvec2id
, colvec2status
, isValid
) where

import qualified Data.KdTree.Static as KT
import           Data.List (transpose)
import           Data.List.Split (chunksOf)
import           Numeric (readHex)
import           Numeric.LinearAlgebra

type Fingerprint = String
type Channel = [Double]

type Color = [Double]  -- (r,g,b) => [r,g,b]
data ColorVector = ColorVector
  { pixel  :: [Color]
  , imgid  :: Int
  , fp     :: Fingerprint
  , status :: CvStatus
  } deriving (Eq, Show)
type ColorVectorMap = KT.KdTree Double ColorVector

data CvStatus = FILED | PEND | DISCARDED | DUPLICATED | INFERIOR deriving (Eq, Show)

-- constants
--   status

-- public functions

makeMap :: [ColorVector] -> ColorVectorMap
makeMap cvs = KT.buildWithDist colvec2points colvecDistance cvs

{-|
  IN:
    r1  : radius for KdTree search
    r2  : difference for each pixel
-}

findSimilarImage :: Double -> Double -> Double -> ColorVectorMap -> ColorVector -> [ColorVector]
findSimilarImage r1 r2 rt cvmap cv = ss2
  where
    ss1 = filter (\s -> colvec2id cv < colvec2id s) $ KT.inRadius cvmap r1 cv
    ss2 = filter (\s -> isNearImage r2 rt (colvec2fp cv) (colvec2fp s)) ss1

{-|

>>> let fp1 = [1.0, 2.0, 1.0, 2.0]
>>> let fp2 = [1.0, 2.0, 1.0, 2.0]
>>> isNearImage 0.0 1.0 fp1 fp2
True
>>> let fp1 = [1.0, 2.0, 1.0, 2.0]
>>> let fp2 = [1.1, 2.0, 1.0, 2.0]
>>> let ds = map abs $ zipWith (-) fp1 fp2
>>> length ds
4
>>> let ds' = filter (<= 0.0) $ ds
>>> length ds'
3
>>> (fromIntegral (length ds') / fromIntegral (length ds))
0.75
>>> isNearImage 0.0 0.5 fp1 fp2
True

-}

isNearImage :: Double -> Double -> [Double] -> [Double] -> Bool
isNearImage rad rate fp1 fp2 = (fromIntegral (length ds') / fromIntegral (length ds) >= rate)
  where
    ds = map abs $ zipWith (-) fp1 fp2
    ds' = filter (<= rad) ds

{-|
  IN:
      oreso - original resolution of x (ex. 8)
      vreso - ColorVector resolution of x (ex. 2)
      imgid - ID of image
      fp    - fingerprint of image (ex. 8 x 8 x 3ch)

>>> let fp = "000000010101020202030303010101020202030303040404020202030303040404050505030303040404050505060606"
>>> fp2colvec 4 2 1 fp "filed"
ColorVector {pixel = [[1.0,1.0,1.0],[3.0,3.0,3.0],[3.0,3.0,3.0],[5.0,5.0,5.0]], imgid = 1, fp = "000000010101020202030303010101020202030303040404020202030303040404050505030303040404050505060606", status = FILED}

-}

fp2colvec :: Int -> Int -> Int -> Fingerprint -> String -> ColorVector
fp2colvec oreso vreso imgid fp st = ColorVector pixel imgid fp st'
  where
    chs = map (matrix oreso) (fp2channels fp)
    d = oreso `div` vreso
    pixel = transpose $ map (summary d vreso) chs
    st' = case st of
      "filed"      -> FILED
      "pending"    -> PEND
      "deleted"    -> DISCARDED
      "duplicated" -> DUPLICATED
      "inferior"   -> INFERIOR
      _            -> DISCARDED

colvec2id :: ColorVector -> Int
colvec2id (ColorVector _ i _ _) = i

colvec2fp :: ColorVector -> [Double]
colvec2fp (ColorVector _ _ fp _) = fp2double fp

colvec2status :: ColorVector -> CvStatus
colvec2status (ColorVector _ _ _ st) = st

isValid :: ColorVector -> Bool
isValid (ColorVector _ _ _ st) = (st == FILED || st == PEND)

-- internal functions

{-|

>>> let cv1 = ColorVector [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0], [10.0, 11.0, 12.0]] 1 "001122334455" FILED
>>> colvec2points cv1
[1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.0,12.0]

-}

colvec2points :: ColorVector -> [Double]
colvec2points (ColorVector points _ _ _) = concat points

{-|

>>> let cv1 = ColorVector [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0], [10.0, 11.0, 12.0]] 1 "001122334455" FILED
>>> let cv2 = ColorVector [[1.0, 1.0, 4.0], [3.0, 6.5, 6.0], [8.2, 8.0, 7.8], [10.0, 12.0, 11.0]] 2 "101122334455" FILED
>>> colvecDistance cv1 cv2
2.5

-}

colvecDistance :: ColorVector -> ColorVector -> Double
colvecDistance cv1 cv2 = maximum ls
  where
    p1s = colvec2points cv1
    p2s = colvec2points cv2
    ls = map sum $ chunksOf 3 $ zipWith (\x y -> abs (x - y)) p1s p2s

{-|

>>> let m = matrix 4 [0,1,2,3,1,2,3,4,2,3,4,5,3,4,5,6]
>>> summary 2 2 m
[1.0,3.0,3.0,5.0]
>>> let m2 = matrix 4 [6,5,4,3,5,4,3,2,4,3,2,1,3,2,1,0]
>>> summary 2 2 m2
[5.0,3.0,3.0,1.0]

-}

summary :: Int -> Int -> Matrix R -> [Double]
summary d vreso m = map (\x -> x / fromIntegral (d*d)) col
  where
    ele = map (*d) [0..(vreso-1)]
    offset = [(x, y)| x <- ele, y <- ele]
    col = map (\o -> sum.toList.flatten $ subMatrix o (d, d) m) offset

{-|

>>> fp2channels "000102030405060708090a0b"
[[0.0,3.0,6.0,9.0],[1.0,4.0,7.0,10.0],[2.0,5.0,8.0,11.0]]

-}

fp2channels :: Fingerprint -> [Channel]
fp2channels fp = splitRgb ds ([], [], [])
  where
    ds = fp2double fp
    splitRgb :: [Double] -> (Channel, Channel, Channel) -> [Channel]
    splitRgb [] (r, g, b) = [reverse r, reverse g, reverse b]
    splitRgb (r:g:b:xs) (rs, gs, bs) = splitRgb xs ((r:rs), (g:gs), (b:bs))

{-|

>>> fp2double "000102030405060708090a0b0c0d0e0f"
[0.0,1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.0,12.0,13.0,14.0,15.0]
>>> fp2double "00102030405060708090a0b0c0d0e0f0"
[0.0,16.0,32.0,48.0,64.0,80.0,96.0,112.0,128.0,144.0,160.0,176.0,192.0,208.0,224.0,240.0]

-}

fp2double :: Fingerprint -> [Double]
fp2double fp = map hex2double $ chunksOf 2 fp

hex2double :: String -> Double
hex2double hex = fromIntegral h
  where
    [(h, _)] = (readHex hex)


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE BangPatterns #-}
module AI.Singularity.Training where
import GHC.TypeLits
import qualified Numeric.LinearAlgebra.Data as LD
import qualified Numeric.LinearAlgebra.Devel as LDev
import qualified Numeric.LinearAlgebra as LA
import qualified Control.Monad.State as S
import Control.Lens
import Data.Proxy
import Control.Monad.ST
import System.Random
import Control.Monad
import Control.Arrow
import qualified Data.Vector.Storable as VS
import Data.Conduit
import System.IO (hPrint, stderr)
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.List   as CL
import Data.Conduit.Internal (zipSources)
import Numeric.AD
import Control.Monad.Trans.Class (lift)

import AI.Singularity.Utils
import AI.Singularity.Data.Matrix
import AI.Singularity.Data.Vector
import AI.Singularity.Data.Network

data TrainConfig = TrainConfig { _learningRate :: !Float,
                                 _lrUpdate     :: Float -> Float,
                                 _batchsize    :: !Int,
                                 _momentumRate :: !Float }
makeLenses ''TrainConfig

newtype Loss = Loss { getLoss :: forall a. Floating a => a -> a -> a}

lossGrad :: Loss -> (forall a. Floating a => a -> a -> a)
lossGrad (Loss f) = f'
  where fList :: forall a. Floating a => [a] -> a
        fList [y,y'] = f y y'
        f' y y'      = head . tail . grad fList $ [y,y']

data TrainState i o = TrainState { _config  :: !TrainConfig,
                                   _network :: !(TrainedNetwork i o),
                                   _iter    :: !Int,
                                   _loss    :: Loss,
                                   _epochs  :: !Int,
                                   _evalSet :: [(i,o)]}

makeLenses ''TrainState

addGrad :: (KnownNat m, KnownNat n) => TrainConfig -> InTrainMat m n -> Identity (InTrainMat m n)
addGrad tr mp = Identity $ GeneralInTrainingMat newWeights newGradient hGr
  where momentumRatio = view momentumRate tr
        lr            = view learningRate tr
        newWeights    = checkNan $ view trMatr mp + (mmap (*lr) $ view trGrad mp) / mmap sqrt (hGr + 0.00000001)
        newGradient   = checkNan . over matr (LD.cmap (*momentumRatio)) $ view trGrad mp
        hGr           = 0.9 * view hGrad mp + 0.1 * (view trGrad mp ^ 2)

applyGrad :: TrainConfig -> TrainedNetwork inp out -> TrainedNetwork inp out
applyGrad tr = runIdentity . changeMatType (addGrad tr)

calcGrads :: forall n m. (KnownNat n, KnownNat m) => Vector n -> Vector m -> TrainedNetwork (Vector n) (Vector m) -> (Matrix (n+1) m, Vector n)
calcGrads !inp !outGrad (FFLayer !mpair !f) = (gradMat, inpGrad)
  where
    m       = view trMatr mpair
    gradMat = dy `outer` consV 1 inp
    dy      = outGrad * diffM
    diffM   = cmap (diff f) . (#>) m . consV 1 $ inp
    inpGrad = tailV $ transposeM m #> dy

trainWithErr :: forall a b. TrainConfig -> a -> b -> TrainedNetwork a b -> (TrainedNetwork a b, a)
trainWithErr tc !inp !outGrad n@(FFLayer !mpair !f) = (FFLayer  m' f, inpGrad)
  where lr = view learningRate tc
        m' = over trGrad (subtract gradMat) mpair
        (gradMat,inpGrad) = calcGrads inp outGrad n
trainWithErr tc !inp !outGrad (FFSeq !n1 !n2) = (FFSeq n1' n2', inpGrad)
  where (n1', inpGrad)  = trainWithErr tc inp inpGrad' n1
        inp'            = conductSignalInTrain inp n1
        (n2', inpGrad') = trainWithErr tc inp' outGrad n2
trainWithErr tc !inp !outGrad (SplitNet n) = (SplitNet *** splitV) unSpl
  where unSpl = trainWithErr tc (uncurry appendV inp) outGrad n
trainWithErr tc !inp !outGrad (FromBinFunc f) = (FromBinFunc f, grads)
  where grads = (\(x1,x2) -> ((!!0) &&& (!!1)). map (*outGrad) . fgrad $ [x1,x2]) inp
        fwrap :: forall a. Fractional a => [a] -> a
        fwrap [x1,x2] = f x1 x2
        fgrad = grad fwrap
trainWithErr tc !inp !outGrad (Split n1 n2) = (Split n1' n2', inGr)
  where (n1', inGr1) = trainWithErr tc inp (fst outGrad) n1
        (n2', inGr2) = trainWithErr tc inp (snd outGrad) n2
        inGr         = inGr1 + inGr2
trainWithErr tc !inp !outGrad (FromFunc f) = (FromFunc f, inGr)
  where inGr = diff f inp * outGrad
trainWithErr tc !inp !outGrad (FFDiv !n1 !n2) = (FFDiv n1' n2', (inGr1, inGr2))
  where (n1', inGr1) = trainWithErr tc (fst inp) (fst outGrad) n1
        (n2', inGr2) = trainWithErr tc (snd inp) (snd outGrad) n2
trainWithErr tc !inp !outGrad (Recurse n)   = if null inp then (Recurse n, []) else (Recurse n', [fst inGrad])
   where inputs      = scanl (\acc sym -> (sym, conductSignalInTrain acc n)) (head inp,0) (tail inp)
         (n',inGrad) = foldr (\input (net, outG) -> trainWithErr tc input (snd outG) net) (n,(undefined, outGrad)) inputs
trainWithErr tc !inp !outGrad (First n)     = (First n', inGrad)
  where (n', in1Grad) = trainWithErr tc (fst inp) outGrad n
        inGrad        = (in1Grad, 0)
trainWithErr tc !inp !outGrad (Second n)     = (Second n', inGrad)
  where (n', in1Grad) = trainWithErr tc (snd inp) outGrad n
        inGrad        = (0, in1Grad)

trainSample :: forall a b. TrainConfig -> [(a,b)] -> (b -> b -> b) -> Int -> Network a b -> Network a b
trainSample tc !corpus gLoss !epochs = fromTrain . (trainHelp tc 1 . concat $ replicate epochs corpus) . toTrain
  where
    trainHelp :: TrainConfig -> Int -> [(a, b)] -> TrainedNetwork a b -> TrainedNetwork a b
    trainHelp tc i ((!x,!y):corp) !net = next . fst $ trainWithErr tc x l net
      where l    = gLoss y (conductSignalInTrain x net)
            next = if i < b
                   then trainHelp tc (i+1) corp
                   else trainHelp nTc 1 corp . applyGrad tc
            nTc  = tc -- over learningRate (view lrUpdate tc) tc
            b    = view batchsize tc
    trainHelp _ _ [] !net = net

eps :: forall n. Floating n => n
eps = 0.00000001

defaultConfig = TrainConfig 0.05 id 1 0

endOfBatch :: TrainState i o -> Bool
endOfBatch = view iter &&& view (config.batchsize) >>> uncurry (>=)

startState conf net = defaultState conf net logloss
  where logloss y y' = negate $ y * log (y' + eps) + (1 - y) * log (1 - y' + eps)

defaultState :: TrainConfig -> Network i o -> (forall a. Floating a => a -> a -> a) -> Int -> [(i,o)] -> TrainState i o
defaultState conf net f ep = TrainState conf (toTrain net) 1 (Loss f) ep

-- applyTrainingExample :: (Floating o, Show o) => TrainState i o -> (i,o) -> IO (TrainState i o)
applyTrainingExample tr (i,o) = endBatch . over iter (+1) . over network (fst . trainWithErr trConf i outLoss) $ tr
  where
    outLoss  = lossGrad (view loss tr) o (conductSignalInTrain i (view network tr))
    trConf   = view config tr
    endBatch = if endOfBatch tr
               then \x -> do
                           let up = view (config.lrUpdate) x
                               nX = over (config.learningRate) up . set iter 1 . over network (applyGrad trConf) $ x
                               s  = view evalSet tr
                               ac = sum . map (uncurry. flip . getLoss . view loss $ tr) . map (first $ flip conductSignalInTrain (view network nX)) $ s
                           unless (null s) . lift $ do
                             hPrint stderr "LOSS:"
                             hPrint stderr (ac/fromIntegral (length s))
                           return nX
               else return

trainC cond conf net ep = trainStateC cond . startState conf net ep

trainStateC cond trState =
  if view epochs trState >= 1
  then
    runConduitRes (cond .| CL.foldM applyTrainingExample trState) >>= (over epochs (subtract 1) >>> trainStateC cond)
  else
    fromTrain . applyGrad (view config trState) . view network <$> runConduitRes (cond .| CL.foldM applyTrainingExample trState)

trainExample :: forall a b. Floating b => TrainConfig -> [(a, b)] -> Int -> Network a b -> Network a b
trainExample conf x = trainSample conf x logloss'
  where
    logloss [y, y'] = y * log (y' + eps) + (1 - y) * log (1 - y' + eps)
    logloss' y y'   = head . tail . grad logloss $ [y, y']


{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}

module Data.Model where

import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M

import Statistics.Distribution
import Statistics.Distribution.Poisson

import Data.Aeson
import Data.String (IsString)

type Hist = [Double]
type PredHist = [PoissonDistribution]
type DataHist = [Int]

newtype ProcName = ProcName { unPN :: String } deriving (Eq, Ord, FromJSON, IsString)
newtype RegName = RegName { unRN :: String } deriving (Eq, Ord, FromJSON, IsString)

instance Show ProcName where
    show = unPN

instance Show RegName where
    show = unRN

rebinH :: Num a => Int -> [a] -> [a]
rebinH _ []            = []
rebinH n h | n <= 1    = h
           | otherwise = let (hs, hs') = splitAt n h
                         in  sum hs : rebinH n hs'

zipWithLen :: (a -> b -> c) -> [a] -> [b] -> [c]
zipWithLen f (a:as) (b:bs) = f a b : zipWithLen f as bs
zipWithLen _ [] []         = []
zipWithLen _ _ _           = error "zipping lists of different lengths"

-- the log likelihood of a prediction histogram given a data histogram
predHistLLH :: DataHist -> PredHist -> Double
predHistLLH dh ph = sum $ zipWithLen logProbability ph dh

-- a process's normalization is consistent across regions
type Process = Map RegName Hist

-- a collection of named processes
type Prediction = Map ProcName Process

type TotalPrediction = Map RegName PredHist

-- data in several regions
type Dataset = Map RegName DataHist

liftP :: (Double -> Double) -> PoissonDistribution -> PoissonDistribution
liftP f = poisson . f . poissonLambda

liftP2 :: (Double -> Double -> Double) -> PoissonDistribution -> PoissonDistribution -> PoissonDistribution
liftP2 f p p' = poisson $ f (poissonLambda p) (poissonLambda p')

toPred :: Hist -> PredHist
toPred = fmap (poisson . cleanBin)

scaleH :: Double -> Hist -> Hist
scaleH n = fmap (*n)

addH :: Hist -> Hist -> Hist
addH = zipWithLen (+)

cleanBin :: Double -> Double
cleanBin x = if x <= 0 then 1e-100 else x

mulH :: Hist -> Hist -> Hist
mulH = zipWithLen (*)

divH :: Hist -> Hist -> Hist
divH = zipWithLen (/)

sumH :: [Hist] -> Hist
sumH = foldl1 addH

subH :: Hist -> Hist -> Hist
subH = zipWithLen (-)

-- a ModelParam alters a model in some particular way and has a prior
-- distribution
data ModelParam = ModelParam { mpName :: String
                             , mpPrior :: Double -> Double
                             , mpAlter :: Double -> Prediction -> Prediction
                             }

-- alter a particular process's histograms uniformly by some function
alterProc :: (Hist -> Hist) -> ProcName -> Prediction -> Prediction
alterProc f = M.adjust (fmap f)

-- process normalization ModelParam
procNormParam :: ContDistr d => d -> Double -> ProcName -> ModelParam
procNormParam prior n name = ModelParam (show name ++ "_norm") (logDensity prior)
                                $ (\x -> alterProc (scaleH (1+n*x)) name)

procShapeParam :: Double -> Map RegName Hist -> Process -> Process
procShapeParam x s p = M.differenceWith f p s
    where f p' s' = Just $ addH (scaleH x s') p'

shapeParam :: ContDistr d => d -> String -> Map ProcName (Map RegName Hist) -> ModelParam
shapeParam prior name hshapes = ModelParam name (logDensity prior) f
    where f x p = M.differenceWith (g x) p hshapes
          g x p' s' = Just $ procShapeParam x s' p'

totalPrediction :: Prediction -> TotalPrediction
totalPrediction = fmap toPred . M.foldr (M.unionWith addH) M.empty

-- TODO!!!
expectedData :: Prediction -> Dataset
expectedData = fmap (fmap $ round . poissonLambda) . totalPrediction

-- the poisson likelihood of a model given the input data
modelPoissonLLH :: Dataset -> Prediction -> Double
modelPoissonLLH ds m = M.foldr (+) 0 $ M.intersectionWith predHistLLH ds
                                     $ totalPrediction m

modelLLH :: Dataset -> Prediction -> [ModelParam] -> [Double] -> Double
modelLLH ds hpred hparams params = priorLLH + poissLLH
    where
        priorLLH = sum $ zipWithLen mpPrior hparams params
        hpred' = foldr ($) hpred (zipWithLen mpAlter hparams params)
        poissLLH = modelPoissonLLH ds hpred'


module Shader ( sampleBRDF, probability, reflectance
              , sampleVec
              , sampleCol
              , sampleVal
              ) where

import Types
import Numeric.Vector
import Numeric.Scalar (fromScalar, scalar)
import Control.Monad.State
import VecUtil (unpack3, unpack2)
import qualified SceneRandom as R
import VecUtil (vec3Of)

-- UTIL STUFF

fmod :: (RealFrac a) => a -> a -> a
fmod n d = n - (fromIntegral $ truncate $ n / d) * d

-- END UTIL STUFF

-- BRDF STUFF

sampleBRDF :: BRDF -> State SceneContext (Maybe Ray)
sampleBRDF BRDFEmpty = return Nothing
sampleBRDF (Emission _ _) = return Nothing
sampleBRDF (Diffuse _) = do
  ctx <- get
  let normal = rh_getNormal $ ss_getHit ctx
  let pos = (rh_getPos (ss_getHit ctx)) + (normal * (vec3Of 0.01))
  dir <- R.randVecHemisphere normal
  return $ Just $ Ray pos dir

probability :: BRDF -> Vec3d -> State SceneContext Double
probability BRDFEmpty _ = return 0
probability (Emission _ _) _ = return 0.15915494309
probability (Diffuse _) _ = return 0.15915494309 -- 1 / (2 * PI)

reflectance :: BRDF -> State SceneContext Vec3d
reflectance BRDFEmpty = return $ vec3 0 0 0

reflectance (Emission col str) = do
  col' <- sampleCol col
  str' <- sampleVal str
  return $ col' * (fromScalar (scalar str'))

reflectance (Diffuse col) = do
  col' <- sampleCol col
  return $ col'

reflectance (Glossy col roughness) = do
  col' <- sampleCol col
  return col'

-- END BRDF STUFF

-- VECTOR STUFF

sampleVec :: SVector -> State SceneContext Vec3d

sampleVec UV = do
  ss <- get
  let (u, v) = unpack2 $ rh_getTexCoord $ ss_getHit ss
  return $ vec3 u v 0

sampleVec Position = do
  ss <- get
  return $ rh_getPos $ ss_getHit ss

sampleVec Normal = do
  ss <- get
  return $ rh_getNormal $ ss_getHit ss

sampleVec (VecConst x y z) = return $ vec3 x y z

sampleVec (VecMath op l r) = do
  l' <- sampleVec l
  r' <- sampleVec r
  case op of
    VAdd -> return $ l' + r'
    VSub -> return $ l' - r'
    VMul -> return $ l' * r'
    VDiv -> return $ l' / r'
    VMod -> let (x,y,z) = unpack3 l'; (i,j,k) = unpack3 r' in return $ vec3 (fmod x i) (fmod y j) (fmod z k)
    VAbs -> return $ abs l'
    -- VDot -> return $ l' · r'
    -- VCross -> return $ l' × r'

sampleVec (CombineXYZ x y z) = do
  x' <- sampleVal x
  y' <- sampleVal y
  z' <- sampleVal z
  return $ vec3 x' y' z'

-- END VECTOR STUFF

-- COLOR STUFF

sampleCol :: SColor -> State SceneContext Vec3d

sampleCol ColVertex = do
  ss <- get
  return $ rh_getColor $ ss_getHit ss

sampleCol (ColVector v) = sampleVec v

sampleCol (ColConst r g b) = return $ vec3 r g b

sampleCol (CombineRGB vr vg vb) = do
  r <- sampleVal vr
  g <- sampleVal vg
  b <- sampleVal vb
  return $ vec3 r g b

-- END COLOR STUFF

-- VALUE STUFF
sampleVal :: SValue -> State SceneContext Double

sampleVal (ValConst x) = return x

sampleVal RayLength = do
  ss <- get
  return $ rh_getDistance $ ss_getHit ss

sampleVal (ValMath op l r) = do
  l' <- sampleVal l
  r' <- sampleVal r
  case op of
    Add -> return $ l' + r'
    Sub -> return $ l' - r'
    Mul -> return $ l' * r'
    Div -> return $ l' / r'
    Abs -> return $ abs l'
    Mod -> return $ fmod l' r'
    LessThan -> return $ if l' < r' then 1 else 0
    GreaterThan -> return $ if l' > r' then 1 else 0

sampleVal (SeparateX v) = sampleVec v >>= (return . unpack3) >>= \(x,_,_) -> return x
sampleVal (SeparateY v) = sampleVec v >>= (return . unpack3) >>= \(_,y,_) -> return y
sampleVal (SeparateZ v) = sampleVec v >>= (return . unpack3) >>= \(_,_,z) -> return z

sampleVal (SeparateR v) = sampleCol v >>= (return . unpack3) >>= \(r,_,_) -> return r
sampleVal (SeparateG v) = sampleCol v >>= (return . unpack3) >>= \(_,g,_) -> return g
sampleVal (SeparateB v) = sampleCol v >>= (return . unpack3) >>= \(_,_,b) -> return b

-- END VALUE STUFF


module Scene.Material.Diffuse where

import Numeric.LinearAlgebra

import Scene.Base
import Math.Color
import Math.Ray
import Scene.Monad

import Control.Applicative
import Control.Monad.Random

doWhileM :: Monad m => (a -> Bool) -> m a -> m a
doWhileM cond prog = prog >>= (\x -> if cond x then return x else doWhileM cond prog)

diffuse :: (Color -> Color) -> Material
diffuse absorb (Ray p normal) (Ray q incidence) = do
    s_unit <- doWhileM (\s -> s <.> s < 1) $ vec3 <$> getRandomR (-1,1) <*> getRandomR (-1,1) <*> getRandomR (-1,1)
    let s = scale (norm_2 normal) s_unit + normal
    ahead <- trace (Ray p s)
    return $! absorb ahead


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module NanoML.Learn where

import Control.Arrow (second)
import Control.Monad
import Data.Aeson
import qualified Data.ByteString.Lazy as LBS
import qualified Data.Csv as CSV
import Data.Function
import Data.List
import qualified Data.List as List
import Data.Maybe
import qualified Data.Set as Set
import qualified Data.Vector as V
import GHC.Generics
-- import Grenade
import Numeric.LinearAlgebra hiding (rank)
import Text.Read

import NanoML.Classify
import NanoML.Monad
import NanoML.Types

rankExprs :: Net -> [Feature] -> Prog -> [(Confidence, Constraint)]
rankExprs net fs p = second getThing <$!> rank net samples
  where
  samples = [ s {getThing = c}
            | s <- concatMap mkfsD tbad
            , c <- maybeToList (List.find (\c -> getThing s == fromJust (constraintSpan c)) cores)
            ]

  (tbad, cores, me) = case runEval stdOpts (typeProg p) of
    Left e -> ([], [], Just e)
    Right (tp, cs) -> (tp, Set.toList (mconcat cs), Nothing)

  mkfsD (TDFun _ _ pes) = mconcat (map (mkTypeOut . snd) pes)
  mkfsD (TDEvl _ e)     = mkTypeOut e
  mkfsD _               = mempty

  -- inSlice s = any (\c -> getSpan s == fromJust (constraintSpan c)) cores

  -- mkTypeOut :: TExpr -> [Sample ()]
  mkTypeOut te = actfold f [] te
    where
    f p e acc = (:acc) $ MkSample
                { getThing = infoSpan (texprInfo e)
                , getSample = vector $ concatMap (\(_,c) -> c p e) fs
                }

stdNet :: IO (Net, [Feature])
stdNet = loadNet "models/op+context+type-hidden-500.json" >>= \case
  Nothing -> fail "could not decode 'models/op+context+type-hidden-500.json'"
  Just net -> return (net, preds_tis ++ map only_ctx preds_tis_ctx ++ preds_tcon_ctx)

data Prediction = Change | NoChange
  deriving (Eq, Show)

-- predict :: Net -> Sample -> Prediction
-- predict = undefined

type Confidence = Double

rank :: Net -> [Sample a] -> [(Confidence, Sample a)]
rank net samples =
  sortBy (flip compare `on` fst)
    [ (r ! 1, s)
    | s <- samples
    , let r = runNet net (getSample s)
    ]

data Net = MkNet
  { hidden :: [Weights]
  , output :: Weights
  } deriving (Eq, Show, Generic)
instance FromJSON Net

data Weights = MkWeights
  { weights :: Matrix Double
  , biases :: Vector Double
  } deriving (Eq, Show, Generic)
instance FromJSON Weights where
  parseJSON = withObject "Weights" $ \v -> MkWeights
    <$> fmap mkmatrix (v .: "weights")
    <*> fmap vector (v .: "biases")
    where
    mkmatrix (ws :: [[Double]]) = matrix (length $ head ws) (concat ws)

loadNet :: FilePath -> IO (Maybe Net)
loadNet file = decode <$> LBS.readFile file

runNet :: Net -> Vector Double -> Vector Double
runNet MkNet{..} features =
  softmax (runLayer (foldl' runHidden features hidden) output)
  where
  runHidden fs ws = relu $ runLayer fs ws
  runLayer fs MkWeights{..} =
    fs <# weights + biases

relu :: Vector Double -> Vector Double
relu v = cmap (\x -> max 0 x) v

softmax :: Vector Double -> Vector Double
softmax v = cmap (\x -> exp x / summed) v
  where
  summed = sumElements (cmap exp v)

data Sample a = MkSample { getThing :: a, getSample :: Vector Double }
  deriving (Eq, Show, Generic)
instance CSV.FromRecord (Sample SrcSpan) where
  parseRecord r = MkSample
    <$> (maybe mzero return . readMaybe =<< CSV.parseField (r V.! 0))
    <*> fmap vector (CSV.parseRecord (V.drop 4 r))

loadSamples :: FilePath -> IO (Either String (V.Vector (Sample SrcSpan)))
loadSamples file = CSV.decode CSV.HasHeader <$> LBS.readFile file

-- net2network :: Net -> Network ts ds
-- net2network MkNet{..} = foldr addLayer NNil (output : reverse hidden)

-- addLayer :: Weights -> Network ts ds
--          -> Network (FullyConnected dx dy : Relu : ts) (d : d : ds)
-- addLayer ws net = mkLayer ws :~> Relu :~> net


{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE DeriveDataTypeable                       #-}
{-# LANGUAGE DeriveGeneric                            #-}
{-# LANGUAGE FlexibleInstances                        #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE KindSignatures                           #-}
{-# LANGUAGE MultiParamTypeClasses                    #-}
{-# LANGUAGE PatternSynonyms                          #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeFamilies                             #-}
{-# LANGUAGE TypeInType                               #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE UndecidableInstances                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module Backprop.Learn.Model.Neural.LSTM (
    -- * Basic LSTM
    LSTM, pattern LSTM
  , LSTMp(..), lstmForget, lstmInput, lstmUpdate, lstmOutput
  , LSTM'(..)
    -- * GRU
  , GRU, pattern GRU
  , GRUp(..), gruMemory, gruUpdate, gruOutput
  , GRU'(..)
  ) where

import           Backprop.Learn.Initialize
import           Backprop.Learn.Model.Class
import           Backprop.Learn.Model.Function
import           Backprop.Learn.Model.Neural
import           Backprop.Learn.Model.Regression
import           Backprop.Learn.Model.State
import           Control.DeepSeq
import           Data.Typeable
import           GHC.Generics                          (Generic)
import           GHC.TypeNats
import           Lens.Micro
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import           Numeric.OneLiner
import           Numeric.Opto.Ref
import           Numeric.Opto.Update
import qualified Data.Binary                           as Bi
import qualified Numeric.LinearAlgebra.Static          as H

-- | "Base" for 'LSTM'.  An 'LSTM' is just an 'LSTM'' where the second half
-- of the input vector is the previous output.
--
-- This is mostly an implementation detail.  It is recommended that you use
-- 'LSTM' and its constructor, instead of directly using this.
data LSTM' (i :: Nat) (o :: Nat) = LSTM'
  deriving (Typeable)

-- TODO: allow parameterize internal activation function?
-- TODO: Peepholes

-- | 'LSTM' layer parmateters
data LSTMp (i :: Nat) (o :: Nat) =
    LSTMp { _lstmForget :: !(FCp (i + o) o)
          , _lstmInput  :: !(FCp (i + o) o)
          , _lstmUpdate :: !(FCp (i + o) o)
          , _lstmOutput :: !(FCp (i + o) o)
          }
  deriving (Generic, Typeable, Show)

instance NFData (LSTMp i o)
instance (KnownNat i, KnownNat o) => Additive (LSTMp i o)
instance (KnownNat i, KnownNat o) => Scaling Double (LSTMp i o)
instance (KnownNat i, KnownNat o) => Metric Double (LSTMp i o)
instance (KnownNat i, KnownNat o, Ref m (LSTMp i o) v) => AdditiveInPlace m v (LSTMp i o)
instance (KnownNat i, KnownNat o, Ref m (LSTMp i o) v) => ScalingInPlace m v Double (LSTMp i o)
instance (KnownNat i, KnownNat o) => Bi.Binary (LSTMp i o)
instance (KnownNat i, KnownNat o) => Backprop (LSTMp i o)

-- | Forget biases initialized to 1
instance (KnownNat i, KnownNat o) => Initialize (LSTMp i o) where
    initialize d g = LSTMp <$> set (mapped . fcBias) 1 (initialize d g)
                           <*> initialize d g
                           <*> initialize d g
                           <*> initialize d g

instance (KnownNat i, KnownNat o) => Num (LSTMp i o) where
    (+)         = gPlus
    (-)         = gMinus
    (*)         = gTimes
    negate      = gNegate
    abs         = gAbs
    signum      = gSignum
    fromInteger = gFromInteger

instance (KnownNat i, KnownNat o) => Fractional (LSTMp i o) where
    (/)          = gDivide
    recip        = gRecip
    fromRational = gFromRational

instance (KnownNat i, KnownNat o) => Floating (LSTMp i o) where
    pi    = gPi
    sqrt  = gSqrt
    exp   = gExp
    log   = gLog
    sin   = gSin
    cos   = gCos
    asin  = gAsin
    acos  = gAcos
    atan  = gAtan
    sinh  = gSinh
    cosh  = gCosh
    asinh = gAsinh
    acosh = gAcosh
    atanh = gAtanh

lstmForget :: Lens' (LSTMp i o) (FCp (i + o) o)
lstmForget f x = (\y -> x { _lstmForget = y }) <$> f (_lstmForget x)

lstmInput :: Lens' (LSTMp i o) (FCp (i + o) o)
lstmInput f x = (\y -> x { _lstmInput = y }) <$> f (_lstmInput x)

lstmUpdate :: Lens' (LSTMp i o) (FCp (i + o) o)
lstmUpdate f x = (\y -> x { _lstmUpdate = y }) <$> f (_lstmUpdate x)

lstmOutput :: Lens' (LSTMp i o) (FCp (i + o) o)
lstmOutput f x = (\y -> x { _lstmOutput = y }) <$> f (_lstmOutput x)

instance (KnownNat i, KnownNat o, KnownNat io, io ~ (i + o)) => Learn (R io) (R o) (LSTM' i o) where
    type LParamMaybe (LSTM' i o) = 'Just (LSTMp i o)
    type LStateMaybe (LSTM' i o) = 'Just (R o)

    runLearn LSTM' (J_ p) x (J_ s) = (h, J_ s')
      where
        forget = logistic $ runLRp (p ^^. lstmForget) x
        input  = logistic $ runLRp (p ^^. lstmInput ) x
        update = tanh     $ runLRp (p ^^. lstmUpdate) x
        s'     = forget * s + input * update
        o      = logistic $ runLRp (p ^^. lstmOutput) x
        h      = o * tanh s'

-- | Long-term short-term memory layer
--
-- <http://colah.github.io/posts/2015-08-Understanding-LSTMs/>
--
-- @
-- instance 'Learn' ('R' i) (R o) ('LSTM' i o) where
--     type 'LParamMaybe' (LSTM i o) = ''Just' ('LSTMp' i o)
--     type 'LStateMaybe' (LSTM i o) = 'Just ('T2' (R o) (R o))     -- cell state, hidden state
-- @
type LSTM i o = Recurrent (R (i + o)) (R i) (R o) (R o) (LSTM' i o)

-- | Construct an 'LSTM'
pattern LSTM :: (KnownNat i, KnownNat o) => LSTM i o
pattern LSTM <- Rec { _recLearn = LSTM' }
  where
    LSTM = Rec
      { _recSplit = H.split
      , _recJoin  = (H.#)
      , _recLoop  = id
      , _recLearn = LSTM'
      }

-- | "Base" for 'GRU'.  An 'GRU' is just an 'GRU'' where the second half
-- of the input vector is the previous output.
--
-- This is mostly an implementation detail.  It is recommended that you use
-- 'GRU' and its constructor, instead of directly using this.
data GRU' (i :: Nat) (o :: Nat) = GRU'
  deriving (Typeable)

-- | 'GRU' layer parmateters
data GRUp (i :: Nat) (o :: Nat) =
    GRUp { _gruMemory :: !(FCp (i + o) o)
         , _gruUpdate :: !(FCp (i + o) o)
         , _gruOutput :: !(FCp (i + o) o)
         }
  deriving (Generic, Typeable, Show)

instance NFData (GRUp i o)
instance (KnownNat i, KnownNat o) => Additive (GRUp i o)
instance (KnownNat i, KnownNat o) => Scaling Double (GRUp i o)
instance (KnownNat i, KnownNat o) => Metric Double (GRUp i o)
instance (KnownNat i, KnownNat o, Ref m (GRUp i o) v) => AdditiveInPlace m v (GRUp i o)
instance (KnownNat i, KnownNat o, Ref m (GRUp i o) v) => ScalingInPlace m v Double (GRUp i o)
instance (KnownNat i, KnownNat o) => Bi.Binary (GRUp i o)
instance (KnownNat i, KnownNat o) => Backprop (GRUp i o)

instance (KnownNat i, KnownNat o) => Initialize (GRUp i o) where
    initialize d g = GRUp <$> initialize d g
                          <*> initialize d g
                          <*> initialize d g

instance (KnownNat i, KnownNat o) => Num (GRUp i o) where
    (+)         = gPlus
    (-)         = gMinus
    (*)         = gTimes
    negate      = gNegate
    abs         = gAbs
    signum      = gSignum
    fromInteger = gFromInteger

instance (KnownNat i, KnownNat o) => Fractional (GRUp i o) where
    (/)          = gDivide
    recip        = gRecip
    fromRational = gFromRational

instance (KnownNat i, KnownNat o) => Floating (GRUp i o) where
    pi    = gPi
    sqrt  = gSqrt
    exp   = gExp
    log   = gLog
    sin   = gSin
    cos   = gCos
    asin  = gAsin
    acos  = gAcos
    atan  = gAtan
    sinh  = gSinh
    cosh  = gCosh
    asinh = gAsinh
    acosh = gAcosh
    atanh = gAtanh

gruMemory :: Lens' (GRUp i o) (FCp (i + o) o)
gruMemory f x = (\y -> x { _gruMemory = y }) <$> f (_gruMemory x)

gruUpdate :: Lens' (GRUp i o) (FCp (i + o) o)
gruUpdate f x = (\y -> x { _gruUpdate = y }) <$> f (_gruUpdate x)

gruOutput :: Lens' (GRUp i o) (FCp (i + o) o)
gruOutput f x = (\y -> x { _gruOutput = y }) <$> f (_gruOutput x)

instance (KnownNat i, KnownNat o, KnownNat io, io ~ (i + o)) => Learn (R io) (R o) (GRU' i o) where
    type LParamMaybe (GRU' i o) = 'Just (GRUp i o)
    type LStateMaybe (GRU' i o) = 'Nothing

    runLearn GRU' (J_ p) = stateless $ \x ->
        let z      = logistic $ runLRp (p ^^. gruMemory) x
            r      = logistic $ runLRp (p ^^. gruUpdate) x
            r'     = 1 # r
            h'     = tanh     $ runLRp (p ^^. gruOutput) (r' * x)
        in  (1 - z) * snd (split @i x) + z * h'

-- | Gated Recurrent Unit
--
-- <http://colah.github.io/posts/2015-08-Understanding-LSTMs/>
--
-- @
-- instance 'Learn' ('R' i) (R o) ('GRU' i o) where
--     type 'LParamMaybe' (GRU i o) = ''Just' ('GRUp' i o)
--     type 'LStateMaybe' (GRU i o) = 'Just (R o)
-- @
type GRU i o = Recurrent (R (i + o)) (R i) (R o) (R o) (GRU' i o)

-- | Construct an 'GRU'
pattern GRU :: (KnownNat i, KnownNat o) => GRU i o
pattern GRU <- Rec { _recLearn = GRU' }
  where
    GRU = Rec
      { _recSplit = H.split
      , _recJoin  = (H.#)
      , _recLoop  = id
      , _recLearn = GRU'
      }


module Test.Test where

import Test.Framework (defaultMainWithArgs, testGroup, Test)
import Test.Framework.Providers.HUnit (testCase)
import Test.Framework.Providers.QuickCheck2 (testProperty)

import Test.QuickCheck (Property, NonNegative(..))
import Test.QuickCheck.Monadic (assert, monadicIO, run)
import Test.HUnit (assertBool)

import qualified Numeric.LinearAlgebra as LA
--import qualified Graphics.Rendering.Plot.HMatrix as G

import qualified Math.FROG.Tools as FR
import Math.FROG.Retrieval
import Math.FROG.Types
--import Test.Types

type PulseFn = Double -> Double -> LA.Complex Double

-----------------
-- Example Pulses
-----------------
singleGaus :: PulseFn
singleGaus sz k = exp ((-pi) / sz * (k - sz/2.0)^2) LA.:+ 0.0

doubleGaus :: PulseFn
doubleGaus sz k = (singleGaus sz (k - (sz / 8.0))) + (singleGaus sz (k + (sz / 8.0)))

mkPulse :: PulseFn -> Int -> ComplexSignal
mkPulse form sz = LA.buildVector sz mkElem
    where
    dsz = (fromIntegral sz) :: Double
    mkElem = form dsz . fromIntegral

caseRetrieval :: ComplexSignal -> IO ()
caseRetrieval field = do
    let sz = LA.dim field
    let dsz = (fromIntegral sz) :: Double
    let trc = FR.mkTrace field field

    o <- retrieve SHG trc

    let rfield = FR.center . FR.flatPolar2Complex $ o
    let mags = (LA.mapVector LA.magnitude) rfield

    putStrLn . show $ calcLoss SHG trc o

    --G.mplot [LA.linspace sz (0::Double, dsz-1), mags]

    --G.imshow trc

    --G.imshow (FR.mkTrace rfield rfield)

-----------------
---- List of Tests
-----------------

tests :: [Test]
tests =
    [ testGroup "Retrieval" [ testCase "Gaussian" $ caseRetrieval (mkPulse singleGaus 32)
                            , testCase "Double Gaussian" $ caseRetrieval (mkPulse doubleGaus 32)
                            ]
    ]

-- convenient whilst in ghci
runAllTests :: IO ()
runAllTests = defaultMainWithArgs tests []

runTests :: String -> IO ()
runTests p = defaultMainWithArgs tests ["--select-tests", p]

runGroup :: Int -> IO ()
runGroup i = defaultMainWithArgs [tests !! i] []


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
--
-- Module      :  AI.Network.RNN.Data
-- Copyright   :  (c) JP Moresmau
-- License     :  BSD3
--
-- Maintainer  :  JP Moresmau <jp@moresmau.fr>
-- Stability   :  experimental
-- Portability :
--
-- | Conversion from and to text data
--
-----------------------------------------------------------------------------

module AI.Network.RNN.Data  where

import qualified Data.Text as T
import qualified Data.Map as DM
import qualified Data.Set as DS
import Data.Maybe
import Data.Tuple
import Data.List
import Control.Monad.Random as R
import Control.Monad
import Numeric.LinearAlgebra.HMatrix  as M
import Data.Foldable as F (toList)

import qualified Data.SDR as SDR
import qualified Data.IntSet as I

import AI.Network.RNN.Types
import AI.Network.RNN.Util
-- import Debug.Trace

-- | Transform text into training data
-- the encoding of each Char is equilateral encoding
textToTrainData :: T.Text -> TrainData ((DM.Map Int Char),Matrix Double)
textToTrainData t =
    let (ids,m) = T.foldl' toIDs ([],DM.empty) t
        sz = DM.size m
        em = equilateralEncoding sz
        is = map (toArr em) $ reverse ids
        charMap = DM.fromList $ map swap $ DM.assocs m
    in
        TrainData (M.fromList (replicate (sz-1) 0): init is) is (sz-1) (charMap,em)
    where
        toIDs :: ([Int],DM.Map Char Int) -> Char -> ([Int],DM.Map Char Int)
        toIDs (ids,m) c =
            let mmyid = DM.lookup c m
            in case mmyid of
                Just myid -> (myid:ids,m)
                Nothing   ->
                    let myid = DM.size m
                    in (myid:ids,DM.insert c myid m)
--        toArr :: Int -> Int -> Vector Double
--        toArr sz idx = M.fromList $ replicate idx 0 ++ [1] ++ replicate (sz-idx-1) 0
        toArr em idx = flatten $ em ? [idx]

-- | Decode the data from the equilateral encoding
dataToText :: ((DM.Map Int Char),Matrix Double) -> [Vector Double] -> T.Text
dataToText (m,em)  = T.pack . F.toList . fmap toC
    where
        toC :: Vector Double -> Char
        -- toC ds = fromJust $ DM.lookup (maxIndex ds) m
        toC ds = fromJust $ DM.lookup (equilateralDecoding em ds) m

-- | Generate random text with the probabilities given by the vector (not used)
randDataToText :: RandomGen g => DM.Map Int Char -> [Vector Double] -> Rand g T.Text
randDataToText m  = liftM T.pack . mapM (toC . M.toList)
    where
        toC :: RandomGen g => [Double] -> Rand g Char
        toC ls = do
            let m1 = normalize $ map (\(ix,d)->(fromJust $ DM.lookup ix m,d)) $ zip [0..] ls
                m2 = map (\(a,b)->(a,toRational b)) m1
                -- s  = sum $ map (exp . snd) m1
                -- m3 = map (\(a,b)-> (a,exp(b)/s)) m1
            R.fromList m2
            --return $ fst $ last $ sortBy (comparing snd) m3

-- | Generate text using equilateral encoding and a given network
generate :: (RandomGen g,RNNEval a)
    => ((DM.Map Int Char),Matrix Double) -- ^ Character map
    -> Int  -- ^ Number of characters to generate
    -> Int -- ^ Size of the vectors for each characters
    -> a -- network
    -> Rand g T.Text
generate m nb sz rnn =generate' m nb sz rnn dataToText

-- | Generate text using sparse encoding and a given network
generateS :: (RandomGen g,RNNEval a)
    => SDR.SDRSet Char -- ^ Character map
    -> Int  -- ^ Number of characters to generate
    -> Int -- ^ Size of the vectors for each characters
    -> a -- network
    -> Rand g T.Text
generateS m nb sz rnn =generate' m nb sz rnn dataToTextS

-- | Generate text given a decoding function
-- current implementation is not random!
generate' :: (RandomGen g,RNNEval a) => m -> Int -> Int -> a ->
    (m -> [Vector Double] -> T.Text) -> Rand g T.Text
generate' m nb sz rnn f =  do
    let (_,_,alls) = foldl' go (rnn,M.fromList $ replicate sz 0,[]) [1..nb]
    --randDataToText m alls
    return $ f m alls
    where
        go :: (RNNEval a) => (a,Vector Double,[Vector Double]) -> Int -> (a,Vector Double,[Vector Double])
        go (rnn1,is1,oss) _ =
            let (rnn2,os) = evalStep rnn1 is1
                os2 = norm os
            in (rnn2,os2,oss ++ [os])
        norm :: Vector Double -> Vector Double
        norm os = let
            mx = maximum $ M.toList os
            --in M.assoc (M.size os) 0 [(mx,1)]
            in M.fromList $ map (\a->if a == mx then 1 else 0) $ M.toList os

-- | Train data using sparse encoding: 2 values in the vector are set to one for each character
-- This was a test to try to reduce the size of the network but is not efficient
textToTrainDataS :: T.Text -> TrainData (SDR.SDRSet Char)
textToTrainDataS t =
    let charSet = T.foldl' (flip DS.insert) DS.empty t
        sz  = DS.size charSet
        markers = 1
        total = sz * markers
        sdr = SDR.build (sz*markers) markers charSet
        is = T.foldr (\c l->enc total sdr c : l) [] t
    in TrainData (M.fromList (replicate total 0): init is) is total sdr
    where
        enc :: Int -> SDR.SDRSet Char -> Char -> Vector Double
        enc total sdr c =
            let is=SDR.encode sdr c
            in M.assoc total 0 $ zip (I.toList is) (repeat 1)

--textToTrainDataS t =
--    let (ids,m) = T.foldl' toIDs ([],DM.empty) t
--        sz = sparseSize $ DM.size m
--        spr = DM.fromList $ zip [0..] $ sparse $ DM.size m
--        is = map (toArr spr sz) $ reverse ids
--        charMap = DM.fromList $ map (\(c,i)->(fromJust $ DM.lookup i spr,c)) $ DM.assocs m
--    in
--        TrainData (M.fromList (replicate sz 0): init is) is sz charMap
--    where
--        toIDs :: ([Int],DM.Map Char Int) -> Char -> ([Int],DM.Map Char Int)
--        toIDs (ids,m) c =
--            let mmyid = DM.lookup c m
--            in case mmyid of
--                Just myid -> (myid:ids,m)
--                Nothing   ->
--                    let myid = DM.size m
--                    in (myid:ids,DM.insert c myid m)
--        toArr :: DM.Map Int (Int,Int) -> Int -> Int -> Vector Double
--        toArr m sz idx = let
--            (i1,i2) = fromJust $ DM.lookup idx m
--            in M.fromList $ replicate i1 0 ++ [1] ++ replicate (i2-i1-1) 0 ++ [1] ++ replicate (sz-i2-1) 0

-- | Generate text from sparse encoding data
dataToTextS :: SDR.SDRSet Char -> [Vector Double] -> T.Text
dataToTextS sdr = T.pack . map b
    where b v =
            let is = I.fromList $ map fst $ filter (\(_,c)->c>=0.5) $ zipWith (\a b->(b,a)) (M.toList v) [0..]
            in SDR.best sdr is
--dataToTextS m  = T.pack . F.toList . fmap toC
--    where
--        toC :: Vector Double -> Char
--        toC ds = let
--            ix1 = maxIndex ds
--            ix2 = maxIndex $ accum ds const [(ix1,0)]
--            tpl = if ix1 < ix2
--                    then (ix1,ix2)
--                    else (ix2,ix1)
--            mc = DM.lookup tpl m
--            in fromMaybe (fromJust $ DM.lookup (0,1) m) mc
--         --fromJust $ DM.lookup (maxIndex ds) m


{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE DataKinds #-}
{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_GHC -fno-warn-unused-imports #-}

module Numeric.LinearAlgebra where

import Numeric.LinearAlgebra.Internal


{-# LANGUAGE OverloadedStrings,FlexibleContexts,TypeFamilies #-}

import           Control.Applicative
import           Control.Arrow                 ((&&&))
import           Control.Lens                  hiding ((|>))
import           Control.Monad                 (liftM,void)
import           Control.Monad.Primitive.Class
import           Data.Csv                      (HasHeader(..))
import           Data.List                     (mapAccumL)
import           Data.Ratio
import           Data.Time
import qualified Data.Vector                   as V (Vector)
import qualified Data.Vector.Generic           as V hiding (Vector)
import           Graphics.EasyPlot
import           Numeric.LinearAlgebra
import           Numeric.LinearAlgebra.Util
import           Pipes
import qualified Pipes.ByteString              as PB
import           Pipes.Csv                     (decode)
import           Pipes.Safe                    (SafeT, runSafeT)
import qualified Pipes.Safe.Prelude            as PS
import qualified Pipes.Vector                  as PV
import           System.IO                     (IOMode (ReadMode))

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

main :: IO ()
main = do
  runLRStrategy "Cumulative sum of " absReturn cmlAbsReturn 0
  runLRStrategy "Cumulative product of " pctReturn cmlPctReturn 0
  runKFStrategy "Cumulative sum of " absReturn cmlAbsReturn 0

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

-- Linear regression strategy inspired by http://businessforecastblog.com/predicting-the-sp-500-or-the-spy-exchange-traded-fund/

runLRStrategy :: String -> Return -> Return -> Matrix Double -> IO ()
runLRStrategy titleprefix freturn fcmlreturn minret = do
  (mx,my,vdtd) <- prepareData 30 freturn
  let
    t = 2484 -- number of training data points
    mm = takeRows t mx <\> takeRows t my
    my' = mx <> mm
    mr = step (my' - minret) * my
    vy = head $ toColumns my
    vr = head $ toColumns mr
    vcy = fcmlreturn vy
    vcr = fcmlreturn vr
    lcy = V.toList vcy
    lcr = V.toList vcr
    ldt = map (fromRational . dayToYear . readDay) $ V.toList vdtd
  void $ plot' [] X11 [
    Data2D [Title $ titleprefix ++ "in sample strategy returns", Style Lines] [] $ take t $ zip ldt lcr,
    Data2D [Title $ titleprefix ++ "out of sample strategy returns", Style Lines] [] $ drop t $ zip ldt lcr,
    Data2D [Title $ titleprefix ++ "base returns", Style Lines] [] $ zip ldt lcy
    ]

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

-- Kalman filter strategy

runKFStrategy :: String -> Return -> Return -> Matrix Double -> IO ()
runKFStrategy titleprefix freturn fcmlreturn minret = do
  (mx,my,vdtd) <- prepareData 30 freturn
  let
    f = diag $ 61 |> repeat 1 :: Matrix Double
    q = diag $ 61 |> repeat 0.2 :: Matrix Double
    r = (1><1) [0.1] :: Matrix Double
    x0 = 61 |> repeat 0
    p0 = q
    s0 = State x0 p0
    (s',y') = mapAccumL (\st (xt,yt) -> (kalman (System f (fromRows [xt]) q r) st yt,(xt * (sX st)) @> 1)) s0 $ zip (toRows mx) (toRows my)
    my' = fromColumns [fromList y']
    mr = (step (my' - minret)) * my
    vy = head $ toColumns my
    vr = head $ toColumns mr
    vcy = fcmlreturn vy
    vcr = fcmlreturn vr
    lcy = V.toList vcy
    lcr = V.toList vcr
    ldt = map (fromRational . dayToYear . readDay) $ V.toList vdtd
  print $ sX s'
  void $ plot' [] X11 [
    Data2D [Title $ titleprefix ++ "strategy returns", Style Lines] [] $ zip ldt lcr,
    Data2D [Title $ titleprefix ++ "base returns", Style Lines] [] $ zip ldt lcy
    ]

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

-- Kalman code from hmatrix examples

data System = System {kF, kH, kQ, kR :: Matrix Double}
data State = State {sX :: Vector Double , sP :: Matrix Double} deriving Show
type Measurement = Vector Double

kalman :: System -> State -> Measurement -> State
kalman (System f h q r) (State x p) z = State x' p' where
    px = f <> x                            -- prediction
    pq = f <> p <> trans f + q             -- its covariance
    y  = z - h <> px                       -- residue
    cy = h <> pq <> trans h + r            -- its covariance
    k  = pq <> trans h <> inv cy           -- kalman gain
    x' = px + k <> y                       -- new state
    p' = (ident (dim x) - k <> h) <> pq    -- its covariance

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

--TODO: streamline inefficient conversion between lists and vectors

prepareData :: Int -> Return -> IO (Matrix Double,Matrix Double,V.Vector String)
prepareData g freturn = do
  vsp' <- V.map ((^._1) &&& (^._7)) `liftM` readCSV "gspc.csv"
  vvx' <- V.map ((^._1) &&& (^._7)) `liftM` readCSV "vix.csv"
  let
    vm' = V.fromList $ mergeEqual (V.toList vsp') (V.toList vvx') :: V.Vector (String,Double,Double)
    vdtd = V.map (^._1) (V.drop g vm') :: V.Vector String
    [vsp'',vvx''] = [V.map] <*> [(^._2),(^._3)] <*> [vm']
    [vsp,vvx] = map V.convert [vsp'',vvx''] :: [Vector Double]
    [vspd,vvxd] = map freturn [vsp,vvx]
    l = V.length vdtd - 1
    [mspds,mvxds] = map (\v -> fromColumns [subVector (g-n) l v | n <- [0..g]]) [vspd,vvxd]
    mx = (konst 1 (l,1)) ! (dropColumns 1 mspds) ! (dropColumns 1 mvxds)
    my = takeColumns 1 mspds
  return (mx,my,vdtd)

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

type PriceData = (String,Double,Double,Double,Double,Double,Double)

type ParsedPriceData = Either String PriceData

readCSV :: FilePath -> IO (V.Vector PriceData)
readCSV f = runSafeT $ runEffect $ PV.runToVectorP $
    (hoist lift (decode HasHeader (PS.withFile f ReadMode (PB.fromHandle))))
    >-> rightP
    >-> PV.toVector

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

mergeEqual :: Ord a => [(a,b)] -> [(a,b)] -> [(a,b,b)]
mergeEqual x [] = []
mergeEqual [] y = []
mergeEqual x@((xi,xv):xs) y@((yi,yv):ys)
  | xi == yi = (xi,xv,yv) : mergeEqual xs ys
  | xi < yi = mergeEqual xs y
  | xi > yi = mergeEqual x ys

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

type Return = Vector Double -> Vector Double

absReturn :: Return
absReturn v = V.zipWith (-) (V.tail v) (V.init v)

pctReturn :: Return
pctReturn v = (V.zipWith (-) (V.tail v) (V.init v))/(V.init v)

cmlAbsReturn :: Return
cmlAbsReturn = V.scanl (+) 0

cmlPctReturn :: Return
cmlPctReturn = V.scanl (\x y -> x * (1+y)) 1

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

readDay :: String -> Day
readDay = read

dayToYear :: Day -> Ratio Integer
dayToYear day =
  let
    (y,_,_) = toGregorian day
    ystart = fromGregorian y 1 1
    yend = fromGregorian y 12 31
    dd = diffDays day ystart
    dy = diffDays yend ystart
  in
    y % 1 + dd % dy

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

rightP :: (Monad m) => Pipe (Either a b) b m r
rightP = for cat $ either (const $ return ()) yield

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

instance MonadPrim m => MonadPrim (SafeT m) where
  type BasePrimMonad (SafeT m) = BasePrimMonad m
  liftPrim = lift . liftPrim