{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# LANGUAGE FlexibleInstances, OverlappingInstances, ScopedTypeVariables,
    ViewPatterns #-}
module Tests.Distribution (tests) where

import Control.Applicative ((<$), (<$>), (<*>))
import Data.Binary (Binary, decode, encode)
import Data.List (find)
import Data.Typeable (Typeable)
import Statistics.Distribution
import Statistics.Distribution.Beta (BetaDistribution, betaDistr)
import Statistics.Distribution.Binomial (BinomialDistribution, binomial)
import Statistics.Distribution.CauchyLorentz
import Statistics.Distribution.ChiSquared (ChiSquared, chiSquared)
import Statistics.Distribution.Exponential (ExponentialDistribution, exponential)
import Statistics.Distribution.FDistribution (FDistribution, fDistribution)
import Statistics.Distribution.Gamma (GammaDistribution, gammaDistr)
import Statistics.Distribution.Geometric
import Statistics.Distribution.Hypergeometric
import Statistics.Distribution.Normal (NormalDistribution, normalDistr)
import Statistics.Distribution.Poisson (PoissonDistribution, poisson)
import Statistics.Distribution.StudentT
import Statistics.Distribution.Transform (LinearTransform, linTransDistr)
import Statistics.Distribution.Uniform (UniformDistribution, uniformDistr)
import Test.Framework (Test, testGroup)
import Test.Framework.Providers.QuickCheck2 (testProperty)
import Test.QuickCheck as QC
import Test.QuickCheck.Monadic as QC
import Tests.ApproxEq (ApproxEq(..))
import Tests.Helpers (T(..), testAssertion, typeName)
import Tests.Helpers (monotonicallyIncreasesIEEE)
import Text.Printf (printf)
import qualified Control.Exception as E
import qualified Numeric.IEEE as IEEE

-- | Tests for all distributions
tests :: Test
tests = testGroup "Tests for all distributions"
  [ contDistrTests (T :: T BetaDistribution        )
  , contDistrTests (T :: T CauchyDistribution      )
  , contDistrTests (T :: T ChiSquared              )
  , contDistrTests (T :: T ExponentialDistribution )
  , contDistrTests (T :: T GammaDistribution       )
  , contDistrTests (T :: T NormalDistribution      )
  , contDistrTests (T :: T UniformDistribution     )
  , contDistrTests (T :: T StudentT                )
  , contDistrTests (T :: T (LinearTransform StudentT) )
  , contDistrTests (T :: T FDistribution           )

  , discreteDistrTests (T :: T BinomialDistribution       )
  , discreteDistrTests (T :: T GeometricDistribution      )
  , discreteDistrTests (T :: T GeometricDistribution0     )
  , discreteDistrTests (T :: T HypergeometricDistribution )
  , discreteDistrTests (T :: T PoissonDistribution        )

  , unitTests
  ]

----------------------------------------------------------------
-- Tests
----------------------------------------------------------------

-- Tests for continous distribution
contDistrTests :: (Param d, ContDistr d, QC.Arbitrary d, Typeable d, Show d, Binary d, Eq d) => T d -> Test
contDistrTests t = testGroup ("Tests for: " ++ typeName t) $
  cdfTests t ++
  [ testProperty "PDF sanity"              $ pdfSanityCheck     t
  , testProperty "Quantile is CDF inverse" $ quantileIsInvCDF   t
  , testProperty "quantile fails p<0||p>1" $ quantileShouldFail t
  , testProperty "log density check"       $ logDensityCheck    t
  ]

-- Tests for discrete distribution
discreteDistrTests :: (Param d, DiscreteDistr d, QC.Arbitrary d, Typeable d, Show d, Binary d, Eq d) => T d -> Test
discreteDistrTests t = testGroup ("Tests for: " ++ typeName t) $
  cdfTests t ++
  [ testProperty "Prob. sanity"         $ probSanityCheck       t
  , testProperty "CDF is sum of prob."  $ discreteCDFcorrect    t
  , testProperty "Discrete CDF is OK"   $ cdfDiscreteIsCorrect  t
  , testProperty "log probabilty check" $ logProbabilityCheck   t
  ]

-- Tests for distributions which have CDF
cdfTests :: (Param d, Distribution d, QC.Arbitrary d, Show d, Binary d, Eq d) => T d -> [Test]
cdfTests t =
  [ testProperty "C.D.F. sanity"        $ cdfSanityCheck         t
  , testProperty "CDF limit at +inf"    $ cdfLimitAtPosInfinity  t
  , testProperty "CDF limit at -inf"    $ cdfLimitAtNegInfinity  t
  , testProperty "CDF at +inf = 1"      $ cdfAtPosInfinity       t
  , testProperty "CDF at -inf = 1"      $ cdfAtNegInfinity       t
  , testProperty "CDF is nondecreasing" $ cdfIsNondecreasing     t
  , testProperty "1-CDF is correct"     $ cdfComplementIsCorrect t
  , testProperty "Binary OK"            $ p_binary t
  ]

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

-- CDF is in [0,1] range
cdfSanityCheck :: (Distribution d) => T d -> d -> Double -> Bool
cdfSanityCheck _ d x = c >= 0 && c <= 1
  where c = cumulative d x

-- CDF never decreases
cdfIsNondecreasing :: (Distribution d) => T d -> d -> Double -> Double -> Bool
cdfIsNondecreasing _ d = monotonicallyIncreasesIEEE $ cumulative d

-- cumulative d +∞ = 1
cdfAtPosInfinity :: (Param d, Distribution d) => T d -> d -> Bool
cdfAtPosInfinity _ d
  = cumulative d (1/0) == 1

-- cumulative d - ∞ = 0
cdfAtNegInfinity :: (Param d, Distribution d) => T d -> d -> Bool
cdfAtNegInfinity _ d
  = cumulative d (-1/0) == 0

-- CDF limit at +∞ is 1
cdfLimitAtPosInfinity :: (Param d, Distribution d) => T d -> d -> Property
cdfLimitAtPosInfinity _ d =
  okForInfLimit d ==> counterexample ("Last elements: " ++ show (drop 990 probs))
                    $ Just 1.0 == (find (>=1) probs)
  where
    probs = take 1000 $ map (cumulative d) $ iterate (*1.4) 1000

-- CDF limit at -∞ is 0
cdfLimitAtNegInfinity :: (Param d, Distribution d) => T d -> d -> Property
cdfLimitAtNegInfinity _ d =
  okForInfLimit d ==> counterexample ("Last elements: " ++ show (drop 990 probs))
                    $ case find (< IEEE.epsilon) probs of
                        Nothing -> False
                        Just p  -> p >= 0
  where
    probs = take 1000 $ map (cumulative d) $ iterate (*1.4) (-1)

-- CDF's complement is implemented correctly
cdfComplementIsCorrect :: (Distribution d) => T d -> d -> Double -> Bool
cdfComplementIsCorrect _ d x = (eq 1e-14) 1 (cumulative d x + complCumulative d x)

-- CDF for discrete distribution uses <= for comparison
cdfDiscreteIsCorrect :: (DiscreteDistr d) => T d -> d -> Property
cdfDiscreteIsCorrect _ d
  = counterexample (unlines badN)
  $ null badN
  where
    -- We are checking that:
    --
    -- > CDF(i) - CDF(i-e) = P(i)
    --
    -- Apporixmate equality is tricky here. Scale is set by maximum
    -- value of CDF and probability. Case when all proabilities are
    -- zero should be trated specially.
    badN = [ printf "N=%3i    p[i]=%g\tp[i+1]=%g\tdP=%g\trelerr=%g" i p p1 dp ((p1-p-dp) / max p1 dp)
           | i <- [0 .. 100]
           , let p      = cumulative d $ fromIntegral i - 1e-6
                 p1     = cumulative d $ fromIntegral i
                 dp     = probability d i
                 relerr = ((p1 - p) - dp) / max p1 dp
           ,  not (p == 0 && p1 == 0 && dp == 0)
           && relerr > 1e-14
           ]

logDensityCheck :: (ContDistr d) => T d -> d -> Double -> Property
logDensityCheck _ d x
  = counterexample (printf "density    = %g" p)
  $ counterexample (printf "logDensity = %g" logP)
  $ counterexample (printf "log p      = %g" (log p))
  $ counterexample (printf "eps        = %g" (abs (logP - log p) / max (abs (log p)) (abs logP)))
  $ or [ p == 0     && logP == (-1/0)
       , p < 1e-308 && logP < 609
       , eq 1e-14 (log p) logP
       ]
  where
    p    = density d x
    logP = logDensity d x

-- PDF is positive
pdfSanityCheck :: (ContDistr d) => T d -> d -> Double -> Bool
pdfSanityCheck _ d x = p >= 0
  where p = density d x

-- Quantile is inverse of CDF
quantileIsInvCDF :: (Param d, ContDistr d) => T d -> d -> Double -> Property
quantileIsInvCDF _ d (snd . properFraction -> p) =
  p > 0 && p < 1  ==> ( counterexample (printf "Quantile     = %g" q )
                      $ counterexample (printf "Probability  = %g" p )
                      $ counterexample (printf "Probability' = %g" p')
                      $ counterexample (printf "Error        = %e" (abs $ p - p'))
                      $ abs (p - p') < invQuantilePrec d
                      )
  where
    q  = quantile   d p
    p' = cumulative d q

-- Test that quantile fails if p<0 or p>1
quantileShouldFail :: (ContDistr d) => T d -> d -> Double -> Property
quantileShouldFail _ d p =
  p < 0 || p > 1 ==> QC.monadicIO $ do r <- QC.run $ E.catch
                                              (False <$ (return $! quantile d p))
                                              (\(_ :: E.SomeException) -> return True)
                                       QC.assert r

-- Probability is in [0,1] range
probSanityCheck :: (DiscreteDistr d) => T d -> d -> Int -> Bool
probSanityCheck _ d x = p >= 0 && p <= 1
  where p = probability d x

-- Check that discrete CDF is correct
discreteCDFcorrect :: (DiscreteDistr d) => T d -> d -> Int -> Int -> Property
discreteCDFcorrect _ d a b
  = counterexample (printf "CDF   = %g" p1)
  $ counterexample (printf "Sum   = %g" p2)
  $ counterexample (printf "Delta = %g" (abs (p1 - p2)))
  $ abs (p1 - p2) < 3e-10
  -- Avoid too large differeneces. Otherwise there is to much to sum
  --
  -- Absolute difference is used guard againist precision loss when
  -- close values of CDF are subtracted
  where
    n  = min a b
    m  = n + (abs (a - b) `mod` 100)
    p1 = cumulative d (fromIntegral m + 0.5) - cumulative d (fromIntegral n - 0.5)
    p2 = sum $ map (probability d) [n .. m]

logProbabilityCheck :: (DiscreteDistr d) => T d -> d -> Int -> Property
logProbabilityCheck _ d x
  = counterexample (printf "probability    = %g" p)
  $ counterexample (printf "logProbability = %g" logP)
  $ counterexample (printf "log p          = %g" (log p))
  $ counterexample (printf "eps            = %g" (abs (logP - log p) / max (abs (log p)) (abs logP)))
  $ or [ p == 0     && logP == (-1/0)
       , p < 1e-308 && logP < 609
       , eq 1e-14 (log p) logP
       ]
  where
    p    = probability d x
    logP = logProbability d x

p_binary :: (Eq a, Show a, Binary a) => T a -> a -> Bool
p_binary _ a = a == (decode . encode) a

----------------------------------------------------------------
-- Arbitrary instances for ditributions
----------------------------------------------------------------

instance QC.Arbitrary BinomialDistribution where
  arbitrary = binomial <$> QC.choose (1,100) <*> QC.choose (0,1)
instance QC.Arbitrary ExponentialDistribution where
  arbitrary = exponential <$> QC.choose (0,100)
instance QC.Arbitrary GammaDistribution where
  arbitrary = gammaDistr <$> QC.choose (0.1,10) <*> QC.choose (0.1,10)
instance QC.Arbitrary BetaDistribution where
  arbitrary = betaDistr <$> QC.choose (1e-3,10) <*> QC.choose (1e-3,10)
instance QC.Arbitrary GeometricDistribution where
  arbitrary = geometric <$> QC.choose (0,1)
instance QC.Arbitrary GeometricDistribution0 where
  arbitrary = geometric0 <$> QC.choose (0,1)
instance QC.Arbitrary HypergeometricDistribution where
  arbitrary = do l <- QC.choose (1,20)
                 m <- QC.choose (0,l)
                 k <- QC.choose (1,l)
                 return $ hypergeometric m l k
instance QC.Arbitrary NormalDistribution where
  arbitrary = normalDistr <$> QC.choose (-100,100) <*> QC.choose (1e-3, 1e3)
instance QC.Arbitrary PoissonDistribution where
  arbitrary = poisson <$> QC.choose (0,1)
instance QC.Arbitrary ChiSquared where
  arbitrary = chiSquared <$> QC.choose (1,100)
instance QC.Arbitrary UniformDistribution where
  arbitrary = do a <- QC.arbitrary
                 b <- QC.arbitrary `suchThat` (/= a)
                 return $ uniformDistr a b
instance QC.Arbitrary CauchyDistribution where
  arbitrary = cauchyDistribution
                <$> arbitrary
                <*> ((abs <$> arbitrary) `suchThat` (> 0))
instance QC.Arbitrary StudentT where
  arbitrary = studentT <$> ((abs <$> arbitrary) `suchThat` (>0))
instance QC.Arbitrary (LinearTransform StudentT) where
  arbitrary = studentTUnstandardized
           <$> ((abs <$> arbitrary) `suchThat` (>0))
           <*> ((abs <$> arbitrary))
           <*> ((abs <$> arbitrary) `suchThat` (>0))
instance QC.Arbitrary FDistribution where
  arbitrary =  fDistribution
           <$> ((abs <$> arbitrary) `suchThat` (>0))
           <*> ((abs <$> arbitrary) `suchThat` (>0))

-- Parameters for distribution testing. Some distribution require
-- relaxing parameters a bit
class Param a where
  -- Precision for quantileIsInvCDF
  invQuantilePrec :: a -> Double
  invQuantilePrec _ = 1e-14
  -- Distribution is OK for testing limits
  okForInfLimit :: a -> Bool
  okForInfLimit _ = True

instance Param a

instance Param StudentT where
  invQuantilePrec _ = 1e-13
  okForInfLimit   d = studentTndf d > 0.75

instance Param (LinearTransform StudentT) where
  invQuantilePrec _ = 1e-13
  okForInfLimit   d = (studentTndf . linTransDistr) d > 0.75

instance Param FDistribution where
  invQuantilePrec _ = 1e-12

----------------------------------------------------------------
-- Unit tests
----------------------------------------------------------------

unitTests :: Test
unitTests = testGroup "Unit tests"
  [ testAssertion "density (gammaDistr 150 1/150) 1 == 4.883311" $
      4.883311418525483 =~ (density (gammaDistr 150 (1/150)) 1)
    -- Student-T
  , testStudentPDF 0.3  1.34  0.0648215  -- PDF
  , testStudentPDF 1    0.42  0.27058
  , testStudentPDF 4.4  0.33  0.352994
  , testStudentCDF 0.3  3.34  0.757146   -- CDF
  , testStudentCDF 1    0.42  0.626569
  , testStudentCDF 4.4  0.33  0.621739
    -- Student-T General
  , testStudentUnstandardizedPDF 0.3    1.2  4      0.45 0.0533456  -- PDF
  , testStudentUnstandardizedPDF 4.3  (-2.4) 3.22 (-0.6) 0.0971141
  , testStudentUnstandardizedPDF 3.8    0.22 7.62   0.14 0.0490523
  , testStudentUnstandardizedCDF 0.3    1.2  4      0.45 0.458035   -- CDF
  , testStudentUnstandardizedCDF 4.3  (-2.4) 3.22 (-0.6) 0.698001
  , testStudentUnstandardizedCDF 3.8    0.22 7.62   0.14 0.496076
    -- F-distribution
  , testFdistrPDF  1  3   3     (1/(6 * pi)) -- PDF
  , testFdistrPDF  2  2   1.2   0.206612
  , testFdistrPDF  10 12  8     0.000385613179281892790166
  , testFdistrCDF  1  3   3     0.81830988618379067153 -- CDF
  , testFdistrCDF  2  2   1.2   0.545455
  , testFdistrCDF  10 12  8     0.99935509863451408041
  ]
  where
    -- Student-T
    testStudentPDF ndf x exact
      = testAssertion (printf "density (studentT %f) %f ~ %f" ndf x exact)
      $ eq 1e-5  exact  (density (studentT ndf) x)
    testStudentCDF ndf x exact
      = testAssertion (printf "cumulative (studentT %f) %f ~ %f" ndf x exact)
      $ eq 1e-5  exact  (cumulative (studentT ndf) x)
    -- Student-T General
    testStudentUnstandardizedPDF ndf mu sigma x exact
      = testAssertion (printf "density (studentTUnstandardized %f %f %f) %f ~ %f" ndf mu sigma x exact)
      $ eq 1e-5  exact  (density (studentTUnstandardized ndf mu sigma) x)
    testStudentUnstandardizedCDF ndf mu sigma x exact
      = testAssertion (printf "cumulative (studentTUnstandardized %f %f %f) %f ~ %f" ndf mu sigma x exact)
      $ eq 1e-5  exact  (cumulative (studentTUnstandardized ndf mu sigma) x)
    -- F-distribution
    testFdistrPDF n m x exact
      = testAssertion (printf "density (fDistribution %i %i) %f ~ %f [got %f]" n m x exact d)
      $ eq 1e-5  exact d
      where d = density (fDistribution n m) x
    testFdistrCDF n m x exact
      = testAssertion (printf "cumulative (fDistribution %i %i) %f ~ %f [got %f]" n m x exact d)
      $ eq 1e-5  exact d
      where d = cumulative (fDistribution n m) x


{-# LANGUAGE FlexibleContexts #-}
-- | Mixed-radix FFT calculation.
--
-- Arbitrary input vector lengths are handled using a mixed-radix
-- Cooley-Tukey decimation in time algorithm with residual prime
-- length vectors being treated using Rader's algorithm or hand-coded
-- codelets for small primes.
module Numeric.FFT
       ( fft, ifft, fftWith, ifftWith
       , plan, planFromFactors, execute
       , Plan (..), Direction (..), BaseTransform (..)
       ) where

import Prelude hiding (length, map, sum, zipWith)
import Data.Vector.Generic
import Data.Complex

import Numeric.FFT.Types
import Numeric.FFT.Plan
import Numeric.FFT.Execute

-- | Forward FFT with embedded plan calculation.  For an input vector
-- /h/ of length /N/, with entries numbered from 0 to /N - 1/,
-- calculates the entries in /H/, the discrete Fourier transform of
-- /h/, as:
--
-- <<doc-formulae/fft-formula.svg>>
fft :: Vector v (Complex Double) =>
       v (Complex Double) -> IO (v (Complex Double))
fft xs = do
  p <- plan $ length xs
  return $ fftWith p xs

-- | Inverse FFT with embedded plan calculation.  For an input vector
-- /H/ of length /N/, with entries numbered from 0 to /N - 1/,
-- representing Fourier amplitudes of a signal, calculates the entries
-- in /h/, the inverse discrete Fourier transform of /H/, as:
--
-- <<doc-formulae/ifft-formula.svg>>
ifft :: Vector v (Complex Double) =>
        v (Complex Double) -> IO (v (Complex Double))
ifft xs = do
  p <- plan $ length xs
  return $ ifftWith p xs

-- | Forward FFT with pre-computed plan.
fftWith :: Vector v (Complex Double) =>
           Plan -> v (Complex Double) -> v (Complex Double)
fftWith p = convert . execute p Forward . convert

-- | Inverse FFT with pre-computed plan.
ifftWith :: Vector v (Complex Double) =>
            Plan -> v (Complex Double) -> v (Complex Double)
ifftWith p = convert . execute p Inverse . convert


{-# LANGUAGE DuplicateRecordFields     #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TemplateHaskell           #-}

module Minimizer(module Minimizer) where

import           Brain
import           Control.Arrow
import           Control.Lens
import           Control.Monad.Random
import           Convenience
import qualified Data.Vector.Storable           as V
import           GHC.TypeLits
import           Numeric.GSL.Minimization
import           Numeric.GSL.SimulatedAnnealing
import           Simulator
import qualified Graphics.Rendering.Chart.Easy  as E
import qualified Graphics.Rendering.Chart.Gtk   as C
import           Numeric.LinearAlgebra.Data
import           System.IO.Unsafe
import qualified Numeric.AD as AD

type Minimizer n = (forall a. Floating a => Weights n a -> a) -> Weights n Double -> Weights n Double

newtype NelderMeadsParams = NelderMeadsParams Int

instance Default NelderMeadsParams where
  auto = NelderMeadsParams 1000

instance Default SimulatedAnnealingParams where
  auto = SimulatedAnnealingParams 100 100 1.0 1.0 100 1.5 10

newtype AutoDiffParams = AutoDiffParams Int

instance Default AutoDiffParams where
  auto = AutoDiffParams 10

data MinSettings = Annealing SimulatedAnnealingParams | NelderMeads NelderMeadsParams | AutoDiff AutoDiffParams

instance Default MinSettings where
  auto = AutoDiff auto

data Settings = Settings {
    _minSettings      :: MinSettings,
    _seed             :: Int,
    _systems          :: Int,
    _groupedBy        :: Int,
    _iterRange        :: (Int,Int),
    _numMinimizations :: Int
}
makeLenses ''Settings

instance Default Settings where
  auto = Settings {
    _minSettings = auto,
    _seed = 4374653543,
    _systems = 3,
    _groupedBy = 30,
    _iterRange = (400,400),
    _numMinimizations = 5
  }

data NeuralSim system enabled all f = NeuralSim {
    _settings   :: Settings,
    _weights    :: Weights enabled f,
    _restorer   :: Restorer enabled all f,
    _setWeights :: forall a . Floating a => Weights enabled a -> system a -> system a
}
makeLenses ''NeuralSim

minimizeS :: (KnownNat n) => NelderMeadsParams -> Minimizer n
minimizeS (NelderMeadsParams iterations) cost xi =
  minimizeV NMSimplex2 0.0001 iterations (V.replicate (ssize xi) 0.2) (fromVec &. cost) (toVec xi)
  -- & chartify
  & fst & fromVec

chartify :: (a, Matrix Double) -> (a, ())
chartify = second (toColumns &. (!!1)
            &. Numeric.LinearAlgebra.Data.toList
            &. zip [1::Int ..]
            &. (:[]) &. E.line "30" &. E.plot
            &. C.toWindow 500 500 &. unsafePerformIO)
            &. (\t -> seq (snd t) t)

annealing :: (KnownNat n) => SimulatedAnnealingParams -> Minimizer n
annealing anParams cost xi = simanSolve 123 (ssize xi) anParams xi cost metricDist stepFunction (Just $ const "")
  where metricDist a1 a2 = sZipWith (\x1 x2 -> (x1-x2)*(x1-x2)) a1 a2 & ssum & sqrt
        stepFunction rands stepSize current = rands & fromVec &> (\x -> x*2*stepSize - stepSize) & sZipWith (+) current

automaticDiff :: (KnownNat n) => AutoDiffParams -> Minimizer n
automaticDiff (AutoDiffParams iterations) cost (Sized weights) = AD.conjugateGradientDescent (Sized &. cost) weights & take iterations & last & Sized

toMinimizer :: KnownNat n => MinSettings -> Minimizer n
toMinimizer (NelderMeads pars) = minimizeS pars
toMinimizer (Annealing pars)   = annealing pars
toMinimizer (AutoDiff pars)    = automaticDiff pars

minimizer :: (KnownNat all,KnownNat enabled,Simulator system) => NeuralSim system enabled all Double -> IO ()
minimizer n = print (evalRand (minimizeRand n) (mkStdGen $ n^.settings.seed))

minimizeRand :: (KnownNat all, KnownNat enabled, Simulator system) =>
                    NeuralSim system enabled all Double -> Rand StdGen ([(Double,Double)], Weights all Double)
minimizeRand n = composeN (n^.settings.numMinimizations) episodic ([], n^.weights)
                 &> second (n^.restorer)
  where
    episodic (costs,weights) =
      episodicM (toMinimizer (n^.settings.minSettings)) (n^.settings) (n^.setWeights) weights
      &> first (:costs)

episodicM :: forall s enabled. (Simulator s)
          => Minimizer enabled
          -> Settings
          -> (forall a . (Floating a) => Weights enabled a -> s a -> s a)
          -> Weights enabled Double
          -> Rand StdGen ((Double,Double), Weights enabled Double)
episodicM optimizer settins weightSetter initialWeights = do
  iters <- getRandomR (settins^.iterRange)
  randSystems <- getRandoms &> take (settins^.systems)

  let episodeSize = settins^.groupedBy

  let simsEpisode :: Floating a => [s a] -> [s a]
      simsEpisode = apply episodeSize simsStep

  let episodes = iters `div` episodeSize

  let episodedSystems :: Floating a => [[s a]]
      episodedSystems = randSystems
                        &> realToFracSim
                        &  iterate simsEpisode
                        &  take episodes

  let costOfWeights :: (Floating a) => Weights enabled a -> a
      costOfWeights ws = episodedSystems
                         &>>weightSetter ws
                         &> simsEpisode
                         &> simsCost
                         &  sum

  let resultWeights = optimizer costOfWeights initialWeights
  return ((costOfWeights initialWeights,costOfWeights resultWeights), resultWeights)

-- inputRecreator :: _ => Brain n n numWeights -> Minimizer numWeights -> g (Weights numWeights)
-- inputRecreator (Brain feed) optimizer =
--   getRandoms &>> randWeights
--              &> (\inputs -> optimizer (\w -> let outputs = inputs &> feed w
--                                              in zipWith hamming inputs outputs & sum)
--                                       (randWeights 23423))
--   where
--     hamming :: Weights n -> Weights n -> Double
--     hamming w1 w2 = sZipWith (\x y -> (x-y)**2) w1 w2 & sum


module STCR2Z2T0S0EdgeBinary where

import           Control.Arrow
import           Control.Monad
import           Data.Array.Repa         as R
import           Data.Binary             (encodeFile, decodeFile)
import           Data.Complex
import           Data.List               as L
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.Edge
import           Image.IO
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types
import           Utils.Array
import           Utils.Parallel

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:numTrailStr:maxTrailStr:thetaFreqsStr:scaleFreqsStr:hollowRadiusStr:cutoffRadiusStr:histFilePath:filterFileFolder:numIterationStr:writeSourceFlagStr:saveEdgeDataFlagStr:loadEdgeDataFlagStr:edgeFilePath:numNoisePointStr:scaleFactorStr:useFFTWWisdomFlagStr:fftwWisdomFileName:batchSizeStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Double
      scaleFreqs = [-scaleFreq .. scaleFreq]
      hollowRadius = read hollowRadiusStr :: Double
      cutoffRadius = read cutoffRadiusStr :: Double
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      saveEdgeDataFlag = read saveEdgeDataFlagStr :: Bool
      loadEdgeDataFlag = read loadEdgeDataFlagStr :: Bool
      numNoisePoint = read numNoisePointStr :: Int
      scaleFactor = read scaleFactorStr :: Double
      useFFTWWisdomFlag = read useFFTWWisdomFlagStr :: Bool
      batchSize = read batchSizeStr :: Int
      numThread = read numThreadStr :: Int
      parallelParams = ParallelParams numThread batchSize
      folderPath = "output/test/STCR2Z2T0S0EdgeBinary"
      fftwWisdomFilePath = folderPath </> fftwWisdomFileName
      filterFileName =
        printf
          "Filter_%.0f_%.0f_%s"
          hollowRadius
          cutoffRadius
          (takeFileName histFilePath)
      filterFilePath = filterFileFolder </> filterFileName
      biasFilePath =
        printf
          "%s/Bias_%s_%d.dat"
          filterFileFolder
          (takeBaseName edgeFilePath)
          numNoisePoint
      eigenVecFilePath =
        printf
          "%s/EigenVec_%s_%d.dat"
          filterFileFolder
          (takeBaseName edgeFilePath)
          numNoisePoint
  createDirectoryIfMissing True folderPath
  createDirectoryIfMissing True filterFileFolder
  plan <-
    makePlanBinary
      emptyPlan
      useFFTWWisdomFlag
      fftwWisdomFilePath
      (L.length thetaFreqs)
      (L.length scaleFreqs)
      numPoint
      numPoint
  filterFlag <- doesFileExist filterFilePath
  flag <-
    if filterFlag
      then do
        size <- getFileSize filterFilePath
        return $
          if size == 0
            then False
            else True
      else return False
  unless
    flag
    (do histFlag <- doesFileExist histFilePath
        radialArr <-
          if histFlag
            then R.map magnitude . getNormalizedHistogramArr <$>
                 decodeFile histFilePath
            else do
              putStrLn
                "Couldn't find a Green's function data. Start simulation..."
              solveMonteCarloR2Z2T0S0Radial
                numThread
                numTrail
                maxTrail
                numPoint
                numPoint
                thetaSigma
                scaleSigma
                maxScale
                tao
                thetaFreqs
                thetaFreqs
                scaleFreqs
                scaleFreqs
                histFilePath
                (emptyHistogram
                   [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
                   , L.length scaleFreqs
                   , L.length thetaFreqs
                   , L.length scaleFreqs
                   , L.length thetaFreqs
                   ]
                   0)
        putStrLn "Write filter to disk..."
        writeDFTPinwheel
          parallelParams
          plan
          radialArr
          hollowRadius
          cutoffRadius
          thetaFreqs
          scaleFreqs
          maxScale
          (numPoint, numPoint)
          filterFilePath)
  (bias, eigenVec) <-
    if loadEdgeDataFlag
      then do
        bias <- readRepaArray biasFilePath
        eigenVec <- readRepaArray eigenVecFilePath
        return (bias, eigenVec)
      else do
        xs <- parseEdgeFile edgeFilePath
        randomPonintSet <-
          generateRandomPointSet numNoisePoint numPoint numPoint
        let (centerX, centerY) =
              join (***) (\x -> div x . L.length $ xs) .
              L.foldl'
                (\(a, b) (R2S1RPPoint (c, d, _, _)) -> (a + c, b + d))
                (0, 0) $
              xs
            ys =
              L.map
                (\(R2S1RPPoint (a, b, c, d)) ->
                   (R2S1RPPoint
                      ( round $ (fromIntegral $ a - centerX) / scaleFactor
                      , round $ (fromIntegral $ b - centerX) / scaleFactor
                      , c
                      , d)))
                xs
            zs =
              L.map
                (\(R2S1RPPoint (a, b, c, d)) ->
                   (R2S1RPPoint (a - center numPoint, b - center numPoint, c, d)))
                randomPonintSet
            points = ys L.++ zs
        let bias =
              computeBiasR2T0S0 numPoint numPoint thetaFreqs scaleFreqs points
            eigenVec =
              computeInitialEigenVectorBinary
                numPoint
                numPoint
                thetaFreqs
                scaleFreqs
                points
        when saveEdgeDataFlag (writeRepaArray biasFilePath bias)
        return (bias, eigenVec)
  powerMethodBinary
    parallelParams
    plan
    folderPath
    numPoint
    numPoint
    numOrientation
    thetaFreqs
    numScale
    scaleFreqs
    maxScale
    filterFilePath
    numIteration
    writeSourceFlag
    (printf
       "_%.2f_%.2f_%d_%d_%d_%d"
       thetaSigma
       scaleSigma
       (round maxScale :: Int)
       (round tao :: Int)
       (round thetaFreq :: Int)
       (round scaleFreq :: Int))
    saveEdgeDataFlag
    eigenVecFilePath
    bias
    eigenVec


{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE InstanceSigs         #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE UndecidableInstances #-}

module AI.Network.SOM
( SOM(..)
, MapDefinition(..)

, randomNeuron
, makeVectors
, reshapeList
, distance
) where

import           AI.Layer
import           AI.Network
import           AI.Neuron

import           Numeric.LinearAlgebra
import           System.Random

-- | The SOM definition is simple, it only contains a 2-dimensional list of weights
data SOM = SOM { neuronMap :: [[Vector Double]] }

-- | A definitution type for the SOM, it contains the dimensions of each layer (x, y)
--   and the dimension of the input vector (dim)
data MapDefinition = MapDefinition { x        :: Int
                                   , y        :: Int
                                   , inputDim :: Int
                                   }

instance Network SOM where
  type Parameters SOM g = MapDefinition

  predict :: Vector Double -> SOM -> Vector Double
  predict inputs network = inputs

  -- | Create a SOM and initialize it with given weights
  createNetwork :: (RandomGen g) => RandomTransform -> g -> Parameters SOM g -> SOM
  createNetwork transformation g def = SOM randomVectors
    where randomVectors = reshapeList (x def) $
                          makeVectors transformation g (inputDim def)
                          (x def * y def)

-- | A helper function to reshape a 1D list into a 2D list
reshapeList :: Int -> [a] -> [[a]]
reshapeList x [] = [[]]
reshapeList x lst = h : reshapeList x t
  where (h, t) = splitAt x lst

-- | Create a random set of weights for a given neuron
randomNeuron :: (RandomGen g) => RandomTransform -> g -> Int -> Vector Double
randomNeuron transform g inputDim = inputDim |> randomList transform g

-- | Make a 1D list of vectors to be used by the SOM in creating a map of weights
makeVectors :: (RandomGen g) => RandomTransform -> g -> Int -> Int -> [Vector Double]
makeVectors transform g inputDim 0 = []
makeVectors transform g inputDim num = randomNeuron transform g' inputDim :
                                       makeVectors transform g'' inputDim (num - 1)
  where (g', g'') = split g

-- | Calculate the distance between a SOM neuron and an input
distance :: Vector Double -> Vector Double -> Double
distance a b = sqrt $ sum $ map (^2) $ zipWith (-) (toList a) (toList b)


{-# LANGUAGE BangPatterns #-}

module School.Types.Decoding
( binToDouble
, binToInt
, getDouble
, getInt
, binToMatrixDouble
, binToListInt
, runGet
) where

import Control.Monad (replicateM)
import Data.ByteString (ByteString, length, splitAt, unpack)
import Data.Int (Int32)
import Data.Serialize.Get (Get, getInt32be, runGet)
import Data.Serialize.IEEE754 (getFloat64be)
import Numeric.LinearAlgebra ((><), Matrix, I, R)
import Prelude hiding (length, splitAt)
import School.Types.DataType (DataType(..))
import School.Types.Error (Error)

getDouble :: Get R
getDouble = getFloat64be

binToDouble :: ByteString -> Either Error R
binToDouble = runGet getDouble

binToMatrixDouble :: DataType
                  -> Int
                  -> Int
                  -> (ByteString -> Either Error
                                           (Matrix R))
binToMatrixDouble DBL64B nRows nCols =
  runGet (getDoubleMatrixDouble nRows nCols)
binToMatrixDouble INT08B nRows nCols =
   Right
 . (nRows >< nCols)
 . map (fromIntegral . fromEnum)
 . unpack
binToMatrixDouble dType _ _ = const . Left $
  "Decoding to Matrix Double undefined for " ++ show dType

getDoubleMatrixDouble :: Int
                      -> Int
                      -> Get (Matrix Double)
getDoubleMatrixDouble nRows nCols = do
  let nElements = nRows * nCols
  list <- replicateM nElements getDouble
  return $ (nRows >< nCols) list

getInt :: Get Int32
getInt = getInt32be

binToInt :: ByteString -> Either Error I
binToInt = fmap fromIntegral . runGet getInt

binToListInt :: DataType
             -> (ByteString -> Either Error [Int])
binToListInt DBL64B = const . Left $
  "Reject conversion from floating point DBL64B to integral"
binToListInt INT32B = loop [] where
  loop acc bytes = let len = length bytes in
    if len < 4
      then return acc
      else do
        let !(next, current) = splitAt (len - 4) bytes
        int <- fromIntegral <$> binToInt current
        loop (int:acc) next
binToListInt INT08B =
    Right
  . map (fromIntegral . fromEnum)
  . unpack


module PCA.QModels where

import PCA.Distribution
import Numeric.LinearAlgebra.Data
       (Matrix, Vector, ident, konst, size, toRows, fromRows, diag, diagl,
        takeDiag)
import qualified Numeric.LinearAlgebra.HMatrix as H

initQtau :: Distr Double Double
initQtau = Gamma 1.0e-3 1.0e-3

initQmu :: Int -> Distr (Vector Double) (Matrix Double)
initQmu d = MNormal (konst 0 d) (ident d)

initQalpha :: Int -> Distr (Vector Double) (Vector Double)
initQalpha d = MGamma (konst 1.0e-3 d) (konst 1.0e-3 d)

initQW :: Int -> Distr (Matrix Double) (Matrix Double)
initQW d = MatrixMNormal (ident d) (ident d)

initQX :: Matrix Double -> Distr (Matrix Double) (Matrix Double)
initQX trainT = let (_, d) = size trainT
                    sigmaX = calcSigmaX d initQtau (initQW d)
                    mX = calcMX trainT initQtau sigmaX (initQW d) (initQmu d)
                in MatrixMNormal mX sigmaX

type QModel = (Matrix Double, Vector Double)  -- ^ Mode of W and \mu

calculateQ :: Matrix Double -> QModel
calculateQ trainT =
    go 5 initQtau (initQmu d) (initQalpha d) (initQW d) (initQX trainT) trainT
  where
    (n,d) = size trainT
    go 0 _ mu _ w _ _ = (mean w, mean mu)
    go k (Gamma aTau bTau) mu (MGamma aAlpha bAlpha) w x t =
        go (k - 1) ntau nmu nalpha nw nx t
      where
        ntau = Gamma (calcAtau n d aTau) (calcBtau bTau t mu w x)
        nalpha = MGamma (calcAalpha d aAlpha) (calcBalpha bAlpha w)
        nSigmaMu = calcSigmaMu n ntau mu
        nmu = MNormal (calcMMu t ntau nSigmaMu w x) nSigmaMu
        nSigmaW = calcSigmaW ntau nalpha x
        nw = MatrixMNormal (calcMW ntau nSigmaW x t nmu) nSigmaW
        nSigmaX = calcSigmaX d ntau nw
        nx = MatrixMNormal (calcMX t ntau nSigmaX nw nmu) nSigmaX

calcSigmaX
    :: Int
    -> Distr Double Double
    -> Distr (Matrix Double) (Matrix Double)
    -> Matrix Double
calcSigmaX d tau w = H.inv (ident d + mean tau `H.scale` wSquareMean)
  where
    mW = mean w
    wSquareMean = variance w + (H.tr mW H.<> mW)

calcMX
    :: Matrix Double
    -> Distr Double Double
    -> Matrix Double
    -> Distr (Matrix Double) (Matrix Double)
    -> Distr (Vector Double) (Matrix Double)
    -> Matrix Double
calcMX trainT tau sigmaX w mu =
    H.tr
        (mean tau `H.scale` sigmaX H.<> H.tr (mean w) H.<>
         H.tr
             (fromRows
                  (map
                       (\r ->
                             r - mean mu)
                       (toRows trainT))))

calcSigmaMu
    :: Int
    -> Distr Double Double
    -> Distr (Vector Double) (Matrix Double)
    -> Matrix Double
calcSigmaMu n tau mu =
    vMu + 1 / (fromIntegral n * mean tau) `H.scale` ident d
  where
    vMu = variance mu
    (_,d) = size vMu

calcMMu
    :: Matrix Double
    -> Distr Double Double
    -> Matrix Double
    -> Distr (Matrix Double) (Matrix Double)
    -> Distr (Matrix Double) (Matrix Double)
    -> Vector Double
calcMMu trainT tau sigmaMu w x =
    (mean tau `H.scale` sigmaMu) H.#>
    foldr
         (flip (+))
         (konst 0 d)
         (toRows (trainT - H.tr (mW H.<> H.tr mX)))
  where
    mW = mean w
    mX = mean x
    (_, d) = size trainT

calcSigmaW
    :: Distr Double Double
    -> Distr (Vector Double) (Vector Double)
    -> Distr (Matrix Double) (Matrix Double)
    -> Matrix Double
calcSigmaW tau alpha x =
    H.inv
        (diag (mean alpha) +
         mean tau `H.scale`
         (fromIntegral n `H.scale` variance x + H.tr mX H.<> mX))
  where
    mX = mean x
    (n,_) = size mX

calcMW
    :: Distr Double Double
    -> Matrix Double
    -> Distr (Matrix Double) (Matrix Double)
    -> Matrix Double
    -> Distr (Vector Double) (Matrix Double)
    -> Matrix Double
calcMW tau sigmaW x trainT mu =
    (mean tau `H.scale` sigmaW) H.<>
    (H.tr mX H.<>
     (fromRows
          (map
               (\tn ->
                     tn - mMu)
               (toRows trainT))))
  where
    mX = mean x
    mMu = mean mu

calcAalpha :: Int -> Vector Double -> Vector Double
calcAalpha d = H.cmap (+ fromIntegral d / 2)

calcBalpha :: Vector Double
           -> Distr (Matrix Double) (Matrix Double)
           -> Vector Double
calcBalpha b w = b + 0.5 `H.scale` (takeDiag (H.tr mW H.<> mW))
  where
    mW = mean w

calcAtau :: Int -> Int -> Double -> Double
calcAtau n d a = a + fromIntegral (n * d) / 2

-- | TODO: Formula is not correct.
-- Skip part Tr (\langle W^TW \rangle \langle x_n x^T_n 'rangle')
-- because i think it is not correct in that result must be number
calcBtau
    :: Double
    -> Matrix Double
    -> Distr (Vector Double) (Matrix Double)
    -> Distr (Matrix Double) (Matrix Double)
    -> Distr (Matrix Double) (Matrix Double)
    -> Double
calcBtau b trainT mu w x =
    b +
    0.5 *
    (fromIntegral n * norm mMu +
     H.sumElements (takeDiag (trainT H.<> H.tr trainT)) +
     2 * (H.sumElements (H.tr (mW H.<> H.tr mX) H.#> mMu)) -
     2 * (H.sumElements (takeDiag (trainT H.<> mW H.<> H.tr mX))) -
     2 * (H.sumElements (trainT H.#> mMu)))
  where
    (n,_) = size trainT
    norm v = v `H.dot` v
    mMu = mean mu
    mW = mean w
    mX = mean x


{-# LANGUAGE DeriveDataTypeable, DeriveGeneric, FlexibleContexts #-}
-- |
-- Module    : Statistics.Sample.KernelDensity.Simple
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Kernel density estimation code, providing non-parametric ways to
-- estimate the probability density function of a sample.
--
-- The techniques used by functions in this module are relatively
-- fast, but they generally give inferior results to the KDE function
-- in the main 'Statistics.KernelDensity' module (due to the
-- oversmoothing documented for 'bandwidth' below).

module Statistics.Sample.KernelDensity.Simple
    {-# DEPRECATED "Use Statistics.Sample.KernelDensity instead." #-}
    (
    -- * Simple entry points
      epanechnikovPDF
    , gaussianPDF
    -- * Building blocks
    -- These functions may be useful if you need to construct a kernel
    -- density function estimator other than the ones provided in this
    -- module.

    -- ** Choosing points from a sample
    , Points(..)
    , choosePoints
    -- ** Bandwidth estimation
    , Bandwidth
    , bandwidth
    , epanechnikovBW
    , gaussianBW
    -- ** Kernels
    , Kernel
    , epanechnikovKernel
    , gaussianKernel
    -- ** Low-level estimation
    , estimatePDF
    , simplePDF
    -- * References
    -- $references
    ) where

import Data.Data (Data, Typeable)
import GHC.Generics (Generic)
import Numeric.MathFunctions.Constants (m_1_sqrt_2, m_2_sqrt_pi)
import Prelude hiding (sum)
import Statistics.Function (minMax)
import Statistics.Sample (stdDev)
import Statistics.Sample.Internal (sum)
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U

-- | Points from the range of a 'Sample'.
newtype Points = Points {
      fromPoints :: U.Vector Double
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

-- | Bandwidth estimator for an Epanechnikov kernel.
epanechnikovBW :: Double -> Bandwidth
epanechnikovBW n = (80 / (n * m_2_sqrt_pi)) ** 0.2

-- | Bandwidth estimator for a Gaussian kernel.
gaussianBW :: Double -> Bandwidth
gaussianBW n = (4 / (n * 3)) ** 0.2

-- | The width of the convolution kernel used.
type Bandwidth = Double

-- | Compute the optimal bandwidth from the observed data for the
-- given kernel.
--
-- This function uses an estimate based on the standard deviation of a
-- sample (due to Deheuvels), which performs reasonably well for
-- unimodal distributions but leads to oversmoothing for more complex
-- ones.
bandwidth :: G.Vector v Double =>
             (Double -> Bandwidth)
          -> v Double
          -> Bandwidth
bandwidth kern values = stdDev values * kern (fromIntegral $ G.length values)

-- | Choose a uniform range of points at which to estimate a sample's
-- probability density function.
--
-- If you are using a Gaussian kernel, multiply the sample's bandwidth
-- by 3 before passing it to this function.
--
-- If this function is passed an empty vector, it returns values of
-- positive and negative infinity.
choosePoints :: G.Vector v Double =>
                Int             -- ^ Number of points to select, /n/
             -> Double          -- ^ Sample bandwidth, /h/
             -> v Double        -- ^ Input data
             -> Points
choosePoints n h sample = Points . U.map f $ U.enumFromTo 0 n'
  where lo     = a - h
        hi     = z + h
        (a, z) = minMax sample
        d      = (hi - lo) / fromIntegral n'
        f i    = lo + fromIntegral i * d
        n'     = n - 1

-- | The convolution kernel.  Its parameters are as follows:
--
-- * Scaling factor, 1\//nh/
--
-- * Bandwidth, /h/
--
-- * A point at which to sample the input, /p/
--
-- * One sample value, /v/
type Kernel =  Double
            -> Double
            -> Double
            -> Double
            -> Double

-- | Epanechnikov kernel for probability density function estimation.
epanechnikovKernel :: Kernel
epanechnikovKernel f h p v
    | abs u <= 1 = f * (1 - u * u)
    | otherwise  = 0
    where u = (v - p) / (h * 0.75)

-- | Gaussian kernel for probability density function estimation.
gaussianKernel :: Kernel
gaussianKernel f h p v = exp (-0.5 * u * u) * g
    where u = (v - p) / h
          g = f * 0.5 * m_2_sqrt_pi * m_1_sqrt_2

-- | Kernel density estimator, providing a non-parametric way of
-- estimating the PDF of a random variable.
estimatePDF :: G.Vector v Double =>
               Kernel           -- ^ Kernel function
            -> Bandwidth        -- ^ Bandwidth, /h/
            -> v Double         -- ^ Sample data
            -> Points           -- ^ Points at which to estimate
            -> U.Vector Double
estimatePDF kernel h sample
    | n < 2     = errorShort "estimatePDF"
    | otherwise = U.map k . fromPoints
  where
    k p = sum . G.map (kernel f h p) $ sample
    f   = 1 / (h * fromIntegral n)
    n   = G.length sample
{-# INLINE estimatePDF #-}

-- | A helper for creating a simple kernel density estimation function
-- with automatically chosen bandwidth and estimation points.
simplePDF :: G.Vector v Double =>
             (Double -> Double) -- ^ Bandwidth function
          -> Kernel             -- ^ Kernel function
          -> Double             -- ^ Bandwidth scaling factor (3 for a Gaussian kernel, 1 for all others)
          -> Int                -- ^ Number of points at which to estimate
          -> v Double           -- ^ sample data
          -> (Points, U.Vector Double)
simplePDF fbw fpdf k numPoints sample =
    (points, estimatePDF fpdf bw sample points)
  where points = choosePoints numPoints (bw*k) sample
        bw     = bandwidth fbw sample
{-# INLINE simplePDF #-}

-- | Simple Epanechnikov kernel density estimator.  Returns the
-- uniformly spaced points from the sample range at which the density
-- function was estimated, and the estimates at those points.
epanechnikovPDF :: G.Vector v Double =>
                   Int          -- ^ Number of points at which to estimate
                -> v Double     -- ^ Data sample
                -> (Points, U.Vector Double)
epanechnikovPDF = simplePDF epanechnikovBW epanechnikovKernel 1

-- | Simple Gaussian kernel density estimator.  Returns the uniformly
-- spaced points from the sample range at which the density function
-- was estimated, and the estimates at those points.
gaussianPDF :: G.Vector v Double =>
               Int              -- ^ Number of points at which to estimate
            -> v Double         -- ^ Data sample
            -> (Points, U.Vector Double)
gaussianPDF = simplePDF gaussianBW gaussianKernel 3

errorShort :: String -> a
errorShort func = error ("Statistics.KernelDensity." ++ func ++
                        ": at least two points required")

-- $references
--
-- * Deheuvels, P. (1977) Estimation non paramétrique de la densité
--   par histogrammes
--   généralisés. Mhttp://archive.numdam.org/article/RSA_1977__25_3_5_0.pdf>


{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveDataTypeable #-}
#if __GLASGOW_HASKELL__ >= 786
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}
#endif

module Data.Complex.Polar (
#if __GLASGOW_HASKELL__ >= 800
    Polar((:<)),
#elif __GLASGOW_HASKELL >= 786
    Polar,
    pattern (:<),
#else
    Polar,
#endif
    fromPolar,
    fromComplex,
    realPart,
    imagPart,
    conjugate,
    mkPolar,
    cis,
    polar,
    magnitude,
    phase
) where

import           Data.Complex (Complex(..))
import qualified Data.Complex as C

import Data.Typeable
#ifdef __GLASGOW_HASKELL__
import Data.Data (Data)
#endif

#ifdef __HUGS__
import Hugs.Prelude(Num(fromInt), Fractional(fromDouble))
#endif

import Text.Read (parens)
import Text.ParserCombinators.ReadPrec (prec, step)
import Text.Read.Lex (Lexeme (Ident))

infix 6 :<<

-- -----------------------------------------------------------------------------
-- The Polar type

-- | Complex numbers are an algebraic type.
--
-- For a complex number @z@, @'abs' z@ is a number with the magnitude of @z@,
-- but oriented in the positive real direction, whereas @'signum' z@
-- has the phase of @z@, but unit magnitude.
data (RealFloat a) => Polar a = !a :<< !a -- ^ forms a complex number from its magnitude
                                          --   and its phase in radians.
#if __GLASGOW_HASKELL__
    deriving (Eq, Data, Typeable)
#else
    deriving Eq
#endif

#if __GLASGOW_HASKELL__ >= 786
infix 6 :<

-- | Smart constructor that canonicalizes the magnitude and phase.
pattern r :< theta <-
  ( \(r :<< theta) -> Just (r, theta) ->
      Just (r, theta)
    )
  where
    r :< theta = mkPolar r theta
#endif

instance (RealFloat a, Show a) => Show (Polar a) where
    {-# SPECIALISE instance Show (Polar Float) #-}
    {-# SPECIALISE instance Show (Polar Double) #-}
    showsPrec d (r :<< theta) =
      showParen (d >= 11) (showString "mkPolar "
                           . showsPrec 11 r
                           . showString " "
                           . showsPrec 11 theta)

instance (RealFloat a, Read a) => Read (Polar a) where
    {-# SPECIALISE instance Read (Polar Float) #-}
    {-# SPECIALISE instance Read (Polar Double) #-}
    readsPrec d = readParen (d > 10)
                        (\s -> do ("mkPolar", s2) <- lex s
                                  (r, s3) <- readsPrec 11 s2
                                  (theta, s4) <- readsPrec 11 s3
                                  return (mkPolar r theta, s4))

-- | Wrap phase back in interval @(-'pi', 'pi']@.
wrap :: (RealFloat a) => a -> a
{-# SPECIALISE wrap :: Float  -> Float   #-}
{-# SPECIALISE wrap :: Double -> Double  #-}
wrap phi | phi <= (-pi) = wrap (phi+2*pi)
wrap phi | phi > pi     = wrap (phi-2*pi)
wrap phi                = phi

-- | Convert to rectangular form.
fromPolar :: (RealFloat a) => Polar a -> Complex a
{-# INLINE fromPolar #-}
fromPolar p = realPart p :+ imagPart p

-- | Convert to polar form.
fromComplex :: (RealFloat a) => Complex a -> Polar a
{-# INLINE fromComplex #-}
fromComplex = uncurry mkPolar_ . C.polar

-- | Extracts the real part of a complex number.
realPart :: (RealFloat a) => Polar a -> a
realPart (r :<< theta) = r * cos theta

-- | Extracts the imaginary part of a complex number.
imagPart :: (RealFloat a) => Polar a -> a
imagPart (r :<< theta) = r * sin theta

-- | The conjugate of a complex number.
conjugate :: (RealFloat a) => Polar a -> Polar a
{-# SPECIALISE conjugate :: Polar Float  -> Polar Float #-}
{-# SPECIALISE conjugate :: Polar Double -> Polar Double #-}
conjugate (r :<< theta) = mkPolar r (negate theta)

-- | Form a complex number from polar components of magnitude and phase.
-- The magnitude and phase are expected to be in canonical form.
mkPolar_ :: RealFloat a => a -> a -> Polar a
{-# SPECIALISE mkPolar_ :: Float  -> Float  -> Polar Float  #-}
{-# SPECIALISE mkPolar_ :: Double -> Double -> Polar Double #-}
mkPolar_ r theta = r :<< theta

-- | Form a complex number from polar components of magnitude and phase.
mkPolar :: RealFloat a => a -> a -> Polar a
{-# SPECIALISE mkPolar :: Float  -> Float  -> Polar Float  #-}
{-# SPECIALISE mkPolar :: Double -> Double -> Polar Double #-}
mkPolar r theta | r == 0 = 0 :<< 0
mkPolar r theta | r < 0 = mkPolar (- r) (theta + pi)
mkPolar r theta = mkPolar_ r (wrap theta)

-- | @'cis' t@ is a complex value with magnitude @1@
-- and phase @t@ (modulo @2*'pi'@).
cis :: (RealFloat a) => a -> Polar a
{-# SPECIALISE cis :: Float  -> Polar Float  #-}
{-# SPECIALISE cis :: Double -> Polar Double #-}
cis theta = mkPolar 1 theta

-- | The function 'polar' takes a complex number and
-- returns a (magnitude, phase) pair in canonical form:
-- the magnitude is nonnegative, and the phase in the range @(-'pi', 'pi']@;
-- if the magnitude is zero, then so is the phase.
polar :: (RealFloat a) => Polar a -> (a,a)
{-# SPECIALISE polar :: Polar Double -> (Double,Double) #-}
{-# SPECIALISE polar :: Polar Float  -> (Float,Float)   #-}
polar (r :<< theta) = (r,theta)

-- | The nonnegative magnitude of a complex number.
magnitude :: (RealFloat a) => Polar a -> a
{-# SPECIALISE magnitude :: Polar Float  -> Float  #-}
{-# SPECIALISE magnitude :: Polar Double -> Double #-}
magnitude (r :<< _) = r

-- | The phase of a complex number, in the range @(-'pi', 'pi']@.
-- If the magnitude is zero, then so is the phase.
phase :: (RealFloat a) => Polar a -> a
{-# SPECIALISE phase :: Polar Float  -> Float  #-}
{-# SPECIALISE phase :: Polar Double -> Double #-}
phase (_ :<< theta) = theta

instance (RealFloat a) => Num (Polar a) where
    {-# SPECIALISE instance Num (Polar Float)  #-}
    {-# SPECIALISE instance Num (Polar Double) #-}
    z + z'           = fromComplex (fromPolar z + fromPolar z')
    z - z'           = fromComplex (fromPolar z - fromPolar z')
    z * z'           = mkPolar  (magnitude z * magnitude z') (phase z + phase z')
    negate z         = mkPolar  (negate (magnitude z)) (phase z)
    abs z            = mkPolar_ (magnitude z) 0
    signum (0 :<< _) = 0
    signum z         = mkPolar_ 1 (phase z)
    fromInteger      = flip mkPolar 0 . fromInteger

instance (RealFloat a) => Fractional (Polar a) where
    {-# SPECIALISE instance Fractional (Polar Float)  #-}
    {-# SPECIALISE instance Fractional (Polar Double) #-}
    z / z'          = mkPolar (magnitude z / magnitude z') (phase z - phase z')
    fromRational r  = mkPolar (fromRational r) 0

instance  (RealFloat a) => Floating (Polar a)  where
    {-# SPECIALISE instance Floating (Polar Float) #-}
    {-# SPECIALISE instance Floating (Polar Double) #-}
    pi             = mkPolar_ pi 0
    exp (r :<< theta) = mkPolar (exp (r * cos theta)) (r * sin theta)
    log (r :<< theta) = fromComplex (log r :+ theta)

    -- sqrt (0 :<< _)       =  0
    -- sqrt z@(r :<< theta) =  u :+ (if y < 0 then -v else v)
    --                           where (u,v) = if x < 0 then (v',u') else (u',v')
    --                                 v'    = abs y / (u'*2)
    --                                 u'    = sqrt ((magnitude z + abs x) / 2)
    sqrt = fromComplex.sqrt.fromPolar

    -- sin (x:+y)     =  sin x * cosh y :+ cos x * sinh y
    sin = fromComplex.sin.fromPolar

    -- cos (x:+y)     =  cos x * cosh y :+ (- sin x * sinh y)
    cos = fromComplex.cos.fromPolar

    -- tan (x:+y)     =  (sinx*coshy:+cosx*sinhy)/(cosx*coshy:+(-sinx*sinhy))
    --                   where sinx  = sin x
    --                         cosx  = cos x
    --                         sinhy = sinh y
    --                         coshy = cosh y
    tan = fromComplex.tan.fromPolar

    -- sinh (x:+y)    =  cos y * sinh x :+ sin  y * cosh x
    sinh = fromComplex.sinh.fromPolar

    -- cosh (x:+y)    =  cos y * cosh x :+ sin y * sinh x
    cosh = fromComplex.cosh.fromPolar

    -- tanh (x:+y)    =  (cosy*sinhx:+siny*coshx)/(cosy*coshx:+siny*sinhx)
    --                   where siny  = sin y
    --                         cosy  = cos y
    --                         sinhx = sinh x
    --                         coshx = cosh x
    tanh = fromComplex.tanh.fromPolar

    -- asin z@(x:+y)  =  y':+(-x')
    --                   where  (x':+y') = log (((-y):+x) + sqrt (1 - z*z))
    asin = fromComplex.asin.fromPolar

    -- acos z         =  y'':+(-x'')
    --                   where (x'':+y'') = log (z + ((-y'):+x'))
    --                         (x':+y')   = sqrt (1 - z*z)
    acos = fromComplex.acos.fromPolar

    -- atan z@(x:+y)  =  y':+(-x')
    --                   where (x':+y') = log (((1-y):+x) / sqrt (1+z*z))
    atan = fromComplex.atan.fromPolar

    asinh z        =  log (z + sqrt (1+z*z))
    acosh z        =  log (z + (z+1) * sqrt ((z-1)/(z+1)))
    atanh z        =  log ((1+z) / sqrt (1-z*z))


module Distributions where

import Control.Monad (replicateM)
import qualified Control.Monad.Bayes.Class as B
import Control.Monad.Bayes.Class (MonadSample, MonadCond)
import Numeric.SpecFunctions (logGamma)

import qualified Data.Vector as V
import qualified Data.Map as M
import qualified Data.MultiSet as S

import Numeric.Log (Log (Exp))

import SymbolicArithmetic
import Util.Numeric (logFact)
import Util.Bijection

type Distr = Distr' Double

data Distr' d a = Distr
  { sample :: forall m. MonadSample m => m a
  , score :: a -> d
  }

mapOutputDiscrete :: Bijection a b -> Distr' d a -> Distr' d b
mapOutputDiscrete (Bijection to from) (Distr sample score) =
  Distr (to <$> sample) (score . from)

factor :: MonadCond m => Double -> m ()
factor = B.score . Numeric.Log.Exp

negInf :: Double
negInf = - 1 / 0

observe :: MonadCond m => Distr a -> a -> m ()
observe d x = factor (score d x)

data ExpFam a = ExpFam
  { dim :: Int
  , suffStats :: [Exp () Double]
  , logPartition :: [Double] -> Double
  , logBaseMeas :: Exp () Double
  , efSample :: forall m. MonadSample m => [Double] -> m a
  , efRepr :: a -> Double
  , mle :: [Double] -> [Double]
  , efDescribe :: [Double] -> String }

bernoulliEF :: ExpFam Bool
bernoulliEF = ExpFam
  { dim = 1
  , suffStats = [Var ()]
  , logPartition = undefined
  , logBaseMeas = undefined
  , efSample = \[p] -> B.bernoulli p
  , efRepr = \x -> if x then 1 else 0
  , mle = id
  , efDescribe = \[eta] -> let z = exp eta in "Bernoulli(" ++ show (z / (1 + z)) ++ ")"
  }

betaEF :: ExpFam Double
betaEF = ExpFam
  { dim = 2
  , suffStats = let x = Var () in  [log x, log (1 - x)]
  , logPartition = \[a, b] -> logBeta a b
  , logBaseMeas = let x = Var () in - log x - log (1 - x)
  , efSample = \[a, b] -> B.beta a b
  , efRepr = id
  , mle = \[lx, lox] -> undefined
  , efDescribe = \[a, b] -> "Beta(" ++ show a ++ ", " ++ show b ++ ")"
  }

logBeta :: Double -> Double -> Double
logBeta a b = logGamma a + logGamma b - logGamma (a + b)

logChoose :: Int -> Int -> Double
logChoose n k = logFact n - logFact k - logFact (n - k)

normalEF :: ExpFam Double
normalEF = ExpFam
  { dim = 2
  , suffStats = let x = Var () in [x, x^2]
  , logPartition = \[eta1, eta2] -> - eta1^2 / (4 * eta2) - 1 / 2 * log(-2 * eta2)
  , logBaseMeas = - 1 / 2 * log (2 * pi)
  , efSample = \eta -> let (mu, sigma2) = conv eta in
      B.normal mu (sqrt sigma2)
  , efRepr = id
  , mle = \[mu, ex2] -> let sigma2 = ex2 - mu^2 in normalNatParams mu sigma2
  , efDescribe = \eta -> let (mu, sigma2) = conv eta in
    "Normal(" ++ show mu ++ ", " ++ show sigma2 ++ ")"
  }
  where
  conv [eta1, eta2] = let sigma2 = - 1 / (2 * eta2) in (sigma2 * eta1, sigma2)
  conv _ = error "conv"

-- indep :: ExpFam a -> ExpFam b -> ExpFam (a, b)
-- indep efa efb = ExpFam
--   { dim = dim efa + dim efb
--   , suffStats = \(a, b) -> suffStats efa a ++ suffStats efb b
--   , logPartition = \eta -> let (eta1, eta2) = split eta in
--       logPartition efa eta1 + logPartition efb eta2
--   , logBaseMeas = \(a, b) -> logBaseMeas efa a + logBaseMeas efb b
--   , efSample = \eta -> let (eta1, eta2) = split eta in
--       (,) <$> efSample efa eta1 <*> efSample efb eta2
--   , mle = \ss -> let (ss1, ss2) = split ss in
--       mle efa ss1 ++ mle efb ss2
--   , efDescribe = \eta -> let (eta1, eta2) = split eta in
--     "Indep(" ++ efDescribe efa eta1 ++ ", " ++ efDescribe efb eta2 ++ ")"
--   }
--   where
--   split = splitAt (dim efa)

normalNatParams :: Fractional d => d -> d -> [d]
normalNatParams mu sigma2 = [mu / sigma2, -1 / (2 * sigma2)]

dirac :: Eq a => a -> Distr a
dirac x = Distr (return x) (\y -> if x == y then 0 else negInf)

dot :: Num d => [d] -> [d] -> d
dot xs ys = sum (zipWith (*) xs ys)

efToDistr :: ExpFam a -> [Double] -> Distr a
efToDistr ef naturalParams =
  Distr (efSample ef naturalParams)
  (\x -> let ev = eval (\_ -> efRepr ef x) in
    naturalParams `dot` map ev (suffStats ef) - logPartition ef naturalParams + ev (logBaseMeas ef))

bernoulli :: Double -> Distr Bool
bernoulli p = Distr (B.bernoulli p) (\b -> bernoulli_ll p (if b then 1 else 0))

bernoulli01 :: Double -> Distr Int
bernoulli01 p = Distr sam sco where
  sam :: MonadSample m => m Int
  sam = do
    b <- B.bernoulli p
    return (if b then 1 else 0)
  sco 1 = log p
  sco 0 = log (1 - p)
  sco _ = negInf

categorical :: [Double] -> Distr Int
categorical ps = Distr (B.categorical (V.fromList ps)) (\i -> log (ps !! i))

categoricalM :: Ord a => M.Map a Double -> Distr a
categoricalM ps = Distr (fmap (\i -> fst (M.elemAt i ps)) (B.categorical (V.fromList (M.elems ps))))
  (\k -> log (ps M.! k))

bernoulli_ll :: Floating a => a -> a -> a
bernoulli_ll p b = b * log p + (1 - b) * log (1 - p)

-- WARNING: NOT NORMALIZED!
beta_ll :: Floating a => a -> a -> a -> a
beta_ll a b p = a * log p + b * log (1 - p)

beta :: Double -> Double -> Distr Double
beta a b = Distr (B.beta a b) (\x -> beta_ll a b x + logPartition betaEF [a, b])

gaussian_ll :: Floating a => a -> a -> a -> a
gaussian_ll mu sigma2 x = - x^2 / sigma2 + 2 * x * mu / sigma2 - mu^2 / sigma2 - log (2 * pi) - log sigma2

gaussian_ll' :: Floating a => a -> a -> a -> a
gaussian_ll' mu sigma2 x = - (x - mu)^2 / sigma2 - log (2 * pi) - log sigma2

gaussian :: Double -> Double -> Distr Double
gaussian mu sigma2 = Distr (B.normal mu (sqrt sigma2)) (gaussian_ll' mu sigma2)

normal = gaussian

poisson_ll :: Double -> Int -> Double
poisson_ll lambda k = fromIntegral k * log lambda - lambda - logGamma (fromIntegral k + 1)

poisson :: Double -> Distr Int
poisson lambda = Distr (B.poisson lambda) (poisson_ll lambda)

uniform :: Double -> Double -> Distr Double
uniform a b = Distr (B.uniform a b) $ \x ->
  if a <= x && x <= b
    then - log range
    else negInf
  where
  range = b - a

replicateNIID :: Int -> Distr a -> Distr [a]
replicateNIID n d = Distr (replicateM n (sample d)) (\obs ->
  if length obs == n
    then sum (map (score d) obs)
    else (-1 / 0))

uniformIntRange :: Int -> Distr Int
uniformIntRange max = Distr (B.uniformD [0 .. max - 1])
  (\obs -> if 0 <= obs && obs < max then - log (fromIntegral max) else -1 / 0)

bind :: Distr a -> (a -> Distr b) -> Distr (a, b)
bind d f = Distr (do {x <- sample d; y <- sample (f x); pure (x, y)})
                 (\(x, y) -> score d x + score (f x) y)

replicateIID :: forall a. Distr Int -> Distr a -> Distr [a]
replicateIID howMany d = Distr sam obs where
  sam :: MonadSample m => m [a]
  sam = do
    n <- sample howMany
    replicateM n (sample d)
  obs xs = score howMany (length xs) + sum (map (score d) xs)

geometric' :: MonadSample m => m Int
geometric' = do
  x <- B.bernoulli (0.5 :: Double)
  if x then return 0 else fmap (+1) geometric'

exponential :: MonadSample m => Double -> m Double
exponential lambda = do
  u <- B.random
  pure (- log u / lambda)

passert :: MonadCond m => Bool -> m ()
passert True = pure ()
passert False = B.score 0

data Some f where
  Some :: f a -> Some f

exponentialFamilies :: [Some ExpFam]
exponentialFamilies = [Some betaEF, Some bernoulliEF, Some normalEF]

shuffleWithRepeats' :: forall a. Ord a => M.Map a (Distr Int) -> Distr [a]
shuffleWithRepeats' countDistrs = Distr sam sco where
  sam :: MonadSample m => m [a]
  sam = do
    counts <- mapM sample countDistrs
    randomlyInterleave [ replicate n k | (k, n) <- M.toList counts ]
  sco :: [a] -> Double
  sco idxes = let ns = mconcat [ S.singleton i | i <- idxes ] in
    sum (M.mapWithKey (\k d -> score d (length (filter (== k) idxes))) countDistrs)
    + randomlyInterleaveLogPDF (map snd (S.toOccurList ns))

shuffleWithRepeats :: [Distr Int] -> Distr [Int]
shuffleWithRepeats countDistrs = Distr sam sco where
  sam :: MonadSample m => m [Int]
  sam = do
    counts <- mapM sample countDistrs
    randomlyInterleave [ replicate n i | (i, n) <- zip [0..] counts ]
  sco :: [Int] -> Double
  sco idxes = let ns = [ length (filter (== i) idxes) | i <- take (length countDistrs) [0..] ] in
    sum (zipWith score countDistrs ns) + randomlyInterleaveLogPDF ns

randomlyInterleave :: MonadSample m => [[a]] -> m [a]
randomlyInterleave xs = if ntot == 0
  then pure []
  else do
    i <- sample (categorical [ fromIntegral n / fromIntegral ntot | n <- ns ])
    let (before, (v : vs) : after) = splitAt i xs
    (v :) <$> randomlyInterleave (before ++ vs : after)
  where
    ns = map length xs
    ntot = sum ns

logNumPermutationsWithRepeats :: [Int] -> Double
logNumPermutationsWithRepeats ns = logFact (sum ns) - sum [ logFact n | n <- ns ]

randomlyInterleaveLogPDF :: [Int] -> Double
randomlyInterleaveLogPDF ns = - logNumPermutationsWithRepeats ns

shuffleListLogPDF :: Int -> Double
shuffleListLogPDF n = - logFact n


{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE TypeApplications    #-}

{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
module Test.Grenade.Batch where

import           Grenade.Types
import           Grenade.Core.Shape
import           Grenade.Core.Network
import           Grenade.Core.Layer
import           Grenade.Layers.FullyConnected
import           Grenade.Layers.Convolution
import           Grenade.Utils.ListStore

import           Numeric.LinearAlgebra hiding (uniformSample, konst, (===))
import           Numeric.LinearAlgebra.Static as H hiding ((===))
import           Numeric.LinearAlgebra.Data as D hiding ((===))

import           Hedgehog
import           GHC.TypeLits

type FFNetwork = Network '[ FullyConnected 3 5, FullyConnected 5 4 ] '[ 'D1 3, 'D1 5, 'D1 4 ]

prop_networkBatchFeedforward = property $ do
  let bias  :: H.R 5              = H.fromList [1..5]
      bias' :: H.R 4              = H.fromList [1..4]
      acts  :: H.L 5 3            = H.fromList [1..15]
      acts' :: H.L 4 5            = H.fromList [1..20]
      fc    :: FullyConnected 3 5 = FullyConnected (FullyConnected' bias  acts)  mkListStore
      fc'   :: FullyConnected 5 4 = FullyConnected (FullyConnected' bias' acts') mkListStore
      ins   :: [S ('D1 3)]        = [S1D (H.fromList [1, 2, 3]), S1D (H.fromList [4, 5, 6])]
      net   :: FFNetwork          = fc :~> fc' :~> NNil
      (_, outs :: [S ('D1 4)]) = runBatchForwards net ins
      outs' = map (\(S1D v) -> (D.toList . H.extract) v) outs
  outs' === [[986, 2312, 3638, 4964], [2336, 5462, 8588, 11714]]

extractFCGrads :: Gradients '[FullyConnected 3 5, FullyConnected 5 4] -> [(Vector RealNum, Matrix RealNum)]
extractFCGrads ((FullyConnected' wB wN) :/> ((FullyConnected' wB' wN') :/> GNil))
  = [(H.extract wB, H.extract wN), (H.extract wB', H.extract wN')]

prop_networkBackpropCalculatesGradients = property $ do
  let bias  :: H.R 5              = H.fromList [1..5]
      bias' :: H.R 4              = H.fromList [1..4]
      acts  :: H.L 5 3            = H.fromList [1..15]
      acts' :: H.L 4 5            = H.fromList [1..20]
      fc    :: FullyConnected 3 5 = FullyConnected (FullyConnected' bias acts) mkListStore
      fc'   :: FullyConnected 5 4 = FullyConnected (FullyConnected' bias' acts') mkListStore
      ins   :: [S ('D1 3)]        = [S1D (H.fromList [1, 2, 3]), S1D (H.fromList [4, 5, 6])]
      net   :: FFNetwork          = fc :~> fc' :~> NNil
      (tapes, outs :: [S ('D1 4)]) = runBatchForwards net ins
      (grads, vs   :: [S ('D1 3)]) = runBatchBackwards net tapes outs
      (grad,  v    :: S ('D1 3))   = runBackwards net (tapes!!0) (outs!!0)
      (grad', v'   :: S ('D1 3))   = runBackwards net (tapes!!1) (outs!!1)
      grads'                       = map extractFCGrads grads
      grads''                      = map extractFCGrads [grad, grad']
      vs'                          = map (\(S1D vec) -> (D.toList . H.extract) vec) vs
      vs''                         = map (\(S1D vec) -> (D.toList . H.extract) vec) [v, v']

  grads' === grads''
  vs'    === vs''

prop_networkAveragesGradients = property $ do
  let bias  :: H.R 5                 = H.fromList [1..5]
      bias' :: H.R 4                 = H.fromList [1..4]
      acts  :: H.L 5 3               = H.fromList [1..15]
      acts' :: H.L 4 5               = H.fromList [1..20]
      fc    :: FullyConnected 3 5    = FullyConnected (FullyConnected' bias acts) mkListStore
      fc'   :: FullyConnected 5 4    = FullyConnected (FullyConnected' bias' acts') mkListStore
      ins   :: [S ('D1 3)]           = [S1D (H.fromList [1, 2, 3]), S1D (H.fromList [4, 5, 6])]
      net   :: FFNetwork             = fc :~> fc' :~> NNil
      (tapes, outs :: [S ('D1 4)])   = runBatchForwards net ins
      (grads, _    :: [S ('D1 3)])   = runBatchBackwards net tapes outs
      (grad,  _    :: S ('D1 3))     = runBackwards net (tapes!!0) (outs!!0)
      (grad', _    :: S ('D1 3))     = runBackwards net (tapes!!1) (outs!!1)
      rgrad                          = extractFCGrads (reduceGradient @FFNetwork grads)
      [(wB, wN), (wB', wN')]         = extractFCGrads grad
      [(wB'', wN''), (wB''', wN''')] = extractFCGrads grad'
      rgrad'                         = (0.5 * (wB + wB''), 0.5 * (wN + wN''))
      rgrad''                        = (0.5 * (wB' + wB'''), 0.5 * (wN' + wN'''))

  rgrad === [rgrad', rgrad'']

prop_convolutionCalculatesOutputOfBatches = property $ do
  let weights :: H.L 25 1 = H.fromList [1..25]
      convLayer :: Convolution 'WithoutBias 'NoPadding 1 1 5 5 2 2 = Convolution weights mkListStore
      ins :: [S ('D2 11 11)] = [S2D (H.fromList [1..121]), S2D (H.fromList [2..122])]
      (_, outs :: [S ('D2 4 4)]) = runBatchForwards convLayer ins
      outs' = map (\(S2D v) -> (concat . D.toLists . H.extract) v) outs
  (take 9 (concat outs')) === [10925, 11575, 12225, 12875, 18075, 18725, 19375, 20025, 25225]

unwrapGradConv :: ( KnownNat c
              , KnownNat f
              , KnownNat kR
              , KnownNat kC
              , KnownNat sR
              , KnownNat sC
              , KnownNat kF
              , kF ~ (kR * kC * c)) => Convolution' 'WithoutBias 'NoPadding c f kR kC sR sC -> Matrix RealNum
unwrapGradConv (Convolution' mat) = H.extract mat

prop_convolutionBackprop = property $ do
  let weights :: H.L 25 1 = H.fromList [1..25]
      convLayer :: Convolution 'WithoutBias 'NoPadding 1 1 5 5 2 2 = Convolution weights mkListStore
      ins :: [S ('D2 11 11)] = [S2D (H.fromList [1..121]), S2D (H.fromList [2..122])]
      (tapes, outs :: [S ('D2 4 4)]) = runBatchForwards convLayer ins
      (grads, vs   :: [S ('D2 11 11)]) = runBatchBackwards convLayer tapes outs
      (grad,  v    :: S ('D2 11 11))   = runBackwards convLayer (tapes!!0) (outs!!0)
      (grad', v'   :: S ('D2 11 11))   = runBackwards convLayer (tapes!!1) (outs!!1)
      grads'                           = map unwrapGradConv grads
      grads''                          = map unwrapGradConv [grad, grad']
      vs'                              = map (\(S2D u) -> (D.toLists . H.extract) u) vs
      vs''                             = map (\(S2D u) -> (D.toLists . H.extract) u) [v, v']

  grads' === grads''
  vs'    === vs''

prop_convolutionAveragesGradients = property $ do
  let weights :: H.L 25 1                  = H.fromList (concat (replicate 5 [1, 2, 3, 4, 5]))
      convLayer :: Convolution 'WithoutBias 'NoPadding 1 1 5 5 2 2 = Convolution weights mkListStore
      ins :: [S ('D2 11 11)]               = [S2D (H.fromList [1..121]), S2D (H.fromList [2..122])]
      (tapes, outs :: [S ('D2 4 4)])       = runBatchForwards convLayer ins
      (grads, _    :: [S ('D2 11 11)])     = runBatchBackwards convLayer tapes outs
      (grad,  _    :: S ('D2 11 11))       = runBackwards convLayer (tapes!!0) (outs!!0)
      (grad', _    :: S ('D2 11 11))       = runBackwards convLayer (tapes!!1) (outs!!1)
      rgrad                                = unwrapGradConv (reduceGradient @(Convolution 'WithoutBias 'NoPadding 1 1 5 5 2 2) grads)
      w                                    = unwrapGradConv grad
      w'                                   = unwrapGradConv grad'
      rgrad'                       = 0.5 * (w + w')

  rgrad === rgrad'

prop_fullyConnectedCalculatesOutputOfBatches = property $ do
  let bias :: H.R 5 = H.fromList [1..5]
      acts :: H.L 5 3 = H.fromList [1..15]
      fc :: FullyConnected 3 5 = FullyConnected (FullyConnected' bias acts) mkListStore
      ins :: [S ('D1 3)] = [S1D (H.fromList [1, 2, 3]), S1D (H.fromList [4, 5, 6])]
      (_, outs :: [S ('D1 5)]) = runBatchForwards fc ins
      outs' = map (\(S1D v) -> (D.toList . H.extract) v) outs
  outs' === [[15, 34, 53, 72, 91], [33, 79, 125, 171, 217]]

prop_fullyConnectedBackprop = property $ do
  let bias :: H.R 5 = H.fromList [1..5]
      acts :: H.L 5 3 = H.fromList [1..15]
      fc :: FullyConnected 3 5 = FullyConnected (FullyConnected' bias acts) mkListStore
      ins :: [S ('D1 3)] = [S1D (H.fromList [1, 2, 3]), S1D (H.fromList [4, 5, 6])]
      (tapes, outs :: [S ('D1 5)]) = runBatchForwards fc ins
      (grads, vs   :: [S ('D1 3)]) = runBatchBackwards fc tapes outs
      (grad,  v    :: S ('D1 3))   = runBackwards fc (tapes!!0) (outs!!0)
      (grad', v'   :: S ('D1 3))   = runBackwards fc (tapes!!1) (outs!!1)
      unwrapGrad                   = \(FullyConnected' wB wN) -> (H.extract wB, H.extract wN)
      grads'                       = map unwrapGrad grads
      grads''                      = map unwrapGrad [grad, grad']
      vs'                          = map (\(S1D vec) -> (D.toList . H.extract) vec) vs
      vs''                         = map (\(S1D vec) -> (D.toList . H.extract) vec) [v, v']

  grads' === grads''
  vs'    === vs''

prop_fullyConnectedAveragesGradients = property $ do
  let bias :: H.R 5 = H.fromList [1..5]
      acts :: H.L 5 3 = H.fromList [1..15]
      fc :: FullyConnected 3 5 = FullyConnected (FullyConnected' bias acts) mkListStore
      ins :: [S ('D1 3)] = [S1D (H.fromList [1, 2, 3]), S1D (H.fromList [4, 5, 6])]
      (tapes, outs :: [S ('D1 5)]) = runBatchForwards fc ins
      (grads, _    :: [S ('D1 3)]) = runBatchBackwards fc tapes outs
      (grad,  _    :: S ('D1 3))   = runBackwards fc (tapes!!0) (outs!!0)
      (grad', _    :: S ('D1 3))   = runBackwards fc (tapes!!1) (outs!!1)
      f                            = \(FullyConnected' bs ns) -> (H.extract bs, H.extract ns)
      rgrad                        = f (reduceGradient @(FullyConnected 3 5) grads)
      (wB, wN)                     = f grad
      (wB', wN')                   = f grad'
      rgrad'                       = (0.5 * (wB + wB'), 0.5 * (wN + wN'))

  rgrad === rgrad'

tests :: IO Bool
tests = checkParallel $$(discover)


module Main where

import MNIST.Prelude
import Control.Monad.Random
import qualified Numeric.LinearAlgebra.Static as SA

import Grenade

type MNIST
  = Network
    '[ FullyConnected 784 300, Logit
     , FullyConnected 300 100, Logit
     , FullyConnected 100   9, Logit
     ]
    '[ 'D1 784, 'D1 300
     , 'D1 300, 'D1 100
     , 'D1 100, 'D1 9
     , 'D1 9
     ]

randomNet :: MonadRandom m => m MNIST
randomNet = randomNetwork

type Input  = S ('D1 784)
type Output = S ('D1   9)
type DataSet = [(Input, Output)]

trainAll :: LearningParameters -> MNIST -> DataSet -> MNIST
trainAll rate net0 dataset = foldl' trainOne net0 dataset
  where
  trainOne :: MNIST -> (Input, Output) -> MNIST
  trainOne !network (i,o) = train rate network i o

netTrain :: MNIST -> LearningParameters -> Int -> IO MNIST
netTrain net0 rate n = do
  inps <- replicateM n randomTrainingData
  let outs = map randomTrainingLabels inps
  return $ trainAll rate net0 (zip inps outs)
  where
    randomTrainingData :: IO Input
    randomTrainingData = do
      s <- getRandom
      return . S1D $ SA.randomVector s SA.Uniform * 2 - 1

    randomTrainingLabels :: Input -> Output
    randomTrainingLabels (S1D v) = S1D . fromIntegral $ fromEnum isTrue
      where
        isTrue :: Bool
        isTrue = (v `inCircle` (fromRational    0.33, 0.33))
              || (v `inCircle` (fromRational (-0.33), 0.33))

    inCircle :: KnownNat n => R n -> (R n, Double) -> Bool
    inCircle v (o, r) = SA.norm_2 (v - o) <= r

netScore :: MNIST -> IO ()
netScore network = putStrLn . unlines $
  (fmap.fmap) (showNorm . go) testIns

  where
    testIns :: [[(Double, Double)]]
    testIns = [ [ (x,y)  | x <- [0..50] ] | y <- [0..20] ]

    go :: (Double, Double) -> Output
    go (x,y) = runNet network (S1D $ SA.vector [x / 25 - 1, y / 10 - 1])

    showNorm :: Output -> Char
    showNorm (S1D r) = render $ SA.mean r

    render :: Double -> Char
    render n | n <= 0.2  = ' '
             | n <= 0.4  = '.'
             | n <= 0.6  = '-'
             | n <= 0.8  = '='
             | otherwise = '#'

main :: IO ()
main = do
  net0 <- randomNet
  net  <- netTrain net0 params examples
  netScore net
  where
    examples :: Int
    examples = 10000

    params :: LearningParameters
    params = LearningParameters
      { learningRate = 0.01
      , learningMomentum = 0.9
      , learningRegulariser = 0.0005
      }


{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ImpredicativeTypes #-}
{-# LANGUAGE RankNTypes #-}
{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}

module Data.Histogram.Util where

import Control.Foldl (fold)
import qualified Control.Foldl as L
import Data.Histogram (Histogram)
import qualified Data.Histogram as H
import Data.Histogram.Fill
import Data.Monoid
import qualified Data.Vector.Unboxed as V
import Statistics.Distribution
import qualified Statistics.Distribution.Combo as OC
import Statistics.Distribution.Exponential
import Statistics.Distribution.Gamma
import qualified Statistics.Distribution.Laplace as OL
import Statistics.Distribution.Normal
import Statistics.Distribution.Poisson

bin :: BinD
bin = binDn (-10.5) 1 10.5

hist :: BinD -> V.Vector Double -> H.Histogram BinD Double
hist b =
  fillBuilderVec (forceDouble -<< mkSimple b)

-- | make a percentile histogram
perc :: Histogram BinD Double -> Histogram BinD Double
perc h = H.map (\x -> x / H.sum h) h

-- reversing out mean and stdev
-- calculating mean using the trick that the Bin boundaries are constructed to be isomorphic to actual values
meanH :: Histogram BinD Double -> Double
meanH h =
  let t = H.sum h
  in H.bfoldl step begin h /
     t
  where begin = 0.0
        step b bv v = b + bv * v

meanSqH :: Histogram BinD Double -> Double
meanSqH h =
  let t = H.sum h
  in H.bfoldl step begin h /
     t
  where begin = 0.0
        step b bv v =
          b +
          (bv ^
           (2 :: Int)) *
          v

type Pdf = BinD -> [Double] -> V.Vector Double

pdfBin :: (Double -> Double) -> BinD -> V.Vector Double
pdfBin cdf h =
  V.map (\(x,y) -> cdf y - cdf x) binsV
  where binsV =
          binsList h :: V.Vector (Double,Double)

pdfBinNormal :: Pdf
pdfBinNormal b [] =
  pdfBin (cumulative $
          normalDistr 0.0 1.0)
         b
pdfBinNormal b [s] =
  pdfBin (cumulative $
          normalDistr 0.0 s')
         b
  where s' = max 1.0e-5 s
pdfBinNormal b (m:s:_) =
  pdfBin (cumulative $
          normalDistr m s')
         b
  where s' = max 1.0e-5 s

pdfBinLaplace :: Pdf
pdfBinLaplace b (r:l:sl:glue:_) =
  pdfBin (cumulative $
          OL.laplace r' l' sl' glue)
         b
  where l' = max 1.0e-5 l
        r' = max 1.0e-5 r
        sl' = max 1.0e-4 sl
pdfBinLaplace _ _ = mempty

pdfBinExponential :: Pdf
pdfBinExponential b (e:_) =
  pdfBin (cumulative $ exponential e') b
  where e' = max 1.0e-5 e
pdfBinExponential _ _ = mempty

pdfBinPoisson :: Pdf
pdfBinPoisson b (e:_) =
  pdfBin (cumulative $ poisson e') b
  where e' = max 1.0e-5 e
pdfBinPoisson _ _ = mempty

pdfBinGamma :: Pdf
pdfBinGamma b (m:k:v:_) =
  V.map (m *) $
  pdfBin (cumulative $
          gammaDistr k' v')
         b
  where k' = max 1.0e-5 k
        v' = max 1.0e-5 v
pdfBinGamma _ _ = mempty

pdfBinCombo :: Pdf
pdfBinCombo b (split:m:s:r:l:sl:_) =
  pdfBin (cumulative $
          OC.combo split' m s' r' l' sl')
         b
  where split' = max 1.0e-5 split
        s' = max 1.0e-5 s
        l' = max 1.0e-5 l
        r' = max 1.0e-5 r
        sl' = max 1.0e-4 sl
pdfBinCombo _ _ = mempty

err :: Histogram BinD Double -> Pdf -> [Double] -> Double
err h pdf p =
  do let pdfT = pdf (H.bins h) p
         pdfE = H.histData (perc h)
         errV =
           zipWith (-)
                   (V.toList pdfE)
                   (V.toList pdfT)
         err' = fold L.sum (fmap abs errV)
     err'


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}

module MLP.Config (
   readConf,
   mkDataSet
) where

import Data.Aeson (eitherDecode)
import Numeric.LinearAlgebra.Data (loadMatrix, (><), Matrix)
import Numeric.LinearAlgebra.Static (create, L)
import Control.Lens ((^.), foldlOf, (%~), (&))
import MLP.DataSet (toPat, normPat)
import MLP.Network (AllCon, Net, Learn, MLP)
import MLP.Types (trainFile,
                  testFile,
                  topology,
                  normalise,
                  topoFold,
                  DataSet(..),
                  Parameters(..),
                  State(..),
                  Topology,
                  MonadFileSystem(..))
import Control.Monad ((=<<))
import System.FilePath (splitFileName, takeFileName, (</>))
import GHC.TypeLits (KnownNat, Nat)
import Control.Monad.Random.Class (MonadRandom)

liftEither :: (MonadFail m) => Either String b -> m b
liftEither = either fail pure

readConf :: MonadFileSystem m => FilePath -> m Parameters
readConf fl = do
   s <- liftEither . eitherDecode =<< readFileM fl
   pure $ s & trainFile %~ absPath
            & testFile %~ absPath

   where
      (base,_) = splitFileName fl
      absPath x = base </> takeFileName x

mkDataSet :: (KnownNat i, KnownNat o, MonadFileSystem m) => Parameters -> m (DataSet i o)
mkDataSet p = do
      trainSet <- mkPat (p ^. trainFile)
      testSet <- mkPat (p ^. testFile)

      return $
         if norm then
            DataSet (normPat trainSet) (normPat testSet)
         else
            DataSet trainSet testSet
   where
      mkPat fl = readMatrixM fl >>= liftEither . toPat
      norm = p ^. normalise


{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE DeriveGeneric             #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE LambdaCase                #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE StandaloneDeriving        #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeFamilies              #-}
{-# LANGUAGE TypeInType                #-}
{-# LANGUAGE TypeSynonymInstances      #-}
{-# LANGUAGE UndecidableInstances      #-}

module Learn.Neural.Layer.Recurrent.FullyConnected (
    FullyConnectedR
  , FullyConnectedR'
  , CommonMap(..)
  , MapFunc(..)
  ) where

import           Data.Kind
import           Data.Proxy
import           Data.Reflection
import           Data.Singletons.Prelude
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import           Learn.Neural.Layer
import           Learn.Neural.Layer.Mapping
import           Numeric.BLAS
import           Numeric.Backprop
import           Numeric.Backprop.Iso           (iso)
import           Numeric.Backprop.Op
import           Statistics.Distribution
import           Statistics.Distribution.Normal
import qualified Generics.SOP                   as SOP

data FullyConnectedR :: Type

deriving instance Generic (CParam FullyConnectedR b '[i] '[o])
instance SOP.Generic (CParam FullyConnectedR b '[i] '[o])

instance (Num (b '[o,o]), Num (b '[o,i]), Num (b '[o]))
      => Num (CParam FullyConnectedR b '[i] '[o]) where
    FCRP wI1 wS1 b1 + FCRP wI2 wS2 b2 = FCRP (wI1 + wI2) (wS1 + wS2) (b1 + b2)
    FCRP wI1 wS1 b1 - FCRP wI2 wS2 b2 = FCRP (wI1 - wI2) (wS1 - wS2) (b1 - b2)
    FCRP wI1 wS1 b1 * FCRP wI2 wS2 b2 = FCRP (wI1 * wI2) (wS1 * wS2) (b1 * b2)
    negate (FCRP wI wS b) = FCRP (negate wI) (negate wS) (negate b)
    signum (FCRP wI wS b) = FCRP (signum wI) (signum wS) (signum b)
    abs    (FCRP wI wS b) = FCRP (abs    wI) (abs    wS) (abs    b)
    fromInteger x = FCRP (fromInteger x) (fromInteger x) (fromInteger x)

instance (Fractional (b '[o,o]), Fractional (b '[o,i]), Fractional (b '[o]))
      => Fractional (CParam FullyConnectedR b '[i] '[o]) where
    FCRP wI1 wS1 b1 / FCRP wI2 wS2 b2 = FCRP (wI1 / wI2) (wS1 / wS2) (b1 / b2)
    recip (FCRP wI wS b) = FCRP (recip wI) (recip wS) (recip b)
    fromRational x       = FCRP (fromRational x) (fromRational x) (fromRational x)

instance (Floating (b '[o,o]), Floating (b '[o,i]), Floating (b '[o]))
      => Floating (CParam FullyConnectedR b '[i] '[o]) where
    sqrt (FCRP wI wS b) = FCRP (sqrt wI) (sqrt wS) (sqrt b)

instance Num (b '[o]) => Num (CState FullyConnectedR b '[i] '[o]) where
    FCRS s1 + FCRS s2 = FCRS (s1 + s2)
    FCRS s1 - FCRS s2 = FCRS (s1 - s2)
    FCRS s1 * FCRS s2 = FCRS (s1 * s2)
    negate (FCRS s) = FCRS (negate s)
    signum (FCRS s) = FCRS (signum s)
    abs    (FCRS s) = FCRS (abs    s)
    fromInteger x  = FCRS (fromInteger x)

instance Fractional (b '[o]) => Fractional (CState FullyConnectedR b '[i] '[o]) where
    FCRS s1 / FCRS s2 = FCRS (s1 / s2)
    recip (FCRS s)    = FCRS (recip s)
    fromRational x    = FCRS (fromRational x)

instance Floating (b '[o]) => Floating (CState FullyConnectedR b '[i] '[o]) where
    sqrt (FCRS s)    = FCRS (sqrt s)

instance ( BLAS b
         , KnownNat i
         , KnownNat o
         , Floating (b '[o])
         , Floating (b '[o,i])
         , Floating (b '[o,o])
         )
        => Component FullyConnectedR b '[i] '[o] where
    data CParam  FullyConnectedR b '[i] '[o] =
            FCRP { _fcrInpWeights   :: !(b '[o,i])
                 , _fcrStateWeights :: !(b '[o,o])
                 , _fcrBiases       :: !(b '[o])
                 }
    data CState  FullyConnectedR b '[i] '[o] = FCRS { _fcrState :: !(b '[o]) }
    type CConstr FullyConnectedR b '[i] '[o] = (Num (b '[o,i]), Num (b '[o,o]))
    data CConf   FullyConnectedR b '[i] '[o] = forall d. ContGen d => FCRC d

    componentOp = bpOp . withInps $ \(x :< p :< s :< Ø) -> do
        wI :< wS :< b :< Ø <- gTuple #<~ p
        s0 <- opIso (iso _fcrState FCRS) ~$ (s :< Ø)
        y  <- matVecOp ~$ (wI :< x  :< Ø)
        s1 <- matVecOp ~$ (wS :< s0 :< Ø)
        let z = y + s1 + b
        s' <- opIso (iso FCRS _fcrState) ~$ (s1 :< Ø)
        return $ z :< s' :< Ø

    defConf = FCRC (normalDistr 0 0.01)
    initParam = \case
      i `SCons` SNil -> \case
        so@(o `SCons` SNil) -> \(FCRC d) g -> do
          wI <- genA (o `SCons` (i `SCons` SNil)) $ \_ ->
            realToFrac <$> genContVar d g
          wS <- genA (o `SCons` (o `SCons` SNil)) $ \_ ->
            realToFrac <$> genContVar d g
          b <- genA so $ \_ ->
            realToFrac <$> genContVar d g
          return $ FCRP wI wS b
        _ -> error "inaccessible"
      _ -> error "inaccessible"
    initState _ so (FCRC d) g =
        FCRS <$> genA so (\_ -> realToFrac <$> genContVar d g)

instance ( BLAS b
         , KnownNat i
         , KnownNat o
         , Floating (b '[o])
         , Floating (b '[o,i])
         , Floating (b '[o,o])
         )
        => ComponentLayer 'Recurrent FullyConnectedR b '[i] '[o] where
    componentRunMode = RMNotFF

data FullyConnectedR' :: k -> Type

deriving instance Generic (CParam (FullyConnectedR' c) b '[i] '[o])
instance SOP.Generic (CParam (FullyConnectedR' c) b '[i] '[o])

instance (Num (b '[o,o]), Num (b '[o,i]), Num (b '[o]))
      => Num (CParam (FullyConnectedR' s) b '[i] '[o]) where
    FCRP' wI1 wS1 b1 + FCRP' wI2 wS2 b2 = FCRP' (wI1 + wI2) (wS1 + wS2) (b1 + b2)
    FCRP' wI1 wS1 b1 - FCRP' wI2 wS2 b2 = FCRP' (wI1 - wI2) (wS1 - wS2) (b1 - b2)
    FCRP' wI1 wS1 b1 * FCRP' wI2 wS2 b2 = FCRP' (wI1 * wI2) (wS1 * wS2) (b1 * b2)
    negate (FCRP' wI wS b) = FCRP' (negate wI) (negate wS) (negate b)
    signum (FCRP' wI wS b) = FCRP' (signum wI) (signum wS) (signum b)
    abs    (FCRP' wI wS b) = FCRP' (abs    wI) (abs    wS) (abs    b)
    fromInteger x = FCRP' (fromInteger x) (fromInteger x) (fromInteger x)

instance (Fractional (b '[o,o]), Fractional (b '[o,i]), Fractional (b '[o]))
      => Fractional (CParam (FullyConnectedR' s) b '[i] '[o]) where
    FCRP' wI1 wS1 b1 / FCRP' wI2 wS2 b2 = FCRP' (wI1 / wI2) (wS1 / wS2) (b1 / b2)
    recip (FCRP' wI wS b) = FCRP' (recip wI) (recip wS) (recip b)
    fromRational x        = FCRP' (fromRational x) (fromRational x) (fromRational x)

instance (Floating (b '[o,o]), Floating (b '[o,i]), Floating (b '[o]))
      => Floating (CParam (FullyConnectedR' s) b '[i] '[o]) where
    sqrt (FCRP' wI wS b) = FCRP' (sqrt wI) (sqrt wS) (sqrt b)

instance Num (b '[o]) => Num (CState (FullyConnectedR' s) b '[i] '[o]) where
    FCRS' s1 + FCRS' s2 = FCRS' (s1 + s2)
    FCRS' s1 - FCRS' s2 = FCRS' (s1 - s2)
    FCRS' s1 * FCRS' s2 = FCRS' (s1 * s2)
    negate (FCRS' s) = FCRS' (negate s)
    signum (FCRS' s) = FCRS' (signum s)
    abs    (FCRS' s) = FCRS' (abs    s)
    fromInteger x  = FCRS' (fromInteger x)

instance Fractional (b '[o]) => Fractional (CState (FullyConnectedR' s) b '[i] '[o]) where
    FCRS' s1 / FCRS' s2 = FCRS' (s1 / s2)
    recip (FCRS' s)     = FCRS' (recip s)
    fromRational x      = FCRS' (fromRational x)

instance Floating (b '[o]) => Floating (CState (FullyConnectedR' s) b '[i] '[o]) where
    sqrt (FCRS' s)     = FCRS' (sqrt s)

instance ( BLAS b
         , KnownNat i
         , KnownNat o
         , Floating (b '[o])
         , Floating (b '[o,i])
         , Floating (b '[o,o])
         , Reifies s MapFunc
         )
      => Component (FullyConnectedR' s) b '[i] '[o] where
    data CParam  (FullyConnectedR' c) b '[i] '[o] =
            FCRP' { _fcrInpWeights'   :: !(b '[o,i])
                  , _fcrStateWeights' :: !(b '[o,o])
                  , _fcrBiases'       :: !(b '[o])
                  }
    data CState  (FullyConnectedR' c) b '[i] '[o] = FCRS' { _fcrState' :: !(b '[o]) }
    type CConstr (FullyConnectedR' c) b '[i] '[o] =
      ( Num (b '[o,i])
      , Num (b '[o,o])
      )
    data CConf   (FullyConnectedR' c) b '[i] '[o] = forall d. ContGen d => FCRC' d

    componentOp = bpOp . withInps $ \(x :< p :< s :< Ø) -> do
        wI :< wS :< b :< Ø <- gTuple #<~ p
        s0 <- opIso (iso _fcrState' FCRS') ~$ (s :< Ø)
        y  <- matVecOp ~$ (wI :< x  :< Ø)
        s1 <- matVecOp ~$ (wS :< s0 :< Ø)
        s2 <- tmapOp (runMapFunc mf) ~$ (s1 :< Ø)
        let z = y + s2 + b
        s' <- opIso (iso FCRS' _fcrState') ~$ (s2 :< Ø)
        return $ z :< s' :< Ø
      where
        mf :: MapFunc
        mf = reflect (Proxy @s)

    defConf = FCRC' (normalDistr 0 0.01)

    initParam = \case
      i `SCons` SNil -> \case
        so@(o `SCons` SNil) -> \(FCRC' d) g -> do
          wI <- genA (o `SCons` (i `SCons` SNil)) $ \_ ->
            realToFrac <$> genContVar d g
          wS <- genA (o `SCons` (o `SCons` SNil)) $ \_ ->
            realToFrac <$> genContVar d g
          b <- genA so $ \_ ->
            realToFrac <$> genContVar d g
          return $ FCRP' wI wS b
        _ -> error "inaccessible"
      _ -> error "inaccessible"

    initState _ so (FCRC' d) g =
        FCRS' <$> genA so (\_ -> realToFrac <$> genContVar d g)

instance ( BLAS b
         , KnownNat i
         , KnownNat o
         , Floating (b '[o])
         , Floating (b '[o,i])
         , Floating (b '[o,o])
         , Reifies s MapFunc
         )
      => ComponentLayer 'Recurrent (FullyConnectedR' s) b '[i] '[o] where
    componentRunMode = RMNotFF


{-# OPTIONS_GHC  -fno-warn-unused-binds -fno-warn-unused-matches -fno-warn-name-shadowing -fno-warn-missing-signatures #-}
{-# LANGUAGE FlexibleInstances, ConstraintKinds, ExistentialQuantification, GADTs, RankNTypes, MultiParamTypeClasses, RankNTypes, UndecidableInstances, FlexibleContexts, TypeSynonymInstances #-}

---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------
-- |
-- | Module : Approximate median
-- | Creator: Xiao Ling
-- | Created: 12/08/2015
-- | TODO   : test standard deviation of alpha, beta, and final version
-- |
---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------

module ApproxMedian (

      MaxSize
    , Store

    , median
    , median'

  ) where

import Control.Monad.State
import Control.Monad.Random
import Control.Monad.Random.Class

import Data.Conduit
import Data.Foldable (toList)
import qualified Data.Sequence as S
import qualified Data.Conduit.List as Cl
import Data.Sequence (Seq,(|>),update,empty,index)

import Test.QuickCheck
import Test.QuickCheck.Random

import Core
import Utils
import qualified Statistics as St

{-----------------------------------------------------------------------------
  Types
------------------------------------------------------------------------------}

type MaxSize = Float
type Store a = (Seq a, MaxSize)

-- * the algorithm uses `some m` equipped with state and randomness
-- * also consider `some'` `a`
type Some  m = (MonadState Counter m, MonadRandom m)
type Some' a = (Floating a          , Ord a        )

{-----------------------------------------------------------------------------
  Approximate Median
------------------------------------------------------------------------------}

-- * @Use: test (tMedian (E 0.05, D 0.05)) xs
-- * compute `tMedian` within accuracy `Eps` and confidence `Delta`
median' :: Some' a => EpsDelta -> Batch a IO a
median' t = median t `using` eval

-- * `tick` up a counter for each item `a` seen and with probability
-- * min(1, s/i) put item `a` into store `t` of maxSize `s`
median :: (Some m, Some' a) => EpsDelta -> Streaming a m (Store a)
median (ED e d) = Cl.foldM step $ store s
  where s = 7/(e^2) * log (2/d)

-- * a concrete `eval`uation of `m`
eval :: Some' a => StateT Counter (Rand StdGen) (Store a) -> IO a
eval m = fmap (St.median . toList) . fmap fst . evalRandIO $ evalStateT m 0

{-----------------------------------------------------------------------------
  Subroutines. Read $|> to understand the actual algorithm
------------------------------------------------------------------------------}

-- * One step of the algorithm
step :: (Some m, Some' a) => Store a -> a -> m (Store a)
step t a = tick >> t $|> a

-- * with probability `s/i` uniformly select an item from the store `t` and
-- * replace it with `a`.
-- * mnemonic: `|>` is insertion, `$f` means do f probabilistically
infixl 7 $|>
($|>) :: Some m => Store a -> a -> m (Store a)
($|>) t@(_,s) a = do
  i <- get
  if s > i then t $> a else do
    h <- toss . coin $ s/i
    if isHead h then t $> a else return t

{-----------------------------------------------------------------------------
  Utils
------------------------------------------------------------------------------}

-- * Construct an empty store given max size `s`
store :: MaxSize -> Store a
store = (,) empty

-- * Probabilistic insertion. If the store is not full, insert item `a` deterministically
-- * else with uniform probability replace some existing item in store with `a`
-- * mnemonic: `|>` is insertion, `$f` means do f probabilistically
($>) :: Some m => Store a -> a -> m (Store a)
(as,s) $> a | length as < round s = return (as |> a, s)
            | otherwise           = (\x -> (update (round x) a as, s))
                                  <$> getRandomR (0,s :: MaxSize)