module Cas.Internal.Utils
    ( module Cas.Internal.Utils
    , module Cas.Internal.Def.Interface
    , module Cas.Internal.Instances.TFunctor
    , module Cas.Internal.Instances.TFoldable ) where

import           PreludeCustom

import           Cas.Internal.Def.Interface
import           Cas.Internal.Instances.TFunctor
import           Cas.Internal.Instances.TFoldable
import           Data.Ratio
import           Data.Complex.Generic
import           Cas.Misc

isNegative :: T -> Bool
isNegative = isF   ?>>> (<0) . tF
         ||> isMul ?>>- uniqueOnDepth 1 ((&&) <$> isF <*> (<0) . tF)
         ||>            False

selectCs, selectFs, selectXs, selectNonFs  :: T   -> [T]
lengthCs, lengthFs, lengthXs, lengthNonFs, lengthOpers :: T   -> Int
selectCs = select isC
selectFs = select isF
selectXs = select isX
selectNonFs = select (not . isF)
selectNonFLeafs = select isNonFLeaf
lengthCs = tLength isC
lengthFs = tLength isF
lengthXs = tLength isX
lengthNonFs = tLength (not . isF)
lengthNonFLeafs = tLength isNonFLeaf
lengthLeafs = tLength isLeaf
lengthChilds = tLengthOnDepth 1 (const True)
lengthOpers = tLength isOper
lengthNonFlats = tLength (not . isFlat)
lengthAddChilds = length . concatMap tAddTs . select isAdd
lengthMulChilds = length . concatMap tMulTs . select isMul
lengthNotPowAddChilds = length
                         . concatMap tAddTs
                         . filter isAdd
                         . ((:) <$> id
                                <*> concatMap tToListChilds
                                  . select (not . isPow) )
lengthLns = tLength isLn
lengthNonOpers = tLength (not . isOper)

subTsAdd :: Add -> [[T]]
subTsAdd Add {..} = fmap (\case (TMul (Mul {..})) -> mulTs; t -> [t]) addTs

subTsMul :: Mul -> [[T]]
subTsMul Mul {..} = fmap (\case (TAdd (Add {..})) -> addTs; t -> [t]) mulTs

subTs :: T   -> [[T]]
subTs TAdd {..}  = subTsAdd tAdd
subTs TMul {..}  = subTsMul tMul
subTs _          = error "subTs called on non Add or Mul term"

addElems, mulElems :: T -> [T]
addElems = isAdd ?>>> tAddTs ||> return
mulElems = isMul ?>>> tMulTs ||> return

addFBin, mulFBin :: F -> F -> F
addFBin w0 w1 = w0+w1

mulFBin w0 w1 = w0*w1

addF, mulF :: [F] -> F
addF = foldr addFBin zeroF
mulF = foldr mulFBin oneF

filterF = fmap tF . filter isF
filterNotF = filter (not . isF)

addCollapse, mulCollapse :: [T] -> T
addCollapse = add
         . filter (not . isZero)
         . ((:) <$> TF .  addF . filterF
                <*> filterNotF)

mulCollapse = any isZero ?>>> const zero
                          ||> mul
                            . filter (not . isOne)
                            . ((:) <$> TF . mulF . filterF
                                   <*> filterNotF)

addBinCollapse, mulBinCollapse, powCollapse :: T -> T -> T
addBinCollapse x y = addCollapse $ addElems x <> addElems y
mulBinCollapse x y = mulCollapse $ mulElems x <> mulElems y

powCollapse _    (TF 0)  = one
powCollapse (TF 0) _     = zero
powCollapse (TF 1) _     = one
powCollapse x    (TF 1)  = x
powCollapse (TF z) (TF w) | imagPart w == 0 && denominator (realPart w) == 1 = f $ z^^(numerator $ realPart w)
powCollapse x    y           = pow x y

mulParts :: T -> (F,T)
mulParts = isF    ?>>>  (,) <$> tF
                            <*> const one
       ||> isMul  ?>>> ((,) <$> mulF . filterF
                            <*> mul . filterNotF ) . tMulTs
       ||>              (,) <$> const oneF
                            <*> id

powParts :: T -> (T,T)
powParts = ((&&) <$> isPow
                 <*> isF . tPowE) ?>>> (,) <$> tPowE <*> tPowB
                                   ||> (,) <$> const one <*> id
-- powParts =  isPow ?>>> (,) <$> tPowE <*> tPowB
--                                    ||> (,) <$> const one <*> id


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RecordWildCards  #-}

module AI.Layer
( LayerDefinition(..)
, Layer(..)
, Connectivity
, RandomTransform
, Randomization

, showableToLayer

, createLayer
, scaleLayer
, randomizeFully
, randomizeLocally
, connectFully
, connectLocally

, randomList
, boxMuller
, normals
, uniforms
, boundedUniforms
) where

import           Data.Binary           (Binary (..), decode, encode)
import           AI.Neuron
import           Numeric.LinearAlgebra
import           System.Random

-- | The LayerDefinition type is an intermediate type initialized by the
--   library user to define the different layers of the network.
data LayerDefinition g = LayerDefinition { neuronDef   :: Neuron
                                         , neuronCount :: Int
                                         , connect     :: Connectivity
                                         , randomize   :: Randomization g
                                         }

-- | The Layer type, which stores the weight matrix, the bias matrix, and
--   a neuron type.
data Layer = Layer { weightMatrix :: Matrix Double
                   , biasVector   :: Vector Double
                   , neuron       :: Neuron
                   } deriving Show

instance Binary Layer where
  put Layer{..} = do put weightMatrix; put biasVector
  get = do weightMatrix <- get; biasVector <- get; return Layer{..}

-- | Connectivity is the type alias for a function that defines the connective
--   matrix for two layers (fully connected, convolutionally connected, etc.)
--   and takes in the number of output and input neurons
type Connectivity = Int -> Int -> Matrix Double

-- | Randomiation is the type alias for a function that defines
--   the initial random values for the weight matrix and bias vector
--   for two layers and takes in a random transformation on an infinite
--   stream of uniformly generated numbers, a source of entropy,
--   the number of output neurons, and the number of input neurons
type Randomization g = g -> RandomTransform -> Int -> Int -> (Matrix Double, Vector Double)

-- | ConvolutionalSettings is a type alias for the receptive field size,
--   stride, zero-padding, the number of filters, the number of dimensions,
--   the weight and height of the input and output fields
type ConvolutionalSettings = Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int

-- | A random transformation type alias. It is a transformation defined on an
--   infinite list of uniformly distributed random numbers, and returns a list
--   distributed on the transforming distribution.
type RandomTransform = [Double] -> [Double]

-- | The createLayer function takes in a random transformation on an infinite
--   stream of uniformly generated numbers, a source of entropy, and two
--   layer definitions, one for the previous layer and one for the next layer.
--   It returns a layer defined by the Layer type -- a weight matrix, a bias
--   vector, and a neuron type.
createLayer :: (RandomGen g)
  => RandomTransform -> g -> LayerDefinition g -> LayerDefinition g -> Layer
createLayer t g layerDef layerDef' =
  Layer (randomMatrix * connectivity i j)
        (randomVector * bias)
        (neuronDef layerDef)
  where (randomMatrix, randomVector) = randomize layerDef' g t i j
        i = neuronCount layerDef'
        j = neuronCount layerDef
        connectivity = connect layerDef'
        bias = i |> repeat 1 -- bias connectivity (full)

scaleLayer :: Double -> Layer -> Layer
scaleLayer factor l =
  Layer (factor `scale` weightMatrix l) (factor `scale` biasVector l) (neuron l)

-- | The randomizeFully function takes in a source of entropy, the number of output
--   neurons, and the number of input neurons, and returns a tuple of the
--   a fully random matrix, and a fully random vector
randomizeFully :: (RandomGen g) => Randomization g
randomizeFully g t i j = (randomMatrix, randomVector)
  where randomMatrix = (i >< j) (randomList t g')
        randomVector = i |> randomList t g''
        (g', g'') = split g

-- | The randomizeLocally function takes in ConvolutionalSettings a source
--   of entropy, a random transform, the number of output and input neurons.
--   It returns a tuple with a random matrix and a random vector
randomizeLocally :: (RandomGen g) => Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Randomization g
randomizeLocally f s p k d w1 h1 w2 h2 g t i j = (randomMatrix, randomVector)
  where randomMatrix = fromLists (locallyRandomList f s p k d w1 h1 w2 h2 g' t i j) :: Matrix Double
        randomVector = i |> randomList t g''
        (g', g'') = split g

-- | The locallyRandomList function takes in ConvolutionalSettings a source
--   of entropy, a random transform, the number of output and input neurons.
--   It returns a list of lists of locally random numbers
locallyRandomList :: (RandomGen g) => Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> g -> RandomTransform -> Int -> Int -> [[Double]]
locallyRandomList f s p k d w1 h1 w2 h2 g t i j =
  if k == 0 then []
  else filterValues ++ nextFilterValues
    where filterValues = [replicate (rowZeroOffset + colZeroOffset) 0
                          ++ take (j - rowZeroOffset - colZeroOffset) (randomList t g')
                          | n <- [0..div i k-1],
                          let rowSize = w1 + 2 * p,
                          let postsynPerFilter = rem n (div i k),
                          let rowZeroOffset = (1 + s) * quot postsynPerFilter w2 * rowSize,
                          let colZeroOffset = (1 + s) * mod n w2]
          nextFilterValues = locallyRandomList f s p (k - 1) k w1 h1 w2 h2 g'' t (i - div i k) j
          (g', g'') = split g

-- | The connectFully function takes the number of input neurons for a layer, i,
--   and the number of output neurons of a layer, j, and returns an i x j
--   connectivity matrix for a fully connected network.
connectFully :: Connectivity
connectFully i j = (i >< j) (repeat 1)

-- | The connectLocally function takes in ConvolutionalSettings and the number
--   of output and input neurons
connectLocally :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Connectivity
connectLocally f s p k d w1 h1 w2 h2 i j =
  repmat (fromLists conn :: Matrix Double) k d
  where conn = [replicate rowZeroOffset 0
                ++ take (f * rowSize) (cycle fieldArea)
                ++ replicate (rowSize * colSize - rowSize * f - rowZeroOffset) 0
                | n <- [0.. quot i k-1],
                  let rowSize = w1 + 2 * p,
                  let colSize = h1 + 2 * p,
                  let rowZeroOffset = (1 + s) * quot n w2 * rowSize,
                  let fieldAreaZeroOffset = (1 + s) * mod n w2,
                  let fieldArea = replicate fieldAreaZeroOffset 0
                                   ++ replicate f 1
                                   ++ replicate (rowSize - f - fieldAreaZeroOffset) 0]

-- | To go from a showable to a layer, we also need a neuron type,
--   which is an unfortunate restriction owed to Haskell's inability to
--   serialize functions.
showableToLayer :: (RandomGen g) => (Layer, LayerDefinition g) -> Layer
showableToLayer (s, d) = Layer (weightMatrix s) (biasVector s) (neuronDef d)

-- | Initialize an infinite random list given a random transform and a source
--   of entroy.
randomList :: RandomGen g => RandomTransform -> g -> [Double]
randomList transform = transform . randoms

-- | Define a transformation on the uniform distribution to generate
--   normally distributed numbers in Haskell (the Box-Muller transform)
boxMuller :: Double -> Double -> (Double, Double)
boxMuller x1 x2 = (z1, z2)
  where z1 = sqrt ((-2) * log x1) * cos (2 * pi * x2)
        z2 = sqrt ((-2) * log x1) * sin (2 * pi * x2)

-- | This is a function of type RandomTransform that transforms a list of
--   uniformly distributed numbers to a list of normally distributed numbers.
normals :: RandomTransform
normals (x1:x2:xs) = z1:z2:normals xs
  where (z1, z2) = boxMuller x1 x2
normals _ = []

-- | A non-transformation to return a list of uniformly distributed numbers
--   from a list of uniformly distributed numbers. It's really a matter of
--   naming consistency. It generates numbers on the range (0, 1]
uniforms :: RandomTransform
uniforms xs = xs

-- | An affine transformation to return a list of uniforms on the range
--   (a, b]
boundedUniforms :: (Double, Double) -> [Double] -> [Double]
boundedUniforms (lower, upper) = map affine
  where affine x = lower + x * (upper - lower)


{-# LANGUAGE DataKinds, TypeFamilies #-}
{-# OPTIONS_GHC -Wno-missing-export-lists #-}
-- | Vector-based (meaning that dual numbers for gradient computation
-- consider vectors, not scalars, as the primitive differentiable type)
-- implementation of fully connected neutral network for classification
-- of MNIST digits. Sports 2 hidden layers.
module HordeAd.Tool.MnistFcnnVector where

import Prelude

import           Control.Exception (assert)
import qualified Data.Array.DynamicS as OT
import qualified Data.Vector.Generic as V
import           GHC.Exts (inline)
import           Numeric.LinearAlgebra (Vector)

import HordeAd.Core.DualNumber
import HordeAd.Core.Engine
import HordeAd.Core.PairOfVectors (DualNumberVariables, var1)
import HordeAd.Tool.MnistData

sumTrainableInputsV
  :: IsScalar d r
  => DualNumber d (Vector r) -> Int -> DualNumberVariables d r -> DualNumber d r
sumTrainableInputsV x offset variables =
  let v = var1 variables offset
  in v <.>! x

sumTrainableInputsL
  :: forall d r. IsScalar d r
  => DualNumber d (Vector r) -> Int -> DualNumberVariables d r -> Int
  -> DualNumber d (Vector r)
sumTrainableInputsL x offset variables width =
  let f :: Int -> DualNumber d r
      f i = sumTrainableInputsV x (offset + i) variables
  in seq1 $ V.generate width f

sumConstantDataV
  :: IsScalar d r
  => Vector r -> Int -> DualNumberVariables d r -> DualNumber d r
sumConstantDataV x offset variables =
  let v = var1 variables offset
  in v <.>!! x

sumConstantDataL
  :: forall d r. IsScalar d r
  => Vector r -> Int -> DualNumberVariables d r -> Int
  -> DualNumber d (Vector r)
sumConstantDataL x offset variables width =
  let f :: Int -> DualNumber d r
      f i = sumConstantDataV x (offset + i) variables
  in seq1 $ V.generate width f

fcnnMnistLen1 :: Int -> Int -> (Int, [Int], [(Int, Int)], [OT.ShapeL])
fcnnMnistLen1 widthHidden widthHidden2 =
  ( 0
  , replicate widthHidden sizeMnistGlyph ++ [widthHidden]
    ++ replicate widthHidden2 widthHidden ++ [widthHidden2]
    ++ replicate sizeMnistLabel widthHidden2 ++ [sizeMnistLabel]
  , []
  , []
  )

-- | Fully connected neural network for the MNIST digit classification task.
-- There are two hidden layers and both use the same activation function.
-- The output layer uses a different activation function.
-- The widths of the hidden layers are @widthHidden@ and @widthHidden2@
-- and from these, the @len*@ functions compute the number and dimensions
-- of scalars (none in this case) and vectors of dual number parameters
-- (variables) to be given to the program.
fcnnMnist1 :: forall d r m. DualMonad d r m
           => (DualNumber d (Vector r) -> m (DualNumber d (Vector r)))
           -> (DualNumber d (Vector r) -> m (DualNumber d (Vector r)))
           -> Int
           -> Int
           -> Vector r
           -> DualNumberVariables d r
           -> m (DualNumber d (Vector r))
fcnnMnist1 factivationHidden factivationOutput widthHidden widthHidden2
          input variables = do
  let !_A = assert (sizeMnistGlyph == V.length input) ()
  let hiddenLayer1 = sumConstantDataL input 0 variables widthHidden
                     + var1 variables widthHidden  -- bias
  nonlinearLayer1 <- factivationHidden hiddenLayer1
  let offsetMiddle = widthHidden + 1
      hiddenLayer2 = sumTrainableInputsL nonlinearLayer1 offsetMiddle
                                         variables widthHidden2
                     + var1 variables (offsetMiddle + widthHidden2)  -- bias
  nonlinearLayer2 <- factivationHidden hiddenLayer2
  let offsetOutput = offsetMiddle + widthHidden2 + 1
      outputLayer = sumTrainableInputsL nonlinearLayer2 offsetOutput
                                        variables sizeMnistLabel
                    + var1 variables (offsetOutput + sizeMnistLabel)  -- bias
  factivationOutput outputLayer

-- | The neural network applied to concrete activation functions
-- and composed with the appropriate loss function.
fcnnMnistLoss1
  :: DualMonad d r m
  => Int -> Int -> MnistData r -> DualNumberVariables d r
  -> m (DualNumber d r)
fcnnMnistLoss1 widthHidden widthHidden2 (input, target) variables = do
  result <- inline fcnnMnist1 logisticAct softMaxActV
                              widthHidden widthHidden2 input variables
  lossCrossEntropyV target result

-- | A function testing the neural network given testing set of inputs
-- and the trained parameters.
fcnnMnistTest1
  :: forall r. IsScalar 'DModeGradient r
  => Int -> Int -> [MnistData r] -> (Domain0 r, Domain1 r) -> r
fcnnMnistTest1 widthHidden widthHidden2 inputs (params0, params1) =
  let matchesLabels :: MnistData r -> Bool
      matchesLabels (glyph, label) =
        let nn = inline fcnnMnist1 logisticAct softMaxActV
                                        widthHidden widthHidden2 glyph
            value = primalValue nn (params0, params1, V.empty, V.empty)
        in V.maxIndex value == V.maxIndex label
  in fromIntegral (length (filter matchesLabels inputs))
     / fromIntegral (length inputs)


{-# LANGUAGE BangPatterns #-}
module LogFourierSeries where

import           Data.Array.Repa     as R
import           Data.Complex
import           Data.List           as L
import           Data.Vector.Unboxed as VU
import           Image.IO
import           System.Environment
import           System.FilePath
import           Utils.Parallel hiding ((.|))
import Control.DeepSeq
import Data.Conduit
import Data.Conduit.List as CL
import Control.Monad.Trans.Resource
import Control.Monad
import           FokkerPlanck.MonteCarlo
import Utils.Array
import            STC.Utils

{-# INLINE gaussian #-}
gaussian :: Double -> Double -> Double -> Double
gaussian !mu !sigma !x =
  (exp $ (x - mu) ^ 2 / (-2 * sigma ^ 2)) / (sqrt $ 2 * pi) / sigma

sink ::
     ParallelParams -> (VU.Vector (Complex Double))
  -> ConduitT (VU.Vector (Complex Double)) Void (ResourceT IO) (VU.Vector Double)
sink !parallelParams !vec = do
  xs <- CL.take . batchSize $ parallelParams
  if (L.null xs)
    then return . VU.map magnitude $ vec
    else sink parallelParams (L.foldl' (VU.zipWith (+)) vec xs)
  -- case x of
  --   Nothing -> return . VU.map magnitude $ vec
  --   Just y -> sink (VU.zipWith (+) vec y)

main = do
  args@(numPointStr:maxThetaFreqStr:maxRFreqStr:deltaStr:maxRStr:stdStr:scaleFactorStr:batchSizeStr:numThreadStr:_) <-
    getArgs
  print args
  let maxThetaFreq = read maxThetaFreqStr :: Double
      maxRFreq = read maxRFreqStr :: Double
      delta = read deltaStr :: Double
      thetaFreqs = [-maxThetaFreq .. maxThetaFreq]
      rFreqs = [-maxRFreq .. maxRFreq]
      phiFreqs = [0] -- thetaFreqs
      muR = 16
      sigmaR = 8
      muTheta = 0
      sigmaTheta = 0.5
      muPhi = pi
      sigmaPhi = 0.5
      deltaPhi = 2 * pi / fromIntegral numOrientation
      alpha = 1
      scaleFactor = read scaleFactorStr :: Double
      numPoint = read numPointStr :: Int
      numOrientation = 72 :: Int
      maxR = read maxRStr :: Double
      std = read stdStr :: Double
      batchSize = read batchSizeStr :: Int
      numThread = read numThreadStr :: Int
      center = div numPoint 2
      folderPath = "output/test/LogFourierSeries"
      gaussianFunc theta r phi =
        (gaussian muTheta sigmaTheta theta) * (gaussian muR sigmaR r) *
        (gaussian muPhi sigmaPhi phi)
      centerArr =
        fromFunction (Z :. (1 :: Int) :. numPoint :. numPoint) $ \(Z :. _ :. i :. j) ->
          if i == center && j == center
            then 1
            else 0
  plotImageRepa (folderPath </> "Center.png") . ImageRepa 8 . computeS $
    centerArr
  -- originalArray <-
  --   computeUnboxedP . fromFunction (Z :. numPoint :. numPoint :. numOrientation) $ \(Z :. i :. j :. k) ->
  --     let !x = (fromIntegral $ i)
  --         !y = (fromIntegral $ j - center)
  --         !theta = atan2 y x
  --         !r = sqrt $ x ^ 2 + y ^ 2
  --         !phi = fromIntegral k * deltaPhi
  --     in gaussianFunc theta r phi
  arrG <-
    solveMonteCarloR2S1
      numThread
      1000000
      1000000
      numPoint
      numPoint
      numOrientation
      0.1
      100
      64
      1
      ""
  print . extent $ arrG
  let originalArray =
        extend (Z :. All :. All :. (1 :: Int)) . sumS $ rotate3D arrG
  plotImageRepa (folderPath </> "Greens3DNorm_1.png") .
    ImageRepa 8 .
    computeS .
    reduceContrast 10 . extend (Z :. (1 :: Int) :. All :. All) . R.sumS $
    originalArray
  plotImageRepa (folderPath </> "Greens3D_1.png") .
    ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) . R.sumS $
    originalArray
  recon <-
    runConduitRes $
    CL.sourceList
      [ (rFreq, thetaFreq, phiFreq)
      | rFreq <- rFreqs
      , thetaFreq <- thetaFreqs
      , phiFreq <- phiFreqs
      ] .|
    parConduit
      (ParallelParams numThread batchSize)
      (\(rFreq, thetaFreq, phiFreq) ->
         let !c =
               sumAllS .
               R.zipWith (\a b -> (a :+ 0) * b) originalArray .
               fromFunction (Z :. numPoint :. numPoint :. (1 :: Int)) $ \(Z :. i :. j :. k) ->
                 let !x = delta * (fromIntegral $ i - center)
                     !y = delta * (fromIntegral $ j - center)
                     !theta = atan2 y x
                     !r = (sqrt $ x ^ 2 + y ^ 2)
                     !phi = fromIntegral k * deltaPhi
                 in if (x ^ 2 + y ^ 2) == 0  -- || r >= maxR -- || pi * r <  (abs thetaFreq)
                      then 0
                           -- exp $ (log r) * (alpha - 1) :+ ((-thetaFreq) * theta - rFreq * (log r))
                      else exp $ (log r) * (-1 + alpha) :+ ((-thetaFreq) * theta - rFreq * (log r) -- * pi / (log maxR)
                                                              )
                           -- cis ((-thetaFreq) * theta - rFreq * (log r) * pi / (log maxR) )
                           -- (x :+ y) ** ((-thetaFreq) :+ 0) *
                           -- ((x ^ 2 + y ^ 2) :+ 0) **
                           -- (((thetaFreq - 2 + alpha) :+ (-rFreq)) / 2)
                           -- (x :+ y) ** ((-thetaFreq) :+ 0) *
                           -- ((x ^ 2 + y ^ 2) :+ 0) **
                           -- (((thetaFreq) :+ (-pi * rFreq / (log maxR))) / 2)
         in toUnboxed .
            computeS .
            -- R.map (* c) .
            R.map
              (* (c *
                  ((exp $ (-thetaFreq ^ 2 - rFreq ^ 2) / (2 * std ^ 2)) :+ 0))) .
            fromFunction (Z :. numPoint :. numPoint :. (1 :: Int)) $ \(Z :. i :. j :. k) ->
              let !x = (fromIntegral $ i - center)
                  !y = (fromIntegral $ j - center)
                  !theta = atan2 y x
                  !r = delta * (sqrt $ x ^ 2 + y ^ 2)
                  !phi = fromIntegral k * deltaPhi
              in if r <= 0 -- || pi * r < (abs thetaFreq)
                   then 0

                   else exp $ (log r) * (-alpha) :+ (thetaFreq * theta + rFreq * (log r) -- * pi / (log maxR)
                                                  )
                        -- cis (thetaFreq * theta + rFreq * (log r) * pi / (log maxR))
                        -- exp $ (log r) * (-alpha) :+ (thetaFreq * theta + rFreq * (log r))
                        -- (x :+ y) ** (thetaFreq :+ 0) *
                        -- ((x ^ 2 + y ^ 2) :+ 0) **
                        -- (((-thetaFreq - alpha) :+ rFreq) / 2)
                        -- (x :+ y) ** (thetaFreq :+ 0) *
                        -- ((x ^ 2 + y ^ 2) :+ 0) **
                        -- (((-thetaFreq) :+ (pi * rFreq / (log maxR))) / 2)
       ) .|
    sink
      (ParallelParams numThread batchSize)
      (VU.replicate (numPoint ^ 2 * 1) 0)
  -- let
      -- !recon =
      --   VU.map magnitude .
      --   L.foldl1' (VU.zipWith (+)) .
      --   parMap
      --     rdeepseq
      --     (\(rFreq, thetaFreq, phiFreq) ->
      --        let !c =
      --              sumAllS .
      --              fromFunction (Z :. numPoint :. numPoint :. numPoint) $ \(Z :. i :. j :. k) ->
      --                let !x = delta * (fromIntegral $ i - center)
      --                    !y = delta * (fromIntegral $ j - center)
      --                    !theta = atan2 y x
      --                    !r = sqrt $ x ^ 2 + y ^ 2
      --                    !phi = fromIntegral k * deltaPhi
      --                in if r < 1 / maxR || r >= maxR
      --                     then 0
      --                     else (gaussianFunc theta r phi :+ 0) *
      --                          (cis $
      --                           (-1) * thetaFreq * theta -
      --                           phiFreq * (phi - theta) -
      --                           pi * rFreq * (r) / (maxR))
      --        in toUnboxed .
      --           computeS .
      --           R.map (* c) .
      --           fromFunction (Z :. numPoint :. numPoint :. numPoint) $ \(Z :. i :. j :. k) ->
      --             let !x = fromIntegral $ i - center
      --                 !y = fromIntegral $ j - center
      --                 !theta = atan2 y x
      --                 !r = sqrt $ x ^ 2 + y ^ 2
      --                 !phi = fromIntegral k * deltaPhi
      --             in if r == 0
      --                  then 0
      --                  else cis $
      --                       thetaFreq * theta + phiFreq * (phi - theta) +
      --                       pi * rFreq * (r) / (maxR)) $
      --   [ (rFreq, thetaFreq, phiFreq)
      --   | rFreq <- rFreqs
      --   , thetaFreq <- thetaFreqs
      --   , phiFreq <- phiFreqs
      --   ]
  plotImageRepa (folderPath </> "Greens3DRecon_2.png") .
    ImageRepa 8 .
    computeS .
    extend (Z :. (1 :: Int) :. All :. All) .
    sumS . fromUnboxed (Z :. numPoint :. numPoint :. (1 :: Int)) $
    recon
  plotImageRepa (folderPath </> "Greens3DReconNorm_2.png") .
    ImageRepa 8 .
    computeS .
    reduceContrast 10 .
    extend (Z :. (1 :: Int) :. All :. All) .
    sumS . fromUnboxed (Z :. numPoint :. numPoint :. (1 :: Int)) $
    recon


{-# LANGUAGE DeriveFunctor, GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DataKinds, TypeFamilies #-}
{-# LANGUAGE FlexibleContexts, FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ConstraintKinds #-}

module DelayedSampling where

import Prelude hiding ((<>))

import Control.Exception.Base (assert)
import Control.Monad (forM_, when)
import Control.Monad.State (get, put)
import Control.Monad.Trans (MonadIO, liftIO)
import Control.Monad.Bayes.Class (MonadSample, MonadCond, MonadInfer)
import Control.Monad.Bayes.Weighted
import Control.Monad.Bayes.Sampler

import Data.Aeson (ToJSON (..), object, (.=))
import Data.List (delete)
import Data.Maybe (isJust)
import Data.Proxy
import Data.Random hiding (normal)

import Debug.Trace

import GHC.TypeLits

import Numeric.AD
import Numeric.AD.Rank1.Tower (Tower)

import Numeric.LinearAlgebra.Static hiding (M)

import Unsafe.Coerce (unsafeCoerce)

import Distributions
import MVDistributions
import Inference
import Util.Ref

data MarginalT = MGaussianT | MBetaT | MBernoulliT | forall (n :: Nat). MMVGaussianT (Proxy n)

type DelayedSampling = MonadState Heap
type DelayedSample m = (MonadState Heap m, MonadSample m)
type DelayedInfer m = (MonadState Heap m, MonadInfer m)

data SMarginalT (m :: MarginalT) where
  SMGaussianT :: SMarginalT MGaussianT
  SMBetaT :: SMarginalT MBetaT
  SMBernoulliT :: SMarginalT MBernoulliT
  SMMVGaussianT :: forall (n :: Nat) (p :: Proxy n). SMarginalT (MMVGaussianT p)

type family MType (m :: MarginalT) where
  MType MGaussianT = Double
  MType MBetaT = Double
  MType MBernoulliT = Bool
  MType (MMVGaussianT (p :: Proxy n)) = R n

data MDistr (m :: MarginalT) where
  MGaussian :: !Double -> !Double -> MDistr MGaussianT
  MBeta :: !Double -> !Double -> MDistr MBetaT
  MBernoulli :: !Double -> MDistr MBernoulliT
  MMVGaussian :: forall (n :: Nat) (p :: Proxy n). KnownNat n => !(R n) -> !(Sym n) -> MDistr (MMVGaussianT p)

deriving instance Show (MDistr m)

instance ToJSON (MDistr a) where
  toJSON (MMVGaussian mu cov) = object ["mean" .= mu, "cov" .= cov]
  toJSON (MGaussian mu cov) = object ["mean" .= mu, "cov" .= cov]
  toJSON (MBeta a b) = object ["alpha" .= a, "beta" .= b]
  toJSON (MBernoulli p) = object ["p" .= p]

data CDistr (m :: MarginalT) (m' :: MarginalT) where
  AffineMeanGaussian :: !Double -> !Double -> !Double -> CDistr MGaussianT MGaussianT
  CBernoulli :: CDistr MBetaT MBernoulliT
  MVAffineMeanGaussian :: forall (n :: Nat) (m :: Nat) (pn :: Proxy n) (pm :: Proxy m).
    (KnownNat n, KnownNat m) => (L m n) -> (R m) -> (Sym m) -> CDistr (MMVGaussianT pn) (MMVGaussianT pm)

data Result a where
  RConst :: a -> Result a
  RMarginal :: MDistr a -> Result (MType a)

deriving instance Show a => Show (Result a)

gaussianConditioning :: Double -> Double -> Double -> Double -> (Double, Double)
gaussianConditioning mu var obs obsvar = (mu', var')
  where
  ivar = recip var
  iobsvar = recip obsvar
  inf = ivar + iobsvar
  var' = recip inf
  mu' = (ivar * mu + iobsvar * obs) / inf

gaussianMeanGaussian :: Double -> CDistr MGaussianT MGaussianT
gaussianMeanGaussian = AffineMeanGaussian 1 0

ascend1 :: Floating d => d -> (d -> d) -> Int -> d -> d
ascend1 lr diff = go where
  go 0 x = x
  go n x = let d = diff x in go (n - 1) (x + (lr * d))

-- XXX: not sure that the factor amount is good
solveLaplaceGaussian :: MDistr MGaussianT
  -> (forall s. AD s (Tower Double) -> AD s (Tower Double))
  -> Double
  -> Int
  -> (MDistr MGaussianT, Double)
solveLaplaceGaussian (MGaussian mu var) likelihood learningRate numIters =
  (MGaussian mu' var', factorAmt - originalFactorAmt)
  where
  fmu' : _ : f''mu' : _ = diffs0 f mu'
  originalFactorAmt = log $ sqrt (pi / var) * exp (gaussian_ll' mu var mu)
  factorAmt = log $ sqrt (- 2 * pi / f''mu') * exp fmu'
  var' :: Double
  var' = -2 / f''mu'
  mu' ::  Double
  mu' = ascend1 learningRate ((!! 1) . diffs0 f) numIters mu
  f :: forall s. AD s (Tower Double) -> AD s (Tower Double)
  f x = gaussian_ll' (auto mu) (auto var) x + likelihood x

mdistrToDistr :: MDistr a -> Distr (MType a)
mdistrToDistr (MGaussian mu var) = normal mu var
mdistrToDistr (MMVGaussian mu var) = mvNormal mu var
mdistrToDistr (MBeta a b) = beta a b
mdistrToDistr (MBernoulli p) = bernoulli p

cdistrToDistr :: CDistr m m' -> MType m -> MDistr m'
cdistrToDistr (AffineMeanGaussian m b obsvar) mu = MGaussian (m * mu + b) obsvar
cdistrToDistr (MVAffineMeanGaussian f b obsvar) mu = MMVGaussian (f #> mu + b) obsvar
cdistrToDistr CBernoulli p = MBernoulli p

makeMarginal :: MDistr a -> CDistr a b -> MDistr b
makeMarginal (MGaussian mu var) (AffineMeanGaussian m b obsvar) =
  MGaussian (m * mu + b) (m^2 * var + obsvar)
makeMarginal (MMVGaussian mu var) (MVAffineMeanGaussian f b obsvar) =
  MMVGaussian (f #> mu + b) (conjugate f var + obsvar)
makeMarginal (MBeta a b) CBernoulli = MBernoulli (a / (a + b))
makeMarginal (MBernoulli _) _ = error "impossible"

makeConditional :: MDistr a -> CDistr a b -> MType b -> MDistr a
makeConditional (MGaussian mu var) (AffineMeanGaussian m b obsvar) obs =
  MGaussian mu' var'
  where (mu', var') = gaussianConditioning mu var ((obs - b) / m) (obsvar / m^2)
makeConditional (MMVGaussian mu var) (MVAffineMeanGaussian f b obsvar) obs =
  MMVGaussian mu' var'
  where (mu', var') = fst (kalmanUpdate obs f obsvar (mu, var))
makeConditional (MBeta a b) CBernoulli tf =
  if tf then MBeta (a + 1) b else MBeta a (b + 1)
makeConditional (MBernoulli _) _ _ = error "impossible"

data State b where
  Initialized :: State b
  Marginalized :: MDistr b -> State b

deriving instance Show (State MGaussianT)
deriving instance Show (State MBetaT)

isMarginalized :: State b -> Bool
isMarginalized (Marginalized _) = True
isMarginalized _ = False

data DSDistr a b where
   UDistr :: !(MDistr b) -> DSDistr a b
   CDistr :: !(Ref (Node z a)) -> !(CDistr a b) -> DSDistr a b

typeOfMDistr :: MDistr b -> SMarginalT b
typeOfMDistr (MGaussian _ _) = SMGaussianT
typeOfMDistr (MBeta _ _) = SMBetaT
typeOfMDistr (MBernoulli _) = SMBernoulliT
typeOfMDistr (MMVGaussian _ _) = SMMVGaussianT

typeOfCDistr :: CDistr m m' -> SMarginalT m'
typeOfCDistr (AffineMeanGaussian _ _ _) = SMGaussianT
typeOfCDistr CBernoulli = SMBernoulliT
typeOfCDistr (MVAffineMeanGaussian _ _ _) = SMMVGaussianT

typeOfDSDistr :: DSDistr a b -> SMarginalT b
typeOfDSDistr (UDistr d) = typeOfMDistr d
typeOfDSDistr (CDistr _ c) = typeOfCDistr c

data Node a b = Node
  { name :: !String
  , children :: [RefNodeFrom b]
  , state :: !(State b)
  , distr :: !(DSDistr a b)
  }
  | RealizedNode (MType b)

data RefNode a b where
  RefNode :: !(Ref (Node a b)) -> RefNode a b

data SomeRefNode where
  SomeRefNode :: !(Ref (Node a b)) -> SomeRefNode

instance Eq SomeRefNode where
  SomeRefNode x == SomeRefNode y = x == unsafeCoerce y

data Path e (a :: MarginalT) (b :: MarginalT) where
  PNil :: Path e a a
  PSnoc :: e b c -> Path e a b -> Path e a c

data PathTo e b where
  PathTo :: !(Path e a b) -> PathTo e b

data RefNodeTo b where
  RefNodeTo :: !(Ref (Node a b)) -> RefNodeTo b

data RefNodeFrom a where
  RefNodeFrom :: !(Ref (Node a b)) -> RefNodeFrom a

instance Eq (RefNodeFrom a) where
  RefNodeFrom x == RefNodeFrom y = x == unsafeCoerce y

instance Eq (RefNodeTo a) where
  RefNodeTo x == RefNodeTo y = x == unsafeCoerce y

ancestry :: DelayedSampling m => Node a b -> m (PathTo RefNode a)
ancestry n = case distr n of
  UDistr _ -> pure $ PathTo PNil
  CDistr par cdistr -> do
    PathTo ancestors <- ancestry =<< readRef par
    pure (PathTo (PSnoc (RefNode par) ancestors))

path_all :: Applicative f => (forall a b. e a b -> f Bool) -> Path e a b -> f Bool
path_all f PNil = pure True
path_all f (PSnoc e es) = (&&) <$> f e <*> path_all f es

isTerminal :: DelayedSampling m =>Node a b -> m Bool
isTerminal n = do
  PathTo ancestors <- ancestry n
  ancestors_all_marginalized <- path_all (\(RefNode n) -> isMarginalized . state <$> readRef n) ancestors
  pure $ isMarginalized (state n) && ancestors_all_marginalized

-- initialize without parent node
constant' :: String -> MDistr a -> Node z a
constant' n d = Node
  { name = n
  , children = []
  , state = Initialized
  , distr = UDistr d
  }

assumeConstant :: DelayedSampling m => String -> MDistr a -> m (Ref (Node z a))
assumeConstant n d = newRef (constant' n d)

-- initialize with parent node
newConditional' :: DelayedSampling m => String -> Ref (Node a b) -> CDistr b c -> m (Ref (Node b c))
newConditional' str par cdistr = do
  childRef <- newRef child
  modifyRef' par $ \n -> n { children = (RefNodeFrom childRef) : children n }
  --childRef' <- mkWeakRef childRef $ putStrLn ("deleted " ++ str)
  pure childRef
  where
  child = Node
    { name = str
    , children = []
    , state = Initialized
    , distr = CDistr par cdistr }

assumeConditional :: DelayedSampling m => String -> Ref (Node a b) -> CDistr b c -> m (Ref (Node b c))
assumeConditional str par cdistr = newConditional' str par cdistr

updater :: DelayedSampling m => (a -> m a) -> Ref a -> m ()
updater f ref = do
  x <- readRef ref
  y <- f x
  writeRef ref y

marginalize :: DelayedSampling m => MDistr a -> Ref (Node a b) -> m (MDistr b)
marginalize parMarginal nref = do
  n <- readRef nref
  writeLog ("marginalize " ++ name n)
  case (state n, distr n) of
    (Initialized, CDistr par cdistr) -> do
      let marginal' = makeMarginal parMarginal cdistr
      writeRef nref $ n { state = Marginalized marginal' }
      pure marginal'
    (state, _) -> error $  "marginalize': " ++ name n

marginalize2' :: MType a -> Node a b -> Node a b
marginalize2' value n =
  case (state n, distr n) of
    (Initialized, CDistr _ cdistr) ->
      let marg = cdistrToDistr cdistr value in
      n { state = Initialized, distr = UDistr marg }
    _ -> error "marginalize2'"

marginalize2 :: DelayedSampling m => MType a -> Ref (Node a b) -> m ()
marginalize2 a = updater (pure . marginalize2' a)

sample :: DelayedSampling m => MonadSample m => Ref (Node a b) -> m ()
sample nref = do
  n <- readRef nref
  ioAssert (isTerminal n)
  writeLog ("sample " ++ name n)
  case state n of
    Marginalized m -> do
      x <- Distributions.sample (mdistrToDistr m)
      realize x nref
    _ -> error "sample"

updateParent :: DelayedSampling m => Node b c -> (forall a. CDistr b c -> Node a b -> m (Node a b)) -> m ()
updateParent n f = case distr n of
  UDistr _ -> pure ()
  CDistr p cdistr -> updater (f cdistr) p

firstSatisfying :: Monad m => (a -> m Bool) -> [a] -> m (Maybe a)
firstSatisfying p [] = pure Nothing
firstSatisfying p (x : xs) = do
  px <- p x
  if px then pure (Just x) else firstSatisfying p xs

-- Invariant 2: A node always has at most one marginal Child
marginalChild :: DelayedSampling m => Node a b -> m (Maybe (RefNodeFrom b))
marginalChild n = firstSatisfying (\(RefNodeFrom x) -> isMarginalized . state <$> readRef x) (children n)

ioAssert :: Applicative m => m Bool -> m ()
ioAssert x = pure () -- do {b <- x; assert b (pure ())}

writeLog :: Applicative m => String -> m ()
writeLog = const $ pure () -- putStrLn

realize :: DelayedSampling m => MType b -> Ref (Node a b) -> m ()
realize val nref = do
  n <- readRef nref
  writeLog ("realize " ++ name n)
  ioAssert (isTerminal n)
  updateParent n $ \cdistr p -> do
    let Marginalized marg = state p
    let marg' = makeConditional marg cdistr val
    pure $ p { state = Marginalized marg', children = delete (RefNodeFrom nref) (children p) }
  forM_ (children n) $ \(RefNodeFrom c) -> do
    marginalize2 val c
  writeRef nref $ RealizedNode val

score' :: DelayedSampling m => MonadCond m => MType b -> Ref (Node a b) -> m Double
score' x nref = do
  n <- readRef nref
  ioAssert (isTerminal n)
  writeLog ("observe " ++ name n)
  pure $ case state n of
    Marginalized marg -> Distributions.score (mdistrToDistr marg) x
    _ -> error "observe'"

observe :: DelayedSampling m => MonadCond m => MType b -> Ref (Node a b) -> m ()
observe x nref = do
  ll <- score' x nref
  factor ll
  realize x nref

prune :: DelayedSampling m => MonadSample m => Ref (Node a b) -> m ()
prune nref = do
  n <- readRef nref
  writeLog ("prune " ++ name n)
  assert (isMarginalized (state n)) $ do
    maybeChild <- marginalChild n
    forM_ maybeChild $ \(RefNodeFrom c) ->
      prune c
    DelayedSampling.sample nref

-- turns `nref` into a terminal node
graft :: DelayedSampling m => MonadSample m => Ref (Node a b) -> m (MDistr b)
graft nref = do
  n <- readRef nref
  writeLog ("graft " ++ name n)
  case state n of
    Marginalized _ -> do
      maybeChild <- marginalChild n
      forM_ maybeChild $ \(RefNodeFrom c) -> prune c
      n'  <- readRef nref
      pure $ case state n' of
        Marginalized m -> m
        _ -> error "graft"
    Initialized -> case distr n of
      UDistr d -> do
        writeRef nref $ n { state = Marginalized d }
        pure d
      CDistr par cdistr -> do
        parMarg <- graft par
        marginalize parMarg nref

obs :: DelayedSampling m => MonadInfer m => MType b -> Ref (Node a b) -> m ()
obs x n = do
  graft n
  DelayedSampling.observe x n

score :: DelayedSampling m => MonadInfer m => MType b -> Ref (Node a b) -> m Double
score x n = withoutModifying $ do
  graft n
  score' x n

withoutModifying :: MonadState s m => m a -> m a
withoutModifying f = do
  s <- get
  x <- f
  put s
  return x

getValue :: DelayedSampling m => MonadSample m => Ref (Node a b) -> m (MType b)
getValue nref = do
  n <- readRef nref
  case n of
    RealizedNode x -> pure x
    _ -> do
      graft nref
      DelayedSampling.sample nref
      getValue nref

printValue :: Show (MType b) => DelayedSampling m => MonadSample m => MonadIO m => Ref (Node a b) -> m ()
printValue n = do
  x <- getValue n
  liftIO $ print x

printState :: Show (State b) => DelayedSampling m => MonadIO m => Ref (Node a b) -> m ()
printState nref = do
  n <- readRef nref
  liftIO $ print (state n)

forget :: DelayedSampling m => Ref (Node a b) -> m ()
forget nref = do
  n <- readRef nref
  case n of
    RealizedNode _ -> pure ()
    _ -> case state n of
      Initialized -> error "forget"
      Marginalized marg -> do
        forM_ (children n) $ \(RefNodeFrom cref) -> do
          modifyRef' cref $ \c ->
            c { distr = case distr c of
              UDistr d -> UDistr d
              CDistr _ _ -> case state c of
                Marginalized marg -> UDistr marg
                _ -> error "forget" }
        case distr n of
          UDistr d -> pure ()
          CDistr par cdistr -> do
            isStale <- stale par
            if isStale
              then error "can't forget: information to incorporate"
              else updateParent n $ \cdistr p -> pure $ p { children = delete (RefNodeFrom nref) (children p) }
        writeRef nref $ n { distr = UDistr marg }
  modifyRef' nref $ \n' -> n' { children = [] }
  freeRef nref

-- sweep :: [SomeRefNode] -> M ()
-- sweep liveVars = mapM_ (\(SomeRefNode n) -> forgetRelated n) liveVars where
--   forgetRelated :: Ref (Node a b) -> M ()
--   forgetRelated nref = undefined
--   forgettable :: Ref (Node a b) -> M Bool
--   forgettable nref =
--     if SomeRefNode nref `elem` liveVars
--       then pure False
--       else case state n of

stale :: DelayedSampling m => Ref (Node a b) -> m Bool
stale nref = do
  n <- readRef nref
  case n of
    RealizedNode _ -> pure False
    _ -> case state n of
      Initialized -> case distr n of
        UDistr _ -> pure False
        CDistr par _ -> stale par
      Marginalized d -> isJust <$> marginalChild n

-- Examples

delay_triplet :: DelayedInfer m => MonadIO m => Double -> m ()
delay_triplet zobs = do
  x <- assumeConstant "x" (MGaussian 0 1)
  writeLog "x created"
  y <- assumeConditional "y"  x (gaussianMeanGaussian 1)
  writeLog "y created"
  z <- assumeConditional "z" y (gaussianMeanGaussian 1)
  writeLog "z created"
  obs zobs z
  writeLog "z observed"
  printValue z
  printValue x
  printValue y

observeConditional :: DelayedInfer m => String -> Ref (Node a b) -> CDistr b c -> MType c -> m ()
observeConditional str nref cdistr observation = do
  y <- assumeConditional str nref cdistr
  obs observation y
  freeRef y

scoreConditional :: DelayedInfer m => String -> Ref (Node a b) -> CDistr b c -> MType c -> m Double
scoreConditional str nref cdistr observation = do
  y <- assumeConditional str nref cdistr
  ll <- DelayedSampling.score observation y
  modifyRef' nref $ \n -> n { children = tail (children n) }
  freeRef y
  pure ll

-- For some reason, likelihoods aren't exactly  the same
delay_iidADF :: DelayedInfer m => MonadIO m => Bool -> [Double] -> m ()
delay_iidADF useADF yobs = do
  x <- assumeConstant "x" (MGaussian 0 1)
  forM_ (zip [1..] yobs) $ \(t, obsyt) -> do
    if useADF
      then observeGeneric x $ \d -> solveLaplaceGaussian d (\mu -> gaussian_ll' mu 1 (auto obsyt)) 0.05 10000
      else observeConditional ("y" ++ show t) x (gaussianMeanGaussian 1) obsyt  -- normal(x, 1)
    printState x
  printValue x

observeGeneric :: DelayedInfer m => Ref (Node a b) -> (MDistr b -> (MDistr b, Double)) -> m ()
observeGeneric nref f = do
  graft nref
  n <- readRef nref
  case n of
    RealizedNode x -> error "observeGeneric: RealizedNode impossible"
    _ -> case state n of
      Marginalized d -> do
        let (d', factorAmt) = f d
        factor factorAmt
        writeRef nref $ n { state = Marginalized d' }
      _ -> error "observeGeneric"

delay_kalman :: DelayedInfer m => MonadIO m =>
  Bool -> Double -> Double -> [Double] -> m ()
delay_kalman shouldForget m b yobs = do
  x1 <- assumeConstant "x1" (MGaussian 0 1)
  observeConditional "y1" x1 (gaussianMeanGaussian 1) (head yobs)
  loop 2 x1 (tail yobs)
  where
  loop :: DelayedInfer m => MonadIO m =>
    Integer -> Ref (Node a MGaussianT) -> [Double] -> m ()
  loop t xpredt [] = pure ()
  loop t xpredt (y : ys) = do
    xt <- assumeConditional ("x" ++ show t) xpredt (AffineMeanGaussian m b 1)
    observeConditional ("y" ++ show t) xt (gaussianMeanGaussian 1) y
    when (t `mod` 100 == 0) $ do
      liftIO $ putStrLn ("t = " ++ show t)
      printState xt
    when shouldForget $ forget xpredt
    loop (t + 1) xt ys


{-# LANGUAGE RankNTypes, BangPatterns, GADTs #-}
{-# OPTIONS -Wall #-}

module Language.Hakaru.Distribution where

import Control.Monad
import Control.Monad.Primitive
import Control.Monad.Loops
import qualified System.Random.MWC as MWC
import Language.Hakaru.Mixture
import Language.Hakaru.Types
import Data.Ix
import Data.Maybe (fromMaybe)
import Data.List (findIndex, foldl')
import Numeric.SpecFunctions
import qualified Data.Map.Strict as M
import qualified Data.Number.LogFloat as LF

mapFst :: (t -> s) -> (t, u) -> (s, u)
mapFst f (a,b) = (f a, b)

dirac :: (Eq a) => a -> Dist a
dirac theta = Dist {logDensity = (\ (Discrete x) -> if x == theta then 0 else log 0),
                    distSample = (\ _ -> return $ Discrete theta)}

bern :: Double -> Dist Bool
bern p = Dist {logDensity = (\ (Discrete x) -> log (if x then p else 1 - p)),
               distSample = (\ g -> do t <- MWC.uniformR (0,1) g
                                       return $ Discrete (t <= p))}

uniform :: Double -> Double -> Dist Double
uniform lo hi =
    let uniformLogDensity lo' hi' x | lo' <= x && x <= hi' = log (recip (hi' - lo'))
        uniformLogDensity _ _ _ = log 0
    in Dist {logDensity = (\ (Lebesgue x) -> uniformLogDensity lo hi x),
             distSample = (\ g -> liftM Lebesgue $ MWC.uniformR (lo, hi) g)}

uniformD :: (Ix a, MWC.Variate a) => a -> a -> Dist a
uniformD lo hi =
    let uniformLogDensity lo' hi' x | lo' <= x && x <= hi' = log density
        uniformLogDensity _ _ _ = log 0
        density = recip (fromInteger (toInteger (rangeSize (lo,hi))))
    in Dist {logDensity = (\ (Discrete x) -> uniformLogDensity lo hi x),
             distSample = (\ g -> liftM Discrete $ MWC.uniformR (lo, hi) g)}

marsaglia :: (MWC.Variate a, Ord a, Floating a, PrimMonad m) => PRNG m -> m (a, a)
marsaglia g = do -- "Marsaglia polar method"
  x <- MWC.uniformR (-1,1) g
  y <- MWC.uniformR (-1,1) g
  let s = x * x + y * y
      q = sqrt ((-2) * log s / s)
  if 1 >= s && s > 0 then return (x * q, y * q) else marsaglia g

choose :: (PrimMonad m) => Mixture k -> PRNG m -> m (k, Prob)
choose (Mixture m) g = do
  let peak = maximum (M.elems m)
      unMix = M.map (LF.fromLogFloat . (/peak)) m
      total = M.foldl' (+) (0::Double) unMix
  p <- MWC.uniformR (0, total) g
  let f !k !v b !p0 = let p1 = p0 + v in if p <= p1 then k else b p1
      err p0 = error ("choose: failure p0=" ++ show p0 ++
                      " total=" ++ show total ++
                      " size=" ++ show (M.size m))
  return $ (M.foldrWithKey f err unMix 0, LF.logFloat total * peak)

chooseIndex :: (PrimMonad m) => [Double] -> PRNG m -> m Int
chooseIndex probs g = do
  p <- MWC.uniform g
  return $ fromMaybe (error ("chooseIndex: failure p=" ++ show p))
           (findIndex (p <=) (scanl1 (+) probs))

normal_rng :: (Real a, Floating a, MWC.Variate a, PrimMonad m) =>
              a -> a -> PRNG m -> m a
normal_rng mu sd g | sd > 0 = do (x, _) <- marsaglia g
                                 return (mu + sd * x)
normal_rng _ _ _ = error "normal: invalid parameters"

normalLogDensity :: Floating a => a -> a -> a -> a
normalLogDensity mu sd x = (-tau * square (x - mu)
                            + log (tau / pi / 2)) / 2
  where square y = y * y
        tau = 1 / square sd

normal :: Double -> Double -> Dist Double
normal mu sd = Dist {logDensity = normalLogDensity mu sd . fromLebesgue,
                     distSample = (\g -> liftM Lebesgue $ normal_rng mu sd g)}

categoricalLogDensity :: (Eq b, Floating a) => [(b, a)] -> b -> a
categoricalLogDensity list x = log $ fromMaybe 0 (lookup x list)

categoricalSample :: (Num b, Ord b, PrimMonad m, MWC.Variate b) =>
    [(t,b)] -> PRNG m -> m t
categoricalSample list g = do
  let total = sum $ map snd list
  p <- MWC.uniformR (0, total) g
  let sumList = scanl1 (\acc (a, b) -> (a, b + snd(acc))) list
      elem' = fst $ head $ filter (\(_,p0) -> p <= p0) sumList
  return elem'

categorical :: Eq a => [(a,Double)] -> Dist a
categorical list = Dist {logDensity = categoricalLogDensity list . fromDiscrete,
                         distSample = (\g -> liftM Discrete $ categoricalSample list g)}

lnFact :: Int -> Double
lnFact = logFactorial

-- Makes use of Atkinson's algorithm as described in:
-- Monte Carlo Statistical Methods pg. 55
--
-- Further discussion at:
-- http://www.johndcook.com/blog/2010/06/14/generating-poisson-random-values/
poisson_rng :: (PrimMonad m) => Double -> PRNG m -> m Int
poisson_rng lambda g' = make_poisson g'
   where smu = sqrt lambda
         b  = 0.931 + 2.53*smu
         a  = -0.059 + 0.02483*b
         vr = 0.9277 - 3.6224/(b - 2)
         arep  = 1.1239 + 1.1368/(b-3.4)
         lnlam = log lambda

         make_poisson :: (PrimMonad m) => PRNG m -> m Int
         make_poisson g = do u <- MWC.uniformR (-0.5,0.5) g
                             v <- MWC.uniformR (0,1) g
                             let us = 0.5 - abs u
                                 k = floor $ (2*a / us + b)*u + lambda + 0.43
                             case () of
                               () | us >= 0.07 && v <= vr -> return k
                               () | k < 0 -> make_poisson g
                               () | us <= 0.013 && v > us -> make_poisson g
                               () | accept_region us v k -> return k
                               _  -> make_poisson g

         accept_region :: Double -> Double -> Int -> Bool
         accept_region us v k = log (v * arep / (a/(us*us)+b)) <=
                                -lambda + (fromIntegral k)*lnlam - lnFact k

poisson :: Double -> Dist Int
poisson l =
    let poissonLogDensity l' x | l' > 0 && x> 0 = (fromIntegral x)*(log l') - lnFact x - l'
        poissonLogDensity l' x | x==0 = -l'
        poissonLogDensity _ _ = log 0
    in Dist {logDensity = poissonLogDensity l . fromDiscrete,
             distSample = (\g -> liftM Discrete $ poisson_rng l g)}

-- Direct implementation of  "A Simple Method for Generating Gamma Variables"
-- by George Marsaglia and Wai Wan Tsang.
gamma_rng :: (PrimMonad m) => Double -> Double -> PRNG m -> m Double
gamma_rng shape _   _ | shape <= 0.0  = error "gamma: got a negative shape paramater"
gamma_rng _     scl _ | scl <= 0.0  = error "gamma: got a negative scale paramater"
gamma_rng shape scl g | shape <  1.0  = do gvar1 <- gamma_rng (shape + 1) scl g
                                           w <- MWC.uniformR (0,1) g
                                           return $ scl * gvar1 * (w ** recip shape)
gamma_rng shape scl g = do
    let d = shape - 1/3
        c = recip $ sqrt $ 9*d
        -- Algorithm recommends inlining normal generator
        -- n = normal_rng 1 c
    v <- iterateUntil (> 0.0) $ normal_rng 1 c g
        -- (v, g2) = until (\y -> fst y > 0.0) (\ (_, g') -> normal_rng 1 c g') (n g)
    let x = (v - 1) / c
        sqr = x * x
        v3 = v * v * v
    u <- MWC.uniformR (0.0, 1.0) g
    let accept = u < 1.0 - 0.0331*(sqr*sqr) || log u < 0.5*sqr + d*(1.0 - v3 + log v3)
    case accept of
      True -> return $ scl*d*v3
      False -> gamma_rng shape scl g

gammaLogDensity :: Double -> Double -> Double -> Double
gammaLogDensity shape scl x | x>= 0 && shape > 0 && scl > 0 =
     scl * log shape - scl * x + (shape - 1) * log x - logGamma shape
gammaLogDensity _ _ _ = log 0

gamma :: Double -> Double -> Dist Double
gamma shape scl = Dist {logDensity = gammaLogDensity shape scl . fromLebesgue,
                        distSample = (\g -> liftM Lebesgue $ gamma_rng shape scl g)}

beta_rng :: (PrimMonad m) => Double -> Double -> PRNG m -> m Double
beta_rng a b g | a <= 1.0 && b <= 1.0 = do
                 u <- MWC.uniformR (0.0, 1.0) g
                 v <- MWC.uniformR (0.0, 1.0) g
                 let x = u ** (recip a)
                     y = v ** (recip b)
                 case (x+y) <= 1.0 of
                   True -> return $ x / (x + y)
                   False -> beta_rng a b g
beta_rng a b g = do ga <- gamma_rng a 1 g
                    gb <- gamma_rng b 1 g
                    return $ ga / (ga + gb)

betaLogDensity :: Double -> Double -> Double -> Double
betaLogDensity _ _ x | x < 0 || x > 1 = error "beta: value must be between 0 and 1"
betaLogDensity a b _ | a <= 0 || b <= 0 = error "beta: parameters must be positve"
betaLogDensity a b x = (logGamma (a + b)
                        - logGamma a
                        - logGamma b
                        + (a - 1) * log x
                        + (b - 1) * log (1 - x))

beta :: Double -> Double -> Dist Double
beta a b = Dist {logDensity = betaLogDensity a b . fromLebesgue,
                 distSample = (\g -> liftM Lebesgue $ beta_rng a b g)}

laplace_rng :: (PrimMonad m) => Double -> Double -> PRNG m -> m Double
laplace_rng mu sd g = MWC.uniformR (0.0, 1.0) g >>= sample
   where sample u = return $ case u < 0.5 of
                               True  -> mu + sd * log (u + u)
                               False -> mu - sd * log (2.0 - u - u)

laplaceLogDensity :: Floating a => a -> a -> a -> a
laplaceLogDensity mu sd x = - log (2 * sd) - abs (x - mu) / sd

laplace :: Double -> Double -> Dist Double
laplace mu sd = Dist {logDensity = laplaceLogDensity mu sd . fromLebesgue,
                      distSample = (\g -> liftM Lebesgue $ laplace_rng mu sd g)}

-- Consider having dirichlet return Vector
-- Note: This is actually symmetric dirichlet
dirichlet_rng :: (PrimMonad m) => Int ->  Double -> PRNG m -> m [Double]
dirichlet_rng n' a g' = liftM normalize $ gammas g' n'
  where gammas _ 0 = return ([], 0)
        gammas g n = do (xs, total) <- gammas g (n-1)
                        x <- gamma_rng a 1 g
                        return ((x : xs), x+total)
        normalize (b, total) = map (/ total) b

dirichletLogDensity :: [Double] -> [Double] -> Double
dirichletLogDensity a x | all (> 0) x = sum' (zipWith logTerm a x) + logGamma (sum a)
  where sum' = foldl' (+) 0
        logTerm b y = (b-1) * log y - logGamma b
dirichletLogDensity _ _ = error "dirichlet: all values must be between 0 and 1"


{-# LANGUAGE BangPatterns #-}

import Data.Char
import Data.List
import Debug.Trace
import Numeric.LinearAlgebra hiding (Matrix, Vector)
import qualified Numeric.LinearAlgebra as LA
import System.Random.Mersenne

type FpType = Double
type Vector = LA.Vector FpType
type Matrix = LA.Matrix FpType

-- Initializer

initVector :: MTGen -> Int -> IO (Vector)
initVector rng n = do
    rs <- randoms rng
    let v = fromList (take n rs)
    return v

initMatrix :: MTGen -> Int -> Int -> IO (Matrix)
initMatrix rng n m = do
    rs <- randoms rng :: IO [FpType]
    let v = (n><m) $ take (n*m) rs
    return v

initConnections :: MTGen -> (Layer, Layer) -> IO (WeightMatrix)
initConnections rng (l₁,l₂) = do
    w <- initMatrix rng (neurons l₂) (neurons l₁)
    return w

initMlp :: MTGen -> [Layer] -> IO (Network)
initMlp rng ls = do
    ws <- mapM (initConnections rng) $ zip ls (tail ls)
    let net = mlp ls ws
    return net

-- Helper

vecToStr :: Vector -> String
vecToStr v = "[ " ++ intercalate "  " (map show (toList v)) ++ " ]"

vecsToStr :: [Vector] -> String
vecsToStr vs = intercalate "\n" (map vecToStr vs)

datasetToStr :: Dataset -> String
datasetToStr d = "dataset: " ++ show l ++ " elements, " ++ show u ++ " targets\n"
    where u = nub (map snd d)
          l = length d

-- Layers

--TODO: allow 2d tensors for activation as well to facilitate softmax
data Layer = Layer {
    ϕ         :: FpType -> FpType,
    ϕ'        :: FpType -> FpType,
    neurons   :: Int,
    layerType :: String
}

instance Show Layer where
    show (Layer { neurons = n, layerType = t }) = t ++ ", " ++ show n ++ " neurons"

linearLayer :: Int -> Layer
linearLayer n = Layer {
    ϕ         = id,
    ϕ'        = id,
    neurons   = n,
    layerType = "linear"
}

hyperbolicLayer :: Int -> Layer
hyperbolicLayer n = Layer {
    ϕ         = tanh,
    ϕ'        = \x -> 1 - tanh² x,
    neurons   = n,
    layerType = "tanh"
} where tanh² = tanh . tanh

sigmoidLayer :: Int -> Layer
sigmoidLayer n = Layer {
    ϕ         = sigmoid,
    ϕ'        = \x -> sigmoid x * (1 - sigmoid x),
    neurons   = n,
    layerType = "sigmoid"
} where sigmoid x = 0.5 * (1 + tanh (0.5 * x))

-- Networks

data Network = Network {
    layers  :: [Layer],
    weights :: [WeightMatrix],
    netType :: String
}

type WeightMatrix = Matrix

--TODO: It would be nicer to enforce #layers == #weights+1 by zipping.
mlp :: [Layer] -> [WeightMatrix] -> Network
mlp ls ws = debug Network { layers = ls, weights = ws, netType = "mlp" }
    where text  = map (\x -> "Layer " ++ show (fst x) ++ ": " ++ show (snd x)) (zip [1..] ls)
          debug = trace $ "Creating MLP\n" ++ intercalate "\n" text ++ "\n"

forward :: Vector -> (Layer, WeightMatrix) -> Vector
forward x (l,w) = mapVector (ϕ l) $ w <> x

activate :: Network -> Vector -> [Vector]
activate (Network { layers = ls, weights = ws }) x = debug result
    where result = scanl forward ϕx $ zip (tail ls) ws
          ϕx     = mapVector ϕ₁ x
          ϕ₁     = ϕ (head ls)
          debug  = trace ("activate input\n" ++ vecToStr x ++ "\n\nactivate output\n" ++ vecsToStr result ++ "\n")

-- Datasets

--TODO: decouple classification/regression datasets properly, add bias s.t. it works flawlessly in all usages
type Dataset = [(Vector, Vector)]

-- Training

data Trainer = Trainer {
    η         :: FpType,
    terminate :: TrainState -> Bool
}

data TrainState = TrainState {
    nEpochs :: Int,
    δs      :: [FpType]
} deriving (Show)

type Error = FpType

--TODO: more useful termination criteria
nEpochTrainer :: FpType -> Int -> Trainer
nEpochTrainer lr n = debug Trainer { η = lr, terminate = (> n) . nEpochs }
    where debug = trace ("Creating nEpochTrainer\n" ++ "η = " ++ show lr ++ "\nmaxEpochs = " ++ show n ++ "\n")

--TODO: actually perform epoch inside ST monad, put actual result
epoch :: Dataset -> Network -> Trainer -> (Network, Error)
epoch d net t = (net, 0)

_train :: TrainState -> Dataset -> Network -> Trainer -> Network
_train s d net t = do
    let (net',δ') = epoch d net t
    let s' = update s where update s = TrainState { nEpochs = 1 + (nEpochs s), δs = δs s ++ [δ'] }
    let statusMessage = "δ_" ++ show (nEpochs s') ++ " = " ++ show δ' ++ "\n" ++ show s' ++ "\n"

    if (terminate t) s'
        then trace ("Finished Training.\n") net'
        else trace statusMessage $ _train s' d net' t

--TODO: support testing set/cross validation
train :: Dataset -> Network -> Trainer -> Network
train = _train s₀ where s₀ = TrainState { nEpochs = 0, δs = [] }

-- Evaluation

--TODO

-- Testing

main :: IO()
main = do
    rng <- newMTGen Nothing

    let l₁ = linearLayer 5
    let l₂ = sigmoidLayer 3
    let l₃ = linearLayer 5
    let layers = [l₁,l₂,l₃]

    !net <- initMlp rng layers
    xs <- mapM (initVector rng) (take 100 $ repeat 5)

    let ys = xs
    let dataset = zip xs ys
    let trainer = nEpochTrainer 0.1 5
    let !trainedNet = train dataset net trainer

    putStrLn $ "Done!\n"


--------------------------------------------------------------------------------
--
--  Copyright (c) 2010 - 2013 Tad Doxsee
--  All rights reserved.
--
--  Author: Tad Doxsee
--
--------------------------------------------------------------------------------

{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE  FlexibleInstances    #-}

-- * base
import Foreign           (ForeignPtr, Storable)
import Foreign.C.Types   (CDouble, CInt, CSize)

-- * vector
import qualified Data.Vector.Storable.Mutable as V

-- * cholmod
import Numeric.LinearAlgebra.CHOLMOD.CholmodXFaceLow
import Numeric.LinearAlgebra.CHOLMOD.CholmodXFace

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

main :: IO ()
main = do

  c <- allocCommon :: IO (ForeignPtr Common)
  startC c

  let n = 5                  :: CSize
      nz = (n+1) * n `div` 2 :: CSize
      nn = fromIntegral n    :: CSize

  at <-allocTriplet n n nz stSquareSymmetricLower xtReal c
       :: IO (Matrix Triplet)

  atNRow  <- getNRow  at :: IO CSize
  atNCol  <- getNCol  at :: IO CSize
  atNZMax <- getNZMax at :: IO CSize

  putStrLn ""
  putStrLn $ "at: (" ++ show atNRow
                     ++ " x " ++ show atNCol ++ ")"
  putStrLn $ "nzmax: " ++ show atNZMax

  iv <- tripletGetRowIndices at :: IO (V.IOVector CInt)
  jv <- tripletGetColIndices at :: IO (V.IOVector CInt)
  xt  <- tripletGetX at

  let nni = fromIntegral nn

  let ij      = [(i, j) | i <- [0 .. nni-1], j <- [0 .. i]] :: [(CInt, CInt)]
      (ia,ja) = unzip ij                                    :: ([CInt], [CInt])

      -- the elements of the matrix are (lower half)
      --
      --  11.0
      --  21.0  22.0
      --  31.0  32.0  33.0
      --  41.0  42.0  43.0  44.0
      --  51.0  52.0  53.0  54.0  55.0

      xp = [11, 21, 22, 31, 32, 33, 41, 42, 43, 44, 51, 52, 53, 54, 55]
           :: [CDouble]

  writev iv ia
  writev jv ja
  writev xt xp

  let nnz = fromIntegral $ Prelude.length xp :: CSize

  tripletSetNNZ at nnz

  as <- tripletToSparse at c :: IO (Matrix Sparse)
  b  <- ones n 1 xtReal c    :: IO (Matrix Dense)
  l  <- analyze as c         :: IO (ForeignPtr Factor)
  factorize as l c
  x  <- solve stA l b c      :: IO (Matrix Dense)

  xNRow <- getNRow x
  xNCol <- getNCol x

  xv <- getX x :: IO (V.IOVector CDouble)

  putStrLn ""
  putStrLn $ "xv: (" ++ show xNRow
                     ++ " x " ++ show xNCol ++ ")"

  xl <- readv xv
  mapM_ (putStrLn . show) xl

  r <- denseCopy b c

  let oneL = [1, 0] :: [CDouble]
      m1L = [-1, 0] :: [CDouble]

  _ <- sdMult as noTranspose m1L oneL x r c

  norm <- denseNorm r infNorm c :: IO CDouble

  putStrLn $ "norm: " ++ show norm

  putStrLn "done"

writev :: (Storable a) => V.IOVector a -> [a] -> IO ()

writev v xs =
  sequence_ [V.write v i x | (i, x) <- zip [0 .. (Prelude.length xs - 1)] xs]

readv :: (Storable a) => V.IOVector a -> IO [a]
readv v = sequence [V.read v i | i <- [0 .. (V.length v) - 1]]


module LatticeSymmetries.ComplexRational
  ( ComplexRational (..),
    ConvertibleToComplexDouble (..),
    realPart,
    imagPart,
    conjugate,
    magnitudeSquared,
    Cscalar,
  )
where

import Data.Complex (Complex (..))
import Foreign.C.Types (CDouble (..))
import Text.PrettyPrint.ANSI.Leijen (Pretty (..))
import qualified Text.PrettyPrint.ANSI.Leijen as Pretty

data ComplexRational = ComplexRational {-# UNPACK #-} !Rational {-# UNPACK #-} !Rational
  deriving stock (Eq, Show)

prettyRational :: Rational -> Pretty.Doc
prettyRational x
  | realToFrac (fromRational x :: Double) == x = Pretty.double (fromRational x :: Double)
  | otherwise = Pretty.rational x

instance Pretty ComplexRational where
  pretty (ComplexRational r i)
    | i == 0 = prettyRational r
    | otherwise = Pretty.parens $ prettyRational r <> Pretty.text " + " <> prettyRational i <> "𝕚"

realPart :: ComplexRational -> Rational
realPart (ComplexRational r _) = r
{-# INLINE realPart #-}

imagPart :: ComplexRational -> Rational
imagPart (ComplexRational _ i) = i
{-# INLINE imagPart #-}

conjugate :: ComplexRational -> ComplexRational
conjugate (ComplexRational r i) = (ComplexRational r (-i))
{-# INLINE conjugate #-}

magnitudeSquared :: ComplexRational -> Rational
magnitudeSquared (ComplexRational r i) = r * r + i * i
{-# INLINE magnitudeSquared #-}

instance Num ComplexRational where
  {-# INLINE (+) #-}
  {-# INLINE (-) #-}
  {-# INLINE (*) #-}
  {-# INLINE fromInteger #-}
  (ComplexRational r i) + (ComplexRational r' i') = ComplexRational (r + r') (i + i')
  (ComplexRational r i) - (ComplexRational r' i') = ComplexRational (r - r') (i - i')
  (ComplexRational r i) * (ComplexRational r' i') = ComplexRational (r * r' - i * i') (r * i' + i * r')
  negate (ComplexRational r i) = ComplexRational (-r) (-i)
  abs _ = error "Num instance of ComplexRational does not implement abs"
  signum _ = error "Num instance of ComplexRational does not implement signum"
  fromInteger n = ComplexRational (fromInteger n) 0

instance Fractional ComplexRational where
  {-# INLINE (/) #-}
  {-# INLINE fromRational #-}
  (ComplexRational r i) / (ComplexRational r' i') =
    ComplexRational ((r * r' + i * i') / d) ((-r * i' + i * r') / d)
    where
      d = r' * r' + i' * i'
  fromRational a = ComplexRational a 0

-- class Fractional a => ComplexFloating a where
--   toComplexDouble :: a -> Complex Double
--   fromComplexDouble :: Complex Double -> a

type Cscalar = Complex CDouble

class Fractional a => ConvertibleToComplexDouble a where
  toComplexDouble :: a -> Cscalar
  fromComplexDouble :: Cscalar -> a

instance ConvertibleToComplexDouble ComplexRational where
  toComplexDouble (ComplexRational r i) = (fromRational r) :+ (fromRational i)
  fromComplexDouble (r :+ i) = ComplexRational (toRational r) (toRational i)

instance ConvertibleToComplexDouble (Complex CDouble) where
  toComplexDouble = id
  fromComplexDouble = id

instance ConvertibleToComplexDouble (Complex Double) where
  toComplexDouble = coerce
  fromComplexDouble = coerce


{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleContexts #-}

module Lib
    ( someFunc
    ) where

import Numeric.GSL.Fitting
import Numeric.LinearAlgebra
import Numeric.AD
import Numeric.AD.Internal.Reverse
import Data.Reflection

someFunc :: IO ()
someFunc = putStrLn "someFunc"

type FitRes = ([(Double, Double)], Matrix Double)
type FitData = [([Double],([Double],Double))]

-- type ManualJacob = ([Double] -> [Double] -> [[Double]])

data JacobType = AutoJacob | ManualJacob ([Double] -> [Double] -> [[Double]])

data FitParams = FitParams { jacob :: JacobType
                           , guess :: [Double]
                           , iter :: Int
                           , absTol :: Double
                           , relTol :: Double
                           }

defaultParams = FitParams { jacob = AutoJacob,
                            guess = repeat 1,
                            iter = 1000,
                            absTol = 1E-20,
                            relTol = 1E-20 }

     -- ([Double] -> [Double] -> [Double]) model
     -- ([Double] -> [Double] -> [[Double]]) model der
-- Only Doubles since hmatrix requires it
fit :: [Double]
     -> [Double]
     -> [Double]
     -> (forall a . Floating a => [a]->[a]->[a])
     -> JacobType
     -> [Double]->  Maybe FitRes

fit xs ys sigma model jacobian guess
    | xeqy && ( seqx || seqone)  = Just $ fitModelScaled 1E-20 1E-20 500 (model, sanitizedJacob jacobian) fitdata guess
    | otherwise  = Nothing
    where
    lx = length xs
    ly = length ys
    ls = length sigma
    xeqy = lx == ly
    seqx = ls == lx
    seqone = ls == 1
    fitdata = unsafeFormatData xs ys sigma
    sanitizedJacob jacob = case jacob of AutoJacob -> mkJac model
                                         ManualJacob f -> f

{-
fit :: [Double]
     -> [Double]
     -> [Double]
     -> (forall a . Floating a => [a]->[a]->[a])
     -> ([Double] -> [Double] -> [[Double]])
     -> [Double]->  Maybe FitRes

fit xs ys sigma model jacobian guess
    | xeqy && ( seqx || seqone)  = Just $ fitModelScaled 1E-20 1E-20 500 (model, mkJac model) fitdata guess
    | otherwise  = Nothing
    where
    lx = length xs
    ly = length ys
    ls = length sigma
    xeqy = lx == ly
    seqx = ls == lx
    seqone = ls == 1
    fitdata = unsafeFormatData xs ys sigma
-}

{-
fit :: [Double]
     -> [Double]
     -> [Double]
     -> ([Double] -> [Double] -> [Double])
     -> ([Double] -> [Double] -> [[Double]])
     -> [Double]->  Maybe FitRes

fit xs ys sigma model jacobian guess
    | xeqy && ( seqx || seqone)  = Just $ fitModelScaled 1E-20 1E-20 500 (model, jacobian) fitdata guess
    | otherwise  = Nothing
    where
    lx = length xs
    ly = length ys
    ls = length sigma
    xeqy = lx == ly
    seqx = ls == lx
    seqone = ls == 1
    fitdata = unsafeFormatData xs ys sigma -}

unsafeFormatData :: [Double] -> [Double] -> [Double] -> FitData
unsafeFormatData xs ys sigma
    |  length sigma == length xs = zip (return <$> xs) $ zip (return <$> ys) sigma
    |  length sigma == 1 = zip (return <$> xs) $ zip (return <$> ys) (repeat (head sigma))

--jac :: (Num a, Traversable f, Functor g) =>
  --     (forall s. Reifies s Tape => f (Reverse s a) -> [Reverse s a] -> g (Reverse s a))
    --   -> f a
      -- -> g (f a)

jac2 :: (Num a, Traversable f1, Traversable f2, Functor g) =>
       (forall s. Reifies s Tape => f1 (Reverse s a) -> f2 (Reverse s a) -> g (Reverse s a))
       -> f1 a
       -> f2 a

       -> g (f1 a)

jac2 f p x = jacobian ((flip f) (fmap auto x)) p

jac3 :: (Floating a, Traversable f1, Traversable f2, Functor g) =>
       (forall s. Reifies s Tape => f1 (Reverse s a) -> f2 (Reverse s a) -> g (Reverse s a))
       -> f1 a
       -> f2 a

       -> g (f1 a)

jac3 f p x = jacobian ((flip f) (fmap auto x)) p

mkJac :: (Num a) => (forall s. Reifies s Tape => [Reverse s a] -> [Reverse s a] -> [Reverse s a])
       -> [a]
       -> [a]

       -> [[a]]

mkJac f p x = jacobian ((flip f) (fmap auto x)) p

-- to fit data
expModel [a,lambda,b] [t] = [a * exp (-lambda * t)+b]

expModelDer [a,lambda,b] [t] = [[exp (-lambda * t), -t * a * exp(-lambda*t) , 1]]

-- to generate data
expModel2 [a,lambda,b] t = a * exp (-lambda * t)+b

expModelDer2 [a,lambda,b] t = [[exp (-lambda * t), -t * a * exp(-lambda*t) , 1]]
-- to go from unbox to box:
-- a = (\p -> \x -> return (expModel2 p (x !! 0)))::Floating a => [a] -> [a] -> [a]

xs = [a | a <- [0..100]]
sigma = [0.1]
ys = (map (expModel2 [5,0.1,1]) xs)


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE InstanceSigs           #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE NamedFieldPuns         #-}
{-# LANGUAGE QuantifiedConstraints  #-}
{-# LANGUAGE RecordWildCards        #-}
{-# LANGUAGE ScopedTypeVariables    #-}

module Q.MonteCarlo where
import           Control.Monad.State
import           Data.RVar
import           Q.Stochastic.Discretize
import           Q.Stochastic.Process
import           Control.Monad
import           Q.ContingentClaim
import Data.Random
import Q.Time
import Data.Time
import Statistics.Distribution (cumulative)
import Statistics.Distribution.Normal (standard)
import Q.ContingentClaim.Options
import Q.Types

type Path b = [(Time, b)]

-- |Summary type class aggregates all priced values of paths
class (PathPricer p)  => Summary m p | m->p where
  -- | Updates summary with given priced pathes
  sSummarize      :: m -> [p] -> m

  -- | Defines a metric, i.e. calculate distance between 2 summaries
  sNorm           :: m -> m -> Double

-- | Path generator is a stochastic path generator
class PathGenerator m where
  pgMkNew         :: m->IO m
  pgGenerate      :: Integer -> m -> Path b

-- | Path pricer provides a price for given path
class PathPricer m where
  ppPrice :: m -> Path b -> m

type MonteCarlo s a = StateT [(Time, s)] RVar a

-- | Generate a single trajectory stopping at each provided time.
trajectory :: forall a b d. (StochasticProcess a b, Discretize d b) =>
             d        -- ^ Discretization scheme
           -> a        -- ^ The stochastic process
           -> b        -- ^ \(S(0)\)
           -> [Time]   -- ^ Stopping points \(\{t_i\}_i^n \) where \(t_i > 0\)
           -> [RVar b] -- ^ \(dW\)s. One for each stopping point.
           -> RVar [b] -- ^ \(S(0) \cup \{S(t_i)\}_i^n \)
trajectory disc p s0 times dws = reverse <$> evalStateT (onePath times dws) initState' where
  initState' :: [(Time, b)]
  initState' = [(0, s0)]

  onePath :: [Time] -> [RVar b] -> MonteCarlo b [b]
  onePath [] _ = do
    s <- get
    return $ map snd s
  onePath (t1:tn) (dw1:dws) = do
    s <- get
    let t0 = head s
    b <- lift $ pEvolve p disc t0 t1 dw1
    put $ (t1, b) : s
    onePath tn dws

-- | Generate multiple trajectories. See 'trajectory'
trajectories:: forall a b d. (StochasticProcess a b, Discretize d b) =>
             Int        -- ^Num of trajectories
           -> d          -- ^Discretization scheme
           -> a          -- ^The stochastic process
           -> b          -- ^\(S(0)\)
           -> [Time]     -- ^Stopping points \(\{t_i\}_i^n \) where \(t_i > 0\)
           -> [RVar b]   -- ^\(dW\)s. One for each stopping point.
           -> RVar [[b]] -- ^\(S(0) \cup \{S(t_i)\}_i^n \)
trajectories n disc p initState times dws = replicateM n $ trajectory disc p initState times dws

observationTimes :: ContingentClaim a -> [Day]
observationTimes = undefined

class Model a b | a -> b where
  discountFactor :: a -> YearFrac -> YearFrac -> RVar Rate
  evolve   :: a -> YearFrac -> StateT (YearFrac, b) RVar Double


-- | see examples/
module DimMat (

   -- * Quasiquotes
   matD,
   blockD,

   -- * Data.Packed.Vector
   (@>),
   -- * Data.Packed.Matrix
   -- ** dimension
   cols, rows,
   colsNT, rowsNT,
   hasRows, hasCols,
   -- (><),
   trans,
   -- reshape, flatten, fromLists, toLists, buildMatrix,
   -- broken
   ToHLists(toHLists), toHList, FromHLists(fromHLists), fromHList,
   (@@>),

   -- asRow, asColumn, fromRows, toRows, fromColumns, toColumns
   -- fromBlocks
#if MIN_VERSION_hmatrix(0,15,0)
   diagBlock,
#endif
   -- toBlocks, toBlocksEvery, repmat, flipud, fliprl
   -- subMatrix, takeRows, dropRows, takeColumns, dropColumns,
   -- extractRows, diagRect, takeDiag, mapMatrix,
   -- mapMatrixWithIndexM, mapMatrixWithIndexM_, liftMatrix,
   -- liftMatrix2, liftMatrix2Auto, fromArray2D,

   ident, -- where to put this?
   -- * Numeric.Container
   -- constant, linspace,
   diag,
   ctrans,
   -- ** Container class
   scalar,
   conj,
   scale, scaleRecip,
   recipMat,
   addConstant,
   add,
   sub,
   mulMat, mulVec,
   divideMat, divideVec,
   equal,
   arctan2,
   hconcat,
   vconcat,
   cmap,
   konst,
   zeroes,
   -- build, atIndex, minIndex, maxIndex, minElement, maxElement,
   -- sumElements, prodElements, step, cond, find, assoc, accum,
   -- Convert
   -- ** Product class
   Dot(..),
   -- absSum, norm1, norm2, normInf,
   -- norm1, normInf,
   pnorm,
   -- optimiseMult, mXm, mXv, vXm, (<.>),
   -- (<>), (<\>), outer, kronecker,
   -- ** Random numbers
   -- ** Element conversion
   -- ** Input/Output
   -- ** Experimental

   -- * Numeric.LinearAlgebra.Algorithms
   -- | incomplete wrapper for "Numeric.LinearAlgebra.Algorithms"

   -- ** Linear Systems
   -- linearSolve, luSolve, cholSolve, linearSolveLS, linearSolveSVD,
   inv,
   PInv(pinv),
   pinvTol,
   det,
   -- invlndet,
   rank,
   -- rcond,
   -- ** Matrix factorizations

   -- *** Singular value decomposition
   -- *** Eigensystems
   -- eigs
   {-
   wrapEig, wrapEigOnly,
   EigV, EigE,
   -- **** eigenvalues and eigenvectors
   eig,
   eigC,
   eigH,
   eigH',
   eigR,
   eigS,
   eigS',
   eigSH,
   eigSH',

   -- **** eigenvalues
   eigOnlyC,
   eigOnlyH,
   eigOnlyR,
   eigOnlyS,
   eigenvalues,
   eigenvaluesSH,
   eigenvaluesSH',
   -}

   -- *** QR
   -- *** Cholesky
   -- *** Hessenberg
   -- *** Schur
   -- *** LU

   -- ** Matrix functions
   -- sqrtm, matFunc
   expm,

   -- ** Nullspace
   -- ** Norms
   -- ** Misc
   -- ** Util

   -- * Automatic Differentiation
   -- ad
#ifdef WITH_Ad
   diff,
#endif

    -- * todo arrange
    DotS,
    MultEq,

    -- * to keep types looking ok
    D,
    module Data.HList.CommonMain,
    Complex,

    -- ** "Numeric.NumType"
    Pos, Neg, Succ, Zero, Neg1,
    ) where

import Numeric.NumType
import DimMat.Internal
import DimMat.QQ
import Data.HList.CommonMain
import Data.Complex


{-# LANGUAGE ScopedTypeVariables #-}

module Statistics.Distribution.Hypergeometric.GenVar.Test where

import Data.Word
import Data.Int

import Statistics.Distribution.Hypergeometric.GenVar

import System.Random.MWC (initialize, GenIO)

import Data.Vector (fromList)

import Test.QuickCheck
import Test.QuickCheck.All
import Test.QuickCheck.Monadic

import Distribution.TestSuite.QuickCheck

testArbitraryRandom :: (GenIO -> IO Bool) -> [Word32] -> Property
testArbitraryRandom f a = monadicIO $ do
  b <- run $ initialize (fromList a) >>= f
  assert b

tests :: IO [Test]
tests = return

  [ testGroup "identities"

    [ testGroup "small"

      [ testProperty "draw nil" $ \p -> testArbitraryRandom $ \g -> do
        let i = (getSmall . getPositive) p :: Int64
        v <- genVar (0, i, i) g
        return $ v == 0

      , testProperty "draw all" $ \p -> testArbitraryRandom $ \g -> do
        let i = (getSmall . getPositive) p :: Int64
        v <- genVar (i, i, i) g
        return $ v == i

      ]

    , testGroup "large"

      [ testProperty "draw nil" $ \p -> testArbitraryRandom $ \g -> do
        let i = (getLarge . getPositive) p :: Int64
        v <- genVar (0, i, i) g
        return $ v == 0

      , testProperty "draw all" $ \p -> testArbitraryRandom $ \g -> do
        let i = (getLarge . getPositive) p :: Int64
        v <- genVar (i, i, i) g
        return $ v == i

      ]

    ]

  ]


{-# LANGUAGE NamedFieldPuns, TemplateHaskell #-}

module School.Train.Test.GradientDescent
( gradientDescentTest ) where

import Conduit ((.|), liftIO, runConduit, sinkNull, yield, yieldMany)
import Data.Either (isLeft)
import Numeric.LinearAlgebra ((><))
import School.TestUtils (addClasses, assertRight, empty, isSorted,
                         randomAffineParams, randomNNInts, randomMatrix,
                         unitCorrect, weight1)
import School.FileIO.AppIO (runAppIO)
import School.FileIO.FileType (FileType, toExtension)
import School.FileIO.FileHeader (FileHeader(..))
import School.FileIO.MatrixSink (matrixIntSink)
import School.Train.AppTrain (AppTrain)
import School.Train.GradientDescent
import School.Train.SimpleDescentUpdate (simpleDescentUpdate)
import School.Train.StoppingCondition (maxIterations)
import School.Train.IterationHandler (storeCost)
import School.Train.TrainState (TrainState(..), HandlerStore(..), def)
import School.Types.PingPong (pingPongSingleton)
import School.Types.DataType (DataType(..))
import School.Types.Error (Error)
import School.Types.FloatEq ((~=))
import School.Unit.CostFunction (CostFunction)
import School.Unit.Affine (affine)
import School.Unit.MultiNoulli (multiNoulli)
import School.Unit.RecLin (recLin)
import School.Unit.UnitParams (UnitParams(..))
import System.Directory (removeFile)
import Test.Tasty (TestTree)
import Test.Tasty.QuickCheck hiding ((><))
import Test.Tasty.TH
import Test.QuickCheck.Monadic (assert, monadicIO)

prop_no_units :: Property
prop_no_units = monadicIO $ do
  result <- liftIO $ gradientDescent (yield empty)
                                     []
                                     weight1
                                     simpleDescentUpdate
                                     mempty
                                     mempty
                                     def
  assert $ isLeft result

prop_iterations :: Positive Int -> Positive Int -> Positive Int -> Property
prop_iterations (Positive n) (Positive b) (Positive f) = monadicIO $ do
  input <- liftIO $ randomMatrix b f
  let source =  yieldMany . repeat $ input
  result <- liftIO $ gradientDescent source
                                    [recLin]
                                    weight1
                                    simpleDescentUpdate
                                    (maxIterations n)
                                    mempty
                                    def
  assertRight ((== n) . iterationCount) result

prop_cost_decline :: Positive Int -> Positive Int -> Positive Int -> Property
prop_cost_decline (Positive b) (Positive f) (Positive o) = monadicIO $ do
  input <- liftIO $ randomMatrix b f
  let source =  yieldMany . repeat $ input
  paramList <- liftIO $ pingPongSingleton <$> randomAffineParams f o
  let initState = def { handlerStore = CostList []
                                  , learningRate = 1e-2
                                  , paramList
                                  }
  result <- liftIO $ gradientDescent source
                                     [affine]
                                     weight1
                                     simpleDescentUpdate
                                     (maxIterations 5)
                                     storeCost
                                     initState
  assertRight ((\(CostList c) -> isSorted c) . handlerStore)
              result

multiSingle :: CostFunction Double (AppTrain a)
multiSingle = multiNoulli Nothing Nothing

prop_multinoulli_single_file :: Positive Int -> Positive Int -> Property
prop_multinoulli_single_file (Positive c) (Positive b) = monadicIO $ do
  activation <- liftIO $ randomMatrix b c
  classes <- liftIO $ randomNNInts (c - 1) b
  let input = addClasses classes activation
  let source =  yieldMany . repeat $ input
  let initState = def { handlerStore = CostList []
                      , paramList = pingPongSingleton EmptyParams
                      }
  result <- liftIO $ gradientDescent source
                                     [unitCorrect classes]
                                     multiSingle
                                     simpleDescentUpdate
                                     (maxIterations 5)
                                     storeCost
                                     initState
  let check = CostList $ replicate 5 (-1)
  assertRight ((~= check) . handlerStore) result

multiTwo :: [Int]
         -> Int
         -> FileType
         -> IO ( Either Error ( FilePath
                              , CostFunction Double (AppTrain a)))
multiTwo classes iterations fType = do
  let fName = "test." ++ (toExtension fType)
  let cols = 1
  let nClasses = length classes
  let rows = iterations * nClasses
  let header = FileHeader { dataType = INT32B, cols, rows }
  let matrix = (rows >< cols) . concat . (replicate iterations)
             . (map fromIntegral) $ classes
  writeRes <- runAppIO . runConduit $ yield matrix
                                   .| matrixIntSink fType header fName
                                   .| sinkNull
  let multi = multiNoulli (Just fName) (Just $ header { rows = nClasses })
  either (return . Left)
         (return . pure . Right $ (fName, multi))
         writeRes

prop_multinoulli_two_files :: Positive Int -> Positive Int -> FileType -> Property
prop_multinoulli_two_files (Positive c) (Positive b) fType = monadicIO $ do
  activation <- liftIO $ randomMatrix b c
  classes <- liftIO $ randomNNInts (c - 1) b
  let input = addClasses classes activation
  let source =  yieldMany . repeat $ input
  let initState = def { handlerStore = CostList []
                      , paramList = pingPongSingleton EmptyParams
                      }
  let iterations = 5
  multiResult <- liftIO $ multiTwo classes iterations fType
  either (\e -> do
            liftIO . putStrLn $ "ERROR " ++ e
            assert False)
         (\(fName, multi) -> do
            result <- liftIO $ gradientDescent source
                                               [unitCorrect classes]
                                               multi
                                               simpleDescentUpdate
                                               (maxIterations iterations)
                                               storeCost
                                               initState
            let check = CostList $ replicate iterations (-1)
            liftIO $ removeFile fName
            assertRight ((~= check) . handlerStore) result
          )
          multiResult

gradientDescentTest :: TestTree
gradientDescentTest = $(testGroupGenerator)


{-# LANGUAGE TemplateHaskell #-}

module RBM where

import           Control.Lens          hiding ((<.>), (|>))
import           Control.Monad.State
import           Data.Foldable
import           Numeric.LinearAlgebra
import           System.Random

type Input = Vector R

toBool :: Functor f => f Double -> f Bool
toBool = fmap (> 0.5)

sign :: (Ord a, Num a, Num t) => a -> t
sign x
  | x >= 0 = 1
  | otherwise = -1

sign01 :: (Ord a, Num a, Num t) => a -> t
sign01 x
  | x >= 0 = 1
  | otherwise = 0

data RBM = RBM { _weights :: Matrix R, _vbias :: Vector R, _hbias :: Vector R } deriving (Show)
makeLenses ''RBM

nrVisible :: Getter RBM Int
nrVisible = to (rows . view weights)

nrHidden :: Getter RBM Int
nrHidden = to (cols . view weights)

nullrbm :: Int -> Int -> RBM
nullrbm visible hidden = RBM ((hidden >< visible) zeros) (visible |> zeros) (hidden |> zeros)
  where
    zeros = repeat 0.0

-- rngrbm :: Int -> Int -> Double -> State StdGen RBM
rngrbm visible hidden rnghi = do
  let nulled = nullrbm visible hidden
      addv = flip addUniformNoise  rnghi
      addm = flip addUniformNoiseM rnghi
  vis <- addv $ nulled ^. vbias
  hid <- addv $ nulled ^. hbias
  wei <- addm $ nulled ^. weights
  return $ nulled & vbias .~ vis & hbias .~ hid & weights .~ wei

sigmoid :: Double -> Double
sigmoid x = 1.0 / (1.0 + exp(-x))

sigmoidvec :: Vector R -> Vector R
sigmoidvec x = 1 / (1 + cmap exp (-x))

nRandoms :: Int -> State StdGen [Double]
nRandoms n = do
  replicateM n getone
  where
    getone = do
      rng <- get
      let (nr, rng') = randomR (0.0, 1.0) rng
      put rng'
      return nr

energy :: RBM -> Vector Double -> Vector Double -> Double
energy rbm v h = - a - b - c
  where
    a = v <.> (rbm ^. vbias)
    b = h <.> (rbm ^. hbias)
    c = (h <# (rbm ^. weights)) <.> v

-- | Requires binary data, i.e. v binary, for mathematical correctness
freeEnergyBin :: RBM -> Vector Double -> Double
freeEnergyBin rbm v = let
  b = rbm ^. vbias
  c = rbm ^. hbias
  w = rbm ^. weights
  lhs = b <.> v
  rhs = sumElements $ 1 + exp (c + w #> v)
  in -lhs -rhs

addUniformNoise :: Vector R -> R -> State StdGen (Vector R)
addUniformNoise v hi = do
  rngs <- nRandoms $ size v
  let rngv = (size v |> rngs) * 2 - 1
  return $ scalar hi * rngv + v

addUniformNoiseM :: Matrix R -> R -> State StdGen (Matrix R)
addUniformNoiseM m hi = do
  let (i,j) = size m
  rngs <- nRandoms $ i * j
  let rngm = ((i><j) rngs) * 2 - 1
  return $ scalar hi * rngm + m

-- | random parameters that deeplearning.net used for initializing weights
goodNoiseParam :: Floating a => RBM -> a
goodNoiseParam rbm = 4 * sqrt (6 / fromIntegral (rbm ^. nrHidden + rbm ^. nrVisible))

-- | Probability of hidden nodes firing given sample v
pHid :: RBM -> Input -> Vector R
pHid rbm v = sigmoidvec $ b' + (w #> v)
  where
    b' = rbm ^. hbias
    w = rbm ^. weights

-- | Probability of visible nodes firing given sample h
pVis :: RBM -> Vector R -> Vector R
pVis rbm h = sigmoidvec $ v' + (h <# w)
  where
    v' = rbm ^. vbias
    w = rbm ^. weights

testfire :: Vector R -> State StdGen (Vector R)
testfire probs = do
  rngs <- nRandoms $ size probs
  return $ cmap sign01 $ probs - vector rngs

inpbook :: [Input]
inpbook = vector <$>
  [[1,1,1,0,0,0],[1,0,1,0,0,0],[1,1,1,0,0,0],[0,0,1,1,1,0], [0,0,1,1,0,0],[0,0,1,1,1,0]]

-- | The algorithm part of the book is sort of half on-line, just
-- generally broken and using weeird notation if they actually mean
-- something that would end up as something like a weight. The code,
-- instead, does a full batch thingy without a single mention of it,
-- since the code would work either way. Anyhow, a weight update
-- scheme that might work for one-input-at-a-time is available at
-- http://image.diku.dk/igel/paper/AItRBM-proof.pdf , which this
-- follows. All three weight diffs for this one input, as (pos - neg),
-- is returned.  I.e: dw (Matrix), dv (vector), dh (vector). No
-- assumptions on h, but that algo seems to imply that it maybe should
-- stay as a probability, and not as sampled. v' is newer than v, etc.
-- Should probably divide this by number of inputs (?).
compCDs :: Vector R -> Vector R -> Vector R -> Vector R -> (Matrix R, Vector R, Vector R)
compCDs v h v' h' = (dw, dv, dh)
  where
    dw = h `outer` v - h' `outer` v'
    dv = v - v'
    dh = h - h'

-- | gibbs starting on v, h from input unused (structured this way for ease of folding)
gibbsv :: RBM -> (Vector R, Vector R) -> State StdGen (Vector R, Vector R)
gibbsv rbm (v, _) = do
  h' <- testfire $ pHid rbm v
  v' <- testfire $ pVis rbm h'
  return (v', h')

getDiffs :: RBM -> Input -> State StdGen (Matrix R, Vector R, Vector R)
getDiffs rbm inp = do
  let v = inp
  h <- testfire $ pHid rbm v
  (v', h') <- foldM (\vh _ -> gibbsv rbm vh) (v, h) [1..1]
  return $ compCDs v h v' h'

train :: R -> RBM -> [Input] -> State StdGen RBM
train eta rbm inputs = do
  diffs <- mapM (getDiffs rbm) inputs
  let (dw, dv, dh) = foldr (\(a, b, c) (a', b', c') -> (a+a', b+b', c+c')) (0,0,0) diffs
      dims = fromIntegral $ length inputs
      scale = eta / dims
      newrbm = rbm & weights %~ (+ (dw * scalar scale))
                   & vbias %~ (+ (dv * scalar scale))
                   & hbias %~ (+ (dh * scalar scale))
  return newrbm

fwd :: RBM -> Input -> State StdGen (Vector R)
fwd rbm inp = do
  -- (v', h') <- gibbsv rbm (inp, inp)
  h' <- testfire $ pHid rbm inp
  return $ pVis rbm h'

test :: IO ()
test = do
  rng <- newStdGen
  let tren = train 0.01
      test = vector [0,0,0,1,1,0]
  print . flip evalState rng $ do
      t0 <- rngrbm 6 4 0.01

      t1 <- tren t0 inpbook
      t2 <- foldM tren t1 (replicate 4000 inpbook)
      t3 <- tren t2 inpbook
      out <- fwd t3 (head inpbook)
      outtest <- fwd t3 test
      return (out, outtest)

-- | We can see in the output of test that the probabilities after
-- forwarding the test vector on the one where it is active is very high,
-- often >0.7, which is completely unlike for (head inpbook), where
-- they are <0.2 (positions 4,5, that is). Similarly the first of the
-- training inputs fires higher probabilities on its activations, and
-- all this very consistently. If we keep gibbs-ing, the probabilities
-- tend towards the relative presence in the training set, which
-- asserts the generative-ness of it. Pos 3 is ~1 and pos 6 ~0 in all
-- cases, which is reasonable for our data.
--
-- Doing it batch-wise would probably be faster, doing persistent CD
-- likewise. Absolutely none of the improvements from hinton's very
-- reasonable practical guide are implemented, and it would improve
-- everything manyfold. (Also, the fold is pointless on the singleton
-- list, but it really seems to be the case that sampling once
-- performs the best)
--
-- As it is though, it seems to do some kind of learning!


--
-- PoolLayer : pooling layer
--

module CNN.PoolLayer (
  poolMax
, depoolMax
, reversePooling
) where

import Numeric.LinearAlgebra

import CNN.Algebra
import CNN.Image
import CNN.LayerType

type Pix = (Double, Double)

{- |
poolMax

  In : pooling size (x = y)
       image

  OUT: updated image
       position of max values

>>> let im = [fromLists [[1.0,2.0,3.0,4.0],[8.0,7.0,6.0,5.0],[1.0,3.0,5.0,7.0],[2.0,4.0,6.0,8.0]], fromLists [[1.0,2.0,3.0,4.0],[2.0,2.0,3.0,4.0],[3.0,3.0,3.0,4.0],[4.0,4.0,4.0,4.0]]]
>>> poolMax 2 im
[[(2><2)
 [ 8.0, 6.0
 , 4.0, 8.0 ],(2><2)
 [ 2.0, 4.0
 , 4.0, 4.0 ]],[(2><2)
 [ 2.0, 2.0
 , 3.0, 3.0 ],(2><2)
 [ 1.0, 1.0
 , 2.0, 1.0 ]]]

-}

poolMax :: Int -> Image -> [Image]
poolMax s im = [os, is]
  where
    pl = head im
    x  = cols pl `div` s
    y  = rows pl `div` s
    ps = [(i*s, j*s) | i <- [0..(x-1)], j <- [0..(y-1)]]
    (os, is) = unzip $ map (toPlain x y . maxPix s ps) im

toPlain :: Int -> Int -> [Pix] -> (Plain, Plain)
toPlain x y pls = ((x><y) op, (x><y) ip)
  where
    (op, ip) = unzip pls

maxPix :: Int -> [(Int, Int)] -> Plain -> [Pix]
maxPix s ps is = map (max' . toPix) ps
  where
    toPix :: (Int, Int) -> [Pix]
    toPix p = zip (concat $ toLists $ subMatrix p (s, s) is) [0.0..]

max' :: [Pix] -> Pix
max' [] = error "empty list!"
max' [x] = x
max' (x:xs) = maximum' x (max' xs)

maximum' :: Pix -> Pix -> Pix
maximum' a@(v1, _) b@(v2, _) = if v1 < v2 then b else a

-- back prop

{- |
depoolMax

>>> let im = [fromLists [[0.0,2.0],[3.0,1.0]]]
>>> let dl = [fromLists [[0.1,0.2],[0.3,0.4]]]
>>> depoolMax 2 im dl
([(4><4)
 [ 0.1, 0.0, 0.0, 0.0
 , 0.0, 0.0, 0.2, 0.0
 , 0.0, 0.0, 0.0, 0.4
 , 0.0, 0.3, 0.0, 0.0 ]],Nothing)

-}

depoolMax :: Int -> Image -> Delta -> (Delta, Maybe Layer)
depoolMax s im d = (zipWith (concatWith (expand s 0.0)) im d, Nothing)
  where
    concatWith :: (Double -> Double -> Matrix R) ->  Matrix R -> Matrix R
               -> Matrix R
    concatWith f i d = fromBlocks $ zipWith (zipWith f) is ds
      where
        is = toLists i
        ds = toLists d

{- |
expand

  IN : size of pooling
       filling value
       positions
       delta values

>>> expand 2 0.0 0 3
(2><2)
 [ 3.0, 0.0
 , 0.0, 0.0 ]
>>> expand 2 0.0 2 4
(2><2)
 [ 0.0, 0.0
 , 4.0, 0.0 ]
>>> expand 3 0.0 5 3
(3><3)
 [ 0.0, 0.0, 0.0
 , 0.0, 0.0, 3.0
 , 0.0, 0.0, 0.0 ]
>>> expand 3 0.0 8 3
(3><3)
 [ 0.0, 0.0, 0.0
 , 0.0, 0.0, 0.0
 , 0.0, 0.0, 3.0 ]

-}

expand :: Int -> Double -> Double -> Double -> Matrix R
expand s r p d = (s><s) $ (replicate p1 r ++ [d] ++ replicate p2 r)
  where
    p1 = truncate p
    p2 = max 0 (s * s - p1 - 1)

-- reverse

reversePooling :: Int -> Layer
reversePooling = MaxPoolLayer


module Main where

import Statistics.Quantile.Bench
import Statistics.Quantile.Exact
import Statistics.Quantile.Util
import Statistics.Quantile.Types

import System.IO

main :: IO ()
main = do
  hSetBuffering stdin LineBuffering
  hSetBuffering stdout LineBuffering
  selectFromHandle median external stdin >>= print


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict           #-}
{-# LANGUAGE StrictData       #-}
module Pinwheel.Transform where

import           Control.DeepSeq
import           Data.Array.IArray          as IA
import           Data.Complex
import           Data.List                  as L
import           Data.Vector.Generic        as VG
import           Data.Vector.Unboxed        as VU
import           FokkerPlanck.FourierSeries (getHarmonics)
import           FokkerPlanck.Histogram
import           Pinwheel.List
import           Sparse.Vector              as SV
import           Utils.List
import           Utils.Parallel             hiding (dot)

data PinwheelTransformData a b =
  PinwheelTransformData a
                        a
                        b

instance (NFData a, NFData b) => NFData (PinwheelTransformData a b) where
  rnf (PinwheelTransformData x y z) = x `seq` y `seq` z `seq` ()

{-# INLINE projectVec #-}
projectVec ::
     ( VG.Vector vector e
     , VG.Vector vector (Complex e)
     , VG.Vector vector Int
     , Num e
     , RealFloat e
     )
  => Int
  -> Int
  -> e
  -> vector (Complex e)
  -> PinwheelTransformData e (SparseVector vector (Complex e))
  -> Complex e
projectVec angularFreq radialFreq sigma pinwheel (PinwheelTransformData logR theta sparseVec) =
  ((exp ((sigma - 1) * logR)) :+ 0) *
  (cis $
   (-1) * (theta * fromIntegral angularFreq + logR * fromIntegral radialFreq)) *
  (sparseVec `dot` pinwheel)

{-# INLINE pinwheelTransform #-}
pinwheelTransform ::
     ( VG.Vector vector e
     , VG.Vector vector (Complex e)
     , VG.Vector vector Int
     , NFData e
     , RealFloat e
     , Unbox e
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> IA.Array (Int, Int) (vector (Complex e))
  -> [PinwheelTransformData e (SparseVector vector (Complex e))]
  -> Histogram (Complex e)
pinwheelTransform numThread maxPhiFreq maxRhoFreq maxThetaFreq maxRFreq sigma pinwheelArr xs =
  let coef =
        parMapChunk
          (ParallelParams numThread undefined)
          rdeepseq
          (\(rFreq, thetaFreq, rhoFreq, phiFreq) ->
             let pinwheel =
                   getHarmonics
                     pinwheelArr
                     (fromIntegral phiFreq)
                     (fromIntegral rhoFreq)
                     (fromIntegral thetaFreq)
                     (fromIntegral rFreq)
             in L.foldl'
                  (\s x -> s + projectVec thetaFreq rFreq sigma pinwheel x)
                  0
                  xs) $
        [ (rFreq, thetaFreq, rhoFreq, phiFreq)
        | rFreq <- [-maxRFreq .. maxRFreq]
        , thetaFreq <- [-maxThetaFreq .. maxThetaFreq]
        , rhoFreq <- [-maxRhoFreq .. maxRhoFreq]
        , phiFreq <- [-maxPhiFreq .. maxPhiFreq]
        ]
  in Histogram
       [ 2 * maxPhiFreq + 1
       , 2 * maxRhoFreq + 1
       , 2 * maxThetaFreq + 1
       , 2 * maxRFreq + 1
       ]
       1 .
     VU.fromList $
     coef


{-# LANGUAGE FlexibleContexts #-}

import           AI.Layer
import           AI.Neuron

import           AI.Network
import           AI.Network.FeedForwardNetwork

import           AI.Trainer
import           AI.Trainer.BackpropTrainer

--import Network.Visualizations
import           Numeric.LinearAlgebra
import           System.IO
import           System.Random

main :: IO ()
main = do

  g <- newStdGen
  let l = LayerDefinition sigmoidNeuron 2 connectFully randomizeFully
  let l' = LayerDefinition sigmoidNeuron 2 connectFully randomizeFully
  let l'' = LayerDefinition sigmoidNeuron 1 connectFully randomizeFully

  let n = createNetwork normals g [l, l', l'']

  let t = BackpropTrainer (3 :: Double) quadraticCost quadraticCost'
  let dat = [(fromList [0, 1], fromList [1]), (fromList [1, 1], fromList [0]), (fromList [1, 0], fromList [1]), (fromList [0, 0], fromList [0])]

  let n' = trainNTimes g n t online dat 1000

  putStrLn "==> XOR predictions: "
  print $ predict (fromList [0, 0]) n'
  print $ predict (fromList [1, 0]) n'
  print $ predict (fromList [0, 1]) n'
  print $ predict (fromList [1, 1]) n'

  saveFeedForwardNetwork "xor.ann" n'

  putStrLn "==> Network saved and reloaded: "
  n'' <- loadFeedForwardNetwork "xor.ann" [l, l', l'']

  print $ predict (fromList [0, 0]) n''
  print $ predict (fromList [1, 0]) n''
  print $ predict (fromList [0, 1]) n''
  print $ predict (fromList [1, 1]) n''

  --networkHistogram "weights.png" weightList n''
  --networkHistogram "biases.png" biasList n''