{-# LANGUAGE AllowAmbiguousTypes                      #-}
{-# LANGUAGE ApplicativeDo                            #-}
{-# LANGUAGE DeriveAnyClass                           #-}
{-# LANGUAGE DeriveDataTypeable                       #-}
{-# LANGUAGE DeriveGeneric                            #-}
{-# LANGUAGE DerivingVia                              #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE FlexibleInstances                        #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE InstanceSigs                             #-}
{-# LANGUAGE KindSignatures                           #-}
{-# LANGUAGE MultiParamTypeClasses                    #-}
{-# LANGUAGE PatternSynonyms                          #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE StandaloneDeriving                       #-}
{-# LANGUAGE TemplateHaskell                          #-}
{-# 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.Regression (
  -- * Linear Regression
    linReg, logReg
  , LRp(..), lrAlpha, lrBeta, runLRp
  -- * Reshape
  , reshapeLRpInput, reshapeLRpOutput
  , expandLRpInput, expandLRpOutput
  , premuteLRpInput, permuteLRpOutput
  -- * ARIMA
  , arima, autoregressive, movingAverage, arma
  , ARIMAp(..), ARIMAs(..)
  , arimaPhi, arimaTheta, arimaConstant, arimaYPred, arimaYHist, arimaEHist
  ) where

import           Backprop.Learn.Initialize
import           Backprop.Learn.Model.Combinator
import           Backprop.Learn.Model.Function
import           Backprop.Learn.Model.Types
import           Backprop.Learn.Regularize
import           Control.DeepSeq
import           Control.Monad.Primitive
import           Data.Finite
import           Data.Kind
import           Data.List
import           Data.Maybe
import           Data.Proxy
import           Data.Type.Equality
import           Data.Typeable
import           GHC.Generics                          (Generic)
import           GHC.TypeLits.Compare
import           GHC.TypeNats
import           Lens.Micro.TH
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import           Numeric.LinearAlgebra.Static.Vector
import           Numeric.OneLiner
import           Numeric.Opto.Ref
import           Numeric.Opto.Update hiding            ((<.>))
import           Statistics.Distribution
import           Unsafe.Coerce
import qualified Data.Binary                           as Bi
import qualified Data.Vector.Generic.Sized             as SVG
import qualified Data.Vector.Sized                     as SV
import qualified Data.Vector.Storable.Sized            as SVS
import qualified Numeric.LinearAlgebra                 as HU
import qualified Numeric.LinearAlgebra.Static          as H
import qualified System.Random.MWC                     as MWC

-- | Linear Regression parameter
data LRp (i :: Nat) (o :: Nat) = LRp
    { _lrAlpha :: !(R o)
    , _lrBeta  :: !(L o i)
    }
  deriving stock     (Generic, Typeable, Show)
  deriving anyclass  (NFData, Linear Double, Metric Double, Bi.Binary, Initialize, Backprop)

deriving via (GNum (LRp i o)) instance (KnownNat i, KnownNat o) => Num (LRp i o)
deriving via (GNum (LRp i o)) instance (KnownNat i, KnownNat o) => Fractional (LRp i o)
deriving via (GNum (LRp i o)) instance (KnownNat i, KnownNat o) => Floating (LRp i o)

makeLenses ''LRp

instance (PrimMonad m, KnownNat i, KnownNat o) => Mutable m (LRp i o) where
    type Ref m (LRp i o) = GRef m (LRp i o)
    thawRef   = gThawRef
    freezeRef = gFreezeRef
    copyRef   = gCopyRef
instance (PrimMonad m, KnownNat i, KnownNat o) => LinearInPlace m Double (LRp i o)

instance (KnownNat i, KnownNat o) => Regularize (LRp i o) where
    rnorm_1 = rnorm_1 . _lrBeta
    rnorm_2 = rnorm_2 . _lrBeta
    lasso r LRp{..} = LRp { _lrAlpha = 0
                          , _lrBeta  = lasso r _lrBeta
                          }
    ridge r LRp{..} = LRp { _lrAlpha = 0
                          , _lrBeta  = ridge r _lrBeta
                          }

instance (KnownNat i, KnownNat o, PrimMonad m) => Learnable m (LRp i o)

runLRp
    :: (KnownNat i, KnownNat o, Reifies s W)
    => BVar s (LRp i o)
    -> BVar s (R i)
    -> BVar s (R o)
runLRp lrp x = (lrp ^^. lrBeta) #> x + (lrp ^^. lrAlpha)

-- | Reshape an 'LRp' to take more or less inputs.  If more, new parameters
-- are initialized randomly according to the given distribution.
reshapeLRpInput
    :: (ContGen d, PrimMonad m, KnownNat i, KnownNat i', KnownNat o)
    => d
    -> MWC.Gen (PrimState m)
    -> LRp i o
    -> m (LRp i' o)
reshapeLRpInput d g (LRp α β) =
    LRp α <$> reshapeLCols d g β

-- | Reshape an 'LRp' to return more or less outputs  If more, new
-- parameters are initialized randomly according to the given distribution.
reshapeLRpOutput
    :: (ContGen d, PrimMonad m, KnownNat i, KnownNat o, KnownNat o')
    => d
    -> MWC.Gen (PrimState m)
    -> LRp i o
    -> m (LRp i o')
reshapeLRpOutput d g (LRp α β) =
    LRp <$> reshapeR d g α
        <*> reshapeLRows d g β

linReg
    :: (KnownNat i, KnownNat o)
    => Model ('Just (LRp i o)) 'Nothing (R i) (R o)
linReg = modelStatelessD (\(PJust p) -> runLRp p)

logReg
    :: (KnownNat i, KnownNat o)
    => Model ('Just (LRp i o)) 'Nothing (R i) (R o)
logReg = funcD logistic <~ linReg

-- | Adjust an 'LRp' to take extra inputs, initialized randomly.
--
-- Initial contributions to each output is randomized.
expandLRpInput
    :: (PrimMonad m, ContGen d, KnownNat i, KnownNat j, KnownNat o)
    => LRp i o
    -> d
    -> MWC.Gen (PrimState m)
    -> m (LRp (i + j) o)
expandLRpInput LRp{..} d g = LRp _lrAlpha . (_lrBeta H.|||) <$> initialize d g

-- | Adjust an 'LRp' to return extra ouputs, initialized randomly
expandLRpOutput
    :: (PrimMonad m, ContGen d, KnownNat i, KnownNat o, KnownNat p)
    => LRp i o
    -> d
    -> MWC.Gen (PrimState m)
    -> m (LRp i (o + p))
expandLRpOutput LRp{..} d g = do
    newAlpha <- initialize d g
    newBeta  <- initialize d g
    pure (LRp (_lrAlpha H.# newAlpha) (_lrBeta H.=== newBeta))

-- | Premute (or remove) inputs
--
-- Removed inputs will simply have their contributions removed from each
-- output.
premuteLRpInput
    :: (KnownNat i, KnownNat o)
    => SV.Vector i' (Finite i)
    -> LRp i o
    -> LRp i' o
premuteLRpInput is p = p { _lrBeta = colsL . fmap (β `SV.index`) $ is }
  where
    β = lCols (_lrBeta p)

-- | Premute (or remove) outputs
permuteLRpOutput
    :: (KnownNat i, KnownNat o)
    => SV.Vector o' (Finite o)
    -> LRp i o
    -> LRp i o'
permuteLRpOutput is LRp{..} =
    LRp { _lrAlpha = vecR . SVG.convert . fmap (α `SVS.index`) $ is
        , _lrBeta  = rowsL . fmap (β `SV.index`) $ is
        }
  where
    α = rVec _lrAlpha
    β = lRows _lrBeta

-- | 'ARIMA' parmaeters
data ARIMAp :: Nat -> Nat -> Type where
    ARIMAp :: { _arimaPhi      :: !(R p)
              , _arimaTheta    :: !(R q)
              , _arimaConstant :: !Double
              }
           -> ARIMAp p q
  deriving stock     (Generic, Typeable, Show)
  deriving anyclass  (NFData, Linear Double, Metric Double, Initialize, Backprop, Bi.Binary)

deriving via (GNum (ARIMAp p q)) instance Num (ARIMAp p q)
deriving via (GNum (ARIMAp p q)) instance Fractional (ARIMAp p q)
deriving via (GNum (ARIMAp p q)) instance Floating (ARIMAp p q)

makeLenses ''ARIMAp

-- | 'ARIMA' state
data ARIMAs :: Nat -> Nat -> Nat -> Type where
    ARIMAs :: { _arimaYPred :: !Double
              , _arimaYHist :: !(R (p + d))
              , _arimaEHist :: !(R q)
              }
          -> ARIMAs p d q
  deriving stock     (Generic, Typeable, Show)
  deriving anyclass  (NFData, Linear Double, Metric Double, Initialize, Backprop, Bi.Binary)

deriving via (NoRegularize (ARIMAs p d q)) instance Regularize (ARIMAs p d q)
deriving via (GNum (ARIMAs p d q)) instance Num (ARIMAs p d q)
deriving via (GNum (ARIMAs p d q)) instance Fractional (ARIMAs p d q)
deriving via (GNum (ARIMAs p d q)) instance Floating (ARIMAs p d q)

makeLenses ''ARIMAs

arima
    :: forall p d q. (KnownNat p, KnownNat d, KnownNat q)
    => Model ('Just (ARIMAp p q)) ('Just (ARIMAs p d q)) Double Double
arima = modelD $ \(PJust p) x (PJust s) ->
    let d :: L p (p + d)
        d = difference
        e  = x - (s ^^. arimaYPred)
        y  = (p ^^. arimaConstant)
           + (p ^^. arimaPhi  ) <.> (constVar d #> (s ^^. arimaYHist))
           + (p ^^. arimaTheta) <.> (s ^^. arimaEHist)
        yHist' = case Proxy @1 %<=? Proxy @(p + d) of
          LE Refl -> single y # constVar dropLast #> (s ^^. arimaYHist)
          NLE _ _ -> 0
        eHist' = case Proxy @1 %<=? Proxy @q of
          LE Refl -> single e # constVar dropLast #> (s ^^. arimaEHist)
          NLE _ _ -> 0
        s' = isoVar3 ARIMAs (\(ARIMAs pr yh eh) -> (pr,yh,eh))
                y
                yHist'
                eHist'
    in  (y, PJust s')

autoregressive
    :: KnownNat p
    => Model ('Just (ARIMAp p 0)) ('Just (ARIMAs p 0 0)) Double Double
autoregressive = arima

movingAverage
    :: KnownNat q
    => Model ('Just (ARIMAp 0 q)) ('Just (ARIMAs 0 0 q)) Double Double
movingAverage = arima

arma
    :: (KnownNat p, KnownNat q)
    => Model ('Just (ARIMAp p q)) ('Just (ARIMAs p 0 q)) Double Double
arma = arima

monosquare :: forall n. (n <=? (n ^ 2)) :~: 'True
monosquare = unsafeCoerce Refl

dropLast :: forall n. (KnownNat n, 1 <= n) => L (n - 1) n
dropLast = case monosquare @n of
    Refl -> vecL . SVS.generate $ \ij ->
      let i :: Finite n
          j :: Finite (n - 1)
          (i, j) = separateProduct ij
      in  if fromIntegral @_ @Int i == fromIntegral j
            then 1
            else 0

single :: Reifies s W => BVar s Double -> BVar s (R 1)
single = konst

difference'
    :: Int                  -- ^ initial
    -> Int                  -- ^ target
    -> HU.Matrix Double     -- ^ target x initial
difference' n m = foldl' go (HU.ident m) [m + 1 .. n]
  where
    go x k = x HU.<> d k
    d k = HU.build (k-1, k) $ \i j ->
        case round @_ @Int (j - i) of
          0 -> 1
          1 -> -1
          _ -> 0

difference :: forall n m. (KnownNat n, KnownNat m) => L n (n + m)
difference = fromJust . H.create $ difference' (n + m) n
  where
    n = fromIntegral $ natVal (Proxy @n)
    m = fromIntegral $ natVal (Proxy @m)

instance (PrimMonad m, KnownNat p, KnownNat q) => Mutable m (ARIMAp p q) where
    type Ref m (ARIMAp p q) = GRef m (ARIMAp p q)
    thawRef = gThawRef
    freezeRef = gFreezeRef
    copyRef = gCopyRef

instance (PrimMonad m, KnownNat p, KnownNat d, KnownNat q) => Mutable m (ARIMAs p d q) where
    type Ref m (ARIMAs p d q) = GRef m (ARIMAs p d q)
    thawRef = gThawRef
    freezeRef = gFreezeRef
    copyRef = gCopyRef

instance (KnownNat p, KnownNat q, PrimMonad m)  => LinearInPlace m Double (ARIMAp p q)
instance (KnownNat p, KnownNat d, KnownNat q, PrimMonad m) => LinearInPlace m Double (ARIMAs p d q)

instance (KnownNat p, KnownNat q) => Regularize (ARIMAp p q) where
    rnorm_1 ARIMAp{..} = rnorm_1 _arimaPhi + rnorm_1 _arimaTheta
    rnorm_2 ARIMAp{..} = rnorm_2 _arimaPhi + rnorm_2 _arimaTheta
    lasso r ARIMAp{..} = ARIMAp { _arimaPhi      = lasso r _arimaPhi
                                , _arimaTheta    = lasso r _arimaTheta
                                , _arimaConstant = 0
                                }
    ridge r ARIMAp{..} = ARIMAp { _arimaPhi      = ridge r _arimaPhi
                                , _arimaTheta    = ridge r _arimaTheta
                                , _arimaConstant = 0
                                }

instance (KnownNat p, KnownNat q, PrimMonad m) => Learnable m (ARIMAp p q)
instance (KnownNat p, KnownNat d, KnownNat q, PrimMonad m) => Learnable m (ARIMAs p d q)


{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE DeriveGeneric #-}
module FokkerPlanck.BrownianMotion
  ( Particle(..)
  , generatePath
  , moveParticle
  , thetaPlus
  , scalePlus
  , generateRandomNumber
  , generatePath'
  ) where

import           Control.DeepSeq
import           Control.Monad
import           Data.DList                          as DL
import           Data.Ix
import           Data.List                           as L
import           GHC.Generics                        (Generic)
import           Statistics.Distribution
import           Statistics.Distribution.Exponential
import           Statistics.Distribution.Normal
import           Statistics.Distribution.Uniform
import           System.Random.MWC
import           Text.Printf
import           Utils.Distribution

data Particle =
  Particle {-# UNPACK #-}!Double -- \phi
           {-# UNPACK #-}!Double -- \rho
           {-# UNPACK #-}!Double -- \theta
           {-# UNPACK #-}!Double -- r
           {-# UNPACK #-}!Double -- weight
  deriving (Show,Generic)

instance NFData Particle

{-# INLINE getParticleRho #-}
getParticleRho :: Particle -> Double
getParticleRho (Particle _ rho _ _ _ ) = rho

thetaCheck :: Double -> Double
thetaCheck theta =
  if theta < -pi
    then thetaCheck (theta + 2 * pi)
    else if theta >= pi
           then thetaCheck (theta - 2 * pi)
           else theta

{-# INLINE scaleCheck #-}
scaleCheck :: Double -> Double -> Double
scaleCheck maxScale scale =
  if scale >= (1 / maxScale) && scale <= maxScale
    then scale
    else error
           (printf
              "scaleCheck: %.2f is out of boundary (0,%.2f)\n"
              scale
              maxScale)

{-# INLINE thetaPlus #-}
thetaPlus :: Double -> Double -> Double
thetaPlus !x !y =
  let z = x + y
   in thetaCheck z

{-# INLINE scalePlus #-}
scalePlus :: Double -> Double -> Double
scalePlus x delta = x * exp delta

{-# INLINE scalePlusPeriodic #-}
scalePlusPeriodic :: Double -> Double -> Double -> Double
scalePlusPeriodic !logMaxScale !delta !x
  | z >= m = exp (z - 2 * m)
  | z < -m = exp (2 * m + z)
  | otherwise = exp z
  where
    !m = logMaxScale
    !z = log x + delta

{-# INLINE rhoCutoff #-}
rhoCutoff :: Double -> Double
rhoCutoff !rho =
  if rho < 0
    then 0
    else rho

{-# INLINE generateRandomNumber #-}
generateRandomNumber ::
     (Distribution d, ContGen d) => GenIO -> Maybe d -> IO Double
generateRandomNumber !gen dist =
  case dist of
    Nothing -> return 0
    Just d  -> genContVar d gen

{-# INLINE moveParticle #-}
moveParticle :: Double -> Particle -> Particle
moveParticle deltaT (Particle phi rho theta r0 v) =
  let !r = r0 * deltaT
      !x = theta - phi
      !cosX = cos x
      !newPhi = phi `thetaPlus` atan2 (r * sin x) (rho + r * cosX)
      !newRho = sqrt (rho * rho + r * r + 2 * r * rho * cosX)
  in Particle newPhi newRho theta r0 v

{-# INLINE diffuseParticle #-}
diffuseParticle :: Double -> Double -> Double -> Particle -> Particle
diffuseParticle !deltaTheta !deltaScale maxScale (Particle phi rho theta r v) =
  Particle
    phi
    rho
    (theta `thetaPlus` deltaTheta)
    (r `scalePlus` deltaScale)
    v

brownianMotion ::
     (Distribution d1, ContGen d1, Distribution d2, ContGen d2, Distribution d3)
  => GenIO
  -> Maybe d1
  -> Maybe d2
  -> Maybe d3
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Particle
  -> DList Particle
  -> IO (DList Particle)
brownianMotion randomGen thetaDist scaleDist poissonDist deltaT maxRho maxR tao stdR2 particle xs = do
  deltaTheta <- generateRandomNumber randomGen thetaDist
  deltaScale <- generateRandomNumber randomGen scaleDist
  let newParticle@(Particle phi rho theta r v) = moveParticle deltaT particle
      ys =
        if r <= maxR && r > 1 / maxR && rho <= maxRho && rho >= 1 / maxRho
          then DL.cons
                 (Particle phi rho theta r (1 - gaussian2DPolar rho stdR2))
                 xs
          else xs
  t <- genContVar (uniformDistr 0 1) randomGen :: IO Double
  if t > exp (1 / (-tao))
    then return ys
    else brownianMotion
           randomGen
           thetaDist
           scaleDist
           poissonDist
           deltaT
           maxRho
           maxR
           tao
           stdR2
           (diffuseParticle deltaTheta deltaScale maxR newParticle)
           ys

{-# INLINE generatePath #-}
generatePath ::
     (Distribution d1, ContGen d1, Distribution d2, ContGen d2, Distribution d3)
  => Maybe d1
  -> Maybe d2
  -> Maybe d3
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> GenIO
  -> IO (DList Particle)
generatePath thetaDist scaleDist poissonDist maxRho maxR tao deltaT stdR2 randomGen =
  brownianMotion
    randomGen
    thetaDist
    scaleDist
    poissonDist
    deltaT
    maxRho
    maxR
    tao
    stdR2
    (Particle 0 1 0 1 1)
    DL.empty

{-# INLINE moveParticle' #-}
moveParticle' ::  Particle -> Particle
moveParticle' (Particle phi rho theta r v) =
  let !x = theta - phi
      !cosX = cos x
      !newPhi = phi `thetaPlus` atan2 (r * sin x) (rho + r * cosX)
      !newRho = sqrt (rho * rho + r * r + 2 * r * rho * cosX)
   in Particle newPhi newRho theta r v

{-# INLINE diffuseParticle' #-}
diffuseParticle' :: (Distribution d, ContGen d)
  => GenIO ->  Double -> Maybe d -> Double -> Particle -> IO Particle
diffuseParticle' randomGen thetaSigma scaleDist deltaT (Particle phi rho theta r v) = do
  deltaScale <- generateRandomNumber randomGen scaleDist
  let r1 = r `scalePlus` deltaScale
  deltaTheta <-
    genContVar (normalDistr 0 (thetaSigma * sqrt (r1 * deltaT))) randomGen
  return $ Particle phi rho (theta `thetaPlus` deltaTheta) r1 v

brownianMotion' ::
     (Distribution d, ContGen d)
  => GenIO
  -> Double
  -> Maybe d
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Particle
  -> DList Particle
  -> IO (DList Particle)
brownianMotion' randomGen thetaSigma scaleDist poissonLambda deltaT maxRho maxR tao stdR2 particle xs = do
  let newParticle@(Particle phi rho theta r v) = moveParticle deltaT particle
      ys =
        if r <= maxR && r > 1 / maxR && rho <= maxRho && rho >= 1 / maxRho
          then DL.cons (Particle phi rho theta r (1 - gaussian2DPolar rho stdR2)) xs
          else xs
  t <- genContVar (uniformDistr 0 1) randomGen :: IO Double
  if t > exp (1 / (-tao / (r * deltaT)))
    then return ys
    else do
      diffusedParticle <-
        diffuseParticle' randomGen thetaSigma scaleDist deltaT newParticle
      brownianMotion'
        randomGen
        thetaSigma
        scaleDist
        poissonLambda
        deltaT
        maxRho
        maxR
        tao
        stdR2
        diffusedParticle
        ys

{-# INLINE generatePath' #-}
generatePath' ::
     (Distribution d, ContGen d)
  => Double
  -> Maybe d
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> GenIO
  -> IO (DList Particle)
generatePath' thetaSigma scaleDist poissonLambda maxRho maxR tao deltaT stdR2 randomGen = do
  brownianMotion'
    randomGen
    thetaSigma
    scaleDist
    poissonLambda
    deltaT
    maxRho
    maxR
    tao
    stdR2
    (Particle 0 1 0 1 1)
    DL.empty


module STCR2Z2T0S0EndModal where

import           Data.Array.Repa         as R
import           Data.Binary             (decodeFile)
import           Data.Complex
import           Data.List               as L
import           DFT.Plan
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           Image.Transform         (normalizeValueRange)
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types
import           Utils.Array

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:scale0FreqsStr:scaleFreqsStr:histFilePath:numIterationStr:writeSourceFlagStr:cutoffRadiusEndPointStr:cutoffRadiusStr:reversalFactorStr:cStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scale0Freq = read scale0FreqsStr :: Double
      scaleFreq = read scaleFreqsStr :: Double
      scale0Freqs = [-scale0Freq .. scale0Freq]
      scaleFreqs = [-scaleFreq .. scaleFreq]
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      cutoffRadiusEndPoint = read cutoffRadiusEndPointStr :: Int
      cutoffRadius = read cutoffRadiusStr :: Int
      reversalFactor = read reversalFactorStr :: Double
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2Z2T0S0EndModal.noindex"
      a = 20
      b = 7
      c = read cStr :: Int
      endPointFilePath =
        folderPath </>
        (printf
           "EndPoint_%d_%d_%d_%d_%d_%d_%.2f_%.2f_%d_%d_%d_%f.dat"
           numPoint
           (round thetaFreq :: Int)
           (round scaleFreq :: Int)
           (round maxScale :: Int)
           (round tao :: Int)
           cutoffRadiusEndPoint
           thetaSigma
           scaleSigma
           a
           b
           c
           reversalFactor)
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z2T0S0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          tao
          theta0Freqs
          thetaFreqs
          scale0Freqs
          scaleFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length scale0Freqs
             , L.length theta0Freqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0)
  arrR2Z2T0S0 <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      (cutoff cutoffRadius radialArr)
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      theta0Freqs
      scale0Freqs
  arrR2Z2T0S0EndPoint <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      (cutoff cutoffRadiusEndPoint radialArr)
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      theta0Freqs
      scale0Freqs
  plan <- makeR2Z2T0S0Plan emptyPlan arrR2Z2T0S0
  let a' = round $ (fromIntegral a) * (sqrt 2) / 2
      b' = round $ (fromIntegral b) * (sqrt 2) / 2
      c' = round $ (fromIntegral c) * (sqrt 2) / 2
      xs =
        ([R2S1RPPoint (i, i, 0, 0) | i <- [a',a' + b' .. c']] L.++
         [R2S1RPPoint (i, -i, 0, 0) | i <- [-a',-(a' + b') .. -c']] L.++
         [R2S1RPPoint (i, i, 0, 0) | i <- [-a',-(a' + b') .. -c']] L.++
         [R2S1RPPoint (i, -i, 0, 0) | i <- [a',a' + b' .. c']] L.++
         [R2S1RPPoint (i, 0, 0, 0) | i <- [a,a + b .. c]] L.++
         [R2S1RPPoint (0, i, 0, 0) | i <- [-a,-(a + b) .. -c]] L.++
         [R2S1RPPoint (i, 0, 0, 0) | i <- [-a,-(a + b) .. -c]] L.++
         [R2S1RPPoint (0, i, 0, 0) | i <- [a,a + b .. c]])
      bias = computeBiasR2T0S0' numPoint numPoint theta0Freqs scale0Freqs xs
      eigenVec =
        computeInitialEigenVectorR2T0S0
          numPoint
          numPoint
          theta0Freqs
          scale0Freqs
          thetaFreqs
          scaleFreqs
          xs
  powerMethodR2Z2T0S0EndModal
    plan
    folderPath
    numPoint
    numPoint
    numOrientation
    thetaFreqs
    theta0Freqs
    numScale
    scaleFreqs
    scale0Freqs
    arrR2Z2T0S0EndPoint
    arrR2Z2T0S0
    numIteration
    writeSourceFlag
    (printf
       "_%d_%d_%d_%d_%d_%d_%.2f_%.2f_%f"
       numPoint
       (round thetaFreq :: Int)
       (round scaleFreq :: Int)
       (round maxScale :: Int)
       (round tao :: Int)
       cutoffRadiusEndPoint
       thetaSigma
       scaleSigma
       reversalFactor)
    reversalFactor
    bias
    eigenVec


{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-}
{-# LANGUAGE FlexibleContexts #-}
module Neural.Training
  ( TrainingExample(..)
  , TestingExample(..)
  , train
  , backpropagation
  )
where

import           Data.List.Split                ( chunksOf )
import           Foreign.Marshal.Utils          ( fromBool )
import           Control.Monad
import           Control.Monad.Writer
import           Control.Monad.State
import           Control.Monad.Random
import           System.Random
import           System.Random.Shuffle          ( shuffleM )
import           Numeric.LinearAlgebra
import           Text.Printf
import           GHC.Generics                   ( Generic )
import           Control.DeepSeq

import           Neural.Activation
import qualified Neural.Network                as N
import           Neural.Network          hiding ( weights
                                                , biases
                                                )
import           Neural.Layer                   ( Layer(..) )

-- | A training example with an input and desired output in vector form.
data TrainingExample a = TrainingExample (Vector a) (Vector a)
    deriving (Show, Generic, NFData)
-- | A training example with an input and desired output as an index.
data TestingExample a  = TestingExample (Vector a) Int
    deriving (Show, Generic, NFData)

-- | The Hadamar product, elementwise multiplication.
(⊙) :: Num (Vector a) => Vector a -> Vector a -> Vector a
u ⊙ v = u * v

-- | Train a neural network over some amount of epochs given some training
-- data and the size of the mini batches.
train
  :: (RandomGen g, Fractional a, Numeric a, Num (Vector a))
  => ActivationFunction a      -- ^ The activation function
  -> ActivationFunction' a     -- ^ The derivative of the activation function
  -> a                         -- ^ The learning rate
  -> [TrainingExample a]       -- ^ Training data
  -> Maybe [TestingExample a]  -- ^ Optional testing data
  -> Int                       -- ^ Number of ephocs
  -> Int                       -- ^ Mini batch size
  -> Network a                 -- ^ The untrained network
  -> StateT g (Writer [String]) (Network a) -- ^ The trained network
train σ σ' η trainingData testData epochs mbs net = do
  let tellProgress epoch net' = case testData of
        Nothing        -> tell ["Epoch " ++ show epoch ++ " complete"]
        Just testData' -> tell
          [printf "Epoch %d: %d / %d (%.2f%%)" epoch evaluation n p]
         where
          evaluation = evaluate σ testData' net'
          n          = length testData'
          p          = fromIntegral evaluation / fromIntegral n * 100 :: Double
  let trainEpoch net' epoch = do
        shuffledTrainingData <- state $ runRand $ shuffleM trainingData
        let miniBatches = chunksOf mbs shuffledTrainingData
        let tmb         = trainMiniBatch σ σ' η
        let trainedNet  = foldr tmb net' miniBatches
        lift $ tellProgress epoch trainedNet
        return trainedNet

  foldM trainEpoch net [1 .. epochs]

-- | Evaluate the network assuming the output is the index of the neuron with
-- the highest activation.
evaluate
  :: (Numeric a, Num (Vector a))
  => ActivationFunction a     -- ^ The activation function
  -> [TestingExample a]       -- ^ Testing data
  -> Network a                -- ^ The untrained network
  -> Int                      -- ^ The amount of tests the network passed
evaluate σ testData net =
  let eval (TestingExample x y) = maxIndex (fst $ feedforward σ x net) == y
      results = map eval testData
  in  foldl (\acc r -> acc + fromBool r) 0 results

-- | Train a neural network for a mini batch of training data.
trainMiniBatch
  :: (Fractional a, Numeric a, Num (Vector a))
  => ActivationFunction a    -- ^ The activation function
  -> ActivationFunction' a   -- ^ The derivative of the activation function
  -> a                       -- ^ The learning rate
  -> [TrainingExample a]     -- ^ A mini batch of training examples
  -> Network a               -- ^ The untrained network
  -> Network a               -- ^ The trained network
trainMiniBatch σ σ' η miniBatch net =
  let bp example = backpropagation σ σ' example net
      η'     = η / fromIntegral (length miniBatch)
      ws0    = map (konst 0 . size) (N.weights net)
      bs0    = map (konst 0 . size) (N.biases net)
      deltas = map (unzip . bp) miniBatch
      f (w, b) (dw, db) = (zipWith (+) w dw, zipWith (+) b db)
      (nabla_w, nabla_b) = foldl f (ws0, bs0) deltas
      ws = zipWith (\w nw -> w - scale η' nw) (N.weights net) nabla_w
      bs = zipWith (\b nb -> b - scale η' nb) (N.biases net) nabla_b
  in  newNetwork ws bs

-- | Find the gradient of the cost function with respects to a networks
-- weights and biases for one training example.
backpropagation
  :: (Numeric a, Num (Vector a))
  => ActivationFunction a    -- ^ The activation function
  -> ActivationFunction' a   -- ^ The derivative of the activation function
  -> TrainingExample a       -- ^ A single training example
  -> Network a               -- ^ The neural network
  -> [(Matrix a, Vector a)]  -- ^ The weight and bias gradient of the
                               -- quadratic cost function
backpropagation σ σ' (TrainingExample x y) net@(Network layers) =
  let (as, zs) = unzip $ feedforwards σ x net   -- Feed forward and save the
                                                -- results
      as_      = x : as
      nabla_aC = last as - y                    -- The gradient of the cost
                                                -- function
      calc_δ (l, z) δ = bp2 σ' (weights l) δ z  -- A function that calculates
                                                -- the error for the previous
                                                -- layer
      δ0      = bp1 σ' nabla_aC (last zs)       -- The error of the last layer
      δs      = scanr calc_δ δ0 (zip (tail layers) zs)
      nabla_b = δs                              -- The 3rd equation
      nabla_w = zipWith outer δs as_            -- The 4th equation
  in  zip nabla_w nabla_b

-- | The first equation of backpropagation.
bp1
  :: (Numeric a, Num (Vector a))
  => ActivationFunction' a    -- ^ The derivative of the activation function
  -> Vector a                 -- ^ The gradient of the cost function
                                -- with respects to the activation
  -> Vector a                 -- ^ The weighted input
  -> Vector a                 -- ^ The error in the output layer
bp1 _σ' nabla_aC z = nabla_aC ⊙ σ' z where σ' = cmap _σ'

-- | The second equation of backpropagation.
bp2
  :: (Numeric a, Num (Vector a))
  => ActivationFunction' a    -- ^ The derivative of the activation function
  -> Matrix a                 -- ^ The weights of the next layer
  -> Vector a                 -- ^ The error of the next layer
  -> Vector a                 -- ^ The weighted input of the current layer
  -> Vector a                 -- ^ The error of the current layer
bp2 _σ' w δ z = (tr' w #> δ) ⊙ σ' z where σ' = cmap _σ'


{-# LANGUAGE OverloadedStrings #-}

module Statistics.Classification.ConfusionMatrix.Binary (
    ConfusionMatrix
  , trueConfMatrix
  , resultCount, classCount, totalCount
  , actualPositive, actualNegative, predPositive, predNegative
  , truePosRate, falseNegRate, falsePosRate, trueNegRate
  , prevalence, accuracy, errorRate, recall, precision, specificity
  , oddsRatio
  , phiCorrelation
  , fScore, f1Score
  , actualEntropy, predEntropy, condActualEntropy, condPredEntropy, totalCondEntropy
  , mutualInformation
  , asBinaryConfusion
  , asMultiConfusion
  ) where

import Data.Monoid (Sum(..))

import Data.Csv ((.:), (.=), DefaultOrdered(..), FromNamedRecord(..), ToNamedRecord(..), header, namedRecord)
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set

import Statistics.Classification.ConfusionMatrix.Multi (MultiConfusionMatrix(..), PosNeg
  , truePositiveCount, falsePositiveCount, falseNegativeCount, trueNegativeCount)
import Statistics.Classification.Types (ClassificationResult(..))

data ConfusionMatrix = ConfusionMatrix {
    truePositive :: {-# UNPACK #-}!(Sum Int)
  , falsePositive :: {-# UNPACK #-}!(Sum Int)
  , falseNegative :: {-# UNPACK #-}!(Sum Int)
  , trueNegative :: {-# UNPACK #-}!(Sum Int)
  } deriving (Eq, Show, Read)

instance Semigroup ConfusionMatrix where
  cm0 <> cm1 = ConfusionMatrix {
      truePositive  = truePositive  cm0 <> truePositive  cm1
    , falsePositive = falsePositive cm0 <> falsePositive cm1
    , falseNegative = falseNegative cm0 <> falseNegative cm1
    , trueNegative  = trueNegative  cm0 <> trueNegative  cm1
    }

instance Monoid ConfusionMatrix where
  mempty = ConfusionMatrix mempty mempty mempty mempty
  mappend = (<>)

instance DefaultOrdered ConfusionMatrix where
  headerOrder _ = header [
      "true_positive"
    , "false_positive"
    , "true_negative"
    , "false_negative"
    ]

instance ToNamedRecord ConfusionMatrix where
  toNamedRecord conf = namedRecord [
      "true_positive"  .= getSum (truePositive conf)
    , "false_positive" .= getSum (falsePositive conf)
    , "true_negative"  .= getSum (trueNegative conf)
    , "false_negative" .= getSum (falseNegative conf)
    ]

instance FromNamedRecord ConfusionMatrix where
  parseNamedRecord r = ConfusionMatrix <$>
        fmap Sum (r .: "true_positive")
    <*> fmap Sum (r .: "false_positive")
    <*> fmap Sum (r .: "false_negative")
    <*> fmap Sum (r .: "true_negative")

-- create confusion matrix given assumed true results
-- i.e., classifiedPredicted is the number predicted correctly
--       classifiedActual    is the number of actual samples in class
-- Give numbers for both true and false labels
trueConfMatrix :: ClassificationResult Int -> ClassificationResult Int -> ConfusionMatrix
trueConfMatrix pos neg = ConfusionMatrix {
    truePositive  = Sum $ classifiedPredicted pos
  , falseNegative = Sum $ classifiedActual pos - classifiedPredicted pos
  , falsePositive = Sum $ classifiedActual neg - classifiedPredicted neg
  , trueNegative  = Sum $ classifiedPredicted neg
  }

resultCount :: ConfusionMatrix -> ClassificationResult Bool -> Int
resultCount cm (ClassificationResult actual prd) = classCount cm prd actual

classCount :: ConfusionMatrix -> Bool -> Bool -> Int
classCount cm prd actual = getSum $ if prd
  then (if actual then  truePositive cm else falsePositive cm)
  else (if actual then falseNegative cm else  trueNegative cm)

totalCount :: ConfusionMatrix -> Int
totalCount cm = getSum $ truePositive cm <> falsePositive cm <> falseNegative cm <> trueNegative cm

actualPositive, actualNegative, predPositive, predNegative :: ConfusionMatrix -> Int
actualPositive cm = getSum $  truePositive cm <> falseNegative cm
actualNegative cm = getSum $ falsePositive cm <>  trueNegative cm
predPositive   cm = getSum $  truePositive cm <> falsePositive cm
predNegative   cm = getSum $ falseNegative cm <>  trueNegative cm

truePosRate, falseNegRate, falsePosRate, trueNegRate :: ConfusionMatrix -> Double
truePosRate = recall
falseNegRate cm = fromIntegral (getSum (falseNegative cm)) / fromIntegral (actualPositive cm)
falsePosRate cm = fromIntegral (getSum (falsePositive cm)) / fromIntegral (actualNegative cm)
trueNegRate = specificity

prevalence, accuracy, errorRate, recall, precision, specificity :: ConfusionMatrix -> Double
prevalence cm = fromIntegral (actualPositive cm) / fromIntegral (totalCount cm)
accuracy cm = fromIntegral (getSum (truePositive cm <> trueNegative cm)) / fromIntegral (totalCount cm)
errorRate cm = fromIntegral (getSum (falsePositive cm <> falseNegative cm)) / fromIntegral (totalCount cm)
recall cm = fromIntegral (getSum (truePositive cm)) / fromIntegral (actualPositive cm)
precision cm = fromIntegral (getSum (truePositive cm)) / fromIntegral (predPositive cm)
specificity cm = fromIntegral (getSum (trueNegative cm)) / fromIntegral (actualNegative cm)

oddsRatio :: ConfusionMatrix -> Double
oddsRatio (ConfusionMatrix (Sum tp) (Sum fp) (Sum fn) (Sum tn)) =
  fromIntegral (tp * tn) / fromIntegral (fn * fp)

phiCorrelation :: ConfusionMatrix -> Double
phiCorrelation (ConfusionMatrix (Sum tp) (Sum fp) (Sum fn) (Sum tn)) =
  fromIntegral (tp * tn - fp * fn) / sqrt (fromIntegral $ (tp + fn) * (tn + fp) * (tp + fp) * (tn + fn))

fScore :: Double -> ConfusionMatrix -> Double
fScore alpha cm = recip $ alpha * recip (precision cm) + (1 - alpha) * recip (recall cm)

f1Score :: ConfusionMatrix -> Double
f1Score = fScore 0.5

-- Entropy calculations
binomEntropy :: Double -> Double
binomEntropy p = p * log p + (1 - p) * log (1 - p)

discreteEntropy :: Int -> Int -> Double
discreteEntropy n tot = binomEntropy (fromIntegral n / fromIntegral tot)

actualEntropy, predEntropy :: ConfusionMatrix -> Double
actualEntropy cm = discreteEntropy (actualPositive cm) (totalCount cm)
predEntropy cm = discreteEntropy (predPositive cm) (totalCount cm)

condActualEntropy :: ConfusionMatrix -> Bool -> Double
condActualEntropy cm plbl = if plbl
  then discreteEntropy (getSum (truePositive cm)) (predPositive cm)
  else discreteEntropy (getSum (trueNegative cm)) (predNegative cm)

condPredEntropy :: ConfusionMatrix -> Bool -> Double
condPredEntropy cm albl = if albl
  then discreteEntropy (getSum (truePositive cm)) (actualPositive cm)
  else discreteEntropy (getSum (trueNegative cm)) (actualNegative cm)

totalCondEntropy :: ConfusionMatrix -> Double
totalCondEntropy cm = ppos * condActualEntropy cm True + (1 - ppos) * condActualEntropy cm False
  where ppos = fromIntegral (predPositive cm) / fromIntegral (totalCount cm)

mutualInformation :: ConfusionMatrix -> Double
mutualInformation cm = actualEntropy cm - totalCondEntropy cm

asBinaryConfusion :: Ord a => PosNeg a -> MultiConfusionMatrix a -> ConfusionMatrix
asBinaryConfusion posneg mcm = ConfusionMatrix {
    truePositive = Sum $ truePositiveCount mcm posneg
  , falsePositive = Sum $ falsePositiveCount mcm posneg
  , falseNegative = Sum $ falseNegativeCount mcm posneg
  , trueNegative = Sum $ trueNegativeCount mcm posneg
  }

asMultiConfusion :: ConfusionMatrix -> MultiConfusionMatrix Bool
asMultiConfusion cm = MultiConfusionMatrix {
    confCounts = Map.fromList [
        (ClassificationResult True  True,  getSum (truePositive cm))
      , (ClassificationResult False True,  getSum (falsePositive cm))
      , (ClassificationResult True  False, getSum (falseNegative cm))
      , (ClassificationResult False False, getSum (trueNegative cm))
      ]
  , confClasses = Set.fromList [True, False]
  }


{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE Strict          #-}

module Optimization.Types where

import Control.Lens
import Numeric.LinearAlgebra            ( Numeric )
import Numeric.LinearAlgebra.Data
import Data.Vector.Storable       as VS ( map
                                        , singleton
                                        , zipWith )

{- |Basic "zippy" Num instance for storable vectors. Note that, while this is defined
    for some types of Vector by hmatrix, it unfortunately does not account for any
    Num (Vector a). -}
instance Numeric a => Num (Vector a) where
  a + b    = VS.zipWith (+)   a b
  a - b    = VS.zipWith (-)   a b
  a * b    = VS.zipWith (*)   a b
  abs a    = VS.map  (abs)    a
  signum a = VS.map  (signum) a
  fromInteger i = VS.singleton (fromInteger i)

{- |Termination criterion for a numerical optimization method.
    We terminate either on maxit iterations or once our guesses are within
    threshold of one another. -}
data TerminationCriterion a = TerminateAt {
  _maxit     :: Maybe Int, -- ^Maximum number of iterations (if any).
  _threshold :: Maybe a    -- ^Threshold used to declare convergence (if any).
  } deriving (Show, Eq)

-- |Step parameters for Nelder-Mead simplex search. a must be Storable and Num.
data NMStepParameters a = NMStepParameters {
  _epsilon :: a, -- ^Offset around initial guess to construct starting simplex.
  _alpha   :: a, -- ^Reflection coefficient for Nelder-Mead algorithm.
  _gamma   :: a, -- ^Expansion coefficient for Nelder-Mead algorithm.
  _rho     :: a, -- ^Contraction coefficient for Nelder-Mead algorithm.
  _sigma   :: a  -- ^Shrink coefficient for Nelder-Mead algorithm.
  } deriving (Show, Eq)

-- |Standard Nelder-Mead parameters.
defaultStep :: NMStepParameters Double
defaultStep = NMStepParameters 1.0 1.0 2.0 0.5 0.5

-- |The Nelder-Mead simplex.
type Simplex a = Matrix a

-- |Create a TerminationCriterion with just a convergence threshold.
terminateAt ::
  a ->                   -- ^Convengence threshold
  TerminationCriterion a -- ^A TerminationCriterion for optimization.
terminateAt t = TerminateAt Nothing (Just t)

-- |Create a TerminationCriterion with just a maximum iteration threshold.
terminateAfter ::
  Int               ->   -- ^Maximum number of iterations.
  TerminationCriterion a -- ^A TerminationCriterion for optimization.
terminateAfter m = TerminateAt (Just m) Nothing

-- |Create a TerminationCriterion with a convergence threshold and maximum iteration
-- |limit.
terminateAtOrAfter ::
  a      ->              -- ^Convergence threshold
  Int    ->              -- ^Maximum number of iterations
  TerminationCriterion a -- ^A TerminationCriterion for optimization.
terminateAtOrAfter t m = TerminateAt (Just m) (Just t)

makeLenses ''TerminationCriterion
makeLenses ''NMStepParameters


{-# LANGUAGE TypeOperators, CPP #-}
{-# LANGUAGE FlexibleInstances  #-}
{-# LANGUAGE FlexibleContexts   #-}
{-# LANGUAGE DefaultSignatures   #-}
{-# LANGUAGE ScopedTypeVariables #-}
----------------------------------------------------------------------
-- |
-- Module      :   Data.AdditiveGroup
-- Copyright   :  (c) Conal Elliott and Andy J Gill 2008
-- License     :  BSD3
--
-- Maintainer  :  conal@conal.net, andygill@ku.edu
-- Stability   :  experimental
--
-- Groups: zero, addition, and negation (additive inverse)
----------------------------------------------------------------------

module Data.AdditiveGroup
  (
    AdditiveGroup(..), sumV
  , Sum(..), inSum, inSum2
  ) where

import Prelude hiding (foldr)

import Control.Applicative
#if !(MIN_VERSION_base(4,8,0))
import Data.Monoid (Monoid(..))
import Data.Foldable (Foldable)
#endif
import Data.Foldable (foldr)
import Data.Complex hiding (magnitude)
import Data.Ratio
#if !(MIN_VERSION_base(4,11,0))
import Data.Semigroup (Semigroup(..))
#endif
import Foreign.C.Types (CSChar, CInt, CShort, CLong, CLLong, CIntMax, CFloat, CDouble)

import Data.MemoTrie

import Data.VectorSpace.Generic
import qualified GHC.Generics as Gnrx
import GHC.Generics (Generic, (:*:)(..))

infixl 6 ^+^, ^-^

-- | Additive group @v@.
class AdditiveGroup v where
  -- | The zero element: identity for '(^+^)'
  zeroV :: v
  default zeroV :: (Generic v, AdditiveGroup (VRep v)) => v
  zeroV = Gnrx.to (zeroV :: VRep v)
  -- | Add vectors
  (^+^) :: v -> v -> v
  default (^+^) :: (Generic v, AdditiveGroup (VRep v)) => v -> v -> v
  v ^+^ v' = Gnrx.to (Gnrx.from v ^+^ Gnrx.from v' :: VRep v)
  -- | Additive inverse
  negateV :: v -> v
  default negateV :: (Generic v, AdditiveGroup (VRep v)) => v -> v
  negateV v = Gnrx.to (negateV $ Gnrx.from v :: VRep v)
  -- | Group subtraction
  (^-^) :: v -> v -> v
  v ^-^ v' = v ^+^ negateV v'

-- | Sum over several vectors
sumV :: (Foldable f, AdditiveGroup v) => f v -> v
sumV = foldr (^+^) zeroV

instance AdditiveGroup () where
  zeroV     = ()
  () ^+^ () = ()
  negateV   = id

-- For 'Num' types:
--
-- instance AdditiveGroup n where {zeroV=0; (^+^) = (+); negateV = negate}

#define ScalarTypeCon(con,t) \
  instance con => AdditiveGroup (t) where {zeroV=0; (^+^) = (+); negateV = negate}

#define ScalarType(t) ScalarTypeCon((),t)

ScalarType(Int)
ScalarType(Integer)
ScalarType(Float)
ScalarType(Double)
ScalarType(CSChar)
ScalarType(CInt)
ScalarType(CShort)
ScalarType(CLong)
ScalarType(CLLong)
ScalarType(CIntMax)
ScalarType(CFloat)
ScalarType(CDouble)
ScalarTypeCon(Integral a,Ratio a)

instance (RealFloat v, AdditiveGroup v) => AdditiveGroup (Complex v) where
  zeroV   = zeroV :+ zeroV
  (^+^)   = (+)
  negateV = negate

-- Hm.  The 'RealFloat' constraint is unfortunate here.  It's due to a
-- questionable decision to place 'RealFloat' into the definition of the
-- 'Complex' /type/, rather than in functions and instances as needed.

instance (AdditiveGroup u,AdditiveGroup v) => AdditiveGroup (u,v) where
  zeroV             = (zeroV,zeroV)
  (u,v) ^+^ (u',v') = (u^+^u',v^+^v')
  negateV (u,v)     = (negateV u,negateV v)

instance (AdditiveGroup u,AdditiveGroup v,AdditiveGroup w)
    => AdditiveGroup (u,v,w) where
  zeroV                  = (zeroV,zeroV,zeroV)
  (u,v,w) ^+^ (u',v',w') = (u^+^u',v^+^v',w^+^w')
  negateV (u,v,w)        = (negateV u,negateV v,negateV w)

instance (AdditiveGroup u,AdditiveGroup v,AdditiveGroup w,AdditiveGroup x)
    => AdditiveGroup (u,v,w,x) where
  zeroV                       = (zeroV,zeroV,zeroV,zeroV)
  (u,v,w,x) ^+^ (u',v',w',x') = (u^+^u',v^+^v',w^+^w',x^+^x')
  negateV (u,v,w,x)           = (negateV u,negateV v,negateV w,negateV x)

-- Standard instance for an applicative functor applied to a vector space.
instance AdditiveGroup v => AdditiveGroup (a -> v) where
  zeroV   = pure   zeroV
  (^+^)   = liftA2 (^+^)
  negateV = fmap   negateV

-- Maybe is handled like the Maybe-of-Sum monoid
instance AdditiveGroup a => AdditiveGroup (Maybe a) where
  zeroV = Nothing
  Nothing ^+^ b'      = b'
  a' ^+^ Nothing      = a'
  Just a' ^+^ Just b' = Just (a' ^+^ b')
  negateV = fmap negateV

{-

Alexey Khudyakov wrote:

  I looked through vector-space package and found lawless instance. Namely Maybe's AdditiveGroup instance

  It's group so following relation is expected to hold. Otherwise it's not a group.
  > x ^+^ negateV x == zeroV

  Here is counterexample:

  > let x = Just 2 in x ^+^ negateV x == zeroV
  False

  I think it's not possible to sensibly define group instance for
  Maybe a at all.

I see that the problem here is in distinguishing 'Just zeroV' from
Nothing. I could fix the Just + Just line to use Nothing instead of Just
zeroV when a' ^+^ b' == zeroV, although doing so would require Eq a and
hence lose some generality. Even so, the abstraction leak would probably
show up elsewhere.

Hm.

-}

-- Memo tries
instance (HasTrie u, AdditiveGroup v) => AdditiveGroup (u :->: v) where
  zeroV   = pure   zeroV
  (^+^)   = liftA2 (^+^)
  negateV = fmap   negateV

-- | Monoid under group addition.  Alternative to the @Sum@ in
-- "Data.Monoid", which uses 'Num' instead of 'AdditiveGroup'.
newtype Sum a = Sum { getSum :: a }
  deriving (Eq, Ord, Read, Show, Bounded)

instance Functor Sum where
  fmap f (Sum a) = Sum (f a)

-- instance Applicative Sum where
--   pure a = Sum a
--   Sum f <*> Sum x = Sum (f x)

instance Applicative Sum where
  pure  = Sum
  (<*>) = inSum2 ($)

instance AdditiveGroup a => Semigroup (Sum a) where
  (<>) = liftA2 (^+^)

instance AdditiveGroup a => Monoid (Sum a) where
  mempty  = Sum zeroV
#if !(MIN_VERSION_base(4,11,0))
  mappend = (<>)
#endif

-- | Application a unary function inside a 'Sum'
inSum :: (a -> b) -> (Sum a -> Sum b)
inSum = getSum ~> Sum

-- | Application a binary function inside a 'Sum'
inSum2 :: (a -> b -> c) -> (Sum a -> Sum b -> Sum c)
inSum2 = getSum ~> inSum

instance AdditiveGroup a => AdditiveGroup (Sum a) where
  zeroV   = mempty
  (^+^)   = mappend
  negateV = inSum negateV

---- to go elsewhere

(~>) :: (a' -> a) -> (b -> b') -> ((a -> b) -> (a' -> b'))
(i ~> o) f = o . f . i

-- result :: (b -> b') -> ((a -> b) -> (a -> b'))
-- result = (.)

-- argument :: (a' -> a) -> ((a -> b) -> (a' -> b))
-- argument = flip (.)

-- g ~> f = result g . argument f

instance AdditiveGroup a => AdditiveGroup (Gnrx.Rec0 a s) where
  zeroV = Gnrx.K1 zeroV
  negateV (Gnrx.K1 v) = Gnrx.K1 $ negateV v
  Gnrx.K1 v ^+^ Gnrx.K1 w = Gnrx.K1 $ v ^+^ w
  Gnrx.K1 v ^-^ Gnrx.K1 w = Gnrx.K1 $ v ^-^ w
instance AdditiveGroup (f p) => AdditiveGroup (Gnrx.M1 i c f p) where
  zeroV = Gnrx.M1 zeroV
  negateV (Gnrx.M1 v) = Gnrx.M1 $ negateV v
  Gnrx.M1 v ^+^ Gnrx.M1 w = Gnrx.M1 $ v ^+^ w
  Gnrx.M1 v ^-^ Gnrx.M1 w = Gnrx.M1 $ v ^-^ w
instance (AdditiveGroup (f p), AdditiveGroup (g p)) => AdditiveGroup ((f :*: g) p) where
  zeroV = zeroV :*: zeroV
  negateV (x:*:y) = negateV x :*: negateV y
  (x:*:y) ^+^ (ξ:*:υ) = (x^+^ξ) :*: (y^+^υ)
  (x:*:y) ^-^ (ξ:*:υ) = (x^-^ξ) :*: (y^-^υ)


{-# LANGUAGE AllowAmbiguousTypes                      #-}
{-# LANGUAGE ConstraintKinds                          #-}
{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE DefaultSignatures                        #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE FlexibleInstances                        #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeFamilies                             #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE UndecidableInstances                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module Backprop.Learn.Initialize (
    Initialize(..)
  , gInitialize
  , initializeNormal
  , initializeSingle
  -- * Reshape
  , reshapeR
  , reshapeLRows
  , reshapeLCols
  ) where

import           Control.Monad.Primitive
import           Data.Complex
import           Data.Proxy
import           Data.Type.Equality
import           Data.Type.Tuple
import           Data.Vinyl
import           GHC.TypeLits.Compare
import           GHC.TypeNats
import           Generics.OneLiner
import           Numeric.LinearAlgebra.Static.Vector
import           Statistics.Distribution
import           Statistics.Distribution.Normal
import qualified Data.Vector.Generic                 as VG
import qualified Data.Vector.Generic.Sized           as SVG
import qualified Numeric.LinearAlgebra.Static        as H
import qualified System.Random.MWC                   as MWC

-- | Class for types that are basically a bunch of 'Double's, which can be
-- initialized with a given identical and independent distribution.
class Initialize p where
    initialize
        :: (ContGen d, PrimMonad m)
        => d
        -> MWC.Gen (PrimState m)
        -> m p

    default initialize
        :: (ADTRecord p, Constraints p Initialize, ContGen d, PrimMonad m)
        => d
        -> MWC.Gen (PrimState m)
        -> m p
    initialize = gInitialize

-- | 'initialize' for any instance of 'Generic'.
gInitialize
    :: (ADTRecord p, Constraints p Initialize, ContGen d, PrimMonad m)
    => d
    -> MWC.Gen (PrimState m)
    -> m p
gInitialize d g = createA' @Initialize (initialize d g)

-- | Helper over 'inititialize' for a gaussian distribution centered around
-- zero.
initializeNormal
    :: (Initialize p, PrimMonad m)
    => Double                               -- ^ standard deviation
    -> MWC.Gen (PrimState m)
    -> m p
initializeNormal = initialize . normalDistr 0

-- | 'initialize' definition if @p@ is a single number.
initializeSingle
    :: (ContGen d, PrimMonad m, Fractional p)
    => d
    -> MWC.Gen (PrimState m)
    -> m p
initializeSingle d = fmap realToFrac . genContVar d

instance Initialize Double where
    initialize = initializeSingle
instance Initialize Float where
    initialize = initializeSingle

-- | Initializes real and imaginary components identically
instance Initialize a => Initialize (Complex a) where

instance Initialize T0
instance Initialize a => Initialize (TF a)
instance (Initialize a, Initialize b) => Initialize (a :# b)

instance RPureConstrained Initialize as => Initialize (T as) where
    initialize d g = rtraverse (fmap TF)
                   $ rpureConstrained @Initialize (initialize d g)

-- instance (Initialize a, ListC (Initialize <$> as), Known Length as) => Initialize (NETup (a ':| as)) where
--     initialize d g = NET <$> initialize d g
--                          <*> initialize d g

instance Initialize ()
instance (Initialize a, Initialize b) => Initialize (a, b)
instance (Initialize a, Initialize b, Initialize c) => Initialize (a, b, c)
instance (Initialize a, Initialize b, Initialize c, Initialize d) => Initialize (a, b, c, d)
instance (Initialize a, Initialize b, Initialize c, Initialize d, Initialize e) => Initialize (a, b, c, d, e)

instance (VG.Vector v a, KnownNat n, Initialize a) => Initialize (SVG.Vector v n a) where
    initialize d = SVG.replicateM . initialize d

instance KnownNat n => Initialize (H.R n) where
    initialize d = fmap vecR . initialize d
instance KnownNat n => Initialize (H.C n) where
    initialize d = fmap vecC . initialize d

instance (KnownNat n, KnownNat m) => Initialize (H.L n m) where
    initialize d = fmap vecL . initialize d
instance (KnownNat n, KnownNat m) => Initialize (H.M n m) where
    initialize d = fmap vecM . initialize d

-- | Reshape a vector to have a different amount of items  If the matrix is
-- grown, new weights are initialized according to the given distribution.
reshapeR
    :: forall i j d m. (ContGen d, PrimMonad m, KnownNat i, KnownNat j)
    => d
    -> MWC.Gen (PrimState m)
    -> H.R i
    -> m (H.R j)
reshapeR d g x = case Proxy @j %<=? Proxy @i of
    LE  Refl      -> pure . vecR . SVG.take @_ @j @(i - j) . rVec $ x
    NLE Refl Refl -> (x H.#) <$> initialize @(H.R (j - i)) d g

-- | Reshape a matrix to have a different amount of rows  If the matrix
-- is grown, new weights are initialized according to the given
-- distribution.
reshapeLRows
    :: forall i j n d m. (ContGen d, PrimMonad m, KnownNat n, KnownNat i, KnownNat j)
    => d
    -> MWC.Gen (PrimState m)
    -> H.L i n
    -> m (H.L j n)
reshapeLRows d g x = case Proxy @j %<=? Proxy @i of
    LE Refl       -> pure . rowsL . SVG.take @_ @j @(i - j) . lRows $ x
    NLE Refl Refl -> (x H.===) <$> initialize @(H.L (j - i) n) d g

-- | Reshape a matrix to have a different amount of columns.  If the matrix
-- is grown, new weights are initialized according to the given
-- distribution.
reshapeLCols
    :: forall i j n d m. (ContGen d, PrimMonad m, KnownNat n, KnownNat i, KnownNat j)
    => d
    -> MWC.Gen (PrimState m)
    -> H.L n i
    -> m (H.L n j)
reshapeLCols d g x = case Proxy @j %<=? Proxy @i of
    LE Refl       -> pure . colsL . SVG.take @_ @j @(i - j) . lCols $ x
    NLE Refl Refl -> (x H.|||) <$> initialize @(H.L n (j - i)) d g


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLabels      #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Layers.LeakyRelu
Description : Rectifying linear unit layer
Copyright   : (c) Manuel Schneckenreither, 2020
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.LeakyRelu (
    LeakyRelu (..)
  ) where

import           Control.DeepSeq                (NFData (..))
import           Data.Proxy                     (Proxy(Proxy))
import           Data.Maybe                     (fromJust)
import           Data.Serialize
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static   as H

import           Grenade.Core
import           Grenade.Onnx
import           Grenade.Types
import           Grenade.Layers.Internal.Activations

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

instance UpdateLayer LeakyRelu where
  type Gradient LeakyRelu = ()
  runUpdate _ x _ = x
  reduceGradient _ = ()

instance RandomLayer LeakyRelu where
  createRandomWith _ _ = return $ LeakyRelu 0.01

instance Serialize LeakyRelu where
  put (LeakyRelu alpha) = put alpha
  get = LeakyRelu <$> get

instance (KnownNat i) => Layer LeakyRelu ('D1 i) ('D1 i) where
  type Tape LeakyRelu ('D1 i) ('D1 i) = S ('D1 i)

  runForwards (LeakyRelu alpha) (S1D y) = (S1D y, (S1D . fromJust . H.create) relu)
    where
      w    = fromIntegral $ natVal (Proxy :: Proxy i)
      relu = leakyRelu1d w alpha (H.extract y)

  runBackwards (LeakyRelu alpha) (S1D y) (S1D dEdy) = ((), S1D (relu' y * dEdy))
    where
      relu' = H.dvmap (\a -> if a < 0 then alpha else 1)

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

  runForwards (LeakyRelu alpha) (S2D y) = (S2D y, (S2D . fromJust . H.create) relu)
    where
      h    = fromIntegral $ natVal (Proxy :: Proxy i)
      w    = fromIntegral $ natVal (Proxy :: Proxy j)
      y'   = H.extract y
      relu = leakyRelu 1 h w alpha y'

  runBackwards (LeakyRelu alpha) (S2D y) (S2D dEdy) = ((), S2D (relu' y * dEdy))
    where
      relu' = H.dmmap (\a -> if a < 0 then alpha else 1)

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

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

  runForwards (LeakyRelu alpha) (S3D y) = (S3D y, (S3D . fromJust . H.create) relu)
    where
      c    = fromIntegral $ natVal (Proxy :: Proxy k)
      h    = fromIntegral $ natVal (Proxy :: Proxy i)
      w    = fromIntegral $ natVal (Proxy :: Proxy j)
      y'   = H.extract y
      relu = leakyRelu c h w alpha y'

  runBackwards (LeakyRelu alpha) (S3D y) (S3D dEdy) = ((), S3D (relu' y * dEdy))
    where
      relu' = H.dmmap (\a -> if a < 0 then alpha else 1)

instance OnnxOperator LeakyRelu where
  onnxOpTypeNames _ = ["LeakyRelu"]

instance OnnxLoadable LeakyRelu where
  loadOnnxNode _ node = do
    alpha <- readFloatAttributeToRealNum "alpha" node

    return $ LeakyRelu alpha


{- | Example of sampling

Two samplers are availables : the 'discreteAncestralSampler' and the 'gibbsSampler'.
Only the 'gibbsSampler' can be used with evidence.

In this example, we have a very simple network.

@
    simple :: (['TDV' Bool],'SBN' 'CPT')
    simple = 'runBN' $ do
        a <- 'variable' \"a\" ('t' :: Bool)
        b <- 'variable' \"b\" ('t' :: Bool)
--
        'proba' a '~~' [0.4,0.6]
        'cpt' b [a] '~~' [0.8,0.2,0.2,0.8]
--
        return [a,b]
@

This network is representing a sensor b. We observe the value of b and we want to infer the value of a.

We use the 'gibbsSampler' for this with an initial period of 200 samples which are dropped. The 'gibbsSampler' is
generate a stream of samples. From this stream, we need to compute a probability distribution. For this, we use
the 'samplingHistograms' histogram function which is generating a list : the probability values of each vertex.

@
    let (vars\@[a,b],exampleG) = simple
    n <- 'runSampling' 5000 200 ('gibbsSampler' exampleG [b '=:' True])
    let h = 'samplingHistograms' n
    print $ h
@

Then, we compare this result with the exact one we get with a junction tree.

@
    let jt = 'createJunctionTree' 'nodeComparisonForTriangulation' exampleG
        jt' = 'changeEvidence' [b '=:' True] jt
    mapM_ (\x -> print . 'posterior' jt' $ [x]) vars
@

We can also use the 'discreteAncestralSampler' to compute the posterior but it is not supporting the use of evidence in this
version. The syntax is similar.

@
    n <- 'runSampling' 500 ('discreteAncestralSampler' exampleG)
@

-}
module Bayes.Examples.Sampling(
    -- * Test function
	  testSampling
) where

import Bayes.Sampling
import System.Random.MWC.CondensedTable
import qualified Data.Vector as V
import Bayes.Factor
import Bayes
import Bayes.FactorElimination
import Data.Maybe(fromJust)
import Bayes.BayesianNetwork
import Bayes.Factor.CPT
import Statistics.Sample.Histogram

simple :: ([TDV Bool],SBN CPT)
simple = runBN $ do
    a <- variable "a" (t :: Bool)
    b <- variable "b" (t :: Bool)

    proba a ~~ [0.4,0.6]
    cpt b [a] ~~ [0.8,0.2,0.2,0.8]

    return [a,b]

testSampling = do
    let (vars@[a,b],exampleG) = simple
        jt = createJunctionTree nodeComparisonForTriangulation exampleG
    --n <- runSampling 500 (discreteAncestralSampler exampleG)
    putStrLn "The bayesian network"
    print exampleG
    putStrLn "\nThe values of the Bayesian network"
    printGraphValues exampleG

    n <- runSampling 5000 200 (gibbsSampler exampleG [b =: True])
    let h = samplingHistograms 2 n
    putStrLn "\nThe result of the sampling"
    print $ h

    n <- runSampling 5000 200 (gibbsMCMCSampler exampleG [b =: True])
    let h = samplingHistograms 2 n
    putStrLn "\nThe result of the sampling with MCMC"
    print $ h

    let jt' = changeEvidence [b =: True] jt
    putStrLn "\nThe result of the junction tree inference"
    mapM_ (\x -> print . posterior jt' $ [x]) vars


module Algebra where

import qualified Data.Text as T
import qualified Data.Map as M
import qualified Data.Vector as V
import Data.Maybe
import Statistics.LinearRegression

import Types

selectKeyData :: T.Text -> [Sample ItemsMap] -> [Sample Int]
selectKeyData key samples = [sample {sampleItems = select sample} | sample <- samples]
  where
    select sample = fromMaybe 0 $ M.lookup key (sampleItems sample)

prepareSystem :: [Sample Int] -> (V.Vector Double, V.Vector Double)
prepareSystem samples =
      (V.fromList $ map sampleTime samples,
       V.fromList $ map (fromIntegral . sampleItems) samples)

growCoefficient :: T.Text -> [Sample ItemsMap] -> Double
growCoefficient key samples =
  let (xs, ys) = prepareSystem $ selectKeyData key samples
      (_, beta) = linearRegression xs ys
  in  beta


{-|
    Module      : Vocoder
    Description : Phase vocoder
    Copyright   : (c) Celina Pawlińska, 2020
                      Marek Materzok, 2021
    License     : BSD2

This module implements the phase vocoder algorithms.
The implementation is designed to be used directly or to be integrated
into some convenient abstraction (streaming or FRP).
-}
module Vocoder (
      Moduli,
      Phase,
      PhaseInc,
      Frame,
      Window,
      HopSize,
      Length,
      STFTFrame,
      FFTOutput,
      VocoderParams,
      vocoderParams,
      vocFrameLength,
      vocInputFrameLength,
      vocFreqFrameLength,
      vocHopSize,
      vocWindow,
      doFFT,
      doIFFT,
      analysisBlock,
      analysisStep,
      analysisStage,
      synthesisBlock,
      synthesisStep,
      synthesisStage,
      zeroPhase,
      volumeCoeff,
      frameFromComplex,
      frameToComplex,
      addFrames
    ) where

import Data.List
import Data.Complex
import Data.Fixed
import Data.Tuple
import Control.Arrow
import Numeric.FFT.Vector.Invertible as FFT
import Numeric.FFT.Vector.Plan as FFTp
import qualified Data.Vector.Storable as V

-- | Complex moduli of FFT frames. Represent signal amplitudes.
type Moduli = V.Vector Double

-- | Complex arguments of FFT frames. Represent signal phases.
type Phase = V.Vector Double

-- | Phase increments. Represent the deviation of the phase difference
-- between successive frames from the expected difference for the center
-- frequencies of the FFT bins.
type PhaseInc = V.Vector Double

-- | Time domain frame.
type Frame = V.Vector Double

-- | Sampled STFT window function.
type Window = Frame

-- | Offset between successive STFT frames, in samples.
type HopSize = Int

-- | Size in samples.
type Length = Int

-- | STFT processing unit.
type STFTFrame = (Moduli, PhaseInc)

-- | Frequency domain frame.
type FFTOutput = V.Vector (Complex Double)

-- | Type of FFT plans for real signals.
type FFTPlan = FFTp.Plan Double (Complex Double)

-- | Type of IFFT plans for real signals.
type IFFTPlan = FFTp.Plan (Complex Double) Double

-- | Configuration parameters for the phase vocoder algorithm.
data VocoderParams = VocoderParams{
    -- | FFT plan used in analysis stage.
    vocFFTPlan  :: FFTPlan,
    -- | FFT plan used in synthesis stage.
    vocIFFTPlan :: IFFTPlan,
    -- | STFT hop size.
    vocHopSize :: HopSize,
    -- | Window function used during analysis and synthesis.
    vocWindow :: Window
    -- TODO thread safety?
}

-- | FFT frequency frame length.
vocFreqFrameLength :: VocoderParams -> Length
vocFreqFrameLength par = planOutputSize $ vocFFTPlan par

-- | FFT frame length. Can be larger than `vocInputFrameLength` for zero-padding.
vocFrameLength :: VocoderParams -> Length
vocFrameLength par = planInputSize $ vocFFTPlan par

-- | STFT frame length.
vocInputFrameLength :: VocoderParams -> Length
vocInputFrameLength par = V.length $ vocWindow par

-- | Create a vocoder configuration.
vocoderParams :: Length -> HopSize -> Window -> VocoderParams
vocoderParams len hs wnd = VocoderParams (plan dftR2C len) (plan dftC2R len) hs wnd

-- | Apply a window function on a time domain frame.
applyWindow :: Window -> Frame -> Frame
applyWindow = V.zipWith (*)

-- | Change the vector indexing so that the sample at the middle has the number 0.
-- This is done so that the FFT of the window has zero phase, and therefore does not
-- introduce phase shifts in the signal.
rewind :: (V.Storable a) => V.Vector a -> V.Vector a
rewind vec = uncurry (V.++) $ swap $ V.splitAt (V.length vec `div` 2) vec

-- | Zero-pad the signal symmetrically from both sides.
addZeroPadding :: Length
    -> Frame
    -> Frame
addZeroPadding len v
    | diff < 0  = error $ "addZeroPadding: input is " ++ (show diff) ++ " samples longer than target length"
    | diff == 0 = v
    | otherwise = res
    where
    l = V.length v
    diff = len - l
    halfdiff = diff - (diff `div` 2)
    res = (V.++) ((V.++) (V.replicate halfdiff 0) v) (V.replicate (diff-halfdiff) 0)

-- | Perform FFT processing, which includes the actual FFT, rewinding, zero-paddding
-- and windowing.
doFFT :: VocoderParams -> Frame -> FFTOutput
doFFT par =
    FFT.execute (vocFFTPlan par) . rewind . addZeroPadding (vocFrameLength par) . applyWindow (vocWindow par)

-- | Perform analysis on a sequence of frames. This consists of FFT processing
-- and performing analysis on frequency domain frames.
analysisStage :: Traversable t => VocoderParams -> Phase -> t Frame -> (Phase, t STFTFrame)
analysisStage par ph = mapAccumL (analysisBlock par) ph

-- | Perform FFT transform and frequency-domain analysis.
analysisBlock :: VocoderParams -> Phase -> Frame -> (Phase, STFTFrame)
analysisBlock par prev_ph vec = analysisStep (vocHopSize par) (vocFrameLength par) prev_ph (doFFT par vec)

-- | Analyze a frequency domain frame. Phase from a previous frame must be supplied.
-- It returns the phase of the analyzed frame and the result.
analysisStep :: HopSize -> Length -> Phase -> FFTOutput -> (Phase, STFTFrame)
analysisStep h eN prev_ph vec =
    (ph,(mag,ph_inc))
    where
    (mag, ph) = frameFromComplex vec
    ph_inc = V.imap (calcPhaseInc eN h) $ V.zipWith (-) ph prev_ph

-- | Wraps an angle (in radians) to the range [-pi : pi].
wrap :: Double -> Double
wrap e = (e+pi) `mod'` (2*pi) - pi

calcPhaseInc :: Length -> HopSize -> Int -> Double -> Double
calcPhaseInc eN hop k ph_diff =
    (omega + wrap (ph_diff - omega)) / fromIntegral hop
    where
    omega = (2*pi*fromIntegral k*fromIntegral hop) / fromIntegral eN

-- | Perform synthesis on a sequence of frames. This consists of performing
-- synthesis and IFFT processing.
synthesisStage :: Traversable t => VocoderParams -> Phase -> t STFTFrame -> (Phase, t Frame)
synthesisStage par ph frs = mapAccumL (synthesisBlock par) ph frs

-- | Perform frequency-domain synthesis and IFFT transform.
synthesisBlock :: VocoderParams -> Phase -> STFTFrame -> (Phase, Frame)
synthesisBlock par ph fr = (id *** doIFFT par) $ synthesisStep (vocHopSize par) ph fr

-- | Synthesize a frequency domain frame. Phase from the previously synthesized frame
-- must be supplied. It returns the phase of the synthesized frame and the result.
synthesisStep :: HopSize -> Phase -> STFTFrame -> (Phase, FFTOutput)
synthesisStep hop ph (mag, ph_inc) =
    (new_ph, frameToComplex (mag, new_ph))
    where
    new_ph = V.zipWith (+) ph $ V.map (* fromIntegral hop) ph_inc

-- | Perform IFFT processing, which includes the actual IFFT, rewinding, removing padding
-- and windowing.
doIFFT :: VocoderParams -> FFTOutput -> Frame
doIFFT par =
    applyWindow (vocWindow par) . cutCenter (vocInputFrameLength par) . rewind . FFT.execute (vocIFFTPlan par)

-- | Cut the center of a time domain frame, discarding zero padding.
cutCenter :: (V.Storable a) => Length -> V.Vector a -> V.Vector a
cutCenter len vec = V.take len $ V.drop ((V.length vec - len) `div` 2) vec

-- | Zero phase for a given vocoder configuration.
-- Can be used to initialize the synthesis stage.
zeroPhase :: VocoderParams -> Phase
zeroPhase par = V.replicate (vocFreqFrameLength par) 0

-- | An amplitude change coefficient for the processing pipeline.
-- Can be used to ensure that the output has the same volume as the input.
volumeCoeff :: VocoderParams -> Double
volumeCoeff par = fromIntegral (vocHopSize par) / V.sum (V.map (**2) $ vocWindow par)

-- | Converts frame representation to complex numbers.
frameToComplex :: STFTFrame -> FFTOutput
frameToComplex = uncurry $ V.zipWith mkPolar

-- | Converts frame representation to magnitude and phase.
frameFromComplex :: FFTOutput -> STFTFrame
frameFromComplex = V.map magnitude &&& V.map phase

-- | Adds STFT frames.
addFrames :: STFTFrame -> STFTFrame -> STFTFrame
addFrames f1 f2 = frameFromComplex $ V.zipWith (+) (frameToComplex f1) (frameToComplex f2)


{-# LANGUAGE GADTs #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Control.Quantum
    ( Amplitude
    , WaveFunction
    , Quantum
    , quantize
    , MonadQuantum
    , runQuantum
    , runQuantum'
    , measure
    , measurements
    )
    where

import qualified Data.Map  as M
import           Data.Map (Map)
import           Data.Monoid
import           Control.Monad
import           Control.Applicative
import           Control.Monad.Random as R
import           Data.Complex

nearZero :: Double -> Bool
nearZero x = abs x <= 1e-12

type Amplitude = Complex Double

-- | The Wave function is the distribution of amplitudes over the
-- possible states of the underlying data type
type WaveFunction a = [(a, Amplitude)]

class (Monad m) => MonadQuantum m where
    quantize  :: (Ord a) => WaveFunction a -> m a
    quantize' ::            WaveFunction a -> m a
    condition :: Bool -> m ()

liftQ  :: (Ord b, MonadQuantum m) => (a -> b) -> m a -> m b
liftQ f q1 = q1 >>= always . f

liftQ2 :: (Ord c, MonadQuantum m) => (a -> b -> c) -> m a -> m b -> m c
liftQ2 f p1 p2 = do x1 <- p1
                    x2 <- p2
                    always (f x1 x2)

liftQ3 :: (Ord d, MonadQuantum m) => (a -> b -> c -> d) -> m a -> m b -> m c -> m d
liftQ3 f p1 p2 p3 = do x1 <- p1
                       x2 <- p2
                       x3 <- p3
                       always (f x1 x2 x3)

always a = quantize [(a, 1.0:+0.0)]
always' a = quantize' [(a, 1.0:+0.0)]

-- | Describes the quantized version of an arbitrary data type.
-- The underlying data type defines the basis that the state is measured in.
data Quantum a where
    Quantum :: Ord a => Map (Maybe a) Amplitude -> Quantum a
    QuantumAny :: [(Maybe a, Amplitude)]  -> Quantum a

noState  :: (Ord a) => Quantum a
noState  = Quantum (M.singleton Nothing (1.0:+0.0))

noState' :: Quantum a
noState' = QuantumAny [(Nothing, 1.0 :+ 0.0)]

deriving instance (Show a) => Show (Quantum a)

instance Functor Quantum where
    fmap = liftM

instance Applicative Quantum where
    pure  = return
    (<*>) = ap

instance Monad Quantum where
    return  = always'

    m >>= f = if unitary then next
                         else error "Non-unitary transformation applied to quantum state!"
        where
            unitary = nearZero $ l2norm (map snd $ toList' next) - 1.0
            next = collect [multAmpl q (go a) | (a, q) <- toList' m]
            go = maybe noState' f

multAmpl :: Amplitude -> Quantum a -> Quantum a
multAmpl q (Quantum x) = Quantum $ M.map (* conjugate q) x
multAmpl q (QuantumAny x) = QuantumAny [ (a, q * r) | (a, r) <- x ]

toList' :: Quantum a -> [(Maybe a, Amplitude)]
toList' (Quantum x) = M.toList x
toList' (QuantumAny x) = x

toList :: (Ord a) => Quantum a -> [(Maybe a, Amplitude)]
toList (Quantum x) = M.toList x
toList (QuantumAny x) = merge x

collect :: [Quantum a] -> Quantum a
collect [ ]        = QuantumAny []
collect [x]        = x
collect (Quantum x:t) = case collect t of
                        Quantum y -> Quantum (M.unionWith (+) x y)
                        QuantumAny y -> Quantum (M.unionWith (+) x (M.fromList y))
collect (QuantumAny x:t) = case collect t of
                        Quantum y -> Quantum (M.unionWith (+) (M.fromList x) y)
                        QuantumAny y -> QuantumAny (x ++ y)

merge :: (Ord a) => WaveFunction a -> WaveFunction a
merge = M.toList . M.fromListWith (+)

instance MonadQuantum Quantum where
    quantize    = Quantum . M.fromListWith (+) . normalize
    quantize'   = QuantumAny . normalize
    condition test = if test then always () else noState

instance (Ord a, Monoid a) => Monoid (Quantum a) where
    mempty  = always mempty
    mappend = liftQ2 mappend

normalize :: WaveFunction a -> [(Maybe a, Amplitude)]
normalize xs = map (\(a, q) -> (Just a, q / total)) xs
    where
        total = sum $ zipWith (*) (map conjugate ampl) ampl
        ampl  = map snd xs

-- | Remove all impossible states from the quantum state and renormalize
collapse :: [(Maybe a, Amplitude)] -> WaveFunction a
collapse xs = [ (x, q / (norm:+0)) | (Just x, q) <- xs ]
    where
        norm = l2norm [ q | (Just x, q) <- xs ]

l2norm :: [Amplitude] -> Double
l2norm xs = sqrt $ sum $ map (\x -> (magnitude x)**2) xs

-- | Extracts the wave function from a quantum state.
runQuantum :: (Ord a) => Quantum a -> WaveFunction a
runQuantum = collapse . toList

runQuantum' :: Quantum a -> WaveFunction a
runQuantum' = collapse . toList'

-- | Measures a quantum state. Returns a single realization of the underlying data type
-- with probability equal to the squared magnitude of the amplitude of that realization.
-- Note that if this Module would be backed by a real quantum processor, this would be
-- the only valid way to extract information from the Monad.
measure :: (Ord a) => Quantum a -> IO a
measure q = do
  x <- evalRandIO $ R.fromList $ map (\(a, b) -> (a, toRational b)) $ measurements q
  return x

-- | Converts a quantum state into a probability state through measurement.
-- Equivalent to performing repeated measurements on equally prepared quantum states
-- and noting the frequency of each possible realization.
measurements :: (Ord a) => Quantum a -> [(a, Double)]
measurements = f . runQuantum
    where f xs = [(x, (magnitude q)**2) | (x, q) <- xs]


{-# LANGUAGE CPP                   #-}
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Trustworthy
           , DefaultSignatures     #-}
#define USE_GHC_GENERICS
#endif
{-# LANGUAGE RankNTypes
           , TypeFamilies
           , KindSignatures
           , DeriveGeneric
           , FlexibleInstances
           , FlexibleContexts
           , BangPatterns
           , MultiParamTypeClasses #-}
{-# OPTIONS_GHC -fno-warn-orphans  #-}

-------------------------------------------------------------------------------
-- |
-- Module      : OGLS.Engine.Math.Vectors
-- description : 3D and 4D vectors arithmetic
-- Copyright   : (c) Adam, 2017
-- License     : MIT
-- Maintainer  : awkure@protonmail.ch
-- Stability   : unstable
-- Portability : POSIX
-------------------------------------------------------------------------------

module OGLS.Engine.Math.Vectors
    ( (^$+), (^$-), (^$*)
    , dot, step
    ) where

#ifdef USE_GHC_GENERICS
import GHC.Generics               ( Generic (..) )
#endif
import Control.Applicative        ( liftA2       )
import Control.Monad              ( liftM2       )
import Control.Parallel           ( pseq         )
import Data.Complex               ( Complex (..) )
import Data.Orphans               (              )
import Control.Lens hiding ( (<.>) )

import OGLS.Engine.Math.Instances

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

default ()

class Vector a v where
    -- not yet implemented

(^$+) :: forall (t :: * -> *) a. (Applicative t, Num a) => t a -> t a -> t a
(^$-) :: forall (t :: * -> *) a. (Applicative t, Num a) => t a -> t a -> t a
(^$*) :: forall (t :: * -> *) a. (Applicative t, Num a) => t a -> t a -> t a

newtype E t = E { el :: forall x. Lens' (t x) x }

data Vec3 a = Vec3 { v3x :: {-# UNPACK #-} !a
                   , v3y :: {-# UNPACK #-} !a
                   , v3z :: {-# UNPACK #-} !a
                   } deriving ( Eq, Ord, Show, Generic )

instance Functor Vec3 where
    fmap f (Vec3 x y z) = Vec3 (f x) (f y) (f z)

instance Applicative Vec3 where
    pure x = Vec3 x x x
    Vec3 f g h <*> Vec3 x y z = Vec3 (f x) (g y) (h z)

data Vec4 a = Vec4 { v4x :: {-# UNPACK #-} !a
                   , v4y :: {-# UNPACK #-} !a
                   , v4z :: {-# UNPACK #-} !a
                   , v4w :: {-# UNPACK #-} !a
                   } deriving ( Eq, Ord, Show )

(^$+) = liftA2 (+)
{-# INLINE (^$+) #-}
(^$-) = liftA2 (-)
{-# INLINE (^$-) #-}
(^$*) = liftA2 (*)
{-# INLINE (^$*) #-}

-- | Left scalar product vector
(*^^) :: forall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
(*^^) a = fmap (a*)
{-# INLINE (*^^) #-}

-- | Right scalar product vector
(^^*) :: forall (f :: * -> *) a. (Functor f, Num a) => f a -> a -> f a
f ^^* a = fmap (*a) f
{-# INLINE (^^*) #-}

{- Needs Additive class and instances
basis :: (Traversable t, Num a) => [t a]
basis = basisFor (zero :: Additive v => v Int)

basisFor :: (Traversable t, Num a) => t b -> [t a]
basisFor = \t ->
   ifoldMapOf traversed ?? t $ \i _ ->
     return $ iover  traversed ?? t $ \j _ ->
         if i == j then 1 else 0
{-# INLINABLE basisFor #-}
-}

-- | Linear interpolation
lerp :: forall (f :: * -> *) a. (Functor f, Applicative f, Num a) => a -> f a -> f a -> f a
lerp alpha u v = alpha *^^ u ^$+ (1 - alpha) *^^ v
{-# INLINE lerp #-}

-- | Dot product
dot :: forall (t :: * -> *) a. (Applicative t, Foldable t, Num a) => t a -> t a -> a
dot v1 v2 = sum $ v1 ^$* v2

step :: forall (t :: * -> *) a. (Applicative t, Num a, Ord a) => t a -> t a -> t a
step = liftA2 (\a b -> if b < a then 0 else 1)

#ifdef HLINT
{-# ANN module "HLint: ignore Redundant lambda" #-}
#endif


module DistanceGeometry
  (
  -- * Steps of algorithm
    generateDistBoundsMatr
  , triangleSmooth
  , randomDistMatr
  , distMatrToMetricMatr
  , largestEigValVec
  , generateCoorFromEigValVec
  , coordMatrToDistMatr
  , updateCoord

  -- * Error Functions
  , distanceErrorFunction
  , chiralErrorFunction
  ) where

import           Control.Lens
import           Control.Monad.State
import           Data.List                   (foldl', zipWith3)
import           Numeric.LinearAlgebra
import           Numeric.LinearAlgebra.Data
import           Numeric.LinearAlgebra.Devel
import           System.Random
import           Types

-- * Steps of algorithm
-- | Generate of a distance bounds matrix
generateDistBoundsMatr :: Molecule -> [Bond] -> (Matrix Double, Matrix Double)
generateDistBoundsMatr molecule bonds =
  let n = views atoms length molecule
      upperDist _ _ = 100
      lowerDist i j = getAtom i molecule ^. avdwrad +
                      getAtom j molecule ^. avdwrad
      fixedDist i j = sqrt $ (xj - xi) ^ 2 + (yj - yi) ^ 2 + (zj - zi) ^ 2
        where
          (Point xi yi zi) = getAtom i molecule ^. acoordin
          (Point xj yj zj) = getAtom j molecule ^. acoordin
      upper = build (n, n) (\i' j' ->
        let (i, j) = (fromEnum i', fromEnum j')
        in if i == j
           then 0
           else if isBonded i j bonds
                then fixedDist i j
                else upperDist i j)
      lower = build (n, n) (\i' j' ->
        let (i, j) = (fromEnum i', fromEnum j')
        in if i == j
           then 0
           else if isBonded i j bonds
                then fixedDist i j
                else lowerDist i j)
   in (upper, lower)

-- | Triangle Inequality Bounds Smoothing
-- triangleSmooth :: Matrix Double -> Either String (Matrix Double)
triangleSmooth :: (Matrix Double, Matrix Double) -> (Matrix Double, Matrix Double)
triangleSmooth (upper, lower) =
  let n = rows upper
      kij = [ (k, i, j)
       | k <- [0 .. n - 1]
       , i <- [0 .. n - 2]
       , j <- [i + 1 .. n - 1]
       ]
      smoothing (u0, l0) (k, i, j) = do
        let (u1, l1) = if e1 > e2 + e3
                       then (changeMatrix u0 (i, j) (e2 + e3), l0)
                       else (u0, l0)
                       where
                        e1 = u0 `atIndex` (i, j)
                        e2 = u0 `atIndex` (i, k)
                        e3 = u0 `atIndex` (k, j)
        let (u2, l2) = if e1 < e2 - e3
                       then (u1, changeMatrix l1 (i, j) (e2 - e3))
                       else (u1, l1)
                       where
                        e1 = l1 `atIndex` (i, j)
                        e2 = l1 `atIndex` (i, k)
                        e3 = u1 `atIndex` (k, j)
        let (u3, l3) = if e1 < e2 - e3
                       then (u2, changeMatrix l2 (i, j) (e2 - e3))
                       else (u2, l2)
                       where
                        e1 = l2 `atIndex` (i, j)
                        e2 = l2 `atIndex` (j, k)
                        e3 = u2 `atIndex` (k, i)
        if l3 ! i ! j > u3 ! i ! j
        then error "Erroneous Bounds"
        else (u3, l3)
      (upper', lower') = foldl' smoothing (upper, lower) kij
      newUpper = build (n, n) (\i' j' ->
        let (i, j) = (fromEnum i', fromEnum j')
        in if i < j
           then upper' `atIndex` (i, j)
           else upper' `atIndex` (j, i))
      newLower = build (n, n) (\i' j' ->
        let (i, j) = (fromEnum i', fromEnum j')
        in if i < j
           then lower' `atIndex` (i, j)
           else lower' `atIndex` (j, i))
   in (newUpper, newLower)

-- | Generation of a distance matrix by random selection
-- of distances between the bounds.
randomDistMatr :: (Matrix Double, Matrix Double) -> IO (Matrix Double)
randomDistMatr (upper, lower) = do
  let n = rows upper
  distanceVector <- sequence [ r
   | i <- [0 .. n - 1]
   , j <- [0 .. n - 1]
   , let a = lower `atIndex` (i, j)
         b = upper `atIndex` (i, j)
         r = if j <= i
             then return 0
             else randomRIO (a, b)
             ]
  let distanceMatrix = matrix n distanceVector
  return $ distanceMatrix `add` tr' distanceMatrix

-- | Improving Random Sampling: Metrization
metrization :: Matrix Double -> Matrix Double
metrization matrix = undefined

-- | Improving Random Sampling: Partial Metrization
partialMetrization :: Matrix Double -> Matrix Double
partialMetrization matrix = undefined

-- | Conversion of the distance matrix to a metric matrix
distMatrToMetricMatr :: Matrix Double -> Matrix Double
distMatrToMetricMatr matr =
  let n = rows matr
      m = fromIntegral n
      uTriangleMatr = build (n, n) (\i' j' ->
        let (i, j) = (fromEnum i', fromEnum j')
        in if j <= i
           then 0
           else matr `atIndex` (i, j))
      d0 i = (1 / m) * sumElements ((matr ! i) ^ 2) -
             (1 / m ^ 2) * sumElements (uTriangleMatr ^ 2)
  in build (n, n) (\i' j' ->
    let (i, j) = (fromEnum i', fromEnum j')
        a = d0 i
        b = d0 j
        c = matr `atIndex` (i, j)
    in (a + b - c ^ 2) / 2)

-- | distanceToMetricMatrix is symmetric matrix => eigenvalues is real
-- Function return @k@ pairs of eigenvalues and eigenvectors (as columns)
-- in descending order
largestEigValVec :: Matrix Double -> (Vector Double, Matrix Double)
largestEigValVec matr =
  if (not . isSymmetric) matr
  then error "The matrix is not symmetric"
  else (subVector 0 k eigval, takeColumns k eigvec)
  where
    (eigval, eigvec) = (eigSH . trustSym) matr
    k = min 3 (size eigval)

-- | Generation of three-dimensional coordinates from
-- eigenvalues and eigenvectors.
generateCoorFromEigValVec ::
  (Vector Double, Matrix Double) -> Matrix Double
generateCoorFromEigValVec (val, vec) =
  vec Numeric.LinearAlgebra.<> (sqrt . diag) val

-- | Convert coordinate matrix to
-- distance matrix.
coordMatrToDistMatr :: Matrix Double -> Matrix Double
coordMatrToDistMatr matr =
  let n = rows matr
      dist i j = sqrt $ (xj - xi) ^ 2 + (yj - yi) ^ 2 + (zj - zi) ^ 2
        where
          [xi, yi, zi] = toList $ matr ! i
          [xj, yj, zj] = toList $ matr ! j
      dmatr = build (n, n) (\i' j' ->
       let (i, j) = (fromEnum i', fromEnum j')
       in if j <= i
          then 0
          else dist i j)
  in dmatr `add` tr' dmatr

-- | Update coordinates
updateCoord :: Matrix Double -> Molecule -> Molecule
updateCoord coord mol = set atoms newatoms mol
  where
    newatoms = zipWith updatec (toRows coord) (view atoms mol)
    updatec coord' =
      execState
        (do acoordin . x .= coord' ! 0
            acoordin . y .= coord' ! 1
            acoordin . z .= coord' ! 2)

-- * Error Functions
-- | Distance error function.
distanceErrorFunction ::
     Matrix Double -> Matrix Double -> Matrix Double -> [Bond] -> Double
distanceErrorFunction coord upper lower bonds =
  let n = rows coord
      dist = coordMatrToDistMatr coord
      [d2, u2, l2] = map (^2) [dist, upper, lower]
      ij' = [(i, j) | i <- [0 .. n - 2], j <- [i + 1 .. n - 1]]
      ij = filter (\(i,j) -> not $ isBonded i j bonds) ij'
      f1 (i, j) =
        (+) $ max 0 (d2 ! i ! j - u2 ! i ! j) ^ 2 +
        max 0 (l2 ! i ! j ^ 2 - d2 ! i ! j) ^ 2
      f2 (i, j) =
        (+) $ max 0 (d2 ! i ! j / u2 ! i ! j - 1) ^ 2 +
        max 0 (l2 ! i ! j / d2 ! i ! j - 1) ^ 2
      f3 (i, j) =
        (+) $ max 0 (d2 ! i ! j / u2 ! i ! j - 1) ^ 2 +
        max 0 (2 * l2 ! i ! j / (l2 ! i ! j + d2 ! i ! j) - 1) ^ 2
   in foldr f1 0 ij + foldr f2 0 ij + foldr f3 0 ij

-- | Chiral error function.
chiralErrorFunction ::
     Matrix Double -> Matrix Double -> Matrix Double -> [Bond] -> Double
chiralErrorFunction coord upper lower bonds =
  let n = rows coord
      dist = coordMatrToDistMatr coord
      ijkm' = [(i,j,k,m) | i <- [0..n-1], j <- [i+1..n-1],
                           k <- [j+1..n-1], m <- [k+1..n-1]]
      ijkm = filter (isTetrahedra dist) ijkm'
        where
          isTetrahedra d (i,j,k,m) = all (isTriangle d) [(i,j,k), (i,j,m), (i,k,m), (j,k,m)] && all (checkAngles d) [(i,j,k,m)]-- (j,k,m,i), (k,m,i,j), (m,i,j,k)]
            where isTriangle d (i,j,k) = and [a + b > c, b + c > a, c + a > b]
                   where a = d ! i ! j
                         b = d ! i ! k
                         c = d ! j ! k
                  checkAngles d (i,j,k,m) = and [b1 + b2 > b3, b1 + b3 > b2, b2 + b3 > b1]
                   where a1 = d ! i ! j; a2 = d ! i ! k
                         a3 = d ! i ! m; a4 = d ! j ! k
                         a5 = d ! j ! m; a6 = d ! k ! m
                         b1 = acos $ (a1^2 + a2^2 - a4^2) / (2*a1*a2)
                         b2 = acos $ (a1^2 + a3^2 - a5^2) / (2*a1*a3)
                         b3 = acos $ (a2^2 + a3^2 - a6^2) / (2*a2*a3)
      vol (i,j,k,m) = (xj - xi) `dot` ((xk - xi) `cross` (xm - xi))
        where [xi,xj,xk,xm] = map (coord !) [i,j,k,m]
      u (i,j,k,m) = (sum . map (atIndex upper)) [(i, j), (j, k), (k, m), (m, i)]
      l (i,j,k,m) = (sum . map (atIndex lower)) [(i, j), (j, k), (k, m), (m, i)]
      f x@(i,j,k,m) = (+) $ max 0 (vol x - u x) ^ 2 + max 0 (l x - vol x)
  in foldr f 0 ijkm

-- * Utils.
-- | Изменяет значение матрицы @m@ по индексам (@i@,@j@) на указанное @v@
changeMatrix :: Matrix Double -> (Int, Int) -> Double -> Matrix Double
changeMatrix m (i, j) v =
  runSTMatrix $ do
    m' <- thawMatrix m
    writeMatrix m' i j v
    return m'

-- | Check the symmetric matrix
isSymmetric :: Matrix Double -> Bool
isSymmetric matr =
  let r = rows matr
      c = cols matr
      s = [matr ! i ! j == matr ! j ! i | i <- [0 .. r - 1], j <- [i .. c - 1]]
  in r == c && and s

-- | Определяет, связаны ли атомы
isBonded :: Serial -> Serial -> [Bond] -> Bool
isBonded n m s = (n, m) `elem` bonds s || (m, n) `elem` bonds s
  where bonds = map (\x -> (view bfid x, view bsid x))

-- | Get atom from molecule
getAtom :: Serial -> Molecule -> Atom
getAtom i = views atoms (!! i)


{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE QuantifiedConstraints  #-}
{-# LANGUAGE RecordWildCards        #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE UndecidableInstances   #-}
module Q.Stochastic.Discretize
        where

import           Data.Functor
import           Numeric.LinearAlgebra
import           Q.Stochastic.Process
-- |Euler discretization of stochastic processes
newtype Euler = Euler { eDt :: Double }
        deriving stock (Show, Eq)

-- | Euler end-point discretization of stochastic processes
newtype EndEuler = EndEuler { eeDt :: Double }
        deriving stock (Show, Eq)

instance Discretize Euler Double where
  dDrift p Euler{..} s0 = pDrift p s0 <&> (* eDt)
  dDiff  p Euler{..} b  = (pDiff p b) <&> (* (sqrt eDt))
  dDt    _ Euler{..} _  = eDt

instance Discretize Euler (Vector Double) where
  dDrift p Euler{..} s0 = pDrift p s0 <&> (scale eDt)
  dDiff  p Euler{..} b = (pDiff p b) <&> (scale (sqrt eDt))
  dDt    _ Euler{..} _  = eDt

instance Discretize EndEuler Double where
  dDrift p EndEuler{..} s0@(t0, x0) = pDrift p (t0 + eeDt, x0) <&> (* eeDt)
  dDiff  p EndEuler{..}  s0@(t0, x0) =  pDiff  p (t0 + eeDt, x0) <&> (* (sqrt eeDt))
  dDt    _ e _   = eeDt e


-- Test of the computational correctness of my current FFT candidate.
--
-- Original author: David Banas <capn.freako@gmail.com>
-- Original date:   October 3, 2015
--
-- Copyright (c) 2015 David Banas; all rights reserved World wide.
--
-- I'm waiting for Conal to find time to look into my latest non-termination
-- failure through his compiler. While I do, I'm attempting, here, to verify
-- the computational correctness of my current FFT candidate.

{-# LANGUAGE GADTs #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}

module Main where

import Prelude hiding ({- id,(.), -}foldl,foldr,sum,product,zipWith,reverse,and,or,scanl,minimum,maximum)

import Control.Applicative
import Control.Arrow
import Control.Monad (forM_, unless)
import Data.Complex
import Data.Foldable (Foldable, sum, foldl')
import Data.Newtypes.PrettyDouble
import System.Exit (exitFailure)
import TypeUnary.Nat (IsNat, Nat(..), nat, N2, N3, N4, N5)  -- , N6)

-- import Test.QuickCheck (choose, vectorOf, elements, collect)
import Test.QuickCheck (choose, vectorOf)
import Test.QuickCheck.Arbitrary
import Test.QuickCheck.All (quickCheckAll)

import Circat.Scan (lproducts, LScan)
import qualified Circat.Pair as P
import qualified Circat.RTree as RT
import Circat.RTree (bottomSplit)

type RTree = RT.Tree

-- FFT, as a class
-- (The LScan constraint comes from the use of 'lproducts', in 'addPhase'.)
class (LScan f) => FFT f a where
    fft  :: f a -> f a  -- Computes the FFT of a functor.

-- Note that this definition of the FFT instance for Pair assumes DIT.
-- How can we eliminate this assumption and make this more general?
instance (RealFloat a, Applicative f, Foldable f, Num (f (Complex a)), FFT f (Complex a)) => FFT P.Pair (f (Complex a)) where
    fft = P.inP (uncurry (+) &&& uncurry (-)) . P.secondP addPhase . fmap fft

instance (IsNat n, RealFloat a) => FFT (RTree n) (Complex a) where
    fft = fft' nat
        where   fft' :: (RealFloat a) => Nat n -> RTree n (Complex a) -> RTree n (Complex a)
                fft' Zero     = id
                fft' (Succ _) = inDIT fft
                    where   inDIT g  = RT.toB . g . bottomSplit

-- Adds the proper phase adjustments to a functor containing Complex RealFloats,
-- and instancing Num.
addPhase :: (Applicative f, Foldable f, LScan f, RealFloat a, Num (f (Complex a))) => f (Complex a) -> f (Complex a)
addPhase = liftA2 (*) id phasor
  where phasor f = fst $ lproducts (pure phaseDelta)
          where phaseDelta = cis ((-pi) / fromIntegral n)
                n          = flen f

-- Gives the "length" (i.e. - number of elements in) of a Foldable.
-- (Soon, to be provided by the Foldable class, as "length".)
flen :: (Foldable f) => f a -> Int
flen = foldl' (flip ((+) . const 1)) 0

-- Test config.
realData :: [[PrettyDouble]]
realData = [  [1.0,   0.0,   0.0,   0.0]  -- Delta
            , [1.0,   1.0,   1.0,   1.0]  -- Constant
            , [1.0,  -1.0,   1.0,  -1.0]  -- Nyquist
            , [1.0,   0.0,  -1.0,   0.0]  -- Fundamental
            , [0.0,   1.0,   0.0,  -1.0]  -- Fundamental w/ 90-deg. phase lag
           ]
complexData :: [[Complex PrettyDouble]]
complexData = map (map (:+ 0.0)) realData

myTree2 :: [a] -> RTree N2 a
myTree2 [w, x, y, z] = RT.tree2 w x y z
myTree2 _            = error "Something went horribly wrong!"

myTree3 :: [a] -> RTree N3 a
myTree3 [a, b, c, d, e, f, g, h] = RT.tree3 a b c d e f g h
myTree3 _            = error "Something went horribly wrong!"

myTree4 :: [a] -> RTree N4 a
myTree4 [a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p] = RT.tree4 a b c d e f g h i j k l m n o p
myTree4 _            = error "Something went horribly wrong!"

myTree5 :: [a] -> RTree N5 a
myTree5 [a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p,
         a', b', c', d', e', f', g', h', i', j', k', l', m', n', o', p'] =
            RT.tree5 a b c d e f g h i j k l m n o p a' b' c' d' e' f' g' h' i' j' k' l' m' n' o' p'
myTree5 _            = error "Something went horribly wrong!"

-- myTree6 :: [a] -> RTree N6 a
-- myTree6 [a1, b1, c1, d1, e1, f1, g1, h1, i1, j1, k1, l1, m1, n1, o1, p1,
--          a2, b2, c2, d2, e2, f2, g2, h2, i2, j2, k2, l2, m2, n2, o2, p2,
--          a3, b3, c3, d3, e3, f3, g3, h3, i3, j3, k3, l3, m3, n3, o3, p3,
--          a4, b4, c4, d4, e4, f4, g4, h4, i4, j4, k4, l4, m4, n4, o4, p4] =
--             RT.tree6 a1 b1 c1 d1 e1 f1 g1 h1 i1 j1 k1 l1 m1 n1 o1 p1
--                      a2 b2 c2 d2 e2 f2 g2 h2 i2 j2 k2 l2 m2 n2 o2 p2
--                      a3 b3 c3 d3 e3 f3 g3 h3 i3 j3 k3 l3 m3 n3 o3 p3
--                      a4 b4 c4 d4 e4 f4 g4 h4 i4 j4 k4 l4 m4 n4 o4 p4
-- myTree6 _            = error "Something went horribly wrong!"

-- Discrete Fourier Transform (DFT) (our "truth" reference)
-- O(n^2)
--
dft :: RealFloat a => [Complex a] -> [Complex a]
dft xs = [ sum [ x * exp((0.0 :+ (-1.0)) * 2 * pi / lenXs * fromIntegral(k * n))
                 | (x, n) <- Prelude.zip xs [0..]
               ]
           | k <- [0..(length xs - 1)]
         ]
    where lenXs = fromIntegral $ length xs

-- QuickCheck types & propositions
newtype FFTTestVal = FFTTestVal {
    getVal :: [Complex PrettyDouble]
} deriving (Show)
instance Arbitrary FFTTestVal where
    arbitrary = do
        xs <- vectorOf 32 $ choose (-1.0::Double, 1.0)
        let zs = map ((:+ 0) . PrettyDouble) xs
        return $ FFTTestVal zs

prop_fft_test_N2 :: FFTTestVal -> Bool
prop_fft_test_N2 testVal = fft (myTree2 zs) == RT.fromList (dft zs)
    where zs = take 4 $ getVal testVal

prop_fft_test_N3 :: FFTTestVal -> Bool
prop_fft_test_N3 testVal = fft (myTree3 zs) == RT.fromList (dft zs)
    where zs = take 8 $ getVal testVal

prop_fft_test_N4 :: FFTTestVal -> Bool
prop_fft_test_N4 testVal = fft (myTree4 zs) == RT.fromList (dft zs)
    where zs = take 16 $ getVal testVal

prop_fft_test_N5 :: FFTTestVal -> Bool
prop_fft_test_N5 testVal = fft (myTree5 zs) == RT.fromList (dft zs)
    where zs = take 32 $ getVal testVal

-- Test definitions & choice
basicTest :: IO ()
basicTest = forM_ complexData (\x -> do
                putStr "\nTesting input: "
                print x
                putStr "Expected output: "
                print $ dft x
                putStr "Actual output:   "
                print $ fft $ myTree2 x
                )

-- This weirdness is required, as of GHC 7.8.
return []

runTests :: IO Bool
runTests = $quickCheckAll
-- End weirdness.

advancedTest :: IO ()
advancedTest = do
    allPass <- runTests -- Run QuickCheck on all prop_ functions
    unless allPass exitFailure

main :: IO ()
-- main = basicTest
main = advancedTest