{-# LANGUAGE DataKinds         #-}
{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE RankNTypes        #-}
{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE TypeFamilies      #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeOperators     #-}
-- {-# LANGUAGE UndecidableInstances #-}

module GrenadeExtras (
  epochTraining,
  trainOnBatchesEpochs,
  binaryNetError,
  normalize,
  hotvector,
  applyUpdates
) where

import           Shuffle (shuffle)
import           Control.Monad.Random (MonadRandom)

import           Data.List (foldl')

import           Lens.Micro (over, both)--Lens', set)

import           Data.Singletons (SingI)
import           Data.Singletons.Prelude (Head, Last)

import           Numeric.LinearAlgebra.Static (extract, R, ℝ, toRows, eye)
import           Numeric.LinearAlgebra.Data (Vector, (!))

import           GHC.TypeLits (KnownNat)

import           Grenade (Network((:~>), NNil), Shape(D1), LearningParameters, runNet, S(S1D),
                          Gradients((:/>)), backPropagate, train, runUpdates, Gradient, applyUpdate)

import           GrenadeExtras.OrphanNum()
import           GrenadeExtras.GradNorm (GradNorm(normSquared))

trainOnBatchesEpochs :: (SingI (Last shapes), MonadRandom m, Num (Gradients layers), GradNorm (Gradients layers))
                 => Network layers shapes
                 -> LearningParameters
                 -> [(S (Head shapes), S (Last shapes))]
                 -> Int
                 -> m [(Double, Network layers shapes)]
trainOnBatchesEpochs net0 rate input_data batchSize =

    foldMeOutList (0, net0) [(1::Int)..] $ \(_,net) _-> do
      shuffledInput <- shuffle input_data
      -- traning net (an epoch) with the input shuffled
      let batches = splitInBatches shuffledInput
          (gradientNorm, newNet) = foldl' trainBatch (0, net) batches
      return (gradientNorm, newNet)

  where
    trainBatch :: (SingI (Last shapes), Num (Gradients layers), GradNorm (Gradients layers))
               => (Double, Network layers shapes)
               -> [(S (Head shapes), S (Last shapes))]
               -> (Double, Network layers shapes)
    trainBatch (accNorm, !network) ios =
      let grads = fmap (uncurry $ backPropagate network) ios
          grad = sum grads
          norm = sqrt $ normSquared grad
       in (accNorm+norm, applyUpdate rate network grad)
       --in applyUpdates rate network grads

    --len = length input_data

    splitInBatches :: [a] -> [[a]]
    splitInBatches [] = []
    splitInBatches xs =
      let (start, finish) = splitAt batchSize xs
       in start : splitInBatches finish

applyUpdates :: LearningParameters
             -> Network layers shapes
             -> [Gradients layers]
             -> Network layers shapes
applyUpdates rate (layer :~> rest) gradients
  = runUpdates rate layer layerGradients :~> applyUpdates rate rest restLayersGradients
    where headTailGrad :: Gradients (layer ': layers) -> (Gradient layer, Gradients layers)
          headTailGrad (gradient :/> grest) = (gradient, grest)
          (layerGradients, restLayersGradients) = unzip $ fmap headTailGrad gradients

applyUpdates _ NNil _
  = NNil

epochTraining :: (SingI (Last shapes), MonadRandom m) =>
                 Network layers shapes
                 -> LearningParameters
                 -> [(S (Head shapes), S (Last shapes))]
                 -> m [Network layers shapes]
epochTraining net0 rate input_data =

    foldMeOutList net0 [(1::Int)..] $ \net _-> do
      shuffledInput <- shuffle input_data
      -- traning net (an epoch) with the input shuffled
      let newNet = foldl' trainEach net shuffledInput
      return newNet

  where
    trainEach !network (i,o) = train rate network i o

foldMeOutList :: Monad m => a -> [b] -> (a -> b -> m a) -> m [a]
foldMeOutList z xs_ op = f' z xs_
  where f' _    []     = return []
        f' zero (x:xs) = do
          zero' <- op zero x
          rec   <- f' zero' xs
          return $ zero':rec

binaryNetError :: (Last shapes ~ 'D1 1, Foldable t) =>
  Network layers shapes -> t (S (Head shapes), S ('D1 1)) -> (Double, Double)

binaryNetError net test = (fromIntegral errors / fromIntegral total, distance)
  where
    total, errors :: Integer
    {-(total, errors) = foldl' step (0,0) test-}
    distance :: Double
    (total, errors, distance) = foldl' step (0,0,0) test

    {-step (t, e) song = (t+1, e')-}
    step (t, e, d) song = (t+1, e', d')
      where
        (label, netOut) = valueFromDataAndNet song
        -- the predictions from the network come with numbers between 0 and 1,
        -- everything above .5 is considered 1 and below 0
        e' = if (label > 0.5) == (netOut > 0.5)
                then e
                else e+1
        d' = d + (label - netOut)^(2::Int)

    valueFromDataAndNet (input, label) = over both sd1toDouble (label, runNet net input)
      where
        sd1toDouble :: S ('D1 1) -> Double
        sd1toDouble (S1D r) = (extract :: R 1 -> Vector ℝ) r ! 0

-- taken from https://en.wikipedia.org/wiki/Normalization_(statistics)
-- normalization method: Student's t-statistic
normalize :: KnownNat n => [R n] -> [R n]
normalize features = fmap (\x -> (x-mean)/stdDeviation ) features
  where
    len = length features

    --mean :: R n
    mean = sum features / fromIntegral len

    --stdDeviation :: R n
    stdDeviation = sqrt $ (sum . fmap (\x-> (x-mean)^(2::Int)) $ features)
                          / fromIntegral (len-1)

hotvector :: KnownNat n => Int -> Maybe (R n)
hotvector m = toRows eye ~!! m
  where
    (~!!) :: [a] -> Int -> Maybe a
    []     ~!! _ = Nothing
    (x:xs) ~!! n
      | n == 0 = Just x
      | n < 0  = Nothing
      | otherwise = xs ~!! (n-1)


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE OverloadedStrings #-}

module Model where

import           Control.Monad.Bayes.Class
import           Control.Monad.Bayes.Sampler
import qualified Data.Text.IO as TIO
import           Control.Monad
import           Control.Monad.Bayes.Inference.PMMH
import           Control.Monad.State
import Control.Monad.Bayes.Weighted
import           Numeric.Log
import           Text.Megaparsec

import           DataParser
import           Statistics
import           Utils

data Params = Params {
    rho :: Double, -- ^ Rate of detection
    beta :: Double, -- ^ Mean contact rate between susceptible and infected people
    gamma :: Double -- ^ Mean recovery rate
} deriving Show

data FixedParams = FixedParams {
    numPop :: Int,
    timeSlices :: Int
}

data LatentState = LatentState {
    sus :: Int, -- ^ Number of people susceptible to infection
    inf :: Int, -- ^ Number of people currently infected
    recov :: Int -- ^ Number of people recovered from infection
}

{-
observation model: Poisson rho * I
-}

-- | Model for how we observe the number of infected people
observationModel :: MonadSample m => Params -> LatentState -> m InfectionCount
observationModel (Params rho _ _) (LatentState _ inf _) =
    poisson (rho * fromIntegral inf)

{-
dN_SI <-
-}
-- | Transition the model a single time slice
transitionModelSingleStep
    :: MonadSample m => FixedParams -> Params -> LatentState -> m LatentState
transitionModelSingleStep (FixedParams numPop timeSlices) (Params rho beta gamma) (LatentState sus inf recov)
    = do
        let dt = 1 / fromIntegral timeSlices
        dN_SI <- binomial
            sus
            (1 - exp ((-beta * dt * (fromIntegral inf)) / (fromIntegral numPop))
            )
        dN_IR <- binomial
            inf
            (1 - exp (-gamma * dt)
            )
        let sus'   = sus - dN_SI
        let inf'   = inf + dN_SI - dN_IR
        let recov' = recov + dN_IR
        return (LatentState sus' inf' recov')

-- | Transition the model for a full step
transitionModel :: MonadSample m => FixedParams -> Params -> LatentState -> m LatentState
transitionModel fixedParams params =
    transitionModelSingleStep fixedParams params

-- | Simulate a single step, returning the new latent state and appending the observed infection count to the state.
simulateStep
    :: (MonadSample m, MonadState [InfectionCount] m)
    => FixedParams
    -> Params
    -> LatentState
    -> m LatentState
simulateStep fixedParams params latent = do
    latent'        <- transitionModel fixedParams params latent
    infectionCount <- observationModel params latent'
    modify (++ [infectionCount])
    return latent'

-- | Simulate nsteps steps of an epidemic with the specified initial state and parameters
simulateEpidemic
    :: MonadSample m => LatentState -> FixedParams -> Params -> Int -> m Epidemic
simulateEpidemic initialState fixedParams params nsteps =
    Epidemic
        <$> execStateT
                (repeatFunction nsteps (simulateStep fixedParams params) initialState)
                []

-- | Execute a single simulation of an epidemic
generateSingleEpidemic :: LatentState -> FixedParams -> Params -> Int -> IO Epidemic
generateSingleEpidemic initialState fixedParams params nsteps =
    sampleIO $ simulateEpidemic initialState fixedParams params nsteps

generateEpidemics :: LatentState -> FixedParams -> Params -> Int -> Int -> IO [Epidemic]
generateEpidemics initialState fixedParams params nsteps nepidemics = replicateM nepidemics (generateSingleEpidemic initialState fixedParams params nsteps)

fixedParams :: FixedParams
fixedParams = FixedParams 763 1

params :: Params
params = Params 0.9 2.0 0.6

initialState :: LatentState
initialState = LatentState 762 1 0

scoreEpidemicToDatum
    :: MonadInfer m
    => FixedParams
    -> Epidemic
    -> LatentState
    -> Params
    ->  m Params
scoreEpidemicToDatum fixedParams dat  x params = do
    let obs lambda y = score (poissonPdf lambda y)
        go [] x = return x
        go (y:ys) x = do
            x' <- transitionModel fixedParams params x
            obs ((fromIntegral $ inf x') * (rho params)) y
            return x'
    (go $! (unwrapEpidemic dat)) $! x
    return params

unwrapEpidemic :: Epidemic -> [Int]
unwrapEpidemic (Epidemic xs) = xs

{-
ddprior <- function(params) {
    dgamma(params[[1]], 0.3, 10, log = TRUE) +
    dgamma(params[[3]], 1, 8, log=TRUE) +
    dbeta(params[[2]], 2,7, log=TRUE)
}
-}
paramsPrior :: MonadSample m => m Params
paramsPrior = do
    pBeta <- Control.Monad.Bayes.Class.gamma 2 1
    pRho <- Control.Monad.Bayes.Class.beta 2.0 7.0
    pGamma <- Control.Monad.Bayes.Class.gamma 1.0  (1 / 8.0)
    return (Params pRho pBeta pGamma )

testInferenceEpidemic :: Int -> Int -> IO [[(Params, Numeric.Log.Log Double)]]
testInferenceEpidemic nsteps nparticles = do
    ys <- parseFromFile epidemicParser "data/datafile"
    case ys of (Left _) -> error "naughty"
               (Right dat) -> sampleIO $ do
                       pmmhRes <- prior $ pmmh nsteps 14 nparticles paramsPrior (scoreEpidemicToDatum fixedParams dat initialState)
                       return pmmhRes

parseFromFile p file = runParser p file <$> TIO.readFile file

extractParams :: (Params -> Double) -> [[(Params, Numeric.Log.Log Double)]] -> [Double]
extractParams project samples = (project . fst . head) <$> samples

--scoreEpidemicToData :: (MonadSample m, MonadState [InfectionCount] m => Epidemic ->  Params -> LatentState -> m Epidemic
--scoreEpidemicToData data params initialState = do
--    let obs lambda y = score (poissonPdf lambda y)
    -- simulate a new x from old x (x is a latent state)
    -- calculate new lambda
    -- score observation at time using lambda
    -- store x in the monad state
    -- repeat

generateSamples :: [Double]
generateSamples = sampleSTfixed $ replicateM 1000 (normal 3 1)


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
module Statistics.BBVI.Examples
  ( genMixture
  , mixtureFit
  , genNormal
  , normalFit
  , genDirichlet
  , dirichletFit
  , genMixedMem
  , mixedMemFit
  )
where

import           Control.Monad.ST
import           Data.Propagator
import qualified Data.Vector                   as V
import           Numeric.AD                     ( grad'
                                                , auto
                                                , diff
                                                )
import           System.Random.MWC              ( create
                                                , uniform
                                                , initialize
                                                , GenST
                                                )
import qualified System.Random.MWC.Distributions
                                               as MWCD
import           Statistics.BBVI

--- helper functions
logSum :: (Floating c, Ord c) => V.Vector c -> V.Vector c -> c
logSum v1 = log . minimumProb . V.sum . V.map exp . V.zipWith (+) v1

minimumProb :: (Floating a, Ord a) => a -> a
minimumProb = max 1e-308

--- useful global params
-- for fit
globalMaxStep :: Int
globalMaxStep = 1000
globalDelta :: Double
globalDelta = 1e-16 -- 0.00001 --

-- for generating data
nDimension :: Int
nDimension = 1
nTime :: Int
nTime = 1000
nState :: Int
nState = 3 -- 10

-----------------
--  mixture
------------------
-- | log joint of a mixture model, with categorical marginalized out
mixtureJoint :: (Floating a, Ord a) => a -> V.Vector a -> a -> V.Vector a -> a
mixtureJoint std theta ob betas = logSum
  (V.map log theta)
  (V.map (\mu -> diffableNormalLogProb mu std ob) betas)

-- | transform a log joint to a gradient propagator, for use with
-- models with a mixture-like plate structure. Updates the mixture
-- components with the reparameterization gradient (Kucukelbir et al
-- 2017).
updateReparam
  :: (Dist b SampleVector, Differentiable c Double) -- Dist a SampleVector,
  => Int
  -> (  forall e
      . (Floating e, Ord e)
     => V.Vector e
     -> e
     -> V.Vector e
     -> e
     )
  -> GenST s
  -> V.Vector (V.Vector Double)
  -> DistInvariant b
  -> DistInvariant c
  -> DistCell b
  -> V.Vector (V.Vector (DistCell c))
  -> ST
       s
       (DistCell b, V.Vector (V.Vector (DistCell c)))
updateReparam nSamp jointF gen obss thetaG betaG thetaN betasN = do
  -- obss        <- V.replicateM nObs (resample xss gen)
  thetaSamp   <- V.replicateM nSamp (resample theta gen)
  betaSamples <- V.replicateM nSamp (epsilon ((betass V.! 0) V.! 0) gen)
  let
    (joints, gradJoints) = V.unzip $ V.zipWith
      (\e th ->
        V.foldl1'
            (\(x1, y1) (x2, y2) -> (x1 + x2, V.zipWith (V.zipWith (+)) y1 y2))
          $ V.map
              (\obs ->
                let (joints', gradss) = V.unzip $ V.zipWith
                      (\ob betas -> grad' (jointF (V.map auto th) (auto ob))
                                          (V.map (\d -> transform d e) betas)
                      )
                      obs
                      betass
                in  (V.sum joints', gradss)
              )
              (obss :: V.Vector (V.Vector Double))
      )
      betaSamples
      thetaSamp
  return
    $ ( gradientScore thetaG thetaN (nFactor, joints, thetaSamp)
      , V.imap
        (\dim v -> V.imap
          (\c d -> gradientReparam
            betaG
            d
            (nFactor, V.map ((V.! c) . (V.! dim)) gradJoints, betaSamples)
          )
          v
        )
        betasN
      )
 where
  nFactor = fromIntegral $ V.length obss
  theta   = dist thetaN
  betass  = V.map (V.map dist) betasN

-- | transform a log joint to a gradient propagator, for use with
-- models with a mixture-like plate structure. Updates the mixture
-- components with the score gradient (Ranganath et al 2014)
updateScore
  :: (Dist a1 c1, Dist a2 c2)
  => Int
  -> (c1 -> a3 -> V.Vector c2 -> Double)
  -> GenST s
  -> V.Vector (V.Vector a3)
  -> DistInvariant a1
  -> DistInvariant a2
  -> DistCell a1
  -> V.Vector (V.Vector (DistCell a2))
  -> ST
       s
       ( DistCell a1
       , V.Vector (V.Vector (DistCell a2))
       )
updateScore nSamp jointF gen obss thetaG betaG thetaN betasN = do
  -- obss        <- V.replicateM nObs (resample xs gen)
  thetaSamp   <- V.replicateM nSamp (resample theta gen)
  betaSamples <- V.replicateM nSamp
                              (V.mapM (V.mapM (\b -> resample b gen)) betass)
  let joints = V.zipWith
        (\theta' betass' ->
          V.foldl1' (V.zipWith (+))
            . V.map (\obs -> V.zipWith (jointF theta') obs betass')
            $ obss
        )
        thetaSamp
        betaSamples
  return
    $ ( gradientScore thetaG thetaN (nFactor, V.map V.sum joints, thetaSamp)
      , V.imap
        (\dim v -> V.imap
          (\c d -> gradientScore
            betaG
            d
            ( nFactor
            , V.map (V.! dim) joints
            , V.map ((V.! c) . (V.! dim)) betaSamples
            )
          )
          v
        )
        betasN
      )
 where
  -- xs     = dist xsN
  nFactor = fromIntegral $ V.length obss
  theta   = dist thetaN
  betass  = V.map (V.map dist) betasN

-- | generate data for testing mixture model
genMixture :: ST s (V.Vector (V.Vector Double))
genMixture = do
  gen <- create
  let theta' = MWCD.categorical (V.replicate nState 1.0) gen
  let std    = 1.0
  let
    mixtures = V.generate
      nState
      (\k ->
        (MWCD.normal
          ((fromIntegral k - (fromIntegral (nState - 1) * 0.5)) * 5)
          std
          gen
        )
      )
  xs <- V.replicateM nTime (V.replicateM nDimension . (mixtures V.!) =<< theta')
  return xs

-- | fit a mixture model, given data
mixtureFit
  :: V.Vector (V.Vector Double) -> (Dirichlet, V.Vector (V.Vector NormalDist))
mixtureFit xs = runST $ do
  genG <- create
  gen1 <- initialize =<< V.replicateM 256 (uniform genG)
  let priorTheta = dirichlet (V.replicate nState 1.0)
  let priorBeta = normalDistr 0 (5 * (fromIntegral nState - 1) :: Double)
  let nSamp      = 10
  let localStep  = (20 :: Int)

  -- initialize distribution cells (both invariant and variants)
  let thetaGrad = DistInvariant (fromIntegral $ V.length xs)
                                priorTheta
                                (rhoKuc defaultKucP) --
  qTheta <- cellWith $ mergeGeneric globalMaxStep globalDelta
  write qTheta $ defaultDistCell (dirichlet (V.replicate nState 1.0))
  let betaGrad = DistInvariant (fromIntegral $ V.length xs)
                               priorBeta
                               (rhoKuc defaultKucP) --
  qBetas <- cellWith $ mergeGenericss globalMaxStep globalDelta
  write qBetas =<< V.replicateM
    nDimension
    (V.generateM
      nState
      (\_i -> do
        mu <- resample priorBeta gen1
        return $ defaultDistCell (normalDistr mu (1.0 :: Double))
      )
    )

  -- attach gradient propagators to cells
  stepTogether
    (updateReparam nSamp (mixtureJoint 1.0) gen1 xs thetaGrad betaGrad)
    qTheta
    qBetas
  -- stepSeparate localStep
  --              globalDelta
  --              (updateScore nSamp (mixtureJoint 1.0) gen1 xs thetaGrad betaGrad)
  --              qTheta
  --              qBetas

  -- pull content (run network to quiesence)
  thetaF <- unsafeContent qTheta
  betaF  <- unsafeContent qBetas
  let betaDists = V.map (V.map dist) betaF
  return (dist thetaF, betaDists)

--------------------------------------
-- fit simple normal distribution mean
--------------------------------------
-- | generate normally distributed data for testing
genNormal :: ST s (V.Vector Double)
genNormal = do
  gen <- create
  xs  <- V.replicateM 1000 (resample (normalDistr (5.0 :: Double) 3.0) gen)
  return xs

-- | log joint of normal distribution
jointNormal :: Floating a => a -> a -> a -> a
jointNormal std ob mu = diffableNormalLogProb mu std ob

-- | fit a gaussian using reparam gradient
normalFit :: Dist (Obs a) SampleDouble => V.Vector a -> (NormalDist, Time)
normalFit xs = runST $ do
  genG <- create
  gen1 <- initialize =<< V.replicateM 256 (uniform genG)
  let nSamp = 100
  xProp <- known $ defaultObs xs
  q     <- cellWith $ mergeGeneric globalMaxStep globalDelta
  write q (defaultDistCell $ normalDistr 0.0 2.0)
  let qInvar = DistInvariant (fromIntegral $ V.length xs)
                             (normalDistr 0.0 2.0)
                             (rhoKuc defaultKucP)

  (\qP xP -> watch qP $ \q' -> with xP $ \xs' ->
      singleUpdateReparam nSamp (jointNormal 1.0) gen1 qInvar xs' q' >>= write q
    )
    q
    xProp
  q' <- unsafeContent q
  return (dist q', time q')

-- | transform a log joint to a gradient propagator for a single
-- distribution cell, using reparameterization graidient. TODO: more
-- work on generalizing these types of functions!
singleUpdateReparam
  :: (Dist b SampleDouble, Differentiable c Double)
  => Int
  -> (forall  e . (Floating e, Ord e) => e -> e -> e)
  -> GenST s
  -> DistInvariant c
  -> DistCell b
  -> DistCell c
  -> ST s (DistCell c)
singleUpdateReparam nSamp joint gen qG xsN qN = do
  obs     <- V.replicateM nSamp (resample xs gen)
  samples <- V.replicateM nSamp (epsilon q gen)
  let gradJoint = V.map
        (\mu -> V.sum $ V.map (\ob -> diff (joint (auto ob)) mu) obs)
        (V.map (transform q) samples)
  return $ gradientReparam qG qN ((fromIntegral nSamp), gradJoint, samples)
 where
  xs = dist xsN
  q  = dist qN

---------------------------
-- simple dirichlet example
---------------------------
-- | generate data for testing the fit to a dirichlet
genDirichlet :: ST s (V.Vector Int)
genDirichlet = do
  gen <- create
  xs  <- V.replicateM
    1000
    (   resample (dirichlet (V.fromList [10.0, 20.0])) gen
    >>= \cat -> MWCD.categorical (cat :: V.Vector Double) gen
    )
  return (xs :: V.Vector Int)

-- | git a dirichlet to samples form a categorical
dirichletFit :: V.Vector Int -> (Dirichlet, Time)
dirichletFit xs = runST $ do
  genG <- create
  gen1 <- initialize =<< V.replicateM 256 (uniform genG)
  let priorTheta = dirichlet (V.fromList [1.0, 1.0])
  let nSamp      = 100
  qTheta <- cellWith $ mergeGeneric globalMaxStep globalDelta
  let qInvar = DistInvariant 1 priorTheta (rhoKuc defaultKucP)
  write qTheta (defaultDistCell priorTheta)
  (\tP -> watch tP $ \theta' -> do
      upTh <- dirichletGradProp nSamp gen1 xs qInvar theta'
      write tP upTh
    )
    qTheta
  thetaF <- unsafeContent qTheta
  return (dist thetaF, time thetaF)

dirichletGradProp
  :: Dist a (V.Vector Double)
  => Int
  -> GenST s
  -> V.Vector Int
  -> DistInvariant a
  -> DistCell a
  -> ST s (DistCell a)
dirichletGradProp nSamp gen xs qInvar diriQ =
  V.replicateM nSamp (resample diri gen) >>= \samples -> return $ gradientScore
    qInvar
    diriQ
    (1, V.map (\z -> V.sum $ V.map (\i -> log (z V.! i)) xs) samples, samples)
  where diri = dist diriQ

----------------------
-- mixed membership
----------------------

-- | transform a log joint to a gradient propagator, for use with
-- models with a mixed-membership-like plate structure. Updates the
-- mixture components with the reparameterization gradient (Kucukelbir
-- et al 2017). TODO : generalize all these joint transformer
-- functions better around the idea of plates
updateMembershipReparam
  :: (Dist b SampleVector, Differentiable c Double) -- Dist a SampleVector,
  => Int
  -> (  forall e
      . (Floating e, Ord e)
     => V.Vector e
     -> e
     -> V.Vector e
     -> e
     )
  -> GenST s
  -> V.Vector (V.Vector Double)
  -> DistInvariant b
  -> DistInvariant c
  -> V.Vector (DistCell b)
  -> V.Vector (V.Vector (DistCell c))
  -> ST
       s
       ( V.Vector (DistCell b)
       , V.Vector (V.Vector (DistCell c))
       )
updateMembershipReparam nSamp jointF gen obss thetaG betaG thetaN betasN = do
  thetaSamples <- V.replicateM nSamp (V.mapM (\th -> resample th gen) thetas)
  betaSamples  <- V.replicateM nSamp (epsilon ((betass V.! 0) V.! 0) gen)
  let (joints, gradJoints) = V.unzip $ V.zipWith
        (\e ths -> V.unzip $ V.zipWith
          (\obs th ->
            let (joints', gradss) = V.unzip $ V.zipWith
                  (\ob betas -> grad' (jointF (V.map auto th) (auto ob))
                                      (V.map (\d -> transform d e) betas)
                  )
                  obs
                  betass
            in  (V.sum joints', gradss)
          )
          (obss :: V.Vector (V.Vector Double))
          ths
        )
        betaSamples
        thetaSamples
  return
    $ ( V.imap
        (\i thN -> gradientScore
          thetaG
          thN
          (nFactor, V.map (V.! i) joints, (V.map (V.! i) thetaSamples))
        )
        thetaN
      , V.imap
        (\dim v -> V.imap
          (\c d -> gradientReparam
            betaG
            d
            ( nFactor
            , V.map (V.sum . V.map ((V.! c) . (V.! dim))) gradJoints
            , betaSamples
            )
          )
          v
        )
        betasN
      )
 where
  nFactor = fromIntegral $ V.length obss
  thetas  = V.map dist thetaN
  betass  = V.map (V.map dist) betasN

-- | transform a log joint to a gradient propagator, for use with
-- models with a mixed membership-like plate structure. Updates the
-- mixture components with the score gradient (Ranganath et al
-- 2014). TODO: generalize these log joint to gradient propagator
-- transformer functions
updateMembershipScore
  :: (Dist a1 c1, Dist a2 c2)
  => Int
  -> (c1 -> a3 -> V.Vector c2 -> Double)
  -> GenST s
  -> V.Vector (V.Vector a3)
  -> DistInvariant a1
  -> DistInvariant a2
  -> V.Vector (DistCell a1)
  -> V.Vector (V.Vector (DistCell a2))
  -> ST
       s
       ( V.Vector (DistCell a1)
       , V.Vector (V.Vector (DistCell a2))
       )
updateMembershipScore nSamp jointF gen obss thetaG betaG thetasN betasN = do
  -- obss        <- V.replicateM nObs (resample xs gen)
  thetaSamples <- V.replicateM nSamp
                               (V.mapM (\theta -> resample theta gen) thetas)
  betaSamples <- V.replicateM nSamp
                              (V.mapM (V.mapM (\b -> resample b gen)) betass)
  let joints = V.zipWith
        (\thetas' betass' ->
          V.zipWith (\obs theta -> V.zipWith (jointF theta) obs betass') obss
            $ thetas'
        )
        thetaSamples
        betaSamples
  return
    $ ( V.imap
        (\i thetaN -> gradientScore
          thetaG
          thetaN
          (nFactor, V.map (V.sum . (V.! i)) joints, V.map (V.! i) thetaSamples)
        )
        thetasN
      , V.imap
        (\dim v -> V.imap
          (\c d -> gradientScore
            betaG
            d
            ( nFactor
            , V.map (V.sum . V.map (V.! dim)) joints
            , V.map ((V.! c) . (V.! dim)) betaSamples
            )
          )
          v
        )
        betasN
      )
 where
  -- xs     = dist xsN
  nFactor = fromIntegral $ V.length obss
  thetas  = V.map dist thetasN
  betass  = V.map (V.map dist) betasN

-- | generate data for testing a Gaussian mixed membership model
genMixedMem :: ST s (V.Vector (V.Vector (Double))) -- Maybe Double
genMixedMem = do
  gen <- create
  -- initialize dirichlet that favors sparse categoricals
  let diriParam  = 0.01
  let thetaParam = MWCD.dirichlet (V.replicate nState diriParam) gen
  thetasTrue <- V.replicateM nTime thetaParam
  let thetas' = V.map (\theta' -> MWCD.categorical theta' gen) thetasTrue
  let betaStd = 0.001
  let betas' = V.generate
        nState
        (\k -> V.replicate
          nDimension
          (MWCD.normal
            ((fromIntegral k - (fromIntegral (nState - 1) * 0.5)) * 5)
            betaStd
            gen
          )
        )
  xs <- V.generateM
    nTime
    (\day -> V.generateM
      nDimension
      (\loc -> (thetas' V.! day) >>= \z -> ((betas' V.! z) V.! loc)) -- Just <$>
    )
  return xs

-- | fit a Gaussian mixed membership model to data
mixedMemFit
  :: V.Vector (V.Vector Double)
  -> (V.Vector Dirichlet, V.Vector (V.Vector NormalDist))
mixedMemFit xs = runST $ do
  genG <- create
  gen1 <- initialize =<< V.replicateM 256 (uniform genG)
  let priorTheta = dirichlet (V.replicate nState 1.0)
  let priorBeta = normalDistr 0 (5 * (fromIntegral nState - 1) :: Double)
  let nSamp      = 10
  -- let localStep  = 20
  let betaGrad = DistInvariant (fromIntegral $ V.length xs)
                               priorBeta
                               (rhoKuc defaultKucP)
  qBetas <- cellWith $ mergeGenericss globalMaxStep globalDelta
  write qBetas =<< V.generateM
    nDimension
    (\_i -> V.replicateM
      nState
      (do
        mu <- resample priorBeta gen1
        return $ defaultDistCell (normalDistr mu 1.0)
      )
    )
  qThetas <- cellWith $ mergeGenerics globalMaxStep globalDelta
  write qThetas $ V.replicate nTime (defaultDistCell priorTheta)
  let thetaGrad = DistInvariant (fromIntegral $ V.length xs)
                                priorTheta
                                (rhoKuc defaultKucP)
  stepTogether
    (updateMembershipReparam nSamp (mixtureJoint 1.0) gen1 xs thetaGrad betaGrad
    )
    qThetas
    qBetas
  -- stepSeparate localStep
  --              globalDelta
  --              (updateMembershipScore nSamp (mixtureJoint 1.0) gen1 xs thetaGrad betaGrad)
  --              qThetas
  --              qBetas
  thetaF <- unsafeContent qThetas
  betaF  <- unsafeContent qBetas
  return (V.map dist thetaF, V.map (V.map dist) betaF)


module Tests.ComplexFromIntegral where
import Data.Complex

runTest = do
  putStrLn $ show $ p' where
    p' = zipWith (*) (map (fromIntegral) [1..3]) p
    p = [1 :+ 0, 2 :+ 0, 3 :+ 0]


{-# LANGUAGE CPP #-}

module Year2021.Day15 where

import Numeric.LinearAlgebra
import Util
import Safe
import Linear.V2
import System.FilePath
import Data.Bifunctor

part1, part2 :: Int
part1 = 1
part2 = 5

main = readFile (replaceExtension __FILE__ ".in") >>= \input ->
  let grid = fromLists . map (map (read . pure)) $ lines input :: Matrix Z
      grid' = cmod 9 (fromBlocks [[grid + scalar (r + c) | c <- [0..4]] | r <- [0..4]] - 1) + 1
      neigh (a, b) = [ (grid' ! x ! y, (x, y))
        | p@(V2 x y) <- (+ V2 a b) <$> adjacent
        , inBounds 0 (uncurry V2 (size grid') - 1) p ]
   in print . fst . findJust ((== size grid') . bimap succ succ . head . snd)
        $ dijkstra neigh (0, 0)


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE RankNTypes #-}
-----------------------------------------------------------------------------
--
-- Module      :  AI.Network.RMM.LSTM
-- Copyright   :  (c) JP Moresmau
-- License     :  BSD3
--
-- Maintainer  :  JP Moresmau <jp@moresmau.fr>
-- Stability   :  experimental
-- Portability :
--
-- | Long Short Term Memory networks
--
-----------------------------------------------------------------------------

module AI.Network.RNN.LSTM where

import Control.DeepSeq

import Data.List

import Numeric.LinearAlgebra.HMatrix

import AI.Network.RNN.Types
import AI.Network.RNN.Util
import AI.Network.RNN.Expr
import Debug.Trace
import qualified Data.IntMap as I

import Numeric.AD

-- | The LSTM data type, with the different weights and states
data LSTMNetwork = LSTMNetwork
    { lstmSize :: !Int
    , lstmWeightsW :: !(Matrix Double)
    , lstmWeightsU :: !(Matrix Double)
    , lstmBias :: !(Vector Double)
    , lstmState :: !(Vector Double)
    , lstmOutput :: !(Vector Double)
    } deriving (Show,Read,Eq)

--instance Show LSTMNetwork where
--     show s = let
--        a= show $ (rnnsize s,toVector s)
--        in trace (show $ length a) a
--
--instance Read LSTMNetwork where
--    readsPrec d s =
--        let [(ok,r)] = readsPrec d s
--            dat = uncurry fromVector ok
--        in [(dat,r)]

-- | Force evaluation instance
instance NFData LSTMNetwork where
    rnf LSTMNetwork{..} = rnf (lstmSize,lstmWeightsW,lstmWeightsU,lstmBias,lstmState,lstmOutput)

-- | Implement the network evaluation and conversions functions
instance RNNEval LSTMNetwork where
    type Size LSTMNetwork = Int
    evalStep n@LSTMNetwork{..} is =
        let z = (lstmWeightsW #> is) + (lstmWeightsU #> lstmOutput) + lstmBias
            [i,f,c1,o] = takesV (replicate 4 lstmSize) z
            c2 = cmap tanh c1
            ns = (c2 * cmap sigmoid i) + (cmap sigmoid f * lstmState)
            no = cmap sigmoid o * cmap tanh ns
        in force (n{lstmState=ns,lstmOutput=no},no)
    fromVector sz vs =
        let
            msize = sz * sz
            [v1,v2,v3,v4,v5] = takesV [msize * 4,msize * 4, sz * 4,sz,sz] vs
            m1 = reshape sz v1
            m2 = reshape sz v2
        in LSTMNetwork sz m1 m2 v3 v4 v5
    toVector LSTMNetwork{..} = force $ vjoin
        [ flatten lstmWeightsW, flatten lstmWeightsU, lstmBias, lstmState, lstmOutput]
    rnnsize = lstmSize
    fullSize = lstmFullSize .lstmSize

-- | Full size of a network
lstmFullSize :: FullSize LSTMNetwork
lstmFullSize sz = (sz * sz) * 8 + sz* 4 + sz + sz

data LSTMIO = LSTMIO
    { lioInput :: Matrix Double
    , lioLstms :: [LSTMNetwork]
    , lioOutput :: Matrix Double
    , lioSer   :: Maybe (Vector Double)
    } deriving (Eq)

instance NFData LSTMIO where
  rnf LSTMIO{..} = rnf (lioInput,lioLstms,lioOutput)

instance RNNEval LSTMIO  where
    type Size LSTMIO = (Int,Int,Int,Int)
    evalStep io@LSTMIO{..} is =
        let inps = cmap sigmoid $ lioInput #> (vjoin [is,scalar 1])
            (lstm2,outs) = evalStep lioLstms inps
        in (io{lioLstms=lstm2,lioSer=Nothing},lioOutput #> outs)
    fromVector (is,lnnSize,lnnNumber,os) vs =
        let [v1,v2,v3] = takesV [(is+1)* lnnSize ,lnnNumber * (lstmFullSize lnnSize),os * lnnSize ] vs
            m1 = reshape (is+1) v1
            m2 = reshape lnnSize v3
        in LSTMIO m1 (fromVector (replicate lnnNumber lnnSize) v2) m2 $ Just vs
    toVector LSTMIO{..} = case lioSer of
        Just vs -> vs
        Nothing -> vjoin [flatten lioInput,toVector lioLstms,flatten lioOutput]
    rnnsize LSTMIO{..} = (cols lioInput - 1, rnnsize $ head lioLstms, length lioLstms,rows lioOutput)
    fullSize =lstmioFullSize . rnnsize
--        let (r1,c1)=size lioInput
--            (r2,c2)=size lioOutput
--        in (r1*c1) + (fullSize lioLstms) + (r2*c2)

instance Show LSTMIO where
    show l = show (rnnsize l,toVector l)

instance Read LSTMIO where
    readsPrec d s =
        let [(ok,r)] = readsPrec d s
            dat = uncurry fromVector ok
        in [(dat,r)]

lstmioFullSize :: FullSize LSTMIO
lstmioFullSize (is,lnnSize,lnnNumber,os) = (is+1) * lnnSize + lnnNumber * (lstmFullSize lnnSize) + os * lnnSize

data LSTMList = LSTMList
    { llSize :: (Int,Int,Int,Int)
    , llData :: [Double]
    } deriving (Show, Read, Eq)

instance NFData LSTMList where
  rnf LSTMList{..} = rnf (llSize,llData)

instance RNNEval LSTMList where
    type Size LSTMList = (Int,Int,Int,Int)
    rnnsize = llSize
    fullSize =lstmioFullSize . rnnsize
    toDList = llData
    fromDList sz = LSTMList sz
    toVector = fromList . toDList
    fromVector sz = fromDList sz . toList
    evalStep l@LSTMList{..} is =
        let lstmio :: LSTMIO = fromDList llSize llData
            (io2,out) = evalStep lstmio is
        in (l{llData=toDList io2},out)

-- | Implementation of the LSTM evaluation step without explicit matrices and vectors
-- just using lists, so we can use AD on it
lstmList :: (Num b,Floating b) => Int -> [b] -> [b] -> ([b],[b])
lstmList sz lstm is = let
    msize = sz * sz
    [mW,mU,vB,vS,vO] = takes [msize * 4,msize * 4, sz * 4,sz,sz] lstm
    z = zipWith3 (\a b c->a+b+c) (listMProd mW is) (listMProd mU vO) vB
    [i,f,c1,o] = takes (replicate 4 sz) z
    c2 = map tanh c1
    ns = zipWith (+) (zipWith (*) c2 (map sigmoid i)) (zipWith (*) (map sigmoid f) vS)
    no = zipWith (*) (map sigmoid o) (map tanh ns)
    in (mW++mU++vB++ns++no,no)

-- | Cost calculation using list representation for AD
cost' :: (Num b,Floating b,Fractional b) => Int -> [[b]] -> [[b]] -> [b] -> b
cost' sz is os lstm = let
    (_,res) = mapAccumL (lstmList sz) lstm is
--    in - (calcMeanList $ three (last res) (last os))
--    in - (calcMeanList $ concat $ zipWith (three) res os)
    in sum $ zipWith err os res
    where
      err :: (Num b,Floating b) => [b] -> [b] -> b
      err a b  = sum (zipWith (\c d -> (c- d)**2 ) a b)
--      one i o = zipWith (*) o (map log i)
--      two i o = zipWith (*) (map (\x->1 - x) o) (map (\x->1 - log x) i)
--      three i o= zipWith (+) (one i o) (two i o)

-- | Gradient descent learning
-- The third parameter is a call back function to monitor progress and stop the learning process if needed
learnGradientDescent :: (Monad m) => LSTMNetwork -> TrainData a -> (LSTMNetwork -> TrainData a -> Int -> m Bool) -> m LSTMNetwork
learnGradientDescent lstm td progressF =  go (toList $ toVector lstm) 0
    -- go2 gds 0
    where
      go2 (ls:lss) gen =  do
        let rnn::LSTMNetwork = fromVector (tdRecSize td) (fromList ls)
        cont <- progressF rnn td gen
        if cont
            then go2 lss (gen+1)
            else return rnn
      go ls gen = do
        let rnn::LSTMNetwork = fromVector (tdRecSize td) (fromList ls)
        cont <- progressF rnn td gen
        if cont
            then do
                let
                    gs= gf ls -- gradients using AD
                    ls2 = zipWith (\o g->o-g*0.1) ls gs
                go ls2 (gen+1)
            else return rnn
      lis = map toList (tdInputs td)
      los = map toList (tdOutputs td)
      gf = grad (cost' (tdRecSize td) (map (map auto) lis) (map (map auto) los))
      gds = gradientDescent (cost' (tdRecSize td) (map (map auto) lis) (map (map auto) los)) (toList $ toVector lstm)

-- | Gradient descent learning using symbolic differentiation
--   The goal was to calculate the derivative once and then just close and eval the result using the current data
--   However it is much slower than normal automatic differentiation
learnGradientDescentSym :: (Monad m) => LSTMNetwork -> TrainData a -> (LSTMNetwork -> TrainData a -> Int -> m Bool) -> m LSTMNetwork
learnGradientDescentSym lstm td progressF =  go (toList $ toVector lstm) 0
    where
      go ls gen = do
        let rnn::LSTMNetwork = fromVector (tdRecSize td) (fromList ls)
        cont <- progressF rnn td gen
        if cont
            then do
                let
                    i = I.fromList $ zip [0..] ls
                    cexpr = map (\g-> close g i) gf
                    gs = map eval cexpr
                    ls2 = zipWith (\o g->o-g*0.1) ls gs
                go ls2 (gen+1)
            else return rnn
      lis = map toList (tdInputs td)
      los = map toList (tdOutputs td)
      ls0 = toList $ toVector lstm
      gf = map fullSimplify $ grad (cost'
            (tdRecSize td)
            (map (map (\x -> autoEval (Lit $ Lit x) I.empty)) lis)
            (map (map (\x -> autoEval (Lit $ Lit x) I.empty)) los))
            (zipWith (\_ i->Var i) ls0 [0..])

-- | RMSProp learning, as far as I can make out
-- The third parameter is a call back function to monitor progress and stop the learning process if needed
learnRMSProp :: (Monad m) => LSTMNetwork -> TrainData a -> (LSTMNetwork -> TrainData a -> Int -> m Bool) -> m LSTMNetwork
learnRMSProp lstm td progressF = go ls0 (replicate myl 0) (replicate myl 0) (replicate myl 0) 0
    where
      go ls rgs rgs2 ugs gen = do
        let rnn::LSTMNetwork = fromVector (tdRecSize td) (fromList ls)
        cont <- progressF rnn td gen
        if cont
            then do
                let
                    gs= gf ls -- gradients using AD
                    rgup = zipWith (\rg g-> 0.95 * rg + 0.05 * g) rgs ls
                    rg2up = zipWith (\rg2 g-> 0.95 * rg2 + 0.05 * (g ** 2)) rgs2 ls
                    ugup = zipWith4 (\ud zg rg rg2 -> 0.9 * ud - 1e-4 * zg / sqrt(rg2 - rg ** 2 + 1e-4)) ugs gs rgup rg2up
                    ls2 = zipWith (+) ls ugup
                go (force ls2) (force rgup) (force rg2up) (force ugup) (gen+1)
            else return rnn
      lis = map toList (tdInputs td)
      los = map toList (tdOutputs td)
      gf = grad (cost' (tdRecSize td) (map (map auto) lis) (map (map auto) los))
      ls0 = toList $ toVector lstm
      myl= length ls0


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE TemplateHaskell   #-}
{-# LANGUAGE UnicodeSyntax     #-}
{-# LANGUAGE BangPatterns      #-}
module NLP.Albemarle.LSA
    (-- * Creating a model
    lsa
    -- * Managing terms
    , rebase
    -- * Managing topics
    , pad
    , trim
    -- * Internal methods (may be subject to change)
    , SVD.sparsify
    , SVD.sparseSvd
    , batchLSA
    -- * Types (may be subject to change)
    , LSAModel(..)
    , termvectors
    , topicweights
    ) where
import NLP.Albemarle
import NLP.Albemarle.Dict (Dict)
import qualified NLP.Albemarle.Dict as Dict
import ClassyPrelude hiding (Vector)
import qualified Data.Vector.Generic as Vec
import qualified Data.Vector.Unboxed as UVec
import qualified Data.Vector.Storable as SVec
import qualified Data.HashMap.Strict as HashMap
import Numeric.LinearAlgebra (Vector, Matrix, tr, diag, (|||), (===))
import qualified Numeric.LinearAlgebra as HMatrix
import qualified Numeric.LinearAlgebra.Devel as HMatrix
import qualified Numeric.LinearAlgebra.SVD.SVDLIBC as SVD
import qualified System.IO.Streams as Streams
import System.IO.Streams (Generator, InputStream, OutputStream)
import Lens.Micro.TH
import Lens.Micro
import Data.Tuple

-- | An LSA model (The singular values and right singular vectors of truncated
--   SVD on a term-document matrix, where documents are rows and terms columns)
data LSAModel = LSAModel {
  _topicweights :: !(HMatrix.Vector Double), -- ^ Topic weights
  _termvectors :: !(Matrix Double) -- ^ Rows are topics, columns are terms
} deriving (Show, Eq)
makeLenses ''LSAModel

-- This instance may be a fib. I'm not sure this is actually associative.
-- Even if it is, I imagine there are practical numerical stability concerns.
-- TODO: In particular, I think we need to weight the left and weight according
-- to how many documents they represent.
instance Monoid LSAModel where
  mempty = LSAModel mempty mempty
  mappend left right
    | left == mempty = right
    | right == mempty = left
    | otherwise = trim target_len $ pad target_len $ LSAModel s v
    where
      target_len = max (topicCount left) (topicCount right)
      (s, v) = HMatrix.rightSV $ combine left === combine right
instance Semigroup LSAModel where
  (<>) = mappend

-- | Truncate an LSAModel to a specific number of topics, if there are too many
trim :: Int -- ^ The maximum number of topics the model should have (inclusive)
     -> LSAModel -> LSAModel
trim count model
  | topicCount model <= count = model
  | otherwise = LSAModel
    (Vec.take count $ model^.topicweights)
    (HMatrix.takeRows count $ model^.termvectors)

-- | Pad an LSAModel with extra empty topics of there are too few
pad :: Int -- ^ The minimum number of topics the model should have (inclusive)
    -> LSAModel -> LSAModel
pad count model
  | topicCount model >= count = model
  | otherwise = let
    (height, width) = HMatrix.size $ model^.termvectors
    missing = count - height
    in LSAModel
      (model^.topicweights <> Vec.replicate missing 0)
      (model^.termvectors === HMatrix.konst 0 (missing, width))

-- | Multiply the vectors and values of an SVD to make one matrix
combine :: LSAModel -> Matrix Double
combine model = diag (model^.topicweights) <> (model^.termvectors)

sparseToCSR :: SparseMatrix -> HMatrix.CSR
sparseToCSR (SparseMatrix width vecs) =
  let
    docs = (Vec.fromList . (\ (SparseVector len wds) -> wds) <$> vecs)
    convI = Vec.map fromIntegral . Vec.convert
    counts = Vec.map snd <$> docs         :: [UVec.Vector Double]
    concatcounts = Vec.concat counts      ::  UVec.Vector Double
    storcounts = Vec.convert concatcounts ::  SVec.Vector Double
  in HMatrix.CSR {
    HMatrix.csrVals = storcounts,
    HMatrix.csrCols = convI
      $ Vec.concat
      $ Vec.map ((+1).fst)
      <$> docs,
    HMatrix.csrRows = convI
      $ Vec.scanl (+) 1
      $ Vec.fromList
      $ Vec.length
      <$> docs,
    HMatrix.csrNCols = width,
    HMatrix.csrNRows = length docs
  }

-- | Generate an LSA Model with N topics, based on a sparse matrix.
--
--   You can make the sparse matrix with a Dict and some documents.
--   See Dict.asSparseMatrix.
lsa :: Int -- ^ Number of vectors/dimensions in the new space
    -> SparseMatrix -- ^ Sparse representation of the term-document matrix
    -> LSAModel -- ^ A dense representation of the term-topic matrix
lsa top_vectors termdoc =
  -- This astounds me, but the number of singular values may not match the
  -- number of singular vectors. We have to balance before padding.
  pad top_vectors $ case compare missing_values 0 of
    LT -> LSAModel (Vec.take top_vectors s) vt
    EQ -> LSAModel s vt
    GT -> LSAModel (s <> Vec.replicate missing_values 0) vt
  where
    (u, s, vt) = SVD.sparseSvd top_vectors $ sparseToCSR termdoc
    missing_values = fst (HMatrix.size vt) - Vec.length s

-- | Rebase an LSA model made with one Dict to one made with another Dict.
--
--   By using this, you can split a corpus into batches, make Dicts and topic
--   models, then merge the Dicts, rebase the Models to the merged Dict, and
--   then merge the models. It allows working in parallel, and limits memory
--   usage.
rebase :: Dict -> Dict -> LSAModel -> LSAModel
rebase d1 d2 model
  | model == mempty = mempty
  | otherwise = LSAModel {
    _termvectors = (HMatrix.¿) (model^.termvectors) $ Dict.select d1 d2,
    _topicweights = model^.topicweights
  }

-- | Get the number of topics in a model
topicCount :: LSAModel -> Int
topicCount model = Vec.length $ model^.topicweights

-- | SVD with some transposes, for convenience and speed
-- Normally (U, Sigma, V^T) = svd A, but then the rows are topics and the cols
-- are the documents/words (in U and V). But we use C-style (row major) matrices
-- which means for large matrices, getting the vector of one word or one doc
-- will require reading the whole model from memory. (A terrible waste of cache)
-- So instead we use (U^T, sigma, V) where the rows are the vector embeddings
-- of documents and words instead.
--
-- Prefer lsa (rather than batchLSA) - batchLSA is experimental and may be
-- removed
batchLSA :: Int -> HMatrix.CSR -> (Matrix Double, HMatrix.Vector Double, Matrix Double)
batchLSA top_vectors csr = (tr u, s, tr vt)
  where (u, s, vt) = SVD.sparseSvd top_vectors csr


-- test for trac #8532

import Data.Complex

main :: IO ()
main = do
    print $ acosh ((-1)::Complex Double)
    print $ acosh ((-1)::Complex Float)


{-# LANGUAGE CPP                        #-}
{-# LANGUAGE PolyKinds                  #-}
{-# LANGUAGE ConstraintKinds            #-}
{-# LANGUAGE DefaultSignatures          #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE RebindableSyntax           #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE TypeFamilies               #-}

module Data.Semimodule.Algebra (
  -- * Algebras
    Algebra(..)
  -- * Unital algebras
  , Unital(..)
  -- * Coalgebras
  , Coalgebra(..)
  -- * Unital coalgebras
  , Counital(..)
  -- * Bialgebras
  , Bialgebra
) where

import Control.Arrow
import Control.Applicative
import Control.Category (Category, (<<<), (>>>))
import Data.Bool
import Data.Complex
import Data.Finite (Finite, finites)
import Data.Fixed
import Data.Functor.Apply
import Data.Functor.Compose
import Data.Functor.Contravariant
import Data.Functor.Product
import Data.Functor.Rep
import Data.Functor.Rep
import Data.Int
import Data.Semiring
import Data.Semimodule
import Data.Sequence hiding (reverse,index)
import Data.Tuple (swap)
import Data.Word
import Foreign.C.Types (CFloat(..),CDouble(..))
import GHC.Real hiding (Fractional(..))
import GHC.TypeNats (KnownNat(..))
import Numeric.Natural
import Prelude (Ord, reverse)
import Prelude (fromInteger)
import Prelude hiding (Num(..), Fractional(..), negate, sum, product)
import qualified Control.Category as C
import qualified Control.Monad as M
import qualified Data.IntSet as IntSet
import qualified Data.Map as Map
import qualified Data.Sequence as Seq
import qualified Data.Set as Set

-------------------------------------------------------------------------------
-- Algebras
-------------------------------------------------------------------------------

-- | An < https://en.wikipedia.org/wiki/Algebra_over_a_field#Generalization:_algebra_over_a_ring algebra > over a semiring.
--
-- Note that the algebra < https://en.wikipedia.org/wiki/Non-associative_algebra needn't be associative >.
--
class Semiring a => Algebra a b where

    -- |
    --
    -- @
    -- 'joined' = 'Data.Semimodule.Free.over' 'Data.Semimodule.Free.diagonal' '.' 'uncurry'
    -- @
    --
    joined :: (b -> b -> a) -> b -> a

-- | A < https://en.wikipedia.org/wiki/Algebra_over_a_field#Unital_algebra unital algebra > over a semiring.
--
class Algebra a b => Unital a b where

    -- | Obtain a vector from the unit of a unital algebra.
    --
    -- @
    -- 'unital' = 'Data.Semimodule.Transform.over' 'initial' '.' 'const'
    -- @
    --
    unital :: a -> b -> a

-------------------------------------------------------------------------------
-- Coalgebras
-------------------------------------------------------------------------------

-- | A coalgebra over a semiring.
--
class Semiring a => Coalgebra a c where

    -- |
    --
    -- @
    -- 'cojoined' = 'curry' '.' 'Data.Semimodule.Free.over' 'Data.Semimodule.Free.codiagonal'
    -- @
    --
    cojoined :: (c -> a) -> c -> c -> a

-- | A counital coalgebra over a semiring.
--
class Coalgebra a c => Counital a c where

    -- |
    --
    -- @
    -- 'Data.Semimodule.Free.Cov' 'counital' = 'Data.Semimodule.Free.coover' 'Data.Semimodule.Free.coinitial' $ pure '()'
    -- @
    --
    counital :: (c -> a) -> a

-------------------------------------------------------------------------------
-- Bialgebras
-------------------------------------------------------------------------------

-- | A < https://en.wikipedia.org/wiki/Bialgebra bialgebra > over a semiring.
--
class (Unital a b, Counital a b) => Bialgebra a b

-------------------------------------------------------------------------------
-- Algebra instances
-------------------------------------------------------------------------------

instance Semiring a => Algebra a () where
  joined f = f ()

instance Semiring a => Unital a () where
  unital r () = r

--TODO: consider separating out n=3 cross product case
instance Semiring a => Algebra a (Finite n) where
  joined = M.join

instance Semiring a => Unital a (Finite n) where
  unital = const

instance (Algebra a b1, Algebra a b2) => Algebra a (b1, b2) where
  joined f (a,b) = joined (\a1 a2 -> joined (\b1 b2 -> f (a1,b1) (a2,b2)) b) a

instance (Unital a b1, Unital a b2) => Unital a (b1, b2) where
  unital r (a,b) = unital r a * unital r b

instance (Algebra a b1, Algebra a b2, Algebra a b3) => Algebra a (b1, b2, b3) where
  joined f (a,b,c) = joined (\a1 a2 -> joined (\b1 b2 -> joined (\c1 c2 -> f (a1,b1,c1) (a2,b2,c2)) c) b) a

instance (Unital a b1, Unital a b2, Unital a b3) => Unital a (b1, b2, b3) where
  unital r (a,b,c) = unital r a * unital r b * unital r c

-- | Tensor algebra on /b/.
--
-- >>> joined (<>) [1..3 :: Int]
-- [1,2,3,1,2,3,1,2,3,1,2,3]
--
-- >>> joined (\f g -> fold (f ++ g)) [1..3] :: Int
-- 24
--
instance Semiring a => Algebra a [b] where
  joined f = go [] where
    go ls rrs@(r:rs) = f (reverse ls) rrs + go (r:ls) rs
    go ls [] = f (reverse ls) []

instance Semiring a => Unital a [b] where
  unital a [] = a
  unital _ _ = zero

instance Semiring a => Algebra a (Seq b) where
  joined f = go Seq.empty where
    go ls s = case viewl s of
       EmptyL -> f ls s
       r :< rs -> f ls s + go (ls |> r) rs

instance Semiring a => Unital a (Seq b) where
  unital a b | Seq.null b = a
             | otherwise = zero

instance (Semiring a, Ord b) => Algebra a (Set.Set b) where
  joined f = go Set.empty where
    go ls s = case Set.minView s of
       Nothing -> f ls s
       Just (r, rs) -> f ls s + go (Set.insert r ls) rs

instance (Semiring a, Ord b) => Unital a (Set.Set b) where
  unital a b | Set.null b = a
           | otherwise = zero

instance Semiring a => Algebra a IntSet.IntSet where
  joined f = go IntSet.empty where
    go ls s = case IntSet.minView s of
       Nothing -> f ls s
       Just (r, rs) -> f ls s + go (IntSet.insert r ls) rs

instance Semiring a => Unital a IntSet.IntSet where
  unital a b | IntSet.null b = a
             | otherwise = zero

---------------------------------------------------------------------
-- Coalgebra instances
---------------------------------------------------------------------

--instance (Representable f, Algebra a (Rep (Co f))) => Coalgebra a (Co f a) where
--cojoined k (Co f) (Co g) = k (index f * index g)
--  cojoined k f g = k (f * g)

instance Semiring a => Coalgebra a () where
  cojoined = const

instance Semiring a => Counital a () where
  counital f = f ()

instance Semiring a => Coalgebra a (Finite n) where
  cojoined f i j = bool zero (f i) $ i == j

instance (KnownNat n, Semiring a) => Counital a (Finite n) where
  counital f = sum . fmap f $ finites

instance Algebra a b => Coalgebra a (b -> a) where
  cojoined k f g = k (f * g)

instance Unital a b => Counital a (b -> a) where
  counital f = f one

instance (Coalgebra a c1, Coalgebra a c2) => Coalgebra a (c1, c2) where
  cojoined f (a1,b1) (a2,b2) = cojoined (\a -> cojoined (\b -> f (a,b)) b1 b2) a1 a2

instance (Counital a c1, Counital a c2) => Counital a (c1, c2) where
  counital k = counital $ \a -> counital $ \b -> k (a,b)

instance (Coalgebra a c1, Coalgebra a c2, Coalgebra a c3) => Coalgebra a (c1, c2, c3) where
  cojoined f (a1,b1,c1) (a2,b2,c2) = cojoined (\a -> cojoined (\b -> cojoined (\c -> f (a,b,c)) c1 c2) b1 b2) a1 a2

instance (Counital a c1, Counital a c2, Counital a c3) => Counital a (c1, c2, c3) where
  counital k = counital $ \a -> counital $ \b -> counital $ \c -> k (a,b,c)

-- | The tensor coalgebra on /c/.
--
instance Semiring a => Coalgebra a [c] where
  cojoined f as bs = f (mappend as bs)

instance Semiring a => Counital a [c] where
  counital f = f []

instance Semiring a => Coalgebra a (Seq c) where
  cojoined f as bs = f (mappend as bs)

instance Semiring a => Counital a (Seq c) where
  counital f = f Seq.empty

-- | The free commutative band coalgebra
instance (Semiring a, Ord c) => Coalgebra a (Set.Set c) where
  cojoined f as bs = f (Set.union as bs)

instance (Semiring a, Ord c) => Counital a (Set.Set c) where
  counital f = f Set.empty

-- | The free commutative band coalgebra over Int
instance Semiring a => Coalgebra a IntSet.IntSet where
  cojoined f as bs = f (IntSet.union as bs)

instance Semiring a => Counital a IntSet.IntSet where
  counital f = f IntSet.empty

-- | The free commutative coalgebra over a set and a given semigroup
instance (Semiring r, Ord a, Semigroup b) => Coalgebra r (Map.Map a b) where
  cojoined f as bs = f (Map.unionWith (<>) as bs)

instance (Semiring r, Ord a, Semigroup b) => Counital r (Map.Map a b) where
  counital f = f Map.empty

{-

-- | The free commutative coalgebra over a set and a given semigroup
instance (Semiring r, Ord a, Commutative b) => Coalgebra r (Map a b) where
  cojoined f as bs = f (Map.unionWith (+) as bs)
  counital k = k (Map.empty)

-- | The free commutative coalgebra over a set and Int
instance (Semiring r, Commutative b) => Coalgebra r (IntMap b) where
  cojoined f as bs = f (IntMap.unionWith (+) as bs)
  counital k = k (IntMap.empty)
-}

---------------------------------------------------------------------
-- Bialgebra instances
---------------------------------------------------------------------

instance Semiring a => Bialgebra a () where

instance (KnownNat n, Semiring a) => Bialgebra a (Finite n) where

instance (Bialgebra a b1, Bialgebra a b2) => Bialgebra a (b1, b2) where

instance (Bialgebra a b1, Bialgebra a b2, Bialgebra a b3) => Bialgebra a (b1, b2, b3) where

instance Semiring a => Bialgebra a [b]

instance Semiring a => Bialgebra a (Seq b)


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
module LSI.FrequencyResponse where

import Data.Complex
import Data.Vector as V hiding (map, zip, foldr)
import Linear.V as V

import LSI.RationalFunction

computeFR :: RealFloat f
          => RationalFunction d f
          -> V d f
          -> Complex f
computeFR rf = (computeFR' rf) . V.toList .V.toVector
  where computeFR' f =
          case f of
            Monomial c es ->
              let fs = map toInteger . V.toList . V.toVector $ es in
                \vs -> {-# SCC monom #-}
                  let factors = map iexp $ zip zs fs
                      iexp (z, e) | e < 0 = 1/(z^(-e))
                      iexp (z, e) = z^e
                      zs = map (\w -> mkPolar 1.0 w) vs in
                    (c :+ 0.0) * (foldr (*) 1 factors)

            Add rf1 rf2 ->
              let fr1 = computeFR' rf1
                  fr2 = computeFR' rf2 in
                \vs -> (fr1 vs) + (fr2 vs)

            Mul rf1 rf2 ->
              let fr1 = computeFR' rf1
                  fr2 = computeFR' rf2 in
                \vs -> (fr1 vs) * (fr2 vs)

            Div rf1 rf2 ->
              let fr1 = computeFR' rf1
                  fr2 = computeFR' rf2 in
                \vs -> (fr1 vs) / (fr2 vs)


{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TupleSections         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE FlexibleContexts      #-}

import           Control.Applicative
import           Control.Monad
import           Control.Monad.Random

import           Codec.Compression.GZip ( decompress )
import           Data.Serialize ( Get )
import qualified Data.Serialize as Serialize
import qualified Data.ByteString.Lazy as B

import           Data.List ( foldl' )
import           Data.List.Split ( chunksOf )
import           Data.Maybe ( fromMaybe )
#if ! MIN_VERSION_base(4,13,0)
import           Data.Semigroup ( (<>) )
#endif
import           Data.Word ( Word32 , Word8 )
import qualified Data.Vector.Storable as V

import           Numeric.LinearAlgebra ( maxIndex )
import qualified Numeric.LinearAlgebra.Static as SA

import           Options.Applicative
import           System.FilePath ( (</>) )

import           Grenade
import           Grenade.Utils.OneHot

#ifdef COMPLEX
-- It's logistic regression!
--
-- This network is used to show how we can embed a Network as a layer in the larger MNIST
-- type.
--
type FL i o =
  Network
    '[ FullyConnected i o, Logit ]
    '[ 'D1 i, 'D1 o, 'D1 o ]

-- The definition of our convolutional neural network.
-- In the type signature, we have a type level list of shapes which are passed between the layers.
-- One can see that the images we are inputing are two dimensional with 28 * 28 pixels.

-- It's important to keep the type signatures, as there's many layers which can "squeeze" into the gaps
-- between the shapes, so inference can't do it all for us.

-- With the mnist data from Kaggle normalised to doubles between 0 and 1, learning rate of 0.01 and 15 iterations,
-- this network should get down to about a 1.3% error rate.
--
-- /NOTE:/ This model is actually too complex for MNIST, and one should use the type given belowe instead.
--         This one is just here to demonstrate Inception layers in use.
--
type MNIST =
  Network
    '[ Reshape,
       Concat ('D3 28 28 1) Trivial ('D3 28 28 14) (InceptionMini 28 28 1 5 9),
       Pooling 2 2 2 2, Relu,
       Concat ('D3 14 14 3) (Convolution 15 3 1 1 1 1) ('D3 14 14 15) (InceptionMini 14 14 15 5 10), Crop 1 1 1 1, Pooling 3 3 3 3, Relu,
       Reshape, FL 288 80, FL 80 10 ]
    '[ 'D2 28 28, 'D3 28 28 1,
       'D3 28 28 15, 'D3 14 14 15, 'D3 14 14 15, 'D3 14 14 18,
       'D3 12 12 18, 'D3 4 4 18, 'D3 4 4 18,
       'D1 288, 'D1 80, 'D1 10 ]

#else

-- The simpler network can just be dropped in without changing any of the other code
--

type MNIST
  = Network
    '[ Convolution 1 10 5 5 1 1, Pooling 2 2 2 2, Relu
     , Convolution 10 16 5 5 1 1, Pooling 2 2 2 2, Reshape, Relu
     , FullyConnected 256 80, Logit, FullyConnected 80 10, Logit]
    '[ 'D2 28 28, 'D3 24 24 10, 'D3 12 12 10, 'D3 12 12 10
     , 'D3 8 8 16, 'D3 4 4 16, 'D1 256, 'D1 256
     , 'D1 80, 'D1 80, 'D1 10, 'D1 10]

-- ... and this is an even simpler network from  https://crypto.stanford.edu/~blynn/haskell/brain.html
{-
type MNIST
  = Network
      '[Reshape, FullyConnected 784 30, Relu, FullyConnected 30 10, Logit]
      '[ 'D2 28 28, 'D1 784, 'D1 30, 'D1 30, 'D1 10, 'D1 10]
-}

#endif

randomMnist :: MonadRandom m => m MNIST
randomMnist = randomNetwork

convTest :: Int -> FilePath -> Maybe Int -> LearningParameters -> IO ()
convTest iterations dataDir nSamples rate = do
  net0         <- randomMnist
  trainData    <- readMNIST (dataDir </> "train-images-idx3-ubyte.gz")
                            (dataDir </> "train-labels-idx1-ubyte.gz")
  validateData <- readMNIST (dataDir </> "t10k-images-idx3-ubyte.gz")
                            (dataDir </> "t10k-labels-idx1-ubyte.gz")

  foldM_ (runIteration (maybe trainData (`take` trainData) nSamples) validateData) net0 [1..iterations]

    where
  trainEach rate' !network (i, o) = train rate' network i o

  runIteration trainRows validateRows net i = do
    let trained' = foldl' (trainEach ( rate { learningRate = learningRate rate * 0.9 ^ i} )) net trainRows
    print trained'

    putStrLn "Checking..."
    let res      = fmap (\(rowP,rowL) -> (rowL,) $ runNet trained' rowP) validateRows
    let res'     = fmap (\(S1D label, S1D prediction) -> (maxIndex (SA.extract label), maxIndex (SA.extract prediction))) res
    let matched   = length $ filter ((==) <$> fst <*> snd) res'
    let total     = length res'
    let matchedpc = fromIntegral matched / fromIntegral total * 100.0 :: Float
    putStrLn $ "Iteration " ++ show i ++ ": matched " ++ show matched ++ " of " ++ show total ++ " (" ++ show matchedpc ++ "%)"
    return trained'

data MnistOpts = MnistOpts FilePath (Maybe Int) Int LearningParameters

mnist' :: Parser MnistOpts
mnist' = MnistOpts <$> argument str (metavar "DATADIR")
                       -- option to reduce the number of training samples used from 60,000
                       -- to avoid running out of memory
                   <*> option (Just <$> auto) (long "limit_samples_to" <> short 'l' <> value Nothing)
                   <*> option auto (long "iterations" <> short 'i' <> value 15)
                   <*> (LearningParameters
                       <$> option auto (long "train_rate" <> short 'r' <> value 0.01)
                       <*> option auto (long "momentum" <> value 0.9)
                       <*> option auto (long "l2" <> value 0.0005)
                       )

main :: IO ()
main = do
    MnistOpts dataDir nSamples iter rate <- execParser (info (mnist' <**> helper) idm)
    putStr "Training convolutional neural network with "
    putStr $ maybe "all" show nSamples
    putStrLn " samples..."

    convTest iter dataDir nSamples rate

-- Adapted from https://github.com/tensorflow/haskell/blob/master/tensorflow-mnist/src/TensorFlow/Examples/MNIST/Parse.hs
-- Could also have used Data.IDX, although that uses a different Vector variant from that need for fromStorable
readMNIST :: FilePath -> FilePath -> IO [(S ( 'D2 28 28), S ( 'D1 10))]
readMNIST iFP lFP = do
  labels  <- readMNISTLabels lFP
  samples <- readMNISTSamples iFP
  return $ zip
    (fmap (fromMaybe (error "bad samples") . fromStorable) samples)
    (fromMaybe (error "bad labels") . oneHot . fromIntegral <$> labels)

-- | Check's the file's endianess, throwing an error if it's not as expected.
checkEndian :: Get ()
checkEndian = do
  magic <- Serialize.getWord32be
  when (magic `notElem` ([2049, 2051] :: [Word32]))
    $ error "Expected big endian, but image file is little endian."

-- | Reads an MNIST file and returns a list of samples.
readMNISTSamples :: FilePath -> IO [V.Vector Double]
readMNISTSamples path = do
  raw <- decompress <$> B.readFile path
  either fail ( return . fmap (V.map normalize) ) $ Serialize.runGetLazy getMNIST raw
 where
  getMNIST :: Get [V.Vector Word8]
  getMNIST = do
    checkEndian
    -- Parse header data.
    cnt    <- fromIntegral <$> Serialize.getWord32be
    rows   <- fromIntegral <$> Serialize.getWord32be
    cols   <- fromIntegral <$> Serialize.getWord32be
    -- Read all of the data, then split into samples.
    pixels <- Serialize.getLazyByteString $ fromIntegral $ cnt * rows * cols
    return $ V.fromList <$> chunksOf (rows * cols) (B.unpack pixels)

  normalize :: Word8 -> Double
  normalize = (/ 255) . fromIntegral
  -- There are other normalization functions in the literature, such as
  -- normalize = (/ 0.3081) . (`subtract` 0.1307) . (/ 255) . fromIntegral
  -- but we need values in the range [0..1] for the showShape' pretty printer

-- | Reads a list of MNIST labels from a file and returns them.
readMNISTLabels :: FilePath -> IO [Word8]
readMNISTLabels path = do
  raw <- decompress <$> B.readFile path
  either fail return $ Serialize.runGetLazy getLabels raw
 where
  getLabels :: Get [Word8]
  getLabels = do
    checkEndian
    -- Parse header data.
    cnt <- fromIntegral <$> Serialize.getWord32be
    -- Read all of the labels.
    B.unpack <$> Serialize.getLazyByteString cnt


{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
module FEECa.PolynomialTest where

import Control.Monad    ( liftM, liftM2 )

import FEECa.Internal.Vector
import FEECa.Internal.Simplex
import FEECa.Internal.Spaces
import FEECa.Polynomial

import FEECa.Utility.Utility
import FEECa.Utility.Combinatorics
import FEECa.Utility.Print

import Properties
import FEECa.Utility.Test
import qualified Test.QuickCheck as Q

-- =======
-- import qualified FEECa.Internal.MultiIndex as MI

--import qualified Numeric.LinearAlgebra.HMatrix as M
-- >>>>>>> efa221f... Cleaned up further tests (definitely shouldn't use `fromDouble`).

------------------------------------------------------------------------------
-- Dimension of the space to be tested. Must be greater than zero.
------------------------------------------------------------------------------

n :: Int
n = 2

------------------------------------------------------------------------------
-- Generate a random polynomial of dimension n as fixed by the above parameter.
-- The Polynomial may be of type Constant or Polynomial. If it is of type
-- Polynomial, degree the degree is chosen between 0 and 10 and the constants
-- are chosen randomly.
------------------------------------------------------------------------------

instance (Ring r, Q.Arbitrary r) => Q.Arbitrary (Polynomial r) where
  arbitrary = Q.frequency [(4,arbitraryPolynomial n), (1,arbitraryConstant)]

arbitraryPolynomial :: (Ring r, Q.Arbitrary r) => Int -> Q.Gen (Polynomial r)
arbitraryPolynomial n = do
  r <- Q.choose (0,10)
  s <- Q.choose (0,15)
  let mis = Q.vectorOf s (arbitraryMI n r)
      cs  = Q.vectorOf s Q.arbitrary
  liftM polynomial (liftM2 zip cs mis)

arbitraryConstant :: (Ring r, Q.Arbitrary r) => Q.Gen (Polynomial r)
arbitraryConstant = liftM constant Q.arbitrary

------------------------------------------------------------------------------
-- Generate random vectors of dimension n, so that they can be used to evaluate
-- the randomly generated polynomials.
------------------------------------------------------------------------------

instance (Field f, Q.Arbitrary f) => Q.Arbitrary (Vector f) where
  arbitrary = arbitraryVector n

------------------------------------------------------------------------------
-- Abstract propreties of arithmetic on polynomials. Addition, Subtraction and
-- multiplication of polynomials must commute with the evaluate operator.
------------------------------------------------------------------------------

propArithmetic :: (EuclideanSpace v, Function f v, VectorSpace f, Ring f,
                    r ~ Scalar v, r ~ Scalar f)
                => (r -> r -> Bool)
                -> f -> f -> v -> r
                -> Bool
propArithmetic eq x y v c =
     homomorphicEv addV add && homomorphicEv mul mul && homomorphicEv sub sub
    && prop_operator_commutativity eq (sclV c) (mul c) atV x
  where homomorphicEv o1 o2 = prop_homomorphism eq o1 o2 atV x y
        atV                 = evaluate v

------------------------------------------------------------------------------
-- Concrete arithmetic properties for polynomials defined over rationals.
------------------------------------------------------------------------------

prop_arithmetic_rf :: Polynomial Rational -> Polynomial Rational
                   -> Vector Rational -> Rational
                   -> Bool
prop_arithmetic_rf = propArithmetic (==)

--------------------------------------------------------------------------------
-- Derivation of Polynomials
--------------------------------------------------------------------------------

-- Linearity
propDerivation_linear :: (EuclideanSpace v, Function f v, VectorSpace f,
                          r ~ Scalar v, r ~ Scalar f)
                       => v -> v -> r -> f -> f
                       -> Bool
propDerivation_linear v1 v2 = prop_linearity (==) (evaluate v2 . derive v2)

-- Product rule
propDerivationProduct :: (EuclideanSpace v, Function f v, Ring f)
                        => v -> v -> f -> f
                        -> Bool
propDerivationProduct v1 v2 f g =
    ev (add (mul g (d f)) (mul f (d g))) == (ev . d) (mul f g)
  where ev = evaluate v1
        d  = derive v2

prop_derivation_product :: Vector Rational -> Vector Rational
                        -> Polynomial Rational -> Polynomial Rational
                        -> Bool
prop_derivation_product = propDerivationProduct

prop_derivation_productD :: Vector Double -> Vector Double
                         -> Polynomial Double -> Polynomial Double
                         -> Bool
prop_derivation_productD v1 v2 f g =
    ev (add (mul g (d f)) (mul f (d g))) `eqNum` (ev . d) (mul f g)
  where ev = evaluate v1
        d  = derive v2

-- Test for polynomials using exact arithmetic.
prop_arithmetic_rational :: Polynomial Rational -> Polynomial Rational
                         -> Vector Rational -> Rational
                         -> Bool
prop_arithmetic_rational = propArithmetic (==)

prop_derivation_linear_rational :: Vector Rational -> Vector Rational
                                -> Rational
                                -> Polynomial Rational -> Polynomial Rational
                                -> Bool
prop_derivation_linear_rational = propDerivation_linear

prop_derivation_product_rational :: Vector Rational -> Vector Rational
                                 -> Rational
                                 -> Polynomial Rational -> Polynomial Rational
                                 -> Bool
prop_derivation_product_rational = propDerivation_linear

--------------------------------------------------------------------------------
-- Barycentric Coordinates
--------------------------------------------------------------------------------

-- | Generate random simplex of dimesion 1 <= n <= 10.
instance (EuclideanSpace v, Q.Arbitrary (Scalar v)) => Q.Arbitrary (Simplex v) where
  arbitrary = do n <- Q.choose (1, 3)
                 k <- Q.choose (1, n)
                 arbitrarySubsimplex k n

-- TODO: Fails for n > 3 apparently due to numerical instability. To investigate
-- further.
prop_barycentric :: Simplex (Vector Double) -> Bool
prop_barycentric t =
    allEq [[evaluate v b | v <- vs] | b <- bs] oneLists
  where allEq l1 l2 = and $ zipWith (\l3 l4 -> (and (zipWith eqNum l3 l4))) l1 l2
        bs          = barycentricCoordinates t
        vs          = vertices t
        k           = topologicalDimension t
        oneLists    = map (map fromInt) (sumRLists (k+1) 1)

--------------------------------------------------------------------------------
-- Gradients of Barycentric Coordinates
--------------------------------------------------------------------------------

-- Helper function to create list of vector with 1.0 as first component
d0vectors :: Field a => Int -> [Vector a]
d0vectors n = [fromList $ [mulId]
                 ++ (replicate (i-1) addId)
                 ++ [addInv mulId]
                 ++ (replicate (n - i) addId) | i <- [1..n]]

-- Local gradients of barycentric coordinates, i.e. taken w.r.t the barycentric
-- coordinates themselves. Here we ensure that for a random simplex, if we move
-- from one corner a to corner b (in barycentric coordinates) the multiplication
-- of the  different vector with jacobian yields the corner b
-- (also in barycentric coordinates). This is equivalent to the vectors
-- [x_0, ..., x_n] with x_0 = 1 and x_i = -1 for any i > 0 being eigenvectors of
-- the gradient matrix.
prop_local_gradients :: Simplex (Vector Double) -> Bool
prop_local_gradients t = all (\v -> and $ zipWith eqNum (toList ((mult grads v)::Vector Double)) (toList v)) vs
  where grads     = localBarycentricGradients t
        vs        = d0vectors n :: [Vector Double]
        mult vs v = fromList [dot w v | w <- vs]
        n         = topologicalDimension t

return []
testPolynomial = $quickCheckWithAll


module Numeric.FFT.Execute ( execute ) where

import           Control.Monad               (when)
import qualified Control.Monad               as CM
import           Control.Monad.ST
import           Data.Complex
import qualified Data.IntMap.Strict          as IM
import           Data.STRef
import qualified Data.Vector                 as V
import           Data.Vector.Unboxed
import qualified Data.Vector.Unboxed.Mutable as MV
import           Prelude                     hiding (concatMap, foldr, length,
                                              map, mapM_, null, reverse, sum,
                                              zip, zipWith)

import           Numeric.FFT.Special
import           Numeric.FFT.Types
import           Numeric.FFT.Utils

-- | Main FFT plan execution driver.
execute :: Plan -> Direction -> VCD -> VCD
execute (Plan dlinfo perm base) dir h
  | n == 1 = h
  | V.null dlinfo = runST $ do
                          mhin <- case perm of
                            Nothing -> thaw h
                            Just p  -> unsafeThaw $ backpermute h p
                          mhout <- MV.replicate n 0
                          applyBase base sign mhin mhout
                          when (dir == Inverse) $ do
                            let s = 1.0 / fromIntegral n :+ 0
                            CM.forM_ [0..n-1] $ \i -> do
                              x <- MV.unsafeRead mhout i
                              MV.unsafeWrite mhout i $ s * x
                          unsafeFreeze mhout
  | otherwise=  fullfft
  where
    n = length h              -- Input vector length.
    bsize = baseSize base     -- Size of base transform.

    -- Root of unity sign.
    sign = case dir of
      Forward -> 1
      Inverse -> -1

    -- Apply Danielson-Lanczos steps and base transform to digit
    -- reversal ordered input vector.
    fullfft = runST $ do
      mhin <- case perm of
            Nothing -> thaw h
            Just p  -> unsafeThaw $ backpermute h p
      mhtmp <- MV.replicate n 0
      multBase mhin mhtmp
      mhr <- newSTRef (mhtmp, mhin)
      V.forM_ dlinfo $ \dlstep -> do
        (mh0, mh1) <- readSTRef mhr
        dl sign dlstep mh0 mh1
        writeSTRef mhr (mh1, mh0)
      mhs <- readSTRef mhr
      let vout = fst mhs
      when (dir == Inverse) $ do
        let s = 1.0 / fromIntegral n :+ 0
        CM.forM_ [0..n-1] $ \i -> do
          x <- MV.unsafeRead vout i
          MV.unsafeWrite vout i $ s * x
      unsafeFreeze vout

    -- Multiple base transform application for "bottom" of algorithm.
    multBase :: MVCD s -> MVCD s -> ST s ()
    multBase xmin xmout =
      V.zipWithM_ (applyBase base sign)
                  (slicemvecs bsize xmin) (slicemvecs bsize xmout)

-- | Monadic FFT plan execution driver -- used by Rader's algorithm
-- for convolutions.
executeM :: Plan -> Direction -> MVCD s -> MVCD s -> ST s ()
executeM (Plan dlinfo perm base) dir hin hout =
  if n == 1
  then MV.copy hout hin
  else do
    htmp <- MV.replicate n 0

    -- Input permutation.
    case perm of
      Nothing -> MV.copy htmp hin
      Just p  -> backpermuteM n p hin htmp

    -- Apply Danielson-Lanczos steps and base transform to digit
    -- reversal ordered input vector.
    multBase htmp hout
    mhr <- newSTRef (hout, htmp)
    V.forM_ dlinfo $ \dlstep -> do
      (mh0, mh1) <- readSTRef mhr
      dl sign dlstep mh0 mh1
      writeSTRef mhr (mh1, mh0)
    when (odd $ V.length dlinfo) $ MV.copy hout htmp

    -- Output scaling for inverse transform.
    when (dir == Inverse) $ do
      let s = 1.0 / fromIntegral n :+ 0
      forM_ (enumFromN 0 n) $ \i -> do
        x <- MV.unsafeRead hout i
        MV.unsafeWrite hout i $ s * x
  where
    n = MV.length hin         -- Input vector length.
    bsize = baseSize base     -- Size of base transform.

    -- Root of unity sign.
    sign = case dir of
      Forward -> 1
      Inverse -> -1

    -- Multiple base transform application for "bottom" of algorithm.
    multBase :: MVCD s -> MVCD s -> ST s ()
    multBase xmin xmout =
      V.zipWithM_ (applyBase base sign)
                  (slicemvecs bsize xmin) (slicemvecs bsize xmout)

-- | Single Danielson-Lanczos step: process all duplicates and
-- concatenate into a single vector.
dl :: Int -> (Int, Int, VVVCD, VVVCD) -> MVCD s -> MVCD s -> ST s ()
dl sign (wfac, split, dmatp, dmatm) mhin mhout =
  V.zipWithM_ doone (slicemvecs wfac mhin) (slicemvecs wfac mhout)
  where
    -- Twiddled diagonal entries in row r, column c (both
    -- zero-indexed), where each row and column if a wfac x wfac
    -- matrix.
    dmat = if sign == 1 then dmatp else dmatm
    d r c = (dmat V.! r) V.! c

    -- Size of each diagonal sub-matrix.
    ns = wfac `div` split

    -- Index vectors.
    nsidxs = enumFromN 0 ns
    splitidxs = enumFromN 1 (split-1)

    -- Process one duplicate by processing all rows and writing the
    -- results into a single output vector.
    doone :: MVCD s -> MVCD s -> ST s ()
    doone vin vout = do
      let vs = (slicemvecs ns vin, slicemvecs ns vout)
      mapM_ (single vs) $ enumFromN 0 split
      where
        -- Multiply a single block by its appropriate diagonal
        -- elements and accumulate the result.
        mult :: VMVCD s -> MVCD s -> Int -> Bool -> Int -> ST s ()
        mult vins vo r first c = do
          let vi = vins V.! c
              dvals = d r c
          forM_ nsidxs $ \i -> do
            xi <- MV.unsafeRead vi i
            xo <- if first then return 0 else MV.unsafeRead vo i
            MV.unsafeWrite vo i (xo + xi * dvals ! i)
        -- Multiply all blocks by the corresponding diagonal
        -- elements in a single row.
        single :: (VMVCD s, VMVCD s) -> Int -> ST s ()
        single (vis, vos) r = do
          mult vis (vos V.! r) r True 0
          mapM_ (mult vis (vos V.! r) r False) splitidxs
        -- single (vis, vos) r =
        --   let m = mult vis (vos V.! r) r
        --   in do
        --     m True 0
        --     mapM_ (m False) splitidxs

-- | Apply a base transform to a single vector.
applyBase :: BaseTransform -> Int -> MVCD s -> MVCD s -> ST s ()

-- Simple DFT algorithm.
applyBase (DFTBase sz wsfwd wsinv) sign mhin mhout = do
  h <- freeze mhin
  forM_ (enumFromN 0 sz) $ \i -> MV.unsafeWrite mhout i (doone h i)
  where ws = if sign == 1 then wsfwd else wsinv
        doone h i = sum $ zipWith (*) h $
                    generate sz (\k -> ws ! (i * k `mod` sz))

-- Special hard-coded cases.
applyBase (SpecialBase sz) sign mhin mhout =
  case IM.lookup sz specialBases of
    Just f  -> f sign mhin mhout
    Nothing -> error "invalid problem size for SpecialBase"

-- Rader prime-length FFT.
applyBase (RaderBase sz outperm bfwd binv csz cplan) sign mhin mhout = do
  -- Padding size.
  let pad = csz - (sz - 1)

  -- Permuted input vector padded to next greater power of two size
  -- for fast convolution.
  apad <- MV.replicate csz 0
  forM_ (enumFromN 0 csz) $ \i -> do
    val <- if i == 0 then MV.unsafeRead mhin 1
           else if i > pad
                then MV.unsafeRead mhin $ i - pad + 1
                else return 0
    MV.unsafeWrite apad i val

  -- FFT-based convolution calculation.
  convtmp <- MV.replicate csz 0
  executeM cplan Forward apad convtmp
  let bmult = if sign == 1 then bfwd else binv
  forM_ (enumFromN 0 csz) $ \i -> do
    x <- MV.unsafeRead convtmp i
    MV.unsafeWrite convtmp i $ x * (bmult ! i)
  executeM cplan Inverse convtmp apad
  conv <- unsafeFreeze apad

  -- Input vector sum.
  sumhref <- newSTRef 0
  forM_ (enumFromN 0 sz) $ \i -> do
    val <- MV.unsafeRead mhin i
    modifySTRef sumhref (+ val)
  sumh <- readSTRef sumhref

  -- Write output based on output generator index ordering.
  h0 <- MV.unsafeRead mhin 0
  forM_ (enumFromN 0 sz) $ \i -> do
    let (idx, val) = case i of
          0 -> (0, sumh)
          _ -> (outperm ! (i - 1), h0 + conv ! (i - 1))
    MV.unsafeWrite mhout idx val


{-# language GADTs #-}
{-# language TypeOperators #-}
{-# language FlexibleInstances #-}
{-# language UndecidableInstances #-}

module Control.Commutative
  ( Commutative
  , CommutativeMonad
  ) where

import Control.Applicative (ZipList, Const)
import Control.Applicative.Backwards
-- import Control.Arrow (Kleisli)
import Control.Comonad
import Control.Comonad.Env
import Control.Comonad.Traced
import Control.Monad.Trans.Identity
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.Reader
import Control.Monad.Trans.Writer.Lazy as Lazy
import Control.Monad.Trans.Writer.Strict as Strict
import Data.Complex
import Data.Functor.Compose
import Data.Functor.Identity
import Data.Functor.Product as Functor
import Data.Functor.Reverse
import Data.Monoid as Monoid
import Data.Proxy
import Data.Semigroup as Semigroup
import Data.Semigroup.Commutative
import Data.Tagged
import GHC.Generics
import Linear
import Linear.Plucker
import Linear.V

-- |
-- @
-- ('<*>') = 'flip' ('<**>')
-- @
class Applicative m => Commutative m

instance (Commutative f, Commutative g) => Commutative (Compose f g)

-- instance (CommutativeMonoid a, CommutativeMonoid b) => Commutative ((,,) a b)
-- instance (CommutativeMonoid a, CommutativeMonoid b, CommutativeMonoid c) => Commutative ((,,,) a b c)
-- instance (CommutativeMonoid a, CommutativeMonoid b, CommutativeMonoid c, CommutativeMonoid d) => Commutative ((,,,,) a b c d)

instance CommutativeMonoid w => Commutative ((,) w)
instance Commutative ((->) e)

-- linear
instance Commutative V4
instance Commutative V3
instance Commutative V2
instance Commutative Linear.V1
instance Commutative V0
instance Commutative Plucker
instance Dim n => Commutative (V n)
instance Commutative Complex
instance Commutative Quaternion

-- @transformers@
instance Commutative m => Commutative (ReaderT e m)
instance Commutative Identity
instance Commutative m => Commutative (IdentityT m)
instance (Commutative m, Commutative n) => Commutative (Functor.Product m n)
instance CommutativeMonad m => Commutative (MaybeT m)
instance Commutative m => Commutative (Reverse m)
instance Commutative m => Commutative (Backwards m)
instance (Commutative m, CommutativeMonoid w) => Commutative (Strict.WriterT w m)
instance (Commutative m, CommutativeMonoid w) => Commutative (Lazy.WriterT w m)

-- @tagged@
instance Commutative (Tagged a)
instance Commutative Proxy

-- Control.Applicative
instance Commutative ZipList
instance CommutativeMonoid a => Commutative (Const a)

-- Control.Arrow
-- instance Commutative m => Commutative (Kleisli m a)

-- Data.Semigroup
instance Commutative Semigroup.Last -- NB: Not CommutativeMonoid!
instance Commutative Semigroup.First -- NB: Not CommutativeMonoid!
instance Commutative Option -- NB: Not CommutativeMonoid!

-- Data.Monoid
instance Commutative Monoid.Last -- NB: Not CommutativeMonoid!
instance Commutative Monoid.First -- NB: Not CommutativeMonoid!
instance Commutative Monoid.Product -- NB: Not CommutativeMonoid!
instance Commutative Monoid.Sum -- NB: Not CommutativeMonoid!
instance Commutative f => Commutative (Alt f)

-- GHC.Generics
instance Commutative m => Commutative (M1 i c m)
instance (Commutative f, Commutative g) => Commutative (f :.: g)
instance (Commutative f, Commutative g) => Commutative (f :*: g)
-- instance CommutativeSemigroup c => Commutative (K1 i c) -- missing in base
instance Commutative f => Commutative (Rec1 f)
instance Commutative Par1
instance Commutative U1

-- @comonads@
instance (Commutative w, CommutativeMonoid e) => Commutative (EnvT e w)
instance Commutative w => Commutative (TracedT e w)
instance Commutative (Cokleisli w a)
-- instance Commutative f => Commutative (Ap f) -- ghc 8.6

-- instance CommutativeSemigroup w => Commutative (Validation w)

class    (Commutative m, Monad m) => CommutativeMonad m
instance (Commutative m, Monad m) => CommutativeMonad m


{-# LANGUAGE FlexibleContexts #-}
module PlotGreensFunctionSTS where

import           Control.Monad.Parallel    as MP
import           Data.Array.Repa           as R
import           Data.Binary               (decodeFile)
import           Data.Complex
import           Data.List                 as L
import           Data.Vector.Unboxed       as VU
import           FokkerPlanck.DomainChange
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           System.Directory
import           System.Environment
import           System.FilePath
import           Types
import           Utils.Array

main = do
  args@(numPointStr:spatialFreqStr:numOrientationStr:thetaSigmaStr:numScaleStr:scaleSigmaStr:maxScaleStr:taoStr:initStr:numTrailStr:maxTrailStr:thetaFreqsStr:scaleFreqsStr:histFileName:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      spatialFreq = read spatialFreqStr :: Double
      spatialFreqs = [-spatialFreq .. spatialFreq]
      numOrientation = read numOrientationStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      numScale = read numScaleStr :: Int
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      tao = read taoStr :: Double
      init@(initX, initY, initTheta, initScale) =
        read initStr :: (Double, Double, Double, Double)
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      thetaFreq = read thetaFreqsStr :: Double
      scaleFreq = read scaleFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scaleFreqs = [-scaleFreq .. scaleFreq]
      numThread = read numThreadStr :: Int
      folderPath = "output/test/PlotGreensFunctionSTS"
      histPath = folderPath </> histFileName
      initArr =
        traverse4
          (fromListUnboxed (Z :. (L.length spatialFreqs)) spatialFreqs)
          (fromListUnboxed (Z :. (L.length spatialFreqs)) spatialFreqs)
          (fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs)
          (fromListUnboxed (Z :. (L.length scaleFreqs)) scaleFreqs)
          (\(Z :. a) (Z :. b) (Z :. c) (Z :. d) ->
             (Z :. a :. b :. c :. d :. c :. d)) $ \fx fy ft fs (Z :. x :. y :. _ :. _ :. t :. s) ->
          exp
            (0 :+
             (-1) *
             (ft (Z :. t) * (initTheta / 180 * pi) +
              2 * pi *
              ((fx (Z :. x) * initX + fy (Z :. y) * initY) /
               (fromIntegral numPoint) +
               fs (Z :. s) * initScale / log maxScale)))
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histPath
  arr <-
    if flag
      then fmap
             (fromListUnboxed
                (Z :. (L.length spatialFreqs) :. (L.length spatialFreqs) :.
                 (L.length thetaFreqs) :.
                 (L.length scaleFreqs) :.
                 (L.length thetaFreqs) :.
                 (L.length scaleFreqs))) .
           decodeFile $
           histPath
      else runMonteCarloSTS
             numThread
             numTrail
             maxTrail
             numPoint
             numPoint
             thetaSigma
             scaleSigma
             maxScale
             tao
             spatialFreqs
             spatialFreqs
             thetaFreqs
             scaleFreqs
             histPath $
           VU.replicate
             (L.length scaleFreqs * L.length thetaFreqs * L.length scaleFreqs *
              L.length thetaFreqs *
              L.length spatialFreqs *
              L.length spatialFreqs)
             0
  transformedArr <- sumP . sumS $ R.zipWith (*) arr initArr
  -- arrR2S1RP <-
  --   stsTor2s1rp
  --     numPoint
  --     (fromIntegral numPoint)
  --     spatialFreqs
  --     numPoint
  --     (fromIntegral numPoint)
  --     spatialFreqs
  --     numOrientation
  --     thetaFreqs
  --     numScale
  --     scaleFreqs
  --     maxScale
  --     transformedArr
  -- plotImageRepaComplex (folderPath </> "Greens.png") .
  --   ImageRepa 8 .
  --   computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.sumS . sumS $
  --   arrR2S1RP
  arrMag <-
    stsTor2'
      numPoint
      (fromIntegral numPoint)
      spatialFreqs
      numPoint
      (fromIntegral numPoint)
      spatialFreqs
      transformedArr
  plotImageRepa (folderPath </> "GreensMag.png") .
    ImageRepa 8 .
    computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.map sqrt $
    arrMag