text:
{-# LANGUAGE DataKinds, TypeFamilies #-}
module TestOutdated (testTrees) where

import Prelude

import           Control.Arrow (first)
import qualified Data.Strict.Vector as Data.Vector
import qualified Data.Vector.Generic as V
import           Foreign.Storable (Storable)
import           Foreign.Storable.Tuple ()
import           Numeric.LinearAlgebra (Vector)
import qualified Numeric.LinearAlgebra as HM
import           System.Random
import           Test.Tasty
import           Test.Tasty.HUnit hiding (assert)

import HordeAd
import HordeAd.Core.OutdatedOptimizer
import HordeAd.Tool.MnistTools
import TestCommon ((+\), (*\))

type DualNumberD = DualNumber 'DModeGradient Double

type DualNumberVariablesD = DualNumberVariables 'DModeGradient Double

testTrees :: [TestTree]
testTrees = [ fitTests
            , fit2Tests
            , fit2TestsL
            , smartFitTests
            , smartFit2Tests
            , smartFit2TestsL
            , smartFit3Tests
            , smartFit3TestsL
            , stochasticFit3Tests
            ]

scaleDual :: DualMonad d r m => r -> DualNumber d r -> m (DualNumber d r)
scaleDual r u = returnLet $ scale r u

-- Inlined to avoid the tiny overhead of calling an unknown function.
-- This operation is needed, because @sumListDual@ doesn't (always) fuse.
sumResultsDual :: forall d m a r. (DualMonad d r m, Storable a)
               => (a -> m (DualNumber d r))
               -> Vector a
               -> m (DualNumber d r)
{-# INLINE sumResultsDual #-}
sumResultsDual f as = do
  let g :: DualNumber d r -> a -> m (DualNumber d r)
      g !acc a = do
        u <- f a
        return $! acc + u
  sumUs <- V.foldM g 0 as
  returnLet sumUs

lengthDualNumber :: Storable r => DualNumberVariables d r -> Int
lengthDualNumber (vValue, _, _, _, _, _, _, _) = V.length vValue

-- This, and other Fit and Fit2 nn operations, have unfused Delta let-bindings
-- (one binding per each subexpression, even when not needed), which is fine,
-- just not enough for comprehensive benchmarks.
hiddenLayerFit :: forall m. DualMonad 'DModeGradient Double m
               => (DualNumberD -> m DualNumberD)
               -> Double
               -> DualNumberVariablesD
               -> Int
               -> m (Data.Vector.Vector DualNumberD)
hiddenLayerFit factivation x vec width = do
  let f :: Int -> m DualNumberD
      f i = do
        let weight = var0 vec (2 * i)
            bias = var0 vec (2 * i + 1)
        sx <- scaleDual x weight
        sxBias <- sx +\ bias
        factivation sxBias
  V.generateM width f

outputLayerFit :: DualMonad 'DModeGradient Double m
               => (DualNumberD -> m DualNumberD)
               -> Data.Vector.Vector DualNumberD
               -> Int
               -> DualNumberVariablesD
               -> m DualNumberD
outputLayerFit factivation hiddenVec offset vec = do
  outSum <- sumTrainableInputs hiddenVec offset vec
  factivation outSum

nnFit :: DualMonad 'DModeGradient Double m
      => (DualNumberD -> m DualNumberD)
      -> (DualNumberD -> m DualNumberD)
      -> Double -> DualNumberVariablesD -> m DualNumberD
nnFit factivationHidden factivationOutput x vec = do
  -- One bias of the outer layer, a list of weights of the outer layer,
  -- a list of the same length of weights and biases of the hidden layer.
  let width = (lengthDualNumber vec - 1) `div` 3
  hiddenVec <- hiddenLayerFit factivationHidden x vec width
  outputLayerFit factivationOutput hiddenVec (2 * width) vec

nnFitLoss :: DualMonad 'DModeGradient Double m
          => (DualNumberD -> m DualNumberD)
          -> (DualNumberD -> m DualNumberD)
          -> Double -> Double -> DualNumberVariablesD -> m DualNumberD
nnFitLoss factivationHidden factivationOutput x targ vec = do
  res <- nnFit factivationHidden factivationOutput x vec
  lossSquared targ res

nnFitLossTotal :: forall m. DualMonad 'DModeGradient Double m
               => (DualNumberD -> m DualNumberD)
               -> (DualNumberD -> m DualNumberD)
               -> Vector (Double, Double)
               -> DualNumberVariablesD
               -> m DualNumberD
nnFitLossTotal factivationHidden factivationOutput samples vec = do
  let f :: (Double, Double) -> m DualNumberD
      f (x, res) = nnFitLoss factivationHidden factivationOutput x res vec
  sumResultsDual f samples
    -- an implementation that doesn't delve into implementation details
    -- of dual numbers, but uses @sumListDual@ instead is up to twice slower,
    -- due to no list fusion happening across monadic operations

-- We will use the samples with fixed known good seeds, so we don't care
-- whether any first element of the pair (nearly) repeats,
-- creating (nearly) unsatisfiable samples.
--
-- Alas, this happens too often and is hard to pinpoint,
-- so instead let's generate floats and covert to doubles,
-- to ensure at least minimal separation. Without this,
-- many tests that worked with floats, don't work with doubles,
-- and now more old tests pass and new tests pass with more CPU
-- usage that couldn't pass with floats no matter what, due to numeric errors.
wsFit :: (Float, Float) -> Int -> Int
      -> Vector (Double, Double)
wsFit range seed k =
  let rolls :: RandomGen g => g -> Vector Double
      rolls = V.unfoldrExactN k (first realToFrac . uniformR range)
      (g1, g2) = split $ mkStdGen seed
  in V.zip (rolls g1) (rolls g2)

-- Here a huge separation is ensured.
wsFitSeparated :: (Double, Double) -> Int -> Int
               -> Vector (Double, Double)
wsFitSeparated range@(low, hi) seed k =
  let rolls :: RandomGen g => g -> Vector Double
      rolls = V.unfoldrExactN k (uniformR range)
      width = hi - low
      steps = V.generate k (\n ->
        low + fromIntegral n * width / (fromIntegral k - 1))
      g = mkStdGen seed
  in V.zip steps (rolls g)

gdSimpleShow :: HasDelta r
             => r
             -> (DualNumberVariables 'DModeGradient r -> DualMonadGradient r (DualNumber 'DModeGradient r))
             -> Domain0 r
             -> Int
             -> ([r], r)
gdSimpleShow gamma f initVec n =
  let (res, _, _, _) = gdSimple gamma f n (initVec, V.empty, V.empty, V.empty)
      (_, value) = dReverse 1 f (res, V.empty, V.empty, V.empty)
  in (V.toList res, value)

gdSimpleTestCase
  :: Num a
  => String
  -> ((a, a) -> Int -> Int
      -> Vector (Double, Double))
  -> ((DualNumberD -> DualMonadGradient Double DualNumberD)
      -> (DualNumberD -> DualMonadGradient Double DualNumberD)
      -> Vector (Double, Double)
      -> DualNumberVariablesD
      -> DualMonadGradient Double DualNumberD)
  -> Int -> Int -> Int -> Double -> Int -> Double
  -> TestTree
gdSimpleTestCase prefix sampleFunction lossFunction
                 seedSamples nSamples nParams0 gamma nIterations expected =
  let samples = sampleFunction (-1, 1) seedSamples nSamples
      vec = V.unfoldrExactN nParams0 (uniformR (-1, 1)) $ mkStdGen 33
      name = prefix ++ " "
             ++ unwords [ show seedSamples, show nSamples
                        , show nParams0, show gamma, show nIterations ]
  in testCase name $
       snd (gdSimpleShow gamma (lossFunction tanhAct tanhAct samples)
                         vec nIterations)
       @?= expected

gdSimpleWsTestCase
  :: ((DualNumberD -> DualMonadGradient Double DualNumberD)
      -> (DualNumberD -> DualMonadGradient Double DualNumberD)
      -> Vector (Double, Double)
      -> DualNumberVariablesD
      -> DualMonadGradient Double DualNumberD)
  -> Int -> Int -> Int -> Double -> Int -> Double
  -> TestTree
gdSimpleWsTestCase = gdSimpleTestCase "gdSimple Ws" wsFit

gdSimpleSeparatedTestCase
  :: ((DualNumberD -> DualMonadGradient Double DualNumberD)
      -> (DualNumberD -> DualMonadGradient Double DualNumberD)
      -> Vector (Double, Double)
      -> DualNumberVariablesD
      -> DualMonadGradient Double DualNumberD)
  -> Int -> Int -> Int -> Double -> Int -> Double
  -> TestTree
gdSimpleSeparatedTestCase = gdSimpleTestCase "gdSimple Separated" wsFitSeparated

lenP :: Int -> Int
lenP width = 3 * width + 1

-- Some tests here fail due to not smart enough gradient descent
-- (overshooting, mostly).
fitTests :: TestTree
fitTests = testGroup "Sample fitting fully connected neural net tests"
  [ testCase "wsFit (-1, 1) 42 20" $
      V.toList (wsFit (-1, 1) 42 20) @?= [(-0.22217941284179688,-0.5148218870162964),(0.25622618198394775,0.42662060260772705),(7.794177532196045e-2,-0.5301129817962646),(0.384537935256958,0.8958269357681274),(-0.6027946472167969,-0.5425337553024292),(0.4734766483306885,0.19495820999145508),(0.3921601474285126,0.8963258266448975),(-2.679157257080078e-2,-0.43389952182769775),(-8.326125144958496e-2,-0.17110145092010498),(-6.933605670928955e-2,-0.6602561473846436),(-0.7554467916488647,0.9077622890472412),(-0.17885446548461914,0.14958932995796204),(-0.49340176582336426,0.13965561985969543),(0.4703446626663208,-0.487585186958313),(-0.37681376934051514,-0.39065873622894287),(-0.9820539951324463,-0.10905027389526367),(0.6628230810165405,0.11808493733406067),(4.337519407272339e-3,-7.50422477722168e-3),(-0.270332932472229,0.9103447198867798),(2.815529704093933e-2,-0.9941539764404297)]
  , gdSimpleWsTestCase
      nnFitLossTotal 42 8 (lenP 10) 0.1 10000 1.6225349272445413e-2
  , gdSimpleWsTestCase
      nnFitLossTotal 42 10 (lenP 10) 0.01 100000 0.11821681957239855
      -- failed; even 61 1000000 was not enough
  , testCase "wsFitSeparated (-1, 1) 42 20" $
      V.toList (wsFitSeparated (-1, 1) 42 20) @?= [(-1.0,0.8617048050432681),(-0.8947368421052632,-0.12944690839124995),(-0.7894736842105263,0.7709385349363602),(-0.6842105263157895,0.7043981517795999),(-0.5789473684210527,0.5002907568304664),(-0.4736842105263158,-0.20067467322001753),(-0.368421052631579,-5.526582421799997e-2),(-0.26315789473684215,0.3006213813725571),(-0.1578947368421053,0.12350686811659489),(-5.2631578947368474e-2,-0.7621608299731257),(5.263157894736836e-2,-3.550743010902346e-2),(0.1578947368421053,-0.32868601453242263),(0.26315789473684204,0.7784360517385773),(0.368421052631579,-0.6715107907491862),(0.4736842105263157,-0.41971965075782536),(0.5789473684210527,-0.4920995297212283),(0.6842105263157894,-0.8809132509345221),(0.7894736842105263,-7.615997455596313e-2),(0.894736842105263,0.36412224491658224),(1.0,-0.31352088018219515)]
  , gdSimpleSeparatedTestCase
      nnFitLossTotal 42 8 (lenP 10) 0.1 10000 0.3884360171054549
  , gdSimpleSeparatedTestCase
      nnFitLossTotal 42 10 (lenP 10) 0.01 100000 1.9817301995554423e-2
  , gdSimpleSeparatedTestCase
      nnFitLossTotal 42 16 (lenP 10) 0.01 100000 6.88603932297595e-2
  ]

-- This connects with only one neuron from the first hidden layer.
-- No wonder the extra hidden layer was not particulary effective
-- compared to wider single hidden layer.
middleLayerFit2 :: forall m. DualMonad 'DModeGradient Double m
                => (DualNumberD -> m DualNumberD)
                -> Data.Vector.Vector DualNumberD
                -> Int
                -> DualNumberVariablesD
                -> m (Data.Vector.Vector DualNumberD)
middleLayerFit2 factivation hiddenVec offset vec = do
  let f :: Int -> DualNumberD -> m DualNumberD
      f i x = do
        let weight = var0 vec (offset + 2 * i)
            bias = var0 vec (offset + 1 + 2 * i)
        sx <- x *\ weight
        sxBias <- sx +\ bias
        factivation sxBias
  V.imapM f hiddenVec

nnFit2 :: DualMonad 'DModeGradient Double m
       => (DualNumberD -> m DualNumberD)
       -> (DualNumberD -> m DualNumberD)
       -> (DualNumberD -> m DualNumberD)
       -> Double -> DualNumberVariablesD -> m DualNumberD
nnFit2 factivationHidden factivationMiddle factivationOutput x vec = do
  -- Due to not being fully connected, the parameters are only:
  -- one bias of the outer layer, a list of weights of the outer layer,
  -- a list of the same length of weights and biases of the first hidden layer
  -- and of the middle hidden layer.
  let width = (lengthDualNumber vec - 1) `div` 5
  hiddenVec <- hiddenLayerFit factivationHidden x vec width
  middleVec <- middleLayerFit2 factivationMiddle hiddenVec (2 * width) vec
  outputLayerFit factivationOutput middleVec (4 * width) vec

nnFit2Loss :: DualMonad 'DModeGradient Double m
           => (DualNumberD -> m DualNumberD)
           -> (DualNumberD -> m DualNumberD)
           -> (DualNumberD -> m DualNumberD)
           -> Double -> Double -> DualNumberVariablesD -> m DualNumberD
nnFit2Loss factivationHidden factivationMiddle factivationOutput x targ vec = do
  res <- nnFit2 factivationHidden factivationMiddle factivationOutput x vec
  lossSquared targ res

nnFit2LossTotal :: forall m. DualMonad 'DModeGradient Double m
                => (DualNumberD -> m DualNumberD)
                -> (DualNumberD -> m DualNumberD)
                -> (DualNumberD -> m DualNumberD)
                -> Vector (Double, Double)
                -> DualNumberVariablesD
                -> m DualNumberD
nnFit2LossTotal factivationHidden factivationMiddle factivationOutput
                samples vec = do
  let f :: (Double, Double) -> m DualNumberD
      f (x, res) =
        nnFit2Loss factivationHidden factivationMiddle factivationOutput
                   x res vec
  sumResultsDual f samples

lenP2 :: Int -> Int
lenP2 width = 5 * width + 1

-- Two layers seem to be an advantage for data with points very close
-- together. Otherwise, having all neurons on one layer is more effective.
fit2Tests :: TestTree
fit2Tests = testGroup "Sample fitting 2 hidden layer not fully connected nn tests"
  [ gdSimpleWsTestCase
      (nnFit2LossTotal tanhAct) 42 8 (lenP2 6) 0.1 10000
      1.2856619684390336e-2
  , gdSimpleWsTestCase
      (nnFit2LossTotal tanhAct) 42 10 (lenP2 6) 0.01 400000
      3.835053990072211e-2
  , gdSimpleSeparatedTestCase
      (nnFit2LossTotal tanhAct) 42 8 (lenP2 6) 0.1 10000
      0.31692351465375723
  , gdSimpleSeparatedTestCase
      (nnFit2LossTotal tanhAct) 42 10 (lenP2 6) 0.01 100000
      1.2308485318049472e-3
  , gdSimpleSeparatedTestCase
      (nnFit2LossTotal tanhAct) 42 16 (lenP2 12) 0.01 100000
      1.9398514673723763e-2
  ]

-- The same tests, but with logisticAct for the first hidden layer instead
-- of tanhAct. Usually worse results.
fit2TestsL :: TestTree
fit2TestsL = testGroup "logisticAct: Sample fitting 2 hidden layer not fully connected nn tests"
  [ gdSimpleWsTestCase
      (nnFit2LossTotal logisticAct) 42 8 (lenP2 6) 0.1 10000
      9.323867115794165e-3
  , gdSimpleWsTestCase
      (nnFit2LossTotal logisticAct) 42 10 (lenP2 6) 0.01 400000
      0.12307345215742066
  , gdSimpleSeparatedTestCase
      (nnFit2LossTotal logisticAct) 42 8 (lenP2 6) 0.1 10000
      0.2978076625110597
  , gdSimpleSeparatedTestCase
      (nnFit2LossTotal logisticAct) 42 10 (lenP2 6) 0.01 100000
      8.707658552473477e-2
  , gdSimpleSeparatedTestCase
      (nnFit2LossTotal logisticAct) 42 16 (lenP2 12) 0.01 100000
      1.2453082870396885
  ]

gdSmartShow :: (DualNumberVariablesD -> DualMonadGradient Double DualNumberD)
            -> Domain0 Double
            -> Int
            -> ([Double], (Double, Double))
gdSmartShow f initVec n =
  let ((res, _, _, _), gamma) = gdSmart f n (initVec, V.empty, V.empty, V.empty)
      (_, value) = dReverse 1 f (res, V.empty, V.empty, V.empty)
  in (V.toList res, (value, gamma))

gradSmartTestCase :: Num a
                  => String
                  -> ((a, a) -> Int -> Int
                      -> Vector (Double, Double))
                  -> ((DualNumberD -> DualMonadGradient Double DualNumberD)
                      -> (DualNumberD -> DualMonadGradient Double DualNumberD)
                      -> Vector (Double, Double)
                      -> DualNumberVariablesD
                      -> DualMonadGradient Double DualNumberD)
                  -> Int -> Int -> Int -> Int -> (Double, Double)
                  -> TestTree
gradSmartTestCase prefix sampleFunction lossFunction
                  seedSamples nSamples nParams0 nIterations expected =
  let samples = sampleFunction (-1, 1) seedSamples nSamples
      vec = V.unfoldrExactN nParams0 (uniformR (-1, 1)) $ mkStdGen 33
      name = prefix ++ " "
             ++ unwords [ show seedSamples, show nSamples
                        , show nParams0, show nIterations ]
  in testCase name $
       snd (gdSmartShow (lossFunction tanhAct tanhAct samples) vec nIterations)
       @?= expected

gradSmartWsTestCase
  :: ((DualNumberD -> DualMonadGradient Double DualNumberD)
      -> (DualNumberD -> DualMonadGradient Double DualNumberD)
      -> Vector (Double, Double)
      -> DualNumberVariablesD
      -> DualMonadGradient Double DualNumberD)
  -> Int -> Int -> Int -> Int -> (Double, Double)
  -> TestTree
gradSmartWsTestCase = gradSmartTestCase "gradSmart Ws" wsFit

gradSmartSeparatedTestCase
  :: ((DualNumberD -> DualMonadGradient Double DualNumberD)
      -> (DualNumberD -> DualMonadGradient Double DualNumberD)
      -> Vector (Double, Double)
      -> DualNumberVariablesD
      -> DualMonadGradient Double DualNumberD)
  -> Int -> Int -> Int -> Int -> (Double, Double)
  -> TestTree
gradSmartSeparatedTestCase =
  gradSmartTestCase "gradSmart Separated" wsFitSeparated

-- With Float, the approximation overshoots all the time and makes smaller
-- and smaller steps, getting nowhere. This probably means
-- approximation with this number of neurons is not possible.
-- However, adding neurons doesn't help (without huge increases of iterations).
-- The fact that results are better when freely overshooting
-- suggests there are local minima, which confirms too low dimensionality.
-- The experiments with separated samples seem to confirm both hypotheses.
--
-- With Double, it scales well to twice as many samples or even more,
-- but it takes too long to verify if/when errors crop in again.
smartFitTests :: TestTree
smartFitTests = testGroup "Smart descent sample fitting fully connected nn tests"
  [ gradSmartWsTestCase
      nnFitLossTotal 42 8 (lenP 10) 10000
      (2.0585450568797953e-3,1.25e-2)
  , gradSmartWsTestCase
      nnFitLossTotal 42 10 (lenP 20) 1000000
      (9.072288039580448e-2,6.25e-3)
        -- 31 not enough, 700000 not enough
  , gradSmartWsTestCase
      nnFitLossTotal 42 16 (lenP 20) 1700000
      (4.8336260347113275e-2,1.5625e-3)
  , gradSmartSeparatedTestCase
      nnFitLossTotal 42 8 (lenP 10) 10000
      (1.5742022677967708e-2,2.5e-2)
  , gradSmartSeparatedTestCase
      nnFitLossTotal 42 10 (lenP 10) 100000
      (4.506881373306206e-10,2.5e-2)
  , gradSmartSeparatedTestCase
      nnFitLossTotal 42 16 (lenP 10) 100000
      (5.197706771219677e-2,6.25e-3)
  , gradSmartSeparatedTestCase
      nnFitLossTotal 42 24 (lenP 33) 700000
      (2.967249104936791e-2,6.25e-3)
        -- 61 1300000 not enough
  , gradSmartSeparatedTestCase
      nnFitLossTotal 42 32 (lenP 20) 1700000
      (3.828456463288314e-2,6.25e-3)
        -- 151 1000000 not enough, despite taking twice longer
  ]

smartFit2Tests :: TestTree
smartFit2Tests =
 testGroup "Smart descent sample fitting 2 hidden layer not fully connected nn tests"
  [ gradSmartWsTestCase
      (nnFit2LossTotal tanhAct) 42 8 (lenP2 6) 10000
      (4.896924209457198e-3,2.5e-2)
  , gradSmartWsTestCase
      (nnFit2LossTotal tanhAct) 42 10 (lenP2 6) 400000
      (8.470989419560765e-2,2.5e-2)
  , gradSmartWsTestCase
      (nnFit2LossTotal tanhAct) 42 16 (lenP2 12) 700000
      (5.149610997592684e-2,3.90625e-4)
        -- 61 1000000 not enough for 20, 101 700000 enough
  , gradSmartSeparatedTestCase
      (nnFit2LossTotal tanhAct) 42 8 (lenP2 6) 10000
      (1.832621758590325e-2,1.25e-2)
  , gradSmartSeparatedTestCase
      (nnFit2LossTotal tanhAct) 42 10 (lenP2 6) 100000
      (2.6495249749522148e-2,3.125e-3)
  , gradSmartSeparatedTestCase
      (nnFit2LossTotal tanhAct) 42 16 (lenP2 12) 100000
      (1.8617700399788891e-3,3.125e-3)
  , gradSmartSeparatedTestCase
      (nnFit2LossTotal tanhAct) 42 24 (lenP2 12) 1300000
      (1.0411445668840221e-2,3.125e-3)
        -- this is faster but less accurate than 101 1000000
        -- 151 700000 is not enough
  ]

-- The same tests, but with logisticAct for the first hidden layer instead
-- of tanhAct. Usually worse results.
smartFit2TestsL :: TestTree
smartFit2TestsL =
 testGroup "logisticAct: Smart descent sample fitting 2 hidden layer not fully connected nn tests"
  [ gradSmartWsTestCase
      (nnFit2LossTotal logisticAct) 42 8 (lenP2 6) 10000
      (5.277048486983709e-3,5.0e-2)
  , gradSmartWsTestCase
      (nnFit2LossTotal logisticAct) 42 10 (lenP2 6) 400000
      (0.12231013820426907,2.5e-2)
  , gradSmartWsTestCase
      (nnFit2LossTotal logisticAct) 42 16 (lenP2 12) 700000
      (0.625388101609215,3.125e-3)
  , gradSmartSeparatedTestCase
      (nnFit2LossTotal logisticAct) 42 8 (lenP2 6) 10000
      (0.486980368199192,2.5e-2)
  , gradSmartSeparatedTestCase
      (nnFit2LossTotal logisticAct) 42 10 (lenP2 6) 100000
      (2.9673291826644015e-2,3.125e-3)
  , gradSmartSeparatedTestCase
      (nnFit2LossTotal logisticAct) 42 16 (lenP2 12) 100000
      (2.1703256232095827,1.25e-2)
  , gradSmartSeparatedTestCase
      (nnFit2LossTotal logisticAct) 42 24 (lenP2 12) 1300000
      (2.225663713307959,3.90625e-4)
  ]

-- This is really fully connected.
middleLayerFit3 :: forall m. DualMonad 'DModeGradient Double m
                => (DualNumberD -> m DualNumberD)
                -> Data.Vector.Vector DualNumberD
                -> Int
                -> DualNumberVariablesD
                -> m (Data.Vector.Vector DualNumberD)
middleLayerFit3 factivation hiddenVec offset vec =
  middleLayerMnist factivation hiddenVec offset vec $ V.length hiddenVec

nnFit3 :: DualMonad 'DModeGradient Double m
       => (DualNumberD -> m DualNumberD)
       -> (DualNumberD -> m DualNumberD)
       -> (DualNumberD -> m DualNumberD)
       -> Double -> DualNumberVariablesD -> m DualNumberD
nnFit3 factivationHidden factivationMiddle factivationOutput x vec = do
  -- This is fully connected, so given width w, the number of parameters is:
  -- one bias of the outer layer, a list of weights of the outer layer
  -- of length w, a list of the same length of weights and biases of the first
  -- hidden layer, w * w weigths in the middle hidden layer and w biases.
  -- In total, #params0 == 1 + w + 2 * w + w^2 + w == w^2 + 4 * w + 1,
  -- so the equation to solve is w^2 + 4 * w + (1 - #params0) = 0.
  -- Length 31 gives almost 3. Length 61 gives exactly 6.
  let len :: Double
      len = fromIntegral $ lengthDualNumber vec  -- #params
      width = floor $ (-4 + sqrt (12 + 4 * len)) / 2
  hiddenVec <- hiddenLayerFit factivationHidden x vec width
  middleVec <- middleLayerFit3 factivationMiddle hiddenVec (2 * width) vec
  outputLayerFit factivationOutput middleVec ((3 + width) * width) vec

nnFit3Loss :: DualMonad 'DModeGradient Double m
           => (DualNumberD -> m DualNumberD)
           -> (DualNumberD -> m DualNumberD)
           -> (DualNumberD -> m DualNumberD)
           -> Double -> Double -> DualNumberVariablesD -> m DualNumberD
nnFit3Loss factivationHidden factivationMiddle factivationOutput x targ vec = do
  res <- nnFit3 factivationHidden factivationMiddle factivationOutput x vec
  lossSquared targ res

nnFit3LossTotal :: forall m. DualMonad 'DModeGradient Double m
                => (DualNumberD -> m DualNumberD)
                -> (DualNumberD -> m DualNumberD)
                -> (DualNumberD -> m DualNumberD)
                -> Vector (Double, Double)
                -> DualNumberVariablesD
                -> m DualNumberD
nnFit3LossTotal factivationHidden factivationMiddle factivationOutput
                samples vec = do
  let f :: (Double, Double) -> m DualNumberD
      f (x, res) =
        nnFit3Loss factivationHidden factivationMiddle factivationOutput
                   x res vec
  sumResultsDual f samples

lenP3 :: Int -> Int
lenP3 width = (4 + width) * width + 1

smartFit3Tests :: TestTree
smartFit3Tests =
 testGroup "Smart descent sample fitting 2 hidden layer really fully connected nn tests"
  -- The same (or slightly lower) number of parameters, but on many layers
  -- gives much worse results, except if there are few samples.
  -- It is also less costly for some reason.
  [ gradSmartWsTestCase
      (nnFit3LossTotal tanhAct) 42 8 (lenP3 4) 10000
      (3.5604072240265575e-3,2.5e-2)
  , gradSmartWsTestCase
      (nnFit3LossTotal tanhAct) 42 10 (lenP3 4) 400000
      (0.10126792765437645,3.125e-3)
        -- failed, because too narrow (4 neurons in each hidden layer)
  , gradSmartWsTestCase
      (nnFit3LossTotal tanhAct) 42 16 (lenP3 6) 700000
      (0.19318876323310907,1.5625e-3)
        -- failed, because too narrow (6 neurons in each hidden layer)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal tanhAct) 42 8 (lenP3 4) 10000
      (6.291648505851797e-3,6.25e-3)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal tanhAct) 42 10 (lenP3 4) 100000
      (4.34890234424764e-7,6.25e-3)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal tanhAct) 42 16 (lenP3 6) 100000
      (0.3434691592146121,1.5625e-3)
        -- failed, because too narrow (6 neurons in each hidden layer)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal tanhAct) 42 24 (lenP3 6) 1300000
      (1.665065359469462,9.765625e-5)
        -- failed, because too narrow (6 neurons in each hidden layer)
  -- The same width of fully connected nn, but with 2 instead of 1 hidden
  -- layer has many more parameters and is usually more costly,
  -- but also much more powerful given enough training:
  , gradSmartWsTestCase
      (nnFit3LossTotal tanhAct) 42 8 (lenP3 6) 10000
      (2.1767148242940303e-3,1.25e-2)
  , gradSmartWsTestCase
      (nnFit3LossTotal tanhAct) 42 10 (lenP3 6) 400000
      (1.3871923964300112e-4,6.25e-3)
  , gradSmartWsTestCase
      (nnFit3LossTotal tanhAct) 42 16 (lenP3 12) 700000
      (2.131955325962636e-5,7.8125e-4)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal tanhAct) 42 8 (lenP3 6) 10000
      (2.0369556559640215e-4,5.0e-2)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal tanhAct) 42 10 (lenP3 6) 100000
      (4.0282344474667426e-16,6.25e-3)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal tanhAct) 42 16 (lenP3 12) 100000
      (1.5819824634756573e-5,3.125e-3)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal tanhAct) 42 24 (lenP3 12) 1300000
      (1.1354796858869852e-6,1.5625e-3)
  ]

-- The same tests, but with logisticAct for the first hidden layer instead
-- of tanhAct. Usually worse results.
smartFit3TestsL :: TestTree
smartFit3TestsL =
 testGroup "logisticAct: Smart descent sample fitting 2 hidden layer really fully connected nn tests"
  [ gradSmartWsTestCase
      (nnFit3LossTotal logisticAct) 42 8 (lenP3 4) 10000
      (4.262690053475572e-3,2.5e-2)
  , gradSmartWsTestCase
      (nnFit3LossTotal logisticAct) 42 10 (lenP3 4) 400000
      (8.133312854575311e-2,0.1)
  , gradSmartWsTestCase
      (nnFit3LossTotal logisticAct) 42 16 (lenP3 6) 700000
      (0.35607225062502207,1.5625e-3)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal logisticAct) 42 8 (lenP3 4) 10000
      (8.041780516044969e-2,1.25e-2)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal logisticAct) 42 10 (lenP3 4) 100000
      (1.0877770623826884e-2,6.25e-3)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal logisticAct) 42 16 (lenP3 6) 100000
      (1.444649714852858,1.25e-2)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal logisticAct) 42 24 (lenP3 6) 1300000
      (1.7847189651694308,1.953125e-4)
  , gradSmartWsTestCase
      (nnFit3LossTotal logisticAct) 42 8 (lenP3 6) 10000
      (5.436829414964154e-3,2.5e-2)
  , gradSmartWsTestCase
      (nnFit3LossTotal logisticAct) 42 10 (lenP3 6) 400000
      (0.11938665766915235,1.25e-2)
  , gradSmartWsTestCase
      (nnFit3LossTotal logisticAct) 42 16 (lenP3 12) 700000
      (0.24271694671964975,6.25e-3)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal logisticAct) 42 8 (lenP3 6) 10000
      (8.700033031145113e-2,1.25e-2)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal logisticAct) 42 10 (lenP3 6) 100000
      (7.225917500882556e-6,1.25e-2)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal logisticAct) 42 16 (lenP3 12) 100000
      (2.3379235156826658e-2,3.125e-3)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotal logisticAct) 42 24 (lenP3 12) 1300000
      (0.6440964543158452,3.125e-3)
  ]

{-
nnFit3LossTotalOutput :: forall m. DualMonad Double m
                      => (DualNumberD -> m DualNumberD)
                      -> (DualNumberD -> m DualNumberD)
                      -> (DualNumberD -> m DualNumberD)
                      -> Vector (Double, Double)
                      -> DualNumberVariablesD
                      -> m DualNumberD
nnFit3LossTotalOutput f1 f2 f3 samples vec =
  nnFit3LossTotal f2 f3 f1 samples vec

-- The same tests, but with logisticAct for the output layer instead of tanhAct.
-- Somewhat better results than with tanhAct, but not worth recording.
-- OTOH, with reluAct (understandable, it rules out any negative values)
-- and with reluLeakyAct (gradient explosion for positive values?)
-- the results are disastrous. Also they are disastrous with either relu
-- in any other position (no idea why, there is one more layer to recover
-- anything that's lost; perhaps the sums overflow due to no normalization?)
-- and mediocre with logisticAct in other positions.
smartFit3TestsL3 :: TestTree
smartFit3TestsL3 =
 testGroup "logisticAct3: Smart descent sample fitting 2 hidden layer really fully connected nn tests"
  [ gradSmartWsTestCase
      (nnFit3LossTotalOutput logisticAct) 42 8 (lenP3 4) 10000
      (3.5604072240265575e-3,2.5e-2)
  , gradSmartWsTestCase
      (nnFit3LossTotalOutput logisticAct) 42 10 (lenP3 4) 400000
      (0.10126792765437645,3.125e-3)
  , gradSmartWsTestCase
      (nnFit3LossTotalOutput logisticAct) 42 16 (lenP3 6) 700000
      (0.19318876323310907,1.5625e-3)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotalOutput logisticAct) 42 8 (lenP3 4) 10000
      (6.291648505851797e-3,6.25e-3)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotalOutput logisticAct) 42 10 (lenP3 4) 100000
      (4.34890234424764e-7,6.25e-3)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotalOutput logisticAct) 42 16 (lenP3 6) 100000
      (0.3434691592146121,1.5625e-3)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotalOutput logisticAct) 42 24 (lenP3 6) 1300000
      (1.665065359469462,9.765625e-5)
  , gradSmartWsTestCase
      (nnFit3LossTotalOutput logisticAct) 42 8 (lenP3 6) 10000
      (2.1767148242940303e-3,1.25e-2)
  , gradSmartWsTestCase
      (nnFit3LossTotalOutput logisticAct) 42 10 (lenP3 6) 400000
      (1.3871923964300112e-4,6.25e-3)
  , gradSmartWsTestCase
      (nnFit3LossTotalOutput logisticAct) 42 16 (lenP3 12) 700000
      (2.131955325962636e-5,7.8125e-4)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotalOutput logisticAct) 42 8 (lenP3 6) 10000
      (2.0369556559640215e-4,5.0e-2)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotalOutput logisticAct) 42 10 (lenP3 6) 100000
      (4.0282344474667426e-16,6.25e-3)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotalOutput logisticAct) 42 16 (lenP3 12) 100000
      (1.5819824634756573e-5,3.125e-3)
  , gradSmartSeparatedTestCase
      (nnFit3LossTotalOutput logisticAct) 42 24 (lenP3 12) 1300000
      (1.1354796858869852e-6,1.5625e-3)
  ]
-}

sgdShow :: HasDelta r
        => r
        -> (a -> DualNumberVariables 'DModeGradient r -> DualMonadGradient r (DualNumber 'DModeGradient r))
        -> [a]  -- ^ training data
        -> Domain0 r  -- ^ initial parameters
        -> ([r], r)
sgdShow gamma f trainData params0Init =
  let (res, _, _, _) =
        fst $ sgd gamma f trainData (params0Init, V.empty, V.empty, V.empty)
      (_, value) = dReverse 1 (f $ head trainData) (res, V.empty, V.empty, V.empty)
  in (V.toList res, value)

sgdTestCase
  :: String
  -> IO [a]
  -> (Int
      -> a
      -> DualNumberVariablesD
      -> DualMonadGradient Double DualNumberD)
  -> Double
  -> Double
  -> TestTree
sgdTestCase prefix trainDataIO trainWithLoss gamma expected =
  let widthHidden = 250
      nParams0 = lenMnist widthHidden
      vec = HM.randomVector 33 HM.Uniform nParams0 - HM.scalar 0.5
      name = prefix ++ " "
             ++ unwords [show widthHidden, show nParams0, show gamma]
  in testCase name $ do
       trainData <- trainDataIO
       snd (sgdShow gamma (trainWithLoss widthHidden) trainData vec)
          @?= expected

lenMnist :: Int -> Int
lenMnist widthHidden =
  widthHidden * (sizeMnistGlyph + 1) + sizeMnistLabel * (widthHidden + 1)

nnFit3ForStochastic :: Int
                    -> (Double, Double)
                    -> DualNumberVariablesD
                    -> DualMonadGradient Double DualNumberD
nnFit3ForStochastic _ (x, res) = nnFit3Loss tanhAct tanhAct tanhAct x res

stochasticFit3Tests :: TestTree
stochasticFit3Tests =
 testGroup "Stochastic gradient descent sample fitting 2 hidden layer really fully connected nn tests"
  [ let trainData = V.toList $ wsFit (-1, 1) 42 2
    in sgdTestCase "Ws 42 2"
         (return trainData) nnFit3ForStochastic 0.1 0.23539780131842747
  , let trainData = take 200 $ cycle $ V.toList $ wsFit (-1, 1) 42 2
    in sgdTestCase "Ws 42 2(200)"
         (return trainData) nnFit3ForStochastic 0.1 0.23539780131842747
  , let trainData = V.toList $ wsFitSeparated (-1, 1) 42 128
    in sgdTestCase "separated 42 128"
         (return trainData) nnFit3ForStochastic 0.1 3.465944781121193
  , let trainData = V.toList $ wsFitSeparated (-1, 1) 42 256
    in sgdTestCase "separated 42 256"
         (return trainData) nnFit3ForStochastic 0.1 1.912556094812049e-2
  , let trainData = take 1280 $ cycle $ V.toList $ wsFitSeparated (-1, 1) 42 128
    in sgdTestCase "separated 42 128(1280)"
         (return trainData) nnFit3ForStochastic 0.1 3.465944781121193
  ]


text:
-- |
-- Module      : Streamly.Statistics
-- Copyright   : (c) 2020 Composewell Technologies
-- License     : Apache-2.0
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- Statistical measures over a stream of data. All operations use numerically
-- stable floating point arithmetic.
--
-- Measurements can be performed over the entire input stream or on a sliding
-- window of fixed or variable size.  Where possible, measures are computed
-- online without buffering the input stream.
--
-- Currently there is no overflow detection.
--
-- References:
--
-- * https://en.wikipedia.org/wiki/Statistics
-- * https://mathworld.wolfram.com/topics/ProbabilityandStatistics.html

-- Resources:
--
-- This may be another useful resource for incremental (non-windowed)
-- computation:
--
-- https://www.researchgate.net/publication/287152055_Incremental_Statistical_Measures
--
-- Sample Statistics
--
-- Terms
--
-- Population: the complete data set from which statistical samples are taken.
--
-- Sample: a subset of the population.
--
-- https://en.wikipedia.org/wiki/Sample_(statistics)
--
-- Estimator:
--
-- Statistical measures can be computed either from the actual population
-- or from samples. Statistical measures computed from the samples provide an
-- estimate of the actual measures of the entire population. Measures computed
-- from samples may not truly reflect the actual measures and may have to be
-- adjusted for biases or errors.
--
-- An "estimator" is a method or function to compute a statistical measure from
-- sampled data. For example, the sample variance is an esitmator of the
-- population variance.
--
-- https://en.wikipedia.org/wiki/Estimator
--
-- Bias:
--
-- The result computed by an estimator may not be centered at the true value as
-- determined by computing the measure for the actual population. Such an
-- estimator is called a biased estimator.  For example, notice how
-- 'sampleVariance' is adjusted for bias.
--
-- https://en.wikipedia.org/wiki/Bias_of_an_estimator
--
-- Consistency:
--
-- https://en.wikipedia.org/wiki/Consistent_estimator

{-# LANGUAGE ScopedTypeVariables #-}
module Streamly.Statistics
    (
    -- * Incremental Folds
    -- | Folds of type @Fold m (a, Maybe a) b@ are incremental sliding window
    -- folds. An input of type @(a, Nothing)@ indicates that the input element
    -- @a@ is being inserted in the window without ejecting an old value
    -- increasing the window size by 1. An input of type @(a, Just a)@
    -- indicates that the first element is being inserted in the window and the
    -- second element is being removed from the window, the window size remains
    -- the same. The window size can only increase and never decrease.
    --
    -- You can compute the statistics over the entire stream using sliding
    -- window folds by keeping the second element of the input tuple as
    -- @Nothing@.
    --
      Window.lmap
    , Window.cumulative

    -- * Summary Statistics
    -- | See https://en.wikipedia.org/wiki/Summary_statistics .

    -- ** Sums
    , Window.length
    , Window.sum
    , Window.sumInt
    , Window.powerSum

    -- ** Location
    -- | See https://en.wikipedia.org/wiki/Location_parameter .
    --
    -- See https://en.wikipedia.org/wiki/Central_tendency .
    , Window.minimum
    , Window.maximum
    , rawMoment
    , rawMomentFrac

    -- Pythagorean means (https://en.wikipedia.org/wiki/Pythagorean_means)
    , mean
    , welfordMean
    , geometricMean
    , harmonicMean

    , quadraticMean

    -- Generalized mean
    , powerMean
    , powerMeanFrac

    -- ** Weighted Means
    -- | Exponential Smoothing.
    , ewma
    , ewmaAfterMean
    , ewmaRampUpSmoothing

    -- ** Spread
    -- | Second order central moment is a statistical measure of dispersion.
    -- The \(k\)th moment about the mean (or \(k\)th central moment) is defined
    -- as:
    --
    -- \(\mu_k = \frac{1}{n}\sum_{i=1}^n {(x_{i}-\mu)}^k\)
    --
    -- See https://mathworld.wolfram.com/CentralMoment.html .
    --
    -- See https://en.wikipedia.org/wiki/Statistical_dispersion .
    , Window.range
    , md
    , variance
    , stdDev

    -- ** Shape
    -- | Third and fourth order central moments are a measure of shape.
    --
    -- See https://en.wikipedia.org/wiki/Shape_parameter .
    --
    -- See https://en.wikipedia.org/wiki/Standardized_moment .
    , skewness
    , kurtosis

    -- XXX Move to Statistics.Sample or Statistics.Estimation module?
    -- ** Estimation
    , sampleVariance
    , sampleStdDev
    , stdErrMean

    -- ** Resampling
    , resample
    , foldResamples
    , jackKnifeMean
    , jackKnifeVariance
    , jackKnifeStdDev

    -- ** Probability Distribution
    , frequency
    , mode

    -- Histograms
    , HistBin (..)
    , binOffsetSize
    , binFromSizeN
    , binFromToN
    , binBoundaries
    , histogram

    -- * Transforms
    , fft
    )
where

import Control.Exception (assert)
import Control.Monad (when)
import Control.Monad.IO.Class (MonadIO(..))
import Data.Functor.Identity (runIdentity, Identity)
import Data.Bits (Bits(complement, shiftL, shiftR, (.&.), (.|.)))
import Data.Complex (Complex ((:+)))
import Data.Map.Strict (Map, foldrWithKey)
import Foreign.Storable (Storable)
import Streamly.Data.Fold.Tee(Tee(..), toFold)
import Streamly.Internal.Control.Concurrent (MonadAsync)
import Streamly.Internal.Data.Array.Foreign.Type
    (Array, length, toStream, unsafeIndexIO)
import Streamly.Internal.Data.Fold.Type (Fold(..), Step(..))
import Streamly.Internal.Data.Stream.IsStream (SerialT)
import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))
import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
import System.Random.MWC (createSystemRandom, uniformRM)

import qualified Streamly.Internal.Data.Array.Foreign as Array
import qualified Streamly.Internal.Data.Array.Foreign.Mut as MA
import qualified Streamly.Internal.Data.Fold as Fold
import qualified Streamly.Internal.Data.Fold.Window as Window
import qualified Streamly.Internal.Data.Stream.IsStream as Stream
import qualified Streamly.Internal.Data.Unfold as Unfold

import Prelude hiding (length, sum, minimum, maximum)

-- TODO: Overflow checks. Would be good if we can directly replace the
-- operations with overflow checked operations.
--
-- See https://hackage.haskell.org/package/safe-numeric
-- See https://hackage.haskell.org/package/safeint
--
-- TODO We have many of these functions in Streamly.Data.Fold as well. Need to
-- think about deduplication.

-------------------------------------------------------------------------------
-- Transforms
-------------------------------------------------------------------------------

-- XXX These utility functions can be moved to streamly-numeric

-- | Test if the given integer value is a power of 2.
{-# INLINE isPower2 #-}
isPower2 :: Int -> Bool
isPower2 n = n .&. (n - 1) == 0

-- | Create a power of 2
--
-- Argument must be less than 64 assuming 64-bit Int size.
--
{-# INLINE _power2 #-}
_power2 :: Int -> Int
_power2 n = shiftL 1 n

-- | Create a bit mask with lower n bits 0 and the rest as 1.
--
-- Argument must be less than 64 assuming 64-bit Int size.
--
{-# INLINE maskLowerN #-}
maskLowerN :: Int -> Int
maskLowerN n = complement (shiftL 1 n - 1)

-- | Compute the base 2 logarithm of the given value.
--
-- Assumes the Int size to be 64-bit.
--
{-# INLINE logBase2 #-}
logBase2 :: Int -> Int
logBase2 v0
    | v0 <= 0   = error $ "logBase2: input must be greater than 0 " ++ show v0
    | otherwise = go 32 0 v0

    where

    go !bits !result !v
        | bits == 0 = result
        | v .&. maskLowerN bits /= 0 =
             go (bits `shiftR` 1) (result .|. bits) (v `shiftR` bits)
        | otherwise = go (bits `shiftR` 1) result v

-- Algo translated from the statistics library.
--
-- XXX We can use a wrapper API that takes an array of Double input instead of
-- array of Complex.
--
-- | Compute fast fourier transform of an array of 'Complex' values.
--
-- Array length must be power of 2.
--
{-# INLINE fft #-}
fft :: MonadIO m => MA.Array (Complex Double) -> m ()
fft marr
    | isPower2 len = bitReverse 0 0
    | otherwise  = error "fft: Array length must be power of 2"

    where

    len = MA.length marr

    halve x = x `shiftR` 1

    twice x = x `shiftL` 1

    inner i j k
        | k <= j  = inner i (j - k) (halve k)
        | otherwise = bitReverse (i + 1) (j + k)

    bitReverse i j
        | i == len - 1 = stage 0 1
        | otherwise = do
            when (i < j) $ MA.unsafeSwapIndices i j marr
            inner i j (halve len)

    log2len = logBase2 len

    stage l !l1
        | l == log2len = return ()
        | otherwise = do
            let !l2 = twice l1
                !e  = -6.283185307179586/fromIntegral l2
                flight j !a | j == l1   = stage (l + 1) l2
                            | otherwise = do
                    let butterfly i | i >= len  = flight (j + 1) (a + e)
                                    | otherwise = do
                            let i1 = i + l1
                            xi1 :+ yi1 <- MA.getIndexUnsafe i1 marr
                            let !c = cos a
                                !s = sin a
                                d  = (c * xi1 - s * yi1) :+ (s * xi1 + c * yi1)
                            ci <- MA.getIndexUnsafe i marr
                            MA.putIndexUnsafe  i1 (ci - d) marr
                            MA.putIndexUnsafe  i (ci + d) marr
                            butterfly (i + l2)
                    butterfly j
            flight 0 0

-------------------------------------------------------------------------------
-- Mean
-------------------------------------------------------------------------------

-- | Arithmetic mean of elements in a sliding window:
--
-- \(\mu = \frac{\sum_{i=1}^n x_{i}}{n}\)
--
-- This is also known as the Simple Moving Average (SMA) when used in the
-- sliding window and Cumulative Moving Avergae (CMA) when used on the entire
-- stream.
--
-- Mean is the same as the first raw moment.
--
-- \(\mu = \mu'_1\)
--
-- >>> mean = rawMoment 1
-- >>> mean = powerMean 1
-- >>> mean = Fold.teeWith (/) sum length
--
-- /Space/: \(\mathcal{O}(1)\)
--
-- /Time/: \(\mathcal{O}(n)\)
{-# INLINE mean #-}
mean :: forall m a. (Monad m, Fractional a) => Fold m (a, Maybe a) a
mean = Window.mean

-- | Recompute mean from old mean when an item is removed from the sample.
{-# INLINE _meanSubtract #-}
_meanSubtract :: Fractional a => Int -> a -> a -> a
_meanSubtract n oldMean oldItem =
    let delta = (oldItem - oldMean) / fromIntegral (n - 1)
     in oldMean - delta

-- | Recompute mean from old mean when an item is added to the sample.
{-# INLINE meanAdd #-}
meanAdd :: Fractional a => Int -> a -> a -> a
meanAdd n oldMean newItem =
    let delta = (newItem - oldMean) / fromIntegral (n + 1)
     in oldMean + delta

-- We do not carry rounding errors, therefore, this would be less numerically
-- stable than the kbn mean.
--
-- | Recompute mean from old mean when an item in the sample is replaced.
{-# INLINE meanReplace #-}
meanReplace :: Fractional a => Int -> a -> a -> a -> a
meanReplace n oldMean oldItem newItem =
    let n1 = fromIntegral n
        -- Compute two deltas instead of a single (newItem - oldItem) because
        -- the latter would be too small causing rounding errors.
        delta1 = (newItem - oldMean) / n1
        delta2 = (oldItem - oldMean) / n1
     in (oldMean + delta1) - delta2

-- | Same as 'mean' but uses Welford's algorithm to compute the mean
-- incrementally.
--
-- It maintains a running mean instead of a running sum and adjusts the mean
-- based on a new value.  This is slower than 'mean' because of using the
-- division operation on each step and it is numerically unstable (as of now).
-- The advantage over 'mean' could be no overflow if the numbers are large,
-- because we do not maintain a sum, but that is a highly unlikely corner case.
--
-- /Internal/
{-# INLINE welfordMean #-}
welfordMean :: forall m a. (Monad m, Fractional a) => Fold m (a, Maybe a) a
welfordMean = Fold step initial extract

    where

    initial =
        return
            $ Partial
            $ Tuple'
                (0 :: a)   -- running mean
                (0 :: Int) -- count of items in the window

    step (Tuple' oldMean w) (new, mOld) =
        return
            $ Partial
            $ case mOld of
                Nothing -> Tuple' (meanAdd w oldMean new) (w + 1)
                Just old -> Tuple' (meanReplace w oldMean old new) w

    extract (Tuple' x _) = return x

-------------------------------------------------------------------------------
-- Moments
-------------------------------------------------------------------------------

-- XXX We may have chances of overflow if the powers are high or the numbers
-- are large. A limited mitigation could be to use welford style avg
-- computation. Do we need an overflow detection?
--
-- | Raw moment is the moment about 0. The \(k\)th raw moment is defined as:
--
-- \(\mu'_k = \frac{\sum_{i=1}^n x_{i}^k}{n}\)
--
-- >>> rawMoment k = Fold.teeWith (/) (powerSum p) length
--
-- See https://en.wikipedia.org/wiki/Moment_(mathematics) .
--
-- /Space/: \(\mathcal{O}(1)\)
--
-- /Time/: \(\mathcal{O}(n)\)
{-# INLINE rawMoment #-}
rawMoment :: (Monad m, Fractional a) => Int -> Fold m (a, Maybe a) a
rawMoment k = Fold.teeWith (/) (Window.powerSum k) Window.length

-- | Like 'rawMoment' but powers can be negative or fractional. This is
-- slower than 'rawMoment' for positive intergal powers.
--
-- >>> rawMomentFrac p = Fold.teeWith (/) (powerSumFrac p) length
--
{-# INLINE rawMomentFrac #-}
rawMomentFrac :: (Monad m, Floating a) => a -> Fold m (a, Maybe a) a
rawMomentFrac k = Fold.teeWith (/) (Window.powerSumFrac k) Window.length

-- XXX Overflow can happen when large powers or large numbers are used. We can
-- keep a running mean instead of running sum but that won't mitigate the
-- overflow possibility by much. The overflow can still happen when computing
-- the mean incrementally.

-- | The \(k\)th power mean of numbers \(x_1, x_2, \ldots, x_n\) is:
--
-- \(M_k = \left( \frac{1}{n} \sum_{i=1}^n x_i^k \right)^{\frac{1}{k}}\)
--
-- \(powerMean(k) = (rawMoment(k))^\frac{1}{k}\)
--
-- >>> powerMean k = (** (1 / fromIntegral k)) <$> rawMoment k
--
-- All other means can be expressed in terms of power mean. It is also known as
-- the generalized mean.
--
-- See https://en.wikipedia.org/wiki/Generalized_mean
--
{-# INLINE powerMean #-}
powerMean :: (Monad m, Floating a) => Int -> Fold m (a, Maybe a) a
powerMean k = (** (1 / fromIntegral k)) <$> rawMoment k

-- | Like 'powerMean' but powers can be negative or fractional. This is
-- slower than 'powerMean' for positive intergal powers.
--
-- >>> powerMeanFrac k = (** (1 / k)) <$> rawMomentFrac k
--
{-# INLINE powerMeanFrac #-}
powerMeanFrac :: (Monad m, Floating a) => a -> Fold m (a, Maybe a) a
powerMeanFrac k = (** (1 / k)) <$> rawMomentFrac k

-- | The harmonic mean of the positive numbers \(x_1, x_2, \ldots, x_n\) is
-- defined as:
--
-- \(HM = \frac{n}{\frac1{x_1} + \frac1{x_2} + \cdots + \frac1{x_n}}\)
--
-- \(HM = \left(\frac{\sum\limits_{i=1}^n x_i^{-1}}{n}\right)^{-1}\)
--
-- >>> harmonicMean = Fold.teeWith (/) length (lmap recip sum)
-- >>> harmonicMean = powerMeanFrac (-1)
--
-- See https://en.wikipedia.org/wiki/Harmonic_mean .
--
{-# INLINE harmonicMean #-}
harmonicMean :: (Monad m, Fractional a) => Fold m (a, Maybe a) a
harmonicMean = Fold.teeWith (/) Window.length (Window.lmap recip Window.sum)

-- | Geometric mean, defined as:
--
-- \(GM = \sqrt[n]{x_1 x_2 \cdots x_n}\)
--
-- \(GM = \left(\prod_{i=1}^n x_i\right)^\frac{1}{n}\)
--
-- or, equivalently, as the arithmetic mean in log space:
--
-- \(GM = e ^{{\frac{\sum_{i=1}^{n}\ln a_i}{n}}}\)
--
-- >>> geometricMean = exp <$> lmap log mean
--
-- See https://en.wikipedia.org/wiki/Geometric_mean .
{-# INLINE geometricMean #-}
geometricMean :: (Monad m, Floating a) => Fold m (a, Maybe a) a
geometricMean = exp <$> Window.lmap log mean

-- | The quadratic mean or root mean square (rms) of the numbers
-- \(x_1, x_2, \ldots, x_n\) is defined as:
--
-- \(RMS = \sqrt{ \frac{1}{n} \left( x_1^2 + x_2^2 + \cdots + x_n^2 \right) }.\)
--
-- >>> quadraticMean = powerMean 2
--
-- See https://en.wikipedia.org/wiki/Root_mean_square .
--
{-# INLINE quadraticMean #-}
quadraticMean :: (Monad m, Floating a) => Fold m (a, Maybe a) a
quadraticMean = powerMean 2

-------------------------------------------------------------------------------
-- Weighted Means
-------------------------------------------------------------------------------

-- XXX Is this numerically stable? We can use the kbn summation here.
-- | ewmaStep smoothing-factor old-value new-value
{-# INLINE ewmaStep #-}
ewmaStep :: Double -> Double -> Double -> Double
ewmaStep k x0 x1 = (1 - k) * x0 + k * x1

-- XXX Compute this in a sliding window?
--
-- | @ewma smoothingFactor@.
--
-- @ewma@ of an empty stream is 0.
--
-- Exponential weighted moving average, \(s_n\), of \(n\) values,
-- \(x_1,\ldots,x_n\), is defined recursively as:
--
-- \(\begin{align} s_0& = x_0\\ s_n & = \alpha x_{n} + (1-\alpha)s_{n-1},\quad n>0 \end{align}\)
--
-- If we expand the recursive term it becomes an exponential series:
--
-- \(s_n = \alpha \left[x_n + (1-\alpha)x_{n-1} + (1-\alpha)^2 x_{n-2} + \cdots + (1-\alpha)^{n-1} x_1 \right] + (1-\alpha)^n x_0\)
--
-- where \(\alpha\), the smoothing factor, is in the range \(0 <\alpha < 1\).
-- More the value of \(\alpha\), the more weight is given to newer values.  As
-- a special case if it is 0 then the weighted sum would always be the same as
-- the oldest value, if it is 1 then the sum would always be the same as the
-- newest value.
--
-- See https://en.wikipedia.org/wiki/Moving_average
--
-- See https://en.wikipedia.org/wiki/Exponential_smoothing
--
{-# INLINE ewma #-}
ewma :: Monad m => Double -> Fold m Double Double
ewma k = extract <$> Fold.foldl' step (Tuple' 0 1)

    where

    step (Tuple' x0 k1) x = Tuple' (ewmaStep k1 x0 x) k

    extract (Tuple' x _) = x

-- XXX It can perhaps perform better if implemented as a custom fold?
--
-- | @ewma n k@ is like 'ewma' but uses the mean of the first @n@ values and
-- then uses that as the initial value for the @ewma@ of the rest of the
-- values.
--
-- This can be used to reduce the effect of volatility of the initial value
-- when k is too small.
--
{-# INLINE ewmaAfterMean #-}
ewmaAfterMean :: Monad m => Int -> Double -> Fold m Double Double
ewmaAfterMean n k =
    Fold.concatMap (\i -> (Fold.foldl' (ewmaStep k) i)) (Fold.take n Fold.mean)

-- | @ewma n k@ is like 'ewma' but uses 1 as the initial smoothing factor and
-- then exponentially smooths it to @k@ using @n@ as the smoothing factor.
--
-- This is significantly faster than 'ewmaAfterMean'.
--
{-# INLINE ewmaRampUpSmoothing #-}
ewmaRampUpSmoothing :: Monad m => Double -> Double -> Fold m Double Double
ewmaRampUpSmoothing n k1 = extract <$> Fold.foldl' step initial

    where

    initial = Tuple' 0 1

    step (Tuple' x0 k0) x1 =
        let x = ewmaStep k0 x0 x1
            k = ewmaStep n k0 k1
        in Tuple' x k

    extract (Tuple' x _) = x

-------------------------------------------------------------------------------
-- Spread/Dispersion
-------------------------------------------------------------------------------

-- | @md n@ computes the mean absolute deviation (or mean deviation) in a
-- sliding window of last @n@ elements in the stream.
--
-- The mean absolute deviation of the numbers \(x_1, x_2, \ldots, x_n\) is:
--
-- \(MD = \frac{1}{n}\sum_{i=1}^n |x_i-\mu|\)
--
-- Note: It is expensive to compute MD in a sliding window. We need to
-- maintain a ring buffer of last n elements and maintain a running mean, when
-- the result is extracted we need to compute the difference of all elements
-- from the mean and get the average. Using standard deviation may be
-- computationally cheaper.
--
-- See https://en.wikipedia.org/wiki/Average_absolute_deviation .
--
-- /Pre-release/
{-# INLINE md #-}
md ::  MonadIO m => Fold m ((Double, Maybe Double), m (MA.Array Double)) Double
md =
    Fold.rmapM computeMD
        $ Fold.tee (Fold.lmap fst mean) (Fold.lmap snd Fold.last)

    where

    computeMD (mn, rng) =
        case rng of
            Just action -> do
                arr <- action
                Stream.fold Fold.mean
                    $ Stream.map (\a -> abs (mn - a))
                    $ Stream.unfold MA.read arr
            Nothing -> return 0.0

-- | The variance \(\sigma^2\) of a population of \(n\) equally likely values
-- is defined as the average of the squares of deviations from the mean
-- \(\mu\). In other words, second moment about the mean:
--
-- \(\sigma^2 = \frac{1}{n}\sum_{i=1}^n {(x_{i}-\mu)}^2\)
--
-- \(\sigma^2 = rawMoment(2) - \mu^2\)
--
-- \(\mu_2 = -(\mu'_1)^2 + \mu'_2\)
--
-- Note that the variance would be biased if applied to estimate the population
-- variance from a sample of the population. See 'sampleVariance'.
--
-- See https://en.wikipedia.org/wiki/Variance.
--
-- /Space/: \(\mathcal{O}(1)\)
--
-- /Time/: \(\mathcal{O}(n)\)
{-# INLINE variance #-}
variance :: (Monad m, Fractional a) => Fold m (a, Maybe a) a
variance = Fold.teeWith (\p2 m -> p2 - m ^ (2 :: Int)) (rawMoment 2) mean

-- | Standard deviation \(\sigma\) is the square root of 'variance'.
--
-- This is the population standard deviation or uncorrected sample standard
-- deviation.
--
-- >>> stdDev = sqrt <$> variance
--
-- See https://en.wikipedia.org/wiki/Standard_deviation .
--
-- /Space/: \(\mathcal{O}(1)\)
--
-- /Time/: \(\mathcal{O}(n)\)
{-# INLINE stdDev #-}
stdDev :: (Monad m, Floating a) => Fold m (a, Maybe a) a
stdDev = sqrt <$> variance

-- | Skewness \(\gamma\) is the standardized third central moment defined as:
--
-- \(\tilde{\mu}_3 = \frac{\mu_3}{\sigma^3}\)
--
-- The third central moment can be computed in terms of raw moments:
--
-- \(\mu_3 = 2(\mu'_1)^3 - 3\mu'_1\mu'_2 + \mu'_3\)
--
-- Substituting \(\mu'_1 = \mu\), and \(\mu'_2 = \mu^2 + \sigma^2\):
--
-- \(\mu_3 = -\mu^3 - 3\mu\sigma^2 + \mu'_3\)
--
-- Skewness is a measure of symmetry of the probability distribution. It is 0
-- for a symmetric distribution, negative for a distribution that is skewed
-- towards left, positive for a distribution skewed towards right.
--
-- For a normal like distribution the median can be found around
-- \(\mu - \frac{\gamma\sigma}{6}\) and the mode can be found around
-- \(\mu - \frac{\gamma \sigma}{2}\).
--
-- See https://en.wikipedia.org/wiki/Skewness .
--
{-# INLINE skewness #-}
skewness :: (Monad m, Floating a) => Fold m (a, Maybe a) a
skewness =
    toFold
        $ (\rm3 sd mu ->
            rm3 / sd ^ (3 :: Int) - 3 * (mu / sd) - (mu / sd) ^ (3 :: Int)
          )
        <$> Tee (rawMoment 3)
        <*> Tee stdDev
        <*> Tee mean

-- XXX We can compute the 2nd, 3rd, 4th raw moments by repeatedly multiplying
-- instead of computing the powers every time.
--
-- | Kurtosis \(\kappa\) is the standardized fourth central moment, defined as:
--
-- \(\tilde{\mu}_4 = \frac{\mu_4}{\sigma^4}\)
--
-- The fourth central moment can be computed in terms of raw moments:
--
-- \(\mu_4 = -3(\mu'_1)^4 + 6(\mu'_1)^2\mu'_2 - 4\mu'_1\mu'_3\ + \mu'_4\)
--
-- Substituting \(\mu'_1 = \mu\), and \(\mu'_2 = \mu^2 + \sigma^2\):
--
-- \(\mu_4 = 3\mu^4 + 6\mu^2\sigma^2 - 4\mu\mu'_3 + \mu'_4\)
--
-- It is always non-negative. It is 0 for a point distribution, low for light
-- tailed (platykurtic) distributions and high for heavy tailed (leptokurtic)
-- distributions.
--
-- \(\kappa >= \gamma^2 + 1\)
--
-- For a normal distribution \(\kappa = 3\sigma^4\).
--
-- See https://en.wikipedia.org/wiki/Kurtosis .
--
{-# INLINE kurtosis #-}
kurtosis :: (Monad m, Floating a) => Fold m (a, Maybe a) a
kurtosis =
    toFold
        $ (\rm4 rm3 sd mu ->
             ( 3 * mu ^ (4 :: Int)
            + 6 * mu ^ (2 :: Int) * sd ^ (2 :: Int)
            - 4 * mu * rm3
            + rm4) / (sd ^ (4 :: Int))
          )
        <$> Tee (rawMoment 4)
        <*> Tee (rawMoment 3)
        <*> Tee stdDev
        <*> Tee mean

-------------------------------------------------------------------------------
-- Estimation
-------------------------------------------------------------------------------

-- | Unbiased sample variance i.e. the variance of a sample corrected to
-- better estimate the variance of the population, defined as:
--
-- \(s^2 = \frac{1}{n - 1}\sum_{i=1}^n {(x_{i}-\mu)}^2\)
--
-- \(s^2 = \frac{n}{n - 1} \times \sigma^2\).
--
-- See https://en.wikipedia.org/wiki/Bessel%27s_correction.
--
{-# INLINE sampleVariance #-}
sampleVariance :: (Monad m, Fractional a) => Fold m (a, Maybe a) a
sampleVariance = Fold.teeWith (\n s2 -> n * s2 / (n - 1)) Window.length variance

-- | Sample standard deviation:
--
-- \(s = \sqrt{sampleVariance}\)
--
-- >>> sampleStdDev = sqrt <$> sampleVariance
--
-- See https://en.wikipedia.org/wiki/Unbiased_estimation_of_standard_deviation
-- .
--
{-# INLINE sampleStdDev #-}
sampleStdDev :: (Monad m, Floating a) => Fold m (a, Maybe a) a
sampleStdDev = sqrt <$> sampleVariance

-- | Standard error of the sample mean (SEM), defined as:
--
-- \( SEM = \frac{sampleStdDev}{\sqrt{n}} \)
--
-- See https://en.wikipedia.org/wiki/Standard_error .
--
-- /Space/: \(\mathcal{O}(1)\)
--
-- /Time/: \(\mathcal{O}(n)\)
{-# INLINE stdErrMean #-}
stdErrMean :: (Monad m, Floating a) => Fold m (a, Maybe a) a
stdErrMean = Fold.teeWith (\sd n -> sd / sqrt n) sampleStdDev Window.length

-------------------------------------------------------------------------------
-- Resampling
-------------------------------------------------------------------------------

{-# INLINE foldArray #-}
foldArray :: Storable a => Fold Identity a b -> Array a -> b
foldArray f = runIdentity . Stream.fold f . toStream

-- XXX Is this numerically stable? Should we keep the rounding error in the sum
-- and take it into account when subtracting?
--
-- | Given an array of @n@ items, compute mean of @(n - 1)@ items at a time,
-- producing a stream of all possible mean values omitting a different item
-- every time.
--
{-# INLINE jackKnifeMean #-}
jackKnifeMean :: (Monad m, Fractional a, Storable a) => Array a -> SerialT m a
jackKnifeMean arr = do
    let len = fromIntegral (length arr - 1)
        s = foldArray Fold.sum arr
     in Stream.map (\b -> (s - b) / len) $ toStream arr

-- | Given an array of @n@ items, compute variance of @(n - 1)@ items at a time,
-- producing a stream of all possible variance values omitting a different item
-- every time.
--
{-# INLINE jackKnifeVariance #-}
jackKnifeVariance :: (Monad m, Fractional a, Storable a) =>
    Array a -> SerialT m a
jackKnifeVariance arr = do
    let len = fromIntegral $ length arr - 1
        foldSums (s, s2) x = (s + x, s2 + x ^ (2 :: Int))
        (sum, sum2) = foldArray (Fold.foldl' foldSums (0.0, 0.0)) arr
        var x = (sum2 - x ^ (2 :: Int)) / len -  ((sum - x) / len) ^ (2::Int)
     in Stream.map var $ toStream arr

-- | Standard deviation computed from 'jackKnifeVariance'.
--
{-# INLINE jackKnifeStdDev #-}
jackKnifeStdDev :: (Monad m, Storable a, Floating a) =>
    Array a -> SerialT m a
jackKnifeStdDev = Stream.map sqrt . jackKnifeVariance

-- XXX This can be made more modular if the replicateM unfold can take count
-- from the seed.
--
-- | Randomly select elements from an array, with replacement, producing
-- a stream of the same size as the original array.
{-# INLINE resample #-}
resample :: (MonadIO m, Storable a) => Unfold m (Array a) a
resample = Unfold step inject

    where

    inject arr = liftIO $ do
        g <- createSystemRandom
        return $ (g, arr, length arr, 0)

    chooseOne g arr len = do
        i <- uniformRM (0, len - 1) g
        unsafeIndexIO i arr

    step (g, arr, len, idx) = liftIO $ do
        if idx >= len
        then return Stop
        else do
            e <- chooseOne g arr len
            return $ Yield e (g, arr, len, idx + 1)

-- | Resample an array multiple times and run the supplied fold on each
-- resampled stream, producing a stream of fold results. The fold is usually an
-- estimator fold.
{-# INLINE foldResamples #-}
foldResamples :: (MonadAsync m, Storable a) =>
       Int          -- ^ Number of resamples to compute.
    -> Array a      -- ^ Original sample.
    -> Fold m a b   -- ^ Estimator fold
    -> SerialT m b
foldResamples n arr fld =
    Stream.replicateM n (Unfold.fold fld resample arr)

-------------------------------------------------------------------------------
-- Probability Distribution
-------------------------------------------------------------------------------

-- | Determine the frequency of each element in the stream.
--
{-# INLINE frequency #-}
frequency :: (Monad m, Ord a) => Fold m a (Map a Int)
frequency = Fold.classifyWith id Fold.length

-- | Find out the most frequently ocurring element in the stream and its
-- frequency.
--
{-# INLINE mode #-}
mode :: (Monad m, Ord a) => Fold m a (Maybe (a, Int))
mode = Fold.rmapM findMax frequency

    where

    fmax k v Nothing = Just (k, v)
    fmax k v old@(Just (_, v1))
        | v > v1 = Just (k, v)
        | otherwise = old

    findMax = return . foldrWithKey fmax Nothing

-------------------------------------------------------------------------------
-- Histograms
-------------------------------------------------------------------------------

-- | @binOffsetSize offset binSize input@. Given an integral input value,
-- return its bin index provided that each bin contains @binSize@ items and the
-- bins are aligned such that the 0 index bin starts at @offset@ from 0. If
-- offset = 0 then the bin with index 0 would have values from 0 to binSize -
-- 1.
--
-- This API does not put a bound on the number of bins, therefore, the number
-- of bins could be potentially large depending on the range of values.
--
{-# INLINE binOffsetSize #-}
binOffsetSize :: Integral a => a -> a -> a -> a
binOffsetSize offset binSize x = (x - offset) `div` binSize

data HistBin a = BelowRange | InRange a | AboveRange deriving (Eq, Show)

instance (Eq a, Ord a) => Ord (HistBin a) where
    compare BelowRange BelowRange = EQ
    compare BelowRange (InRange _) = LT
    compare BelowRange AboveRange = LT

    compare (InRange _) BelowRange = GT
    compare (InRange x) (InRange y)= x `compare` y
    compare (InRange _) AboveRange = LT

    compare AboveRange BelowRange = GT
    compare AboveRange (InRange _) = GT
    compare AboveRange AboveRange = EQ

-- | @binFromSizeN low binSize nbins input@. Classify @input@ into bins
-- specified by a @low@ limit, @binSize@ and @nbins@. Inputs below the lower
-- limit are classified into 'BelowRange' and inputs above the highest bin are
-- classified into 'AboveRange'. 'InRange' inputs are classified into bins
-- starting from bin index 0.
--
{-# INLINE binFromSizeN #-}
binFromSizeN :: Integral a => a -> a -> a -> a -> HistBin a
binFromSizeN low binSize nbins x =
    let high = low + binSize * nbins
     in if x < low
        then BelowRange
        else if x >= high
             then AboveRange
             else InRange ((x - low) `div` binSize)

-- | @binFromToN low high nbins input@. Like @binFromSizeN@ except that a range
-- of lower and higher limit is specified. @binSize@ is computed using the
-- range and @nbins@. @nbins@ is rounded to the range @0 < nbins < (high - low
-- + 1)@. @high >= low@ must hold.
--
{-# INLINE binFromToN #-}
binFromToN :: Integral a => a -> a -> a -> a -> HistBin a
binFromToN low high n x =
    let count = high - low + 1
        n1 = max n 1
        n2 = min n1 count
        binSize = count `div` n2
        nbins =
            if binSize * n2 < count
            then n2 + 1
            else n2
     in assert (high >= low) (binFromSizeN low binSize nbins x)

-- Use binary search to find the bin
--
-- | Classify an input value to bins using the bin boundaries specified in an
-- array.
--
-- /Unimplemented/
--
{-# INLINE binBoundaries #-}
binBoundaries :: -- Integral a =>
    Array.Array a -> a -> HistBin a
binBoundaries = undefined

-- | Given a bin classifier function and a stream of values, generate a
-- histogram map from indices of bins to the number of items in the bin.
--
-- >>> Stream.fold (histogram (binOffsetSize 0 3)) $ Stream.fromList [1..15]
-- fromList [(0,2),(1,3),(2,3),(3,3),(4,3),(5,1)]
--
{-# INLINE histogram #-}
histogram :: (Monad m, Ord k) => (a -> k) -> Fold m a (Map k Int)
histogram bin = Fold.classifyWith bin Fold.length

ext: hs

max_stars_repo_name: composewell/streamly-statistics

max_issues_repo_name: composewell/streamly-statistics

max_forks_repo_name: composewell/streamly-statistics


text:
-----------------------------------------------------------------------------
-- |
-- Module      :  Main
-- Copyright   :
-- License     :  AllRightsReserved
--
-- Maintainer  :
-- Stability   :
--
-----------------------------------------------------------------------------

module Main (
    main
) where

--import Data.List
import qualified Data.ByteString.Lazy as B
import qualified Data.ByteString.Lazy.Char8 as BC
--import qualified Data.Double.Conversion.ByteString as DC
--import Workflow
import Test.WorkflowTest
--import Workflow.Toolset
--import Data.Array.Unboxed
--import Data.Complex
--import qualified Data.Map as Map

--import Test.DocTest
--
--
--
--main :: IO ()
--main = doctest ["-isrc", "src/Main.hs"]

--main = do
--    size <- readLn :: IO Int
--    let x1 = xs 100000
--        s1 = x1 `seq` if size > 0 then elems . RTESystem.Toolset.sort $ x1
--                                  else Data.List.sort . elems $ x1
--    writeCSV "data.csv" $ map (\x -> [x]) s1

main :: IO ()
main = do
    size <- readLn :: IO Int
    r <- test size 5.0
    writeCSV "data.csv" . map (\(x,y,z,u) -> [x,y,z,u]) $ r

--main = do
--    size <- readLn :: IO Int
--    let (bnds,wf) = calculateCDF wx 3.0 size
--        --st = flip (Map.!) 1 . nodes . task $ wf
--        x = fst lsp
--        t = snd lsp
--        qu = quantiles wf
--        qu' = charToDist lsp . waves $ wf
--        wa = waves wf
--        wa' = distToChar lsp . quantiles $ wf
--            lsp = linSpace size bnds
--    writeCSV "dataQ.csv" $ zipWith3 (\x a b -> [x,a,b]) x qu qu'
--    writeCSV "dataW.csv" $ zipWith3 (\t a b -> [show t,sh a,sh b]) t wa wa'
--        where sh (x:+y) = show x ++ (if y < 0 then " - " ++ show (-y) else " + " ++ show y) ++ "i"
    --g <- newStdGen --getStdGen
    --let vals = take size (randoms g :: [Double])
    --let vals = take size (randoms g :: [Double])
    --print vals
    --let distr = zipWith3 (\a b c -> [a,b,c]) vals (genGumbel g 2.45 6.23 size) (genWeibull g 3.0 1.8 size)
    --writeCSV "data.csv" distr
    --B.writeFile "data.csv" $ B.tail . B.tail . B.concat . map (BC.pack . (:) ';' . (:) '\n' . show) $ vals

writeCSV :: (Show a) => FilePath -> [[a]] -> IO()
writeCSV path [[]] = B.writeFile path B.empty
writeCSV path [] = B.writeFile path B.empty
writeCSV path (start:rest) = B.writeFile path $
    B.concat $ (:) (BC.pack . createline $ start) (map (BC.pack . (:) '\n' . createline) rest)
  where
    createline (x:xs) = show x ++ (concat . map ((:) ',' . show) $ xs)
    createline []     = ""

max_stars_repo_name: achirkin/haskell-estimatewf

max_issues_repo_name: achirkin/haskell-estimatewf

max_forks_repo_name: achirkin/haskell-estimatewf


text:
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}

{-# OPTIONS_GHC -fno-warn-missing-signatures #-}

-----------------------------------------------------------------------------
{- |
Module      :  Numeric.LinearAlgebra
Copyright   :  (c) Alberto Ruiz 2006-15
License     :  BSD3
Maintainer  :  Alberto Ruiz
Stability   :  provisional

-}
-----------------------------------------------------------------------------
module Numeric.LinearAlgebra (

    -- * Basic types and data manipulation
    -- | This package works with 2D ('Matrix') and 1D ('Vector')
    -- arrays of real ('R') or complex ('C') double precision numbers.
    -- Single precision and machine integers are also supported for
    -- basic arithmetic and data manipulation.
    module Numeric.LinearAlgebra.Data,

    -- * Numeric classes
    -- |
    -- The standard numeric classes are defined elementwise (commonly referred to
    -- as the Hadamard product or the Schur product):
    --
    -- >>>  vector [1,2,3] * vector [3,0,-2]
    -- [3.0,0.0,-6.0]
    -- it :: Vector R
    --
    -- >>> matrix 3 [1..9] * ident 3
    -- (3><3)
    --  [ 1.0, 0.0, 0.0
    --  , 0.0, 5.0, 0.0
    --  , 0.0, 0.0, 9.0 ]

    -- * Autoconformable dimensions
    -- |
    -- In most operations, single-element vectors and matrices
    -- (created from numeric literals or using 'scalar'), and matrices
    -- with just one row or column, automatically
    -- expand to match the dimensions of the other operand:
    --
    -- >>> 5 + 2*ident 3 :: Matrix Double
    -- (3><3)
    --  [ 7.0, 5.0, 5.0
    --  , 5.0, 7.0, 5.0
    --  , 5.0, 5.0, 7.0 ]
    --
    -- >>> (4><3) [1..] + row [10,20,30]
    -- (4><3)
    --  [ 11.0, 22.0, 33.0
    --  , 14.0, 25.0, 36.0
    --  , 17.0, 28.0, 39.0
    --  , 20.0, 31.0, 42.0 ]
    --

    -- * Products
    -- ** Dot
    dot, (<.>),
    -- ** Matrix-vector
    (#>), (<#), (!#>),
    -- ** Matrix-matrix
    (<>),
    -- | The matrix product is also implemented in the "Data.Monoid" instance, where
    -- single-element matrices (created from numeric literals or using 'scalar')
    -- are used for scaling.
    --
    -- >>> import Data.Monoid as M
    -- >>>  let m = matrix 3 [1..6]
    -- >>> m M.<> 2 M.<> diagl[0.5,1,0]
    -- (2><3)
    --  [ 1.0,  4.0, 0.0
    --  , 4.0, 10.0, 0.0 ]
    --
    -- 'mconcat' uses 'optimiseMult' to get the optimal association order.

    -- ** Other
    outer, kronecker, cross,
    scale, add,
    sumElements, prodElements,

    -- * Linear systems
    -- ** General
    (<\>),
    linearSolveLS,
    linearSolveSVD,
    -- ** Determined
    linearSolve,
    luSolve, luPacked,
    luSolve', luPacked',
    -- ** Symmetric indefinite
    ldlSolve, ldlPacked,
    -- ** Positive definite
    cholSolve,
    -- ** Triangular
    UpLo(..),
    triSolve,
    -- ** Tridiagonal
    triDiagSolve,
    -- ** Sparse
    cgSolve,
    cgSolve',

    -- * Inverse and pseudoinverse
    inv, pinv, pinvTol,

    -- * Determinant and rank
    rcond, rank,
    det, invlndet,

    -- * Norms
    Normed(..),
    norm_Frob, norm_nuclear,

    -- * Nullspace and range
    orth,
    nullspace, null1, null1sym,

    -- * Singular value decomposition
    svd,
    thinSVD,
    compactSVD,
    compactSVDTol,
    singularValues,
    leftSV, rightSV,

    -- * Eigendecomposition
    eig, geig, eigSH,
    eigenvalues, geigenvalues, eigenvaluesSH,
    geigSH,

    -- * QR
    qr, thinQR, rq, thinRQ, qrRaw, qrgr,

    -- * Cholesky
    chol, mbChol,

    -- * LU
    lu, luFact,

    -- * Hessenberg
    hess,

    -- * Schur
    schur,

    -- * Matrix functions
    expm,
    sqrtm,
    matFunc,

    -- * Correlation and convolution
    corr, conv, corrMin, corr2, conv2,

    -- * Random arrays

    Seed, RandDist(..), randomVector, rand, randn, gaussianSample, uniformSample,

    -- * Misc
    meanCov, rowOuters, pairwiseD2, normalize, peps, relativeError, magnit,
    haussholder, optimiseMult, udot, nullspaceSVD, orthSVD, ranksv,
    iC, sym, mTm, trustSym, unSym,
    -- * Auxiliary classes
    Element, Container, Product, Numeric, LSDiv, Herm,
    Complexable, RealElement,
    RealOf, ComplexOf, SingleOf, DoubleOf,
    IndexOf,
    Field, Linear(), Additive(),
    Transposable,
    LU(..),
    LDL(..),
    QR(..),
    CGState(..),
    Testable(..)
) where

import Numeric.LinearAlgebra.Data

import Numeric.Matrix()
import Numeric.Vector()
import Internal.Matrix
import Internal.Container hiding ((<>))
import Internal.Numeric hiding (mul)
import Internal.Algorithms hiding (linearSolve,Normed,orth,luPacked',linearSolve',luSolve',ldlPacked')
import qualified Internal.Algorithms as A
import Internal.Util
import Internal.Random
import Internal.Sparse((!#>))
import Internal.CG
import Internal.Conversion
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif

{- | dense matrix product

>>> let a = (3><5) [1..]
>>> a
(3><5)
 [  1.0,  2.0,  3.0,  4.0,  5.0
 ,  6.0,  7.0,  8.0,  9.0, 10.0
 , 11.0, 12.0, 13.0, 14.0, 15.0 ]

>>> let b = (5><2) [1,3, 0,2, -1,5, 7,7, 6,0]
>>> b
(5><2)
 [  1.0, 3.0
 ,  0.0, 2.0
 , -1.0, 5.0
 ,  7.0, 7.0
 ,  6.0, 0.0 ]

>>> a <> b
(3><2)
 [  56.0,  50.0
 , 121.0, 135.0
 , 186.0, 220.0 ]

-}
(<>) :: Numeric t => Matrix t -> Matrix t -> Matrix t
(<>) = mXm
infixr 8 <>

{- | Solve a linear system (for square coefficient matrix and several right-hand sides) using the LU decomposition, returning Nothing for a singular system. For underconstrained or overconstrained systems use 'linearSolveLS' or 'linearSolveSVD'.

@
a = (2><2)
 [ 1.0, 2.0
 , 3.0, 5.0 ]
@

@
b = (2><3)
 [  6.0, 1.0, 10.0
 , 15.0, 3.0, 26.0 ]
@

>>> linearSolve a b
Just (2><3)
 [ -1.4802973661668753e-15,     0.9999999999999997, 1.999999999999997
 ,       3.000000000000001, 1.6653345369377348e-16, 4.000000000000002 ]

>>> let Just x = it
>>> disp 5 x
2x3
-0.00000  1.00000  2.00000
 3.00000  0.00000  4.00000

>>> a <> x
(2><3)
 [  6.0, 1.0, 10.0
 , 15.0, 3.0, 26.0 ]

-}
linearSolve m b = A.mbLinearSolve m b

-- | return an orthonormal basis of the null space of a matrix. See also 'nullspaceSVD'.
nullspace m = nullspaceSVD (Left (1*eps)) m (rightSV m)

-- | return an orthonormal basis of the range space of a matrix. See also 'orthSVD'.
orth m = orthSVD (Left (1*eps)) m (leftSV m)

max_stars_repo_name: schnecki/hmatrix-float

max_issues_repo_name: schnecki/hmatrix-float

max_forks_repo_name: schnecki/hmatrix-float


text:
{-# language BangPatterns #-}
{-# language DataKinds #-}
{-# language FlexibleContexts #-}
{-# language KindSignatures #-}
{-# language LambdaCase #-}
{-# language PolyKinds #-}
{-# language StandaloneDeriving #-}
{-# language TypeFamilies #-}
{-# language TypeFamilyDependencies #-}
{-# language UndecidableInstances #-}

module Qprob where

import Control.Monad (join)
import Data.Coerce (coerce)
import Data.Kind (Type)
import Linear.Epsilon (nearZero)
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Complex as Complex

data Setting = Classical | Quantum

--type family Setting' s = (s' :: Type) | s' -> s

class V (s :: Setting) where
  type Setting' s = (s' :: Type) | s' -> s
  oneS :: Setting' s
  plusS :: Setting' s -> Setting' s -> Setting' s
  zeroS :: Setting' s
  timesS :: Setting' s -> Setting' s -> Setting' s
  negateS :: Setting' s -> Setting' s
  nearZeroS :: Setting' s -> Bool

instance V 'Classical where
  type Setting' 'Classical = Double
  oneS = one; plusS = plus; zeroS = zero; timesS = times; negateS = negate; nearZeroS = nearZero;
  {-# inline oneS #-}
  {-# inline plusS #-}
  {-# inline zeroS #-}
  {-# inline timesS #-}
  {-# inline negateS #-}
  {-# inline nearZeroS #-}

instance V 'Quantum where
  type Setting' 'Quantum = Complex Double
  oneS = one; plusS = plus; zeroS = zero; timesS = times; negateS = negate; nearZeroS = nearZero;
  {-# inline oneS #-}
  {-# inline plusS #-}
  {-# inline zeroS #-}
  {-# inline timesS #-}
  {-# inline negateS #-}
  {-# inline nearZeroS #-}

data Space (s :: Setting) a = Space !a !(Setting' s)

deriving instance (Show a, Show (Setting' s)) => Show (Space s a)
deriving instance (Eq a, Eq (Setting' s)) => Eq (Space s a)

mapSpace :: (Setting' s -> Setting' s) -> Space s a -> Space s a
mapSpace f (Space a s) = Space a (f s)

-- | @W s a@ is a vector space whose basis elements are labelled
--   by objects of type @a@ and where the coefficients are of type @Settings'' s@.
--
--   This is very similar to standard probability monads except that we
--   allow probabilities to be types other than 'Double'.
newtype W (s :: Setting) a = W { runW :: [Space s a] }

deriving instance (Show a, Show (Setting' s)) => Show (W s a)
deriving instance (Eq a, Eq (Setting' s)) => Eq (W s a)

-- | Transform the probabilities inside of 'W'.
mapW :: (Setting' s -> Setting' s) -> W s a -> W s a
mapW f (W l) = W (List.map (mapSpace f) l)

instance Semigroup (W s a) where
  W x <> W y = W (x <> y)
  {-# inline (<>) #-}

instance Monoid (W s a) where
  mempty = W mempty
  {-# inline mempty #-}

instance Functor (W s) where
  fmap f w = fmapW f w
  {-# inline fmap #-}

fmapW :: (a -> b) -> W s a -> W s b
fmapW f (W l) = W (List.map (\(Space a p) -> Space (f a) p) l)
{-# NOINLINE [1] fmapW #-}

{-# RULES "fmapW/coerce" fmapW coerce = coerce #-}

instance V s => Applicative (W s) where
  pure x = W [Space x oneS]
  {-# inline pure #-}
  W fs <*> W xs = W
    ( do Space f a <- fs
         Space x b <- xs
         pure (Space (f x) (timesS a b))
    )

infixl 7 .*
(.*) :: V s => Setting' s -> W s a -> W s a
a .* b = mapW (a `timesS`) b

type P a = W 'Classical a
type Q a = W 'Quantum   a

star :: Q a -> Q a
star = mapW Complex.conjugate

kron :: V s => W s a -> W s c -> W s (a,c)
kron (W x) (W y) = W
  ( do Space a r1 <- x
       Space c r2 <- y
       pure (Space (a,c) (r1 `timesS` r2))
  )

collect :: (V s, Ord a) => W s a -> W s a
collect (W l) = W $ toList . fromListWith plusS $ l
  where
    toList = Map.foldrWithKey (\k x xs -> (Space k x):xs) []
    fromListWith f xs = fromListWithKey (\_ x y -> f x y) xs
    fromListWithKey f xs =
      let ins t (Space a s) = Map.insertWithKey f a s t
      in List.foldl' ins mempty xs

instance V s => Monad (W s) where
  l >>= f = W $ List.concatMap (\(Space (W d) p) -> List.map (\(Space x q) -> (Space x (p `timesS` q))) d) (runW $ fmap f l)

-- | When we come to observe the state of a quantum system, the quantum state becomes an ordinary probablistic one.
observe :: Ord a => Q a -> P a
observe = W . List.map (\(Space a w) -> Space a (Complex.magnitude (w `times` w))) . runW . collect

rotate :: Double -> Bool -> Q Bool
rotate theta = \case
  True ->
    let theta' = theta :+ 0
    in subtract (cos (theta' / 2) .* pure True) (sin (theta' / 2) .* pure False)
  False ->
    let theta' = theta :+ 0
    in (cos (theta' / 2) .* pure False) <> (sin (theta' / 2) .* pure True)

repeat :: Int -> (a -> a) -> (a -> a)
repeat 0 _ = id
repeat n f = repeat (n - 1) f . f

repeatM :: Monad m => Int -> (a -> m a) -> m a -> m a
repeatM n f = repeat n (>>= f)

snot :: Bool -> Q Bool
snot = rotate (pi / 2)

snot1 :: Int -> P Bool
snot1 n = pure True & repeatM n snot & observe

subtract :: V s => W s a -> W s a -> W s a
subtract a b = a <> ((negateS oneS) .* b)

-- | Quantum Zeno effect.
--
-- A watched pot never boils.
zeno1 :: Int -> P Bool
zeno1 n = pure True & repeatM n (rotate (pi / fromIntegral n)) & collect & observe

ffor :: Functor f => f a -> (a -> b) -> f b
ffor = flip fmap

liftBind :: (Functor f, Monad g) => f (g a) -> (a -> g b) -> f (g b)
liftBind fga f = fmap (>>= f) fga

zeno2 :: Int -> P Bool
zeno2 n = pure True & repeat n
  (\x -> x `ffor` pure `liftBind` rotate (pi / fromIntegral n) `ffor` observe & join
  ) & collect

type MixedState a = P (Q a)

data Experimenter = Experimenter
  { experimenterMemory :: [Bool]
  , experimenterState :: !Bool
  }
  deriving (Eq, Ord)

zeno3 :: Int -> P Bool
zeno3 n = pure (Experimenter [] True) & repeatM n
  (\(Experimenter m s) -> do
    s' <- rotate (pi / fromIntegral n) s
    pure $ Experimenter (s:m) s'
  ) & observe & fmap experimenterState & collect

trimZero :: V s => W s a -> W s a
trimZero = W . List.filter (\(Space _ v) -> not $ nearZeroS v) . runW

simplify :: Ord a => Q a -> Q a
simplify = trimZero . collect

max_stars_repo_name: chessai/qprob

max_issues_repo_name: chessai/qprob

max_forks_repo_name: chessai/qprob


text:
{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Trustworthy #-}

-- |
-- Module      : Criterion.Report
-- Copyright   : (c) 2009-2014 Bryan O'Sullivan
--
-- License     : BSD-style
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : GHC
--
-- Reporting functions.

module Criterion.Report
    (
      formatReport
    , report
    , tidyTails
    -- * Rendering helper functions
    , TemplateException(..)
    , loadTemplate
    , includeFile
    , getTemplateDir
    , vector
    , vector2
    ) where

import Control.Exception (Exception, IOException, throwIO)
import Control.Monad (mplus, unless)
import Control.Monad.IO.Class (MonadIO(liftIO))
import Control.Monad.Reader (ask)
import Criterion.Monad (Criterion)
import Criterion.Types
import Data.Aeson (ToJSON (..), Value(..), object, (.=), Value, encode)
import Data.Data (Data, Typeable)
import Data.Foldable (forM_)
import GHC.Generics (Generic)
import Paths_criterion (getDataFileName)
import Statistics.Function (minMax)
import Statistics.Types (confidenceInterval, confidenceLevel, confIntCL, estError)
import System.Directory (doesFileExist)
import System.FilePath ((</>), (<.>), isPathSeparator)
import System.IO (hPutStrLn, stderr)
import Text.Microstache (Key (..), MustacheWarning (..), Node (..), Template (..),
                         compileMustacheText, displayMustacheWarning, renderMustacheW)
import Prelude ()
import Prelude.Compat
import qualified Control.Exception as E
import qualified Data.Text as T
import qualified Data.Text.Lazy.Encoding as TLE
import qualified Data.Text.IO as T
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.IO as TL
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U

#if defined(EMBED)
import Criterion.EmbeddedData (dataFiles, jQueryContents, flotContents,
                               flotErrorbarsContents, flotNavigateContents)
import qualified Data.ByteString.Lazy as BL
import qualified Data.Text.Encoding as TE
#else
import qualified Language.Javascript.Flot as Flot
import qualified Language.Javascript.JQuery as JQuery
#endif

-- | Trim long flat tails from a KDE plot.
tidyTails :: KDE -> KDE
tidyTails KDE{..} = KDE { kdeType   = kdeType
                        , kdeValues = G.slice front winSize kdeValues
                        , kdePDF    = G.slice front winSize kdePDF
                        }
  where tiny     = uncurry subtract (minMax kdePDF) * 0.005
        omitTiny = G.length . G.takeWhile ((<= tiny) . abs)
        front    = omitTiny kdePDF
        back     = omitTiny . G.reverse $ kdePDF
        winSize  = G.length kdePDF - front - back

-- | Return the path to the template and other files used for
-- generating reports.
--
-- When the @-fembed-data-files@ @Cabal@ flag is enabled, this simply
-- returns the empty path.
getTemplateDir :: IO FilePath
#if defined(EMBED)
getTemplateDir = pure ""
#else
getTemplateDir = getDataFileName "templates"
#endif

-- | Write out a series of 'Report' values to a single file, if
-- configured to do so.
report :: [Report] -> Criterion ()
report reports = do
  Config{..} <- ask
  forM_ reportFile $ \name -> liftIO $ do
    td <- getTemplateDir
    tpl <- loadTemplate [td,"."] template
    TL.writeFile name =<< formatReport reports tpl

-- | Format a series of 'Report' values using the given Mustache template.
formatReport :: [Report]
             -> TL.Text    -- ^ Mustache template.
             -> IO TL.Text
formatReport reports templateName = do
    template0 <- case compileMustacheText "tpl" templateName of
        Left err -> fail (show err) -- TODO: throw a template exception?
        Right x -> return x

    jQuery            <- jQueryFileContents
    flot              <- flotFileContents
    flotErrorbars     <- flotErrorbarsFileContents
    flotNavigate      <- flotNavigateFileContents
    jQueryCriterionJS <- readDataFile ("js" </> "jquery.criterion.js")
    criterionCSS      <- readDataFile "criterion.css"

    -- includes, only top level
    templates <- getTemplateDir
    template <- includeTemplate (includeFile [templates]) template0

    let context = object
            [ "json"                .= reports
            , "report"              .= map inner reports
            , "js-jquery"           .= jQuery
            , "js-flot"             .= flot
            , "js-flot-errorbars"   .= flotErrorbars
            , "js-flot-navigate"    .= flotNavigate
            , "jquery-criterion-js" .= jQueryCriterionJS
            , "criterion-css"       .= criterionCSS
            ]

    let (warnings, formatted) = renderMustacheW template context
    -- If there were any issues during mustache template rendering, make sure
    -- to inform the user. See #127.
    forM_ warnings $ \warning -> do
      -- The one thing we choose not to warn about is substituting in the `json`
      -- key. The reason is that `json` is used in:
      --
      --   var reports = {{{json}}};
      --
      -- So `json` represents a raw JavaScript array. This is a bit skeevy by
      -- mustache conventions, but redesigning the template to avoid this
      -- warning would be more work than just substituting the array directly.
      unless (warning == MustacheDirectlyRenderedValue (Key ["json"])) $
        mapM_ (hPutStrLn stderr)
         [ "criterion: warning:"
         , "  " ++ displayMustacheWarning warning
         , ""
         ]
    return formatted
  where
    jQueryFileContents, flotFileContents :: IO T.Text
#if defined(EMBED)
    jQueryFileContents        = pure $ TE.decodeUtf8 jQueryContents
    flotFileContents          = pure $ TE.decodeUtf8 flotContents
    flotErrorbarsFileContents = pure $ TE.decodeUtf8 flotErrorbarsContents
    flotNavigateFileContents  = pure $ TE.decodeUtf8 flotNavigateContents
#else
    jQueryFileContents        = T.readFile =<< JQuery.file
    flotFileContents          = T.readFile =<< Flot.file Flot.Flot
    flotErrorbarsFileContents = T.readFile =<< Flot.file Flot.FlotErrorbars
    flotNavigateFileContents  = T.readFile =<< Flot.file Flot.FlotNavigate
#endif

    readDataFile :: FilePath -> IO T.Text
    readDataFile fp =
      (T.readFile =<< getDataFileName ("templates" </> fp))
#if defined(EMBED)
      `E.catch` \(e :: IOException) ->
        maybe (throwIO e)
              (pure . TE.decodeUtf8)
              (lookup fp dataFiles)
#endif

    includeTemplate :: (FilePath -> IO T.Text) -> Template -> IO Template
    includeTemplate f Template {..} = fmap
        (Template templateActual)
        (traverse (traverse (includeNode f)) templateCache)

    includeNode :: (FilePath -> IO T.Text) -> Node -> IO Node
    includeNode f (Section (Key ["include"]) [TextBlock fp]) =
        fmap TextBlock (f (T.unpack fp))
    includeNode _ n = return n

    -- Merge Report with it's analysis and outliers
    merge :: ToJSON a => a -> Value -> Value
    merge x y = case toJSON x of
        Object x' -> case y of
            Object y' -> Object (x' <> y')
            _         -> y
        _         -> y

    inner r@Report {..} = merge reportAnalysis $ merge reportOutliers $ object
        [ "name"                  .= reportName
        , "json"                  .= TLE.decodeUtf8 (encode r)
        , "number"                .= reportNumber
        , "iters"                 .= vector "x" iters
        , "times"                 .= vector "x" times
        , "cycles"                .= vector "x" cycles
        , "kdetimes"              .= vector "x" kdeValues
        , "kdepdf"                .= vector "x" kdePDF
        , "kde"                   .= vector2 "time" "pdf" kdeValues kdePDF
        , "anMeanConfidenceLevel" .= anMeanConfidenceLevel
        , "anMeanLowerBound"      .= anMeanLowerBound
        , "anMeanUpperBound"      .= anMeanUpperBound
        , "anStdDevLowerBound"    .= anStdDevLowerBound
        , "anStdDevUpperBound"    .= anStdDevUpperBound
        ]
      where
        [KDE{..}]          = reportKDEs
        SampleAnalysis{..} = reportAnalysis

        iters  = measure measIters reportMeasured
        times  = measure measTime reportMeasured
        cycles = measure measCycles reportMeasured
        anMeanConfidenceLevel
               = confidenceLevel $ confIntCL $ estError anMean
        (anMeanLowerBound, anMeanUpperBound)
               = confidenceInterval anMean
        (anStdDevLowerBound, anStdDevUpperBound)
               = confidenceInterval anStdDev

-- | Render the elements of a vector.
--
-- It will substitute each value in the vector for @x@ in the
-- following Mustache template:
--
-- > {{#foo}}
-- >  {{x}}
-- > {{/foo}}
vector :: (G.Vector v a, ToJSON a) =>
          T.Text                -- ^ Name to use when substituting.
       -> v a
       -> Value
{-# SPECIALIZE vector :: T.Text -> U.Vector Double -> Value #-}
vector name v = toJSON . map val . G.toList $ v where
    val i = object [ name .= i ]

-- | Render the elements of two vectors.
vector2 :: (G.Vector v a, G.Vector v b, ToJSON a, ToJSON b) =>
           T.Text               -- ^ Name for elements from the first vector.
        -> T.Text               -- ^ Name for elements from the second vector.
        -> v a                  -- ^ First vector.
        -> v b                  -- ^ Second vector.
        -> Value
{-# SPECIALIZE vector2 :: T.Text -> T.Text -> U.Vector Double -> U.Vector Double
                       -> Value #-}
vector2 name1 name2 v1 v2 = toJSON $ zipWith val (G.toList v1) (G.toList v2) where
    val i j = object
        [ name1 .= i
        , name2 .= j
        ]

-- | Attempt to include the contents of a file based on a search path.
-- Returns 'B.empty' if the search fails or the file could not be read.
--
-- Intended for preprocessing Mustache files, e.g. replacing sections
--
-- @
-- {{#include}}file.txt{{/include}
-- @
--
-- with file contents.
includeFile :: (MonadIO m) =>
               [FilePath]       -- ^ Directories to search.
            -> FilePath         -- ^ Name of the file to search for.
            -> m T.Text
{-# SPECIALIZE includeFile :: [FilePath] -> FilePath -> IO T.Text #-}
includeFile searchPath name = liftIO $ foldr go (return T.empty) searchPath
    where go dir next = do
            let path = dir </> name
            T.readFile path `E.catch` \(_::IOException) -> next

-- | A problem arose with a template.
data TemplateException =
    TemplateNotFound FilePath   -- ^ The template could not be found.
    deriving (Eq, Read, Show, Typeable, Data, Generic)

instance Exception TemplateException

-- | Load a Mustache template file.
--
-- If the name is an absolute or relative path, the search path is
-- /not/ used, and the name is treated as a literal path.
--
-- If the @-fembed-data-files@ @Cabal@ flag is enabled, this also checks
-- the embedded @data-files@ from @criterion.cabal@.
--
-- This function throws a 'TemplateException' if the template could
-- not be found, or an 'IOException' if no template could be loaded.
loadTemplate :: [FilePath]      -- ^ Search path.
             -> FilePath        -- ^ Name of template file.
             -> IO TL.Text
loadTemplate paths name
    | any isPathSeparator name = readFileCheckEmbedded name
    | otherwise                = go Nothing paths
  where go me (p:ps) = do
          let cur = p </> name <.> "tpl"
          x <- doesFileExist' cur
          if x
            then readFileCheckEmbedded cur `E.catch` \e -> go (me `mplus` Just e) ps
            else go me ps
        go (Just e) _ = throwIO (e::IOException)
        go _        _ = throwIO . TemplateNotFound $ name

        doesFileExist' :: FilePath -> IO Bool
        doesFileExist' fp = do
          e <- doesFileExist fp
          pure $ e
#if defined(EMBED)
                 || (fp `elem` map fst dataFiles)
#endif

-- A version of 'readFile' that falls back on the embedded 'dataFiles'
-- from @criterion.cabal@.
readFileCheckEmbedded :: FilePath -> IO TL.Text
readFileCheckEmbedded fp =
  TL.readFile fp
#if defined(EMBED)
  `E.catch` \(e :: IOException) ->
    maybe (throwIO e)
          (pure . TLE.decodeUtf8 . BL.fromStrict)
          (lookup fp dataFiles)
#endif

max_stars_repo_name: rubenpieters/criterion

max_issues_repo_name: rubenpieters/criterion

max_forks_repo_name: rubenpieters/criterion


text:
module Lexer.TypeConversionSpec where

import Lexer.TypeConversion
import Lexer.Token
import Data.Complex
import Common.Number

import Text.Parsec (parse, Parsec, try, ParseError)

import Control.Exception (evaluate)
import Test.Hspec

spec :: Spec
spec = do

  describe "parseBoolean" $ do
    it "can parse short true" $ do
      parseBoolean "#t" `shouldBe` (Bool True)
    it "can parse short false" $ do
      parseBoolean "#f" `shouldBe` (Bool False)
    it "can parse long true" $ do
      parseBoolean "#true" `shouldBe` (Bool True)
    it "can parse long false" $ do
      parseBoolean "#false" `shouldBe` (Bool False)
    it "rejects non-boolean" $ do
      evaluate (parseBoolean "#blah") `shouldThrow` anyErrorCall

  describe "parseCharacter" $ do
    it "can parse single character" $ do
      parseCharacter "#\\e" `shouldBe` (Char 'e')

    it "can parse single unicode character" $ do
      parseCharacter "#\\\3619" `shouldBe` (Char '\3619')

    it "can parse single x" $ do
      parseCharacter "#\\x" `shouldBe` (Char 'x')

    it "can parse alarm" $ do
      parseCharacter "#\\alarm" `shouldBe` (Char '\a')

    it "can parse backspace" $ do
      parseCharacter "#\\backspace" `shouldBe` (Char '\b')

    it "can parse delete" $ do
      parseCharacter "#\\delete" `shouldBe` (Char '\DEL')

    it "can parse escape" $ do
      parseCharacter "#\\escape" `shouldBe` (Char '\ESC')

    it "can parse newline" $ do
      parseCharacter "#\\newline" `shouldBe` (Char '\n')

    it "can parse null" $ do
      parseCharacter "#\\null" `shouldBe` (Char '\0')

    it "can parse return" $ do
      parseCharacter "#\\return" `shouldBe` (Char '\r')

    it "can parse space" $ do
      parseCharacter "#\\space" `shouldBe` (Char ' ')

    it "can parse tab" $ do
      parseCharacter "#\\tab" `shouldBe` (Char '\t')

    it "can parse unicode escape" $ do
      parseCharacter "#\\x3a" `shouldBe` (Char ':')

    it "can parse chinese unicode escape" $ do
      parseCharacter "#\\x3AED" `shouldBe` (Char '\15085')

    it "can parse unicode escape with leading zero" $ do
      parseCharacter "#\\x003a" `shouldBe` (Char ':')

  describe "parseNumber" $ do
    it "can parse integer" $ do
      parseNumber "32\0" `shouldBe` (Number . Exact $ Integer 32)

    it "can parse negative integer" $ do
      parseNumber "-32\0" `shouldBe` (Number . Exact $ Integer (-32))

    it "can parse inexact integer" $ do
      parseNumber "#i3\0" `shouldBe` (Number . Inexact $ Integer 3)

    it "can parse exact integer" $ do
      parseNumber "#e3\0" `shouldBe` (Number . Exact $ Integer 3)

    it "can parse inexact decimal real" $ do
      parseNumber "#i3.141\0" `shouldBe` (Number . Inexact . Real $ InfReal 3.141)

    it "can parse decimal" $ do
      parseNumber "141e2\0" `shouldBe` (Number . Inexact . Real $ InfReal 14100)

    it "can parse decimal point only" $ do
      parseNumber ".131\0" `shouldBe` (Number . Inexact . Real $ InfReal 0.131)

    it "can parse decimal point only suffix" $ do
      parseNumber ".131e1\0" `shouldBe` (Number . Inexact . Real $ InfReal 1.31)

    it "can parse decimal point" $ do
      parseNumber "3.131\0" `shouldBe` (Number . Inexact . Real $ InfReal 3.131)

    it "can parse inexact decimal point" $ do
      parseNumber "#i3.131\0" `shouldBe` (Number . Inexact . Real $ InfReal 3.131)

    it "can parse inexact decimal point suffix" $ do
      parseNumber "#i3.131e2\0" `shouldSatisfy` (\(Number (Inexact (Real (InfReal n)))) -> abs(n - 313.1) < 0.000001)

    it "can parse rational" $ do
      parseNumber "1/2\0" `shouldBe` (Number . Exact $ Rational 1 2)

    it "can parse inexact rational" $ do
      parseNumber "#i1/2\0" `shouldBe` (Number . Inexact $ Rational 1 2)

  describe "parseString" $ do

    it "can parse string" $ do
      parseString "\"hello world\"" `shouldBe` String "hello world"

    it "can parse string with escaped quote" $ do
      parseString "\"hello\\\" world\"" `shouldBe` String "hello\" world"

    it "can parse string with escaped backslash" $ do
      parseString "\"hello\\\\ world\"" `shouldBe` String "hello\\ world"

    it "can parse mnemonic escape" $ do
      parseString "\"\\a\\b\\t\\n\\r\"" `shouldBe` String "\a\b\t\n\r"

    it "can parse unicode escape" $ do
      parseString "\"a\\x003a;\"" `shouldBe` String "a:"

    it "can parse multiline string" $ do
      parseString "\"ab\\  \n  c\"" `shouldBe` String "abc"

  describe "parseEscapedIdentifier" $ do

    it "can parse escaped identifier" $ do
      parseEscapedIdentifier "|hello abc \t\\| \\x003a; \\a \\b |" `shouldBe` Identifier "hello abc \t| : \a \b "

max_stars_repo_name: scheme-metamorph/metamorph

max_issues_repo_name: scheme-metamorph/metamorph

max_forks_repo_name: scheme-metamorph/metamorph