{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE ViewPatterns          #-}
-- |
-- Monoids for calculating various statistics in constant space
module Statistics.Monoid.Numeric (
    -- * Mean & Variance
    -- ** Number of elements
    CountG(..)
  , Count
  , asCount
    -- ** Mean
  , Mean
  , asMean
  , MeanKBN(..)
  , asMeanKBN
    -- ** Variance
  , Variance(..)
  , asVariance
    -- * Maximum and minimum
  , Max(..)
  , Min(..)
  , MaxD(..)
  , MinD(..)
    -- * Binomial trials
  , BinomAcc(..)
  , asBinomAcc
    -- * Accessors
  , CalcCount(..)
  , CalcMean(..)
  , CalcVariance(..)
  , calcStddev
  , calcStddevML
    -- * References
    -- $references
  ) where

import Data.Semigroup               (Semigroup(..))
import Data.Monoid                  (Monoid(..))
import Data.Data                    (Typeable,Data)
import Data.Vector.Unboxed          (Unbox)
import Data.Vector.Unboxed.Deriving (derivingUnbox)
import Numeric.Sum
import GHC.Generics                 (Generic)

import Statistics.Monoid.Class
-- COMPAT
import qualified Data.Vector.Generic         -- Needed for GHC7.4
import qualified Data.Vector.Generic.Mutable -- Needed for GHC7.4

----------------------------------------------------------------
-- Statistical monoids
----------------------------------------------------------------

-- | Calculate number of elements in the sample.
newtype CountG a = CountG { calcCountN :: a }
                  deriving (Show,Eq,Ord,Typeable)

type Count = CountG Int

-- | Type restricted 'id'
asCount :: CountG a -> CountG a
asCount = id

instance Integral a => Semigroup (CountG a) where
  (<>) = mappend
  {-# INLINE (<>) #-}

instance Integral a => Monoid (CountG a) where
  mempty                      = CountG 0
  CountG i `mappend` CountG j = CountG (i + j)
  {-# INLINE mempty  #-}
  {-# INLINE mappend #-}

instance (Integral a) => StatMonoid (CountG a) b where
  singletonMonoid _            = CountG 1
  addValue        (CountG n) _ = CountG (n + 1)
  {-# INLINE singletonMonoid #-}
  {-# INLINE addValue        #-}

instance CalcCount (CountG Int) where
  calcCount = calcCountN
  {-# INLINE calcCount #-}

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

-- | Type alias for currently recommended algorithms for calculation
--   of mean. It should be default choice
type Mean = MeanKBN

asMean :: Mean -> Mean
asMean = id

-- | Incremental calculation of mean. Sum of elements is calculated
--   using Kahan-Babuška-Neumaier summation.
data MeanKBN = MeanKBN !Int !KBNSum
             deriving (Show,Eq,Typeable,Data,Generic)

asMeanKBN :: MeanKBN -> MeanKBN
asMeanKBN = id

instance Semigroup MeanKBN where
  (<>) = mappend
  {-# INLINE (<>) #-}

instance Monoid MeanKBN where
  mempty = MeanKBN 0 mempty
  MeanKBN 0  _  `mappend` m             = m
  m             `mappend` MeanKBN 0  _  = m
  MeanKBN n1 s1 `mappend` MeanKBN n2 s2 = MeanKBN (n1+n2) (s1 `mappend` s2)

instance Real a => StatMonoid MeanKBN a where
  addValue (MeanKBN n m) x = MeanKBN (n+1) (addValue m x)

instance CalcCount MeanKBN where
  calcCount (MeanKBN n _) = n
instance CalcMean MeanKBN where
  calcMean (MeanKBN 0 _) = Nothing
  calcMean (MeanKBN n s) = Just (kbn s / fromIntegral n)

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

-- | This is algorithm for estimation of mean and variance of sample
--   which uses modified Welford algorithm. It uses KBN summation and
--   provides approximately 2 additional decimal digits
data VarWelfordKBN = VarWelfordKBN
  {-# UNPACK #-} !Int    --  Number of elements in the sample
  {-# UNPACK #-} !KBNSum -- Current sum of elements of sample
  {-# UNPACK #-} !KBNSum -- Current sum of squares of deviations from current mean

asVarWelfordKBN :: VarWelfordKBN -> VarWelfordKBN
asVarWelfordKBN = id

-- | Incremental algorithms for calculation the standard deviation.
data Variance = Variance {-# UNPACK #-} !Int    --  Number of elements in the sample
                         {-# UNPACK #-} !Double -- Current sum of elements of sample
                         {-# UNPACK #-} !Double -- Current sum of squares of deviations from current mean
                deriving (Show,Eq,Typeable)

-- | Type restricted 'id '
asVariance :: Variance -> Variance
asVariance = id
{-# INLINE asVariance #-}

instance Semigroup Variance where
  (<>) = mappend
  {-# INLINE (<>) #-}

-- | Iterative algorithm for calculation of variance [Chan1979]
instance Monoid Variance where
  mempty = Variance 0 0 0
  mappend (Variance n1 ta sa) (Variance n2 tb sb)
    = Variance (n1+n2) (ta+tb) sumsq
    where
      na = fromIntegral n1
      nb = fromIntegral n2
      nom = sqr (ta * nb - tb * na)
      sumsq | n1 == 0   = sb
            | n2 == 0   = sa
            | otherwise = sa + sb + nom / ((na + nb) * na * nb)
  {-# INLINE mempty #-}
  {-# INLINE mappend #-}

instance Real a => StatMonoid Variance a where
  addValue (Variance 0 _ _) x = singletonMonoid x
  addValue (Variance n t s) (realToFrac -> x)
    = Variance (n + 1) (t + x) (s + sqr (t  - n' * x) / (n' * (n'+1)))
    where
      n' = fromIntegral n
  singletonMonoid x = Variance 1 (realToFrac x) 0
  {-# INLINE singletonMonoid #-}

instance CalcCount Variance where
  calcCount (Variance n _ _) = n

instance CalcMean Variance where
  calcMean (Variance 0 _ _) = Nothing
  calcMean (Variance n s _) = Just (s / fromIntegral n)

instance CalcVariance Variance where
  calcVariance (Variance n _ s)
    | n < 2     = Nothing
    | otherwise = Just $! s / fromIntegral (n - 1)
  calcVarianceML (Variance n _ s)
    | n < 1     = Nothing
    | otherwise = Just $! s / fromIntegral n

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

-- | Calculate minimum of sample
newtype Min a = Min { calcMin :: Maybe a }
              deriving (Show,Eq,Ord,Typeable,Data,Generic)

instance Ord a => Semigroup (Min a) where
  (<>) = mappend
  {-# INLINE (<>) #-}

instance Ord a => Monoid (Min a) where
  mempty = Min Nothing
  Min (Just a) `mappend` Min (Just b) = Min (Just $! min a b)
  Min a        `mappend` Min Nothing  = Min a
  Min Nothing  `mappend` Min b        = Min b

instance (Ord a, a ~ a') => StatMonoid (Min a) a' where
  singletonMonoid a = Min (Just a)

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

-- | Calculate maximum of sample
newtype Max a = Max { calcMax :: Maybe a }
              deriving (Show,Eq,Ord,Typeable,Data,Generic)

instance Ord a => Semigroup (Max a) where
  (<>) = mappend
  {-# INLINE (<>) #-}

instance Ord a => Monoid (Max a) where
  mempty = Max Nothing
  Max (Just a) `mappend` Max (Just b) = Max (Just $! min a b)
  Max a        `mappend` Max Nothing  = Max a
  Max Nothing  `mappend` Max b        = Max b

instance (Ord a, a ~ a') => StatMonoid (Max a) a' where
  singletonMonoid a = Max (Just a)

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

-- | Calculate minimum of sample of Doubles. For empty sample returns NaN. Any
--   NaN encountered will be ignored.
newtype MinD = MinD { calcMinD :: Double }
              deriving (Show,Typeable,Data,Generic)

instance Eq MinD where
  MinD a == MinD b
    | isNaN a && isNaN b = True
    | otherwise          = a == b

instance Semigroup MinD where
  (<>) = mappend
  {-# INLINE (<>) #-}

-- N.B. forall (x :: Double) (x <= NaN) == False
instance Monoid MinD where
  mempty = MinD (0/0)
  mappend (MinD x) (MinD y)
    | isNaN x   = MinD y
    | isNaN y   = MinD x
    | otherwise = MinD (min x y)
  {-# INLINE mempty  #-}
  {-# INLINE mappend #-}

instance a ~ Double => StatMonoid MinD a where
  singletonMonoid = MinD

-- | Calculate maximum of sample. For empty sample returns NaN. Any
--   NaN encountered will be ignored.
newtype MaxD = MaxD { calcMaxD :: Double }
              deriving (Show,Typeable,Data,Generic)

instance Eq MaxD where
  MaxD a == MaxD b
    | isNaN a && isNaN b = True
    | otherwise          = a == b

instance Semigroup MaxD where
  (<>) = mappend
  {-# INLINE (<>) #-}

instance Monoid MaxD where
  mempty = MaxD (0/0)
  mappend (MaxD x) (MaxD y)
    | isNaN x   = MaxD y
    | isNaN y   = MaxD x
    | otherwise = MaxD (max x y)
  {-# INLINE mempty  #-}
  {-# INLINE mappend #-}

instance a ~ Double => StatMonoid MaxD a where
  singletonMonoid = MaxD

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

-- | Accumulator for binomial trials.
data BinomAcc = BinomAcc { binomAccSuccess :: !Int
                         , binomAccTotal   :: !Int
                         }
  deriving (Show,Eq,Ord,Typeable,Data,Generic)

-- | Type restricted 'id'
asBinomAcc :: BinomAcc -> BinomAcc
asBinomAcc = id

instance Semigroup BinomAcc where
  (<>) = mappend
  {-# INLINE (<>) #-}

instance Monoid BinomAcc where
  mempty = BinomAcc 0 0
  mappend (BinomAcc n1 m1) (BinomAcc n2 m2) = BinomAcc (n1+n2) (m1+m2)

instance StatMonoid BinomAcc Bool where
  addValue (BinomAcc nS nT) True  = BinomAcc (nS+1) (nT+1)
  addValue (BinomAcc nS nT) False = BinomAcc  nS    (nT+1)

----------------------------------------------------------------
-- Helpers
----------------------------------------------------------------

sqr :: Double -> Double
sqr x = x * x
{-# INLINE sqr #-}

----------------------------------------------------------------
-- Unboxed instances
----------------------------------------------------------------

derivingUnbox "CountG"
  [t| forall a. Unbox a => CountG a -> a |]
  [| calcCountN |]
  [| CountG     |]

derivingUnbox "MeanKBN"
  [t| MeanKBN -> (Int,Double,Double) |]
  [| \(MeanKBN a (KBNSum b c)) -> (a,b,c)   |]
  [| \(a,b,c) -> MeanKBN a (KBNSum b c) |]

derivingUnbox "Variance"
  [t| Variance -> (Int,Double,Double) |]
  [| \(Variance a b c) -> (a,b,c)  |]
  [| \(a,b,c) -> Variance a b c    |]

derivingUnbox "MinD"
  [t| MinD -> Double |]
  [| calcMinD |]
  [| MinD     |]

derivingUnbox "MaxD"
  [t| MaxD -> Double |]
  [| calcMaxD |]
  [| MaxD     |]

-- $references
--
-- * [Welford1962] Welford, B.P. (1962) Note on a method for
--   calculating corrected sums of squares and
--   products. /Technometrics/
--   4(3):419-420. <http://www.jstor.org/stable/1266577>
--
-- * [Chan1979] Chan, Tony F.; Golub, Gene H.; LeVeque, Randall
--   J. (1979), Updating Formulae and a Pairwise Algorithm for
--   Computing Sample Variances., Technical Report STAN-CS-79-773,
--   Department of Computer Science, Stanford University. Page 4.


{-# LANGUAGE FlexibleContexts, Rank2Types, NoMonomorphismRestriction,
             StandaloneDeriving #-}

module Numeric.MaxEnt.Linear where

import Control.Applicative

import Data.List (transpose)
import qualified Data.Vector.Storable as S

import Numeric.MaxEnt.ConjugateGradient (minimize, dot)
import Numeric.Optimization.Algorithms.HagerZhang05 (Result, Statistics)
import Numeric.AD

multMV :: (Num a) => [[a]] -> [a] -> [a]
multMV mat vec = map (\row -> dot row vec) mat

probs :: (Floating a) => [[a]] -> [a] -> [a]
probs matrix ls = result where
    norm = partitionFunc matrix ls
    result = map (\x -> exp x / norm ) $ (transpose matrix) `multMV` ls

partitionFunc :: (Floating a) => [[a]] -> [a] -> a
partitionFunc matrix ws = sum . map exp . multMV (transpose matrix) $ ws

-- This is almost the sam as the objectiveFunc

objectiveFunc :: (Floating a) => [[a]] -> [a] -> [a] -> a
objectiveFunc as moments ls = log $ partitionFunc as ls - dot ls moments

data LinearConstraints = LC
  { unLC :: forall a. (Floating a) => ([[a]], [a]) }

-- These instances default the underlying numeric type of `LC` to `Double`,
-- which may be problematic for some usages.
deriving instance Eq LinearConstraints
deriving instance Show LinearConstraints

-- | This is for the linear case Ax = b
--   @x@ in this situation is the vector of probablities.
--
--  Consider the 1 dimensional circular convolution using hmatrix.
--
--  >>> import Numeric.LinearAlgebra
--  >>> fromLists [[0.68, 0.22, 0.1], [0.1, 0.68, 0.22], [0.22, 0.1, 0.68]] <> fromLists [[0.2], [0.5], [0.3]]
--  (3><1) [0.276, 0.426, 0.298]
--
--   Now if we were given just the convolution and the output, we can use 'linear' to infer the input.
--
--   >>> linear 3.0e-17 $ LC ([[0.68, 0.22, 0.1], [0.1, 0.68, 0.22], [0.22, 0.1, 0.68]], [0.276, 0.426, 0.298])
--   Right (fromList [0.2000000000000001,0.49999999999999983,0.30000000000000004])
--
--   I fell compelled to point out that we could also just invert the original
--   convolution matrix. Supposedly using maxent can reduce errors from noise if
--   the convolution matrix is not properly estimated.
linear :: Double
       -- ^ Tolerance for the numerical solver
       -> LinearConstraints
       -- ^ The matrix A and column vector b
       -> Either (Result, Statistics) (S.Vector Double)
       -- ^ Either a description of what went wrong or the probability
       --   distribution
linear tolerance constraints  =
    let (matrix, output) = unLC constraints
        obj = objectiveFunc matrix output
        n = length output
    in (S.fromList . probs matrix . S.toList) <$> minimize tolerance n obj

--------------------------------------------------------------------------------
-- I updated everything below to work with the new types, but it's not clear to
-- me what it's for.  -- EP
--------------------------------------------------------------------------------

linear' :: (Floating a, Ord a)
        => LinearConstraints
        -- ^ The matrix A and column vector b
        -> [[a]]
        -- ^ Either a description of what went wrong or the probability
        --   distribution
linear' constraints =
    let (matrix, output) = unLC constraints
        obj = objectiveFunc matrix output
        guess = 1 : replicate (length output - 1) 0
    in map (probs matrix) . gradientDescent obj $ guess

linear'' :: (Floating a, Ord a)
         => LinearConstraints
         -- ^ The matrix A and column vector b
         -> [[a]]
         -- ^ Either a description of what went wrong or the probability
         --   distribution
linear'' constraints =
    let (matrix, output) = unLC constraints
        obj = objectiveFunc matrix output
        guess = 1 : replicate (length output - 1) 0
    in map (probs matrix) . conjugateGradientDescent obj $ guess

--test1 = LC ( [ [0.892532,0.003851,0.063870,0.001593,0.038155]
--             , [0.237713,0.111149,0.326964,0.271535,0.052639]
--             , [0.133708,0.788233,0.051543,0.003976,0.022539]
--             , [0.238064,0.263171,0.112279,0.270452,0.116034]
--             , [0.844155,0.011312,0.001470,0.001826,0.141237]
--             ]
--           ,
--             [0.246323,0.235600,0.071699,0.211339,0.238439]
--           )


{-# LANGUAGE RecordWildCards #-}
module Statistics.BBVI.StepSize
  ( KucP(..)
  , rhoKuc
  , defaultKucP
  , Rho
  )
where

import qualified Data.Vector                   as V
import           Statistics.BBVI.Propagator     ( DistCell(..)
                                                , Gradient
                                                , Memory
                                                )

-- | step size is a vector the same length as n parameters in
-- the distribution
type Rho = V.Vector Double

-- | calculates a step size following Kucukelbir et al 2017
rhoKuc
  :: KucP -- ^ parameters
  -> Gradient -- ^ gradient at time t
  -> DistCell a -- ^ distribution cell (needed for memory data
                -- e.g. \( s_{t-1} \) in Kucukelbir et al 2017)
  -> (Memory, Rho) -- ^ change in memory, stepsize
rhoKuc KucP {..} gra (Node time memory _dist) =
  ( deltaM
  , V.zipWith
    (\ds s ->
      eta
        *  (fromIntegral time)
        ** (negate 0.5 + eps)
        *  (1.0 / (tau + sqrt (s + ds)))
    )
    deltaM
    memory
  )
 where
  deltaM = V.zipWith (\g s -> alpha * g ^ (2 :: Int) - alpha * s) gra memory
rhoKuc _kp _gra (U{}) = error "called step size cell on an update cell!"

-- | parameters for 'rhoKuc'
data KucP = KucP
  { alpha :: Double
  , eta :: Double
  , tau :: Double
  , eps :: Double
  } deriving (Show, Eq, Ord, Read)

-- | default parameters for 'rhoKuc.' eta is what you probably want to
-- tune: kucukelbir et al try 0.01 0.1 1 10 100
defaultKucP :: KucP
defaultKucP = KucP { alpha = 0.1, eta = 0.1, tau = 1.0, eps = 1e-16 }


{-# LANGUAGE FlexibleInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Foreign.LAPACK
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Types with a Linear Algebra PACKage.
--
module Foreign.LAPACK (
    -- * LAPACK typeclass
    LAPACK(..),
    -- * Enums
    module Foreign.LAPACK.Types,
    ) where

import Control.Monad( forM_, when )
import Control.Exception( assert )
import Data.Complex( Complex )
import Foreign
import Foreign.BLAS
import Foreign.LAPACK.Types
import Foreign.LAPACK.Double
import Foreign.LAPACK.Zomplex

-- | Types with LAPACK operations.
class (BLAS3 e) => LAPACK e where
    geqrf :: Int -> Int -> Ptr e -> Int -> Ptr e -> IO ()
    gelqf :: Int -> Int -> Ptr e -> Int -> Ptr e -> IO ()

    heevr :: EigJob -> EigRange -> Uplo -> Int -> Ptr e -> Int -> Double
          -> Ptr Double -> Ptr e -> Int -> Ptr Int -> IO Int

    larfg :: Int -> Ptr e -> Ptr e -> Int -> IO e
    potrf :: Uplo -> Int -> Ptr e -> Int -> IO Int
    potrs :: Uplo -> Int -> Int -> Ptr e -> Int -> Ptr e -> Int -> IO ()
    pptrf :: Uplo -> Int -> Ptr e -> IO Int
    pptrs :: Uplo -> Int -> Int -> Ptr e -> Ptr e -> Int -> IO ()

    unmqr :: Side -> Trans -> Int -> Int -> Int -> Ptr e -> Int -> Ptr e -> Ptr e -> Int -> IO ()
    unmlq :: Side -> Trans -> Int -> Int -> Int -> Ptr e -> Int -> Ptr e -> Ptr e -> Int -> IO ()

callWithWork :: (Storable e) => (Ptr e -> Ptr LAInt -> Ptr LAInt -> IO ()) -> IO LAInt
callWithWork call =
    alloca $ \pquery ->
    with (-1) $ \pnone ->
    alloca $ \pinfo -> do
        call pquery pnone pinfo
        lwork <- (ceiling . max 1) `fmap` (peek (castPtr pquery) :: IO Double)
        allocaArray lwork $ \pwork ->
            withEnum lwork $ \plwork -> do
                call pwork plwork pinfo
                peek pinfo

callWithWorkIWork :: (Storable e)
    => (Ptr e -> Ptr LAInt -> Ptr LAInt -> Ptr LAInt -> Ptr LAInt -> IO ())
    -> IO LAInt
callWithWorkIWork call =
    alloca $ \pquery ->
    alloca $ \piquery ->
    with (-1) $ \pnone ->
    alloca $ \pinfo -> do
        call pquery pnone piquery pnone pinfo
        lwork <- (ceiling . max 1) `fmap` (peek (castPtr pquery) :: IO Double)
        liwork <- max 1 `fmap` peek piquery
        allocaArray lwork $ \pwork ->
            withEnum lwork $ \plwork ->
                allocaArray (fromEnum liwork) $ \piwork ->
                    with liwork $ \pliwork -> do
                        call pwork plwork piwork pliwork pinfo
                        peek pinfo

callWithWorkRWorkIWork :: (Storable e)
    => (Ptr e -> Ptr LAInt -> Ptr Double -> Ptr LAInt -> Ptr LAInt -> Ptr LAInt -> Ptr LAInt -> IO ())
    -> IO LAInt
callWithWorkRWorkIWork call =
    alloca $ \pquery ->
    alloca $ \prquery ->
    alloca $ \piquery ->
    with (-1) $ \pnone ->
    alloca $ \pinfo -> do
        call pquery pnone prquery pnone piquery pnone pinfo
        lwork <- (ceiling . max 1) `fmap` (peek (castPtr pquery) :: IO Double)
        lrwork <- (ceiling . max 1) `fmap` peek prquery
        liwork <- max 1 `fmap` peek piquery
        allocaArray lwork $ \pwork ->
            withEnum lwork $ \plwork ->
                allocaArray lrwork $ \prwork ->
                    withEnum lrwork $ \plrwork ->
                        allocaArray (fromEnum liwork) $ \piwork ->
                            with liwork $ \pliwork -> do
                                call pwork plwork prwork plrwork piwork pliwork pinfo
                                peek pinfo

checkInfo :: (Num info) => info -> IO ()
checkInfo info = assert (info == 0) $ return ()

withEnum :: (Enum a, Storable a) => Int -> (Ptr a -> IO b) -> IO b
withEnum = with . toEnum
{-# INLINE withEnum #-}

instance LAPACK Double where
    geqrf m n pa lda ptau =
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
            checkInfo =<< callWithWork (dgeqrf pm pn pa plda ptau)

    gelqf m n pa lda ptau =
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
            checkInfo =<< callWithWork (dgelqf pm pn pa plda ptau)

    heevr jobz range uplo n pa lda abstol pw pz ldz psuppz =
        withEigJob jobz $ \pjobz ->
        withEigRange range $ \prange pvl pvu pil piu ->
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
        with abstol $ \pabstol ->
        alloca $ \pm ->
        withEnum ldz $ \pldz ->
        allocaArray nsuppz $ \psuppz' -> do
            checkInfo =<<
                (callWithWorkIWork $
                    dsyevr pjobz prange puplo pn pa plda pvl pvu pil piu
                           pabstol pm pw pz pldz psuppz')
            m <- fromEnum `fmap` peek pm
            when (psuppz /= nullPtr) $
                forM_ [ 0..2*m - 1 ] $ \i -> do
                    s <- peekElemOff psuppz' i
                    pokeElemOff psuppz i (fromEnum s)
            return m
      where
        nsuppz = if psuppz == nullPtr then 0 else 2 * max 1 n

    unmqr side trans m n k pa lda ptau pc ldc =
        withSide side $ \pside ->
        withTrans trans $ \ptrans ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        withEnum lda $ \plda ->
        withEnum ldc $ \pldc ->
            checkInfo =<< callWithWork (dormqr pside ptrans pm pn pk pa plda ptau pc pldc)

    unmlq side trans m n k pa lda ptau pc ldc =
        withSide side $ \pside ->
        withTrans trans $ \ptrans ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        withEnum lda $ \plda ->
        withEnum ldc $ \pldc ->
            checkInfo =<< callWithWork (dormlq pside ptrans pm pn pk pa plda ptau pc pldc)

    larfg n palpha px incx =
        withEnum n $ \pn ->
        withEnum incx $ \pincx ->
        alloca $ \ptau -> do
            dlarfg pn palpha px pincx ptau
            peek ptau

    potrf uplo n pa lda =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
        alloca $ \pinfo -> do
            dpotrf puplo pn pa plda pinfo
            info <- fromEnum `fmap` peek pinfo
            assert (info >= 0) $ return info

    potrs uplo n nrhs pa lda pb ldb =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum nrhs $ \pnrhs ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        alloca $ \pinfo -> do
            dpotrs puplo pn pnrhs pa plda pb pldb pinfo
            checkInfo =<< peek pinfo

    pptrf uplo n pa =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        alloca $ \pinfo -> do
            dpptrf puplo pn pa pinfo
            info <- fromEnum `fmap` peek pinfo
            assert (info >= 0) $ return info

    pptrs uplo n nrhs pa pb ldb =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum nrhs $ \pnrhs ->
        withEnum ldb $ \pldb ->
        alloca $ \pinfo -> do
            dpptrs puplo pn pnrhs pa pb pldb pinfo
            checkInfo =<< peek pinfo

instance LAPACK (Complex Double) where
    geqrf m n pa lda ptau =
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
            checkInfo =<< callWithWork (zgeqrf pm pn pa plda ptau)

    gelqf m n pa lda ptau =
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
            checkInfo =<< callWithWork (zgelqf pm pn pa plda ptau)

    heevr jobz range uplo n pa lda abstol pw pz ldz psuppz =
        withEigJob jobz $ \pjobz ->
        withEigRange range $ \prange pvl pvu pil piu ->
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
        with abstol $ \pabstol ->
        alloca $ \pm ->
        withEnum ldz $ \pldz ->
        allocaArray nsuppz $ \psuppz' -> do
            checkInfo =<<
                (callWithWorkRWorkIWork $
                    zheevr pjobz prange puplo pn pa plda pvl pvu pil piu
                           pabstol pm pw pz pldz psuppz')
            m <- fromEnum `fmap` peek pm
            when (psuppz /= nullPtr) $
                forM_ [ 0..2*m - 1 ] $ \i -> do
                    s <- peekElemOff psuppz' i
                    pokeElemOff psuppz i (fromEnum s)
            return m
      where
        nsuppz = if psuppz == nullPtr then 0 else 2 * max 1 n

    unmqr side trans m n k pa lda ptau pc ldc =
        withSide side $ \pside ->
        withTrans trans $ \ptrans ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        withEnum lda $ \plda ->
        withEnum ldc $ \pldc ->
            checkInfo =<< callWithWork (zunmqr pside ptrans pm pn pk pa plda ptau pc pldc)

    unmlq side trans m n k pa lda ptau pc ldc =
        withSide side $ \pside ->
        withTrans trans $ \ptrans ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        withEnum lda $ \plda ->
        withEnum ldc $ \pldc ->
            checkInfo =<< callWithWork (zunmlq pside ptrans pm pn pk pa plda ptau pc pldc)

    larfg n palpha px incx =
        withEnum n $ \pn ->
        withEnum incx $ \pincx ->
        alloca $ \ptau -> do
            zlarfg pn palpha px pincx ptau
            peek ptau

    potrf uplo n pa lda =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
        alloca $ \pinfo -> do
            zpotrf puplo pn pa plda pinfo
            info <- fromEnum `fmap` peek pinfo
            assert (info >= 0) $ return info

    potrs uplo n nrhs pa lda pb ldb =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum nrhs $ \pnrhs ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        alloca $ \pinfo -> do
            zpotrs puplo pn pnrhs pa plda pb pldb pinfo
            checkInfo =<< peek pinfo

    pptrf uplo n pa =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        alloca $ \pinfo -> do
            zpptrf puplo pn pa pinfo
            info <- fromEnum `fmap` peek pinfo
            assert (info >= 0) $ return info

    pptrs uplo n nrhs pa pb ldb =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum nrhs $ \pnrhs ->
        withEnum ldb $ \pldb ->
        alloca $ \pinfo -> do
            zpptrs puplo pn pnrhs pa pb pldb pinfo
            checkInfo =<< peek pinfo


{-# LANGUAGE FlexibleContexts #-}
module Filter.Pinwheel
  ( module Filter.Utils
  , PinwheelParams(..)
  , pinwheel
  , PinwheelType(..)
  , pinwheelFunc
  , pinwheelFilter
  , convolvePinwheel
  , frequencyDomainMultiply
  ) where

import           Data.Array.Repa     as R
import           Data.Complex
import           Data.List           as L (length, reverse)
import           Data.Vector.Generic as VG (convert)
import           DFT.Plan
import           Filter.Utils
import           Utils.Coordinates

data PinwheelParams = PinwheelParams
  { pinwheelParamsRows         :: !Int
  , pinwheelParamsCols         :: !Int
  , pinwheelParamsAlpha        :: !Double
  , pinwheelParamsRadialPeriod :: !Double
  , pinwheelParamsAngularFreqs :: ![Double]
  , pinwheelParamsRadialFreqs  :: ![Double]
  }

data PinwheelType
  = Pinwheel
  | PinwheelHollow0 !Double
  | PinwheelHollow1 !Double
  deriving (Read, Show)

{-# INLINE rFunc #-}
rFunc :: Int -> Int -> Double
rFunc x y = sqrt . fromIntegral $ x ^ (2 :: Int) + y ^ (2 :: Int)

{-# INLINE thetaFunc #-}
thetaFunc :: Int -> Int -> Double
thetaFunc x y = angleFunctionRad (fromIntegral x) (fromIntegral y)

{-# INLINE pinwheel #-}
pinwheel :: Double -> Double -> Double -> Double -> Int -> Int -> Complex Double
pinwheel af rf maxR alpha x y
  | r == 0 && af == 0 && rf == 0 = 1
  | r == 0 = 0
  | otherwise =
    (r :+ 0) ** (alpha :+ rf * 2 * pi / log maxR) *
    exp (0 :+ af * (thetaFunc x y))
  where
    r = rFunc x y

{-# INCLUDE pinwheelFunc #-}
pinwheelFunc ::
     PinwheelType
  -> Double
  -> Double
  -> Double
  -> Double
  -> Int
  -> Int
  -> Complex Double
pinwheelFunc Pinwheel af rf maxR alpha x y
  | r == 0 && af == 0 && rf == 0 = 1
  | r == 0 = 0
  | otherwise =
    (r :+ 0) ** (alpha :+ rf * 2 * pi / logMaxR) *
    exp (0 :+ af * (thetaFunc x y))
  where
    r = rFunc x y
    logMaxR =
      if maxR == 1
        then 1
        else log maxR
pinwheelFunc (PinwheelHollow0 radius) af rf maxR alpha x y
  | r < radius = 0.0
  | otherwise =
    (r :+ 0) ** (alpha :+ rf * 2 * pi / logMaxR) *
    exp (0 :+ af * (thetaFunc x y))
  where
    r = rFunc x y
    logMaxR =
      if maxR == 1
        then 1
        else log maxR
pinwheelFunc (PinwheelHollow1 radius) af rf maxR alpha x y
  | r == 0 && af == 0 && rf == 0 = 1
  | r < radius = 0.0
  | otherwise =
    (r :+ 0) ** (alpha :+ rf * 2 * pi / logMaxR) *
    exp (0 :+ af * (thetaFunc x y))
  where
    r = rFunc x y
    logMaxR =
      if maxR == 1
        then 1
        else log maxR

{-# INLINE pinwheelPI #-}
pinwheelPI :: Double -> Double -> Double -> Double -> Int -> Int -> Complex Double
pinwheelPI af rf maxR alpha x y
  | r == 0 = 0
  | otherwise =
    (((r) :+ 0) ** (alpha :+ (rf * 2 * pi / (log maxR)))) *
    exp (0 :+ ((af) * theta))
  where
    r = (sqrt . fromIntegral $ x ^ (2 :: Int) + y ^ (2 :: Int))
    theta = pi  + angleFunctionRad (fromIntegral x) (fromIntegral y)

pinwheelFilter ::
     DFTPlan
  -> PinwheelParams
  -> IO (DFTPlan, Array U DIM4 (Complex Double), Array U DIM4 (Complex Double))
pinwheelFilter plan (PinwheelParams rows cols alpha radialPeriod angularFreqs radialFreqs) = do
  let numAngularFreq = L.length angularFreqs
      numRadialFreq = L.length radialFreqs
      filterVec =
        VG.convert . toUnboxed . computeS . makeFilter2D $
        R.traverse2
          (fromListUnboxed (Z :. L.length angularFreqs) angularFreqs)
          (fromListUnboxed (Z :. L.length radialFreqs) radialFreqs)
          (\(Z :. numAngularFreq') (Z :. numRadialFreq') ->
             (Z :. numAngularFreq' :. numRadialFreq' :. cols :. rows))
          (\f1 f2 (Z :. af :. rf :. i :. j) ->
             pinwheel
               (f1 (Z :. af))
               (f2 (Z :. rf))
               radialPeriod
               alpha
               (i - center cols)
               (j - center rows))
      filterPIVec =
        VG.convert . toUnboxed . computeS . makeFilter2D $
        R.traverse2
          (fromListUnboxed (Z :. L.length angularFreqs) angularFreqs)
          (fromListUnboxed (Z :. L.length radialFreqs) radialFreqs)
          (\(Z :. numAngularFreq') (Z :. numRadialFreq') ->
             (Z :. numAngularFreq' :. numRadialFreq' :. cols :. rows))
          (\f1 f2 (Z :. af :. rf :. i :. j) ->
             pinwheelPI
               (f1 (Z :. af))
               (f2 (Z :. rf))
               radialPeriod
               alpha
               (i - div cols 2)
               (j - div rows 2))
      forwardPlanID =
        DFTPlanID DFT1DG [numAngularFreq, numRadialFreq, cols, rows] [2, 3]
      backwardPlanID =
        DFTPlanID IDFT1DG [numAngularFreq, numRadialFreq, cols, rows] [2, 3]
  lock <- getFFTWLock
  (plan1, filterVecF) <-
    dft1dGPlan
      lock
      plan
      [numAngularFreq, numRadialFreq, cols, rows]
      [2, 3]
      filterVec
  (newPlan, _) <-
    idft1dGPlan
      lock
      plan1
      [numAngularFreq, numRadialFreq, cols, rows]
      [2, 3]
      filterVecF
  filterPIVecF <-
    dftExecute
      newPlan
      (DFTPlanID DFT1DG [numAngularFreq, numRadialFreq, cols, rows] [2, 3])
      filterPIVec
  return
    ( newPlan
    , fromUnboxed (Z :. numAngularFreq :. numRadialFreq :. cols :. rows) .
      VG.convert $
      filterVecF
    , fromUnboxed (Z :. numAngularFreq :. numRadialFreq :. cols :. rows) .
      VG.convert $
      filterPIVecF)

{-# INLINE frequencyDomainMultiply #-}
frequencyDomainMultiply ::
     (Source r (Complex Double))
  => Array U DIM4 (Complex Double)
  -> Array r DIM2 (Complex Double)
  -> Array D DIM4 (Complex Double)
frequencyDomainMultiply filterF input =
  R.traverse2
    filterF
    input
    const
    (\f1 f2 idx@(Z :. af :. rf :. i :. j) -> f1 idx * f2 (Z :. i :. j))

{-# INLINE convolvePinwheel #-}
convolvePinwheel ::
     (Source r (Complex Double))
  => DFTPlan
  -> Array U DIM4 (Complex Double)
  -> Array r DIM2 (Complex Double)
  -> IO (Array U DIM4 (Complex Double))
convolvePinwheel plan filterF input =
  fmap (fromUnboxed (extent filterF) . VG.convert) .
  dftExecute
    plan
    (DFTPlanID IDFT1DG (L.reverse . listOfShape . extent $ filterF) [2, 3]) .
  VG.convert . toUnboxed . computeUnboxedS $
  frequencyDomainMultiply filterF input


{-# LANGUAGE NamedFieldPuns #-}

module ParticleProcessingTests
  ( testParticleHandling
  ) where

import Control.Applicative (liftA2)
import Data.Complex (Complex(..))
import Data.List (sort)
import qualified Data.Vector as V
import HABQTlib.Data
import HABQTlib.Data.Particle
import HABQTlib.RandomStates
import qualified Numeric.LinearAlgebra as LA
import StateGenTests (arbPropDensityMatrix)
import qualified Streaming.Prelude as S
import qualified System.Random.MWC as MWC
import qualified Test.QuickCheck as QC
import Test.QuickCheck
  ( Args(..)
  , Positive(..)
  , Property
  , (.&&.)
  , arbitrary
  , quickCheckWith
  , stdArgs
  , suchThat
  )
import TestHelpers

lik :: Double -> Double
lik p = p / (1 - p)

loglik :: Double -> Double
loglik = log . lik

likProp :: PureStateVector -> DensityMatrix -> Property
likProp sv dm =
  let l = pureStateLikelihood sv dm
   in QC.property (l >= 0) .&&. QC.property (l <= 1)

testLikelihood :: IO ()
testLikelihood = do
  putStrLn "Testing likelihood computation for density matrices:"
  quickCheckWith stdArgs {maxSuccess = 1000} likProp

expProp :: QC.Blind Particles -> Property
expProp (QC.Blind pts) =
  let dm = snd . getWDM . V.foldl1' (<+>) . ptsParticles $ pts
      dim = LA.rows . getDensityMatrix $ dm
      z = (dim LA.>< dim) $ repeat 0
      scaleDM :: Weight -> DensityMatrix -> LA.Matrix (Complex Double)
      scaleDM w (DensityMatrix dm'') = LA.scale (w :+ 0) dm''
      dm' = DensityMatrix $ foldOverPts id scaleDM (+) z pts
      matchProp = QC.property $ dm <==> dm'
      validProp = arbPropDensityMatrix dm'
   in matchProp .&&. validProp

testExpectation :: IO ()
testExpectation = do
  putStrLn "Testing particle expectation function:"
  quickCheckWith stdArgs {maxSuccess = 500} expProp

particleUpdateFidR :: PureStateVector -> Particles -> Double
particleUpdateFidR testState@(PureStateVector sv) particles =
  let testPure = DensityMatrix $ sv LA.<> LA.tr sv
      originalEstimate = snd . getWDM $ reduceParticlesToMean particles
      originalFid = fidelityDM testPure originalEstimate
      updatedParticles = updateParticles testState particles
      updatedEstimate = snd . getWDM $ reduceParticlesToMean updatedParticles
      updatedFid = fidelityDM testPure updatedEstimate
   in loglik updatedFid - loglik originalFid

updateProp :: Property
updateProp =
  let gen = do
        (dim, rank) <- arbDimRank
        Positive num <- arbitrary :: QC.Gen (Positive NumberOfParticles)
        stateVector <- QC.resize dim arbitrary
        particles <- arbParticles dim rank (fromIntegral num)
        return (stateVector, particles)
      prop (sv, pts) =
        let wtsf Particles {ptsParticles} =
              V.foldl1' (+) . V.map (fst . getWDM) $ ptsParticles
            ws = wtsf pts
            pts' = updateParticles sv pts
            ws' = wtsf pts'
            wsProp = abs (ws - 1) < 1e-10
            ws'Prop = abs (ws' - 1) < 1e-10
         in wsProp .&&. ws'Prop
      msg = "of updates move the estimate closer"
   in QC.forAll
        gen
        (\(sv, pts) ->
           QC.classify (particleUpdateFidR sv pts >= 0) msg (prop (sv, pts)))

ecdfProp :: Particles -> Property
ecdfProp Particles {ptsParticles} =
  let cdf = ecdf ptsParticles
      hp = V.head cdf > 0
      tp = abs (V.last cdf - 1) < 1e-10
      cdfSorted = V.fromList . sort . V.toList $ cdf
      sorted = cdf == cdfSorted
   in hp .&&. tp .&&. sorted

icdfProp :: Particles -> Property
icdfProp Particles {ptsParticles} =
  QC.ioProperty $ do
    gen <- MWC.createSystemRandom
    x <- MWC.uniform gen
    let cdf = ecdf ptsParticles
        idx = icdf cdf x
        Just idx' = V.findIndex (> x) cdf
    return (idx == idx')

hDistProp :: ParticleHierarchy -> PureStateVector -> Double
hDistProp ph sv =
  let dm = DensityMatrix $ getStateVector sv LA.<> LA.tr (getStateVector sv)
      originalEstimate = getMixedEstimate ph
      updatedEstimate = getMixedEstimate $ updateParticleHierarchy sv ph
      originalFid = fidelityDM dm originalEstimate
      updatedFid = fidelityDM dm updatedEstimate
   in loglik updatedFid - loglik originalFid

hierarchyBatchProp ::
     Int -> Positive Dim -> Positive NumberOfParticles -> IO Property
hierarchyBatchProp batchSize (Positive dim) (Positive num) = do
  ph <- initialiseParticleHierarchy dim num
  passed <-
    S.length_ . S.filter (>= 0) . S.replicateM batchSize $
    hDistProp ph <$> genPureSV dim
  return . QC.property $
    fromIntegral passed / (fromIntegral batchSize :: Double) >= 0.95

hierarchyUpdateProp :: Int -> Property
hierarchyUpdateProp batchSize =
  let gen = do
        dim <- suchThat arbitrary (\x -> x > 1 && x < 20)
        num <- suchThat arbitrary (> 100)
        return (Positive dim, Positive num)
   in QC.forAll
        gen
        (\(dim, num) -> QC.ioProperty $ hierarchyBatchProp batchSize dim num)

effectiveSampleSizeProp :: Property
effectiveSampleSizeProp =
  let genP = do
        (dim, rank) <- arbDimRank
        Positive num <- arbitrary :: QC.Gen (Positive NumberOfParticles)
        arbParticles dim rank (fromIntegral num)
      sizeP ps@(Particles _ _ n _) =
        QC.classify
          (es <= fromIntegral n / 2)
          "of effective sizes are less than half of total particle number"
          ((es <= fromIntegral n) && (es >= 0))
        where
          es = effectiveSize ps
   in QC.forAll genP sizeP

nudgeProp :: Property
nudgeProp =
  let gen = do
        (dim, rank) <- arbDimRank
        dm <- arbDM dim rank
        return (dim, rank, dm)
      nudged (dim, rank, dm) =
        truncateRank rank <$> nudgeParticle dim 1e-2 (mkWDM1 dm)
      prop (dim, rank, dm) =
        QC.ioProperty $ do
          WeighedDensityMatrix (w, ndm) <- nudged (dim, rank, dm)
          let validDM = arbPropDensityMatrix ndm
              preservesRank = getRank dm == getRank ndm
              preservesWeight = w == 1
          return $
            QC.classify
              (fidelityDM ndm dm > 0.99)
              "fidelities exceed 99%"
              (validDM .&&. preservesRank .&&. preservesWeight)
   in QC.forAll gen prop

type ResamplingFunIO = MWC.GenIO -> Particles -> IO Particles

resampleProp :: MWC.GenIO -> ResamplingFunIO -> Property
resampleProp gen resamplingFun =
  let genA = do
        dim <- suchThat (arbitrary :: QC.Gen Dim) (liftA2 (&&) (< 10) (> 1))
        Positive rank <-
          suchThat (arbitrary :: QC.Gen (Positive Rank)) (< Positive dim)
        num <- suchThat arbitrary (> 1000)
        QC.Blind <$> arbParticles dim rank num
      preservesMeanIO :: QC.Blind Particles -> IO Property
      preservesMeanIO (QC.Blind pts) = do
        rpts <- resamplingFun gen pts
        let WeighedDensityMatrix (w1, dm1) = reduceParticlesToMean pts
            WeighedDensityMatrix (w2, dm2) = reduceParticlesToMean rpts
            sumw wacc (WeighedDensityMatrix (w', _)) = wacc + w'
            getpw Particles {ptsParticles = ps} = V.foldl' sumw 0 ps
            w1' = getpw pts
            w2' = getpw rpts
            propWeight = QC.property $ 2 * abs (w1 - w2) / (w1 + w2) < 1e-12
            propWeight' =
              QC.property $ 2 * abs (w1' - w2') / (w1' + w2') < 1e-12
            Particles _ _ num _ = rpts
            propESabs =
              QC.property $ abs (effectiveSize rpts - fromIntegral num) < 1
            propESincrease =
              QC.property $ effectiveSize rpts > effectiveSize pts
            propAll =
              propWeight .&&. propWeight' .&&. propESincrease .&&. propESabs
        return $
          QC.classify
            (fidelityDM dm1 dm2 > 0.99)
            "fidelities between pre- and post-resampling means exceed 99%"
            propAll
   in QC.forAll genA (QC.ioProperty . preservesMeanIO)

testParticleResampling :: IO ()
testParticleResampling = do
  gen <- MWC.createSystemRandom
  putStrLn "Testing multinomial particle resampling without checks:"
  quickCheckWith stdArgs {maxSuccess = 100} $ resampleProp gen resampleMultinom

testParticleNudging :: IO ()
testParticleNudging = do
  putStrLn "Testing particle nudging:"
  quickCheckWith stdArgs {maxSuccess = 1000} nudgeProp

testEffectiveSampleSize :: IO ()
testEffectiveSampleSize = do
  putStrLn "Testing calculation of effective sample size:"
  quickCheckWith stdArgs {maxSuccess = 1000} effectiveSampleSizeProp

testParticleUpdate :: IO ()
testParticleUpdate = do
  putStrLn "Testing update of particles:"
  quickCheckWith stdArgs {maxSuccess = 1000} updateProp

testEcdf :: IO ()
testEcdf = do
  putStrLn "Testing ecdf calculation:"
  quickCheckWith stdArgs {maxSuccess = 10000} ecdfProp
  putStrLn "Testing iecdf calculation:"
  quickCheckWith stdArgs {maxSuccess = 10000} icdfProp

testHierarchyUpdate :: IO ()
testHierarchyUpdate = do
  putStrLn "Testing update of particle hierarchies:"
  quickCheckWith stdArgs {maxSuccess = 50} (hierarchyUpdateProp 100)

testParticleHandling :: IO ()
testParticleHandling = do
  testLikelihood
  testExpectation
  testParticleUpdate
  testEcdf
  testHierarchyUpdate
  testEffectiveSampleSize
  testParticleNudging
  testParticleResampling


{-# LANGUAGE GADTs, Rank2Types, DataKinds, TypeOperators #-}

module Lib.Utility.Mnist.IO
(
    MnistTrainingSize, MnistTestSize, MnistData, TrainingData,
    getMnistData
) where

import Network.HTTP as HTTP
import Data.ByteString.Lazy as Lazy (ByteString, readFile, writeFile, empty)
import Data.ByteString.Lazy.Char8 as Char8 (pack)
import Codec.Compression.GZip as GZip
import qualified Data.IDX as IDX
import Control.Monad.Trans.Maybe
import Data.Vector.Unboxed (Vector, toList)
import qualified Numeric.LinearAlgebra as HMatrix
import qualified Numeric.LinearAlgebra.Static as Static
import System.IO.Error
import qualified Data.Serialize as Serialize
import Control.Arrow
import Control.Applicative
import GHC.TypeLits
import GHC.Generics (Generic)
import Data.Proxy

type TrainingData n = (Static.L n 784, Static.L n 10)

type MnistTrainingSize = 60000
type MnistTestSize = 10000
type MnistData = (TrainingData MnistTrainingSize, TrainingData MnistTestSize)

type SerializedMnistData = ([(Int, [Int])], [(Int, [Int])])
type LabeledIntData = [(Int, Vector Int)]

getMnistData :: MaybeT IO MnistData
getMnistData = do
    loadResult <- loadMnistData
    (tr, te) <- case loadResult of
        Just x -> return x
        Nothing -> do
            () <- MaybeT (Just <$> putStrLn "Try fetch for load fail")
            x <- fetchMnistData
            () <- saveMnistData x
            MaybeT $ return $ parseMatrix $ fromLabeledIntData2Matrix x
    return (tr, te)

loadMnistData :: MaybeT IO (Maybe MnistData)
loadMnistData = MaybeT $ do
    putStrLn $ "Loading from " ++ fpath
    matrixMaybe <- HMatrix.loadMatrix' fpath -- 70000><785
    let loadResult = matrixMaybe >>= parseMatrix
    return $ Just loadResult

-- toOneHot :: HMatrix.Vector Double -> HMatrix.Matrix Double
toOneHot :: forall n. KnownNat n => Static.L n 1 -> Static.L n 10
toOneHot label = Static.build f
    where m0 = Static.extract $ (Static.konst 0 :: Static.L 10 10)
          f r c = if c == (atIndex r 0) then 1.0 else 0
          atIndex r c = Static.unwrap label `HMatrix.atIndex` (truncate r, truncate c)

parseMatrix :: HMatrix.Matrix Double -> Maybe MnistData
parseMatrix matrix = do
        matrix' <- Static.create matrix :: Maybe (Static.L (MnistTrainingSize + MnistTestSize) 785)
        let (label, image) = Static.splitCols matrix'
        let (trL, teL) = Static.splitRows label :: (Static.L MnistTrainingSize 1, Static.L MnistTestSize 1)
        let (trD, teD) = Static.splitRows $ image / 255
        return ((trD, toOneHot trL), (teD, toOneHot teL))

fromLabeledIntData2Matrix :: (LabeledIntData, LabeledIntData) -> HMatrix.Matrix Double
fromLabeledIntData2Matrix (tr, te) = matrix
    where transform = map (second toList)
          tr' = transform tr
          te' = transform te
          join = map (uncurry (:))
          matrix = HMatrix.fromLists $ map (map realToFrac) $ join tr' ++ join te'

saveMnistData :: (LabeledIntData, LabeledIntData) -> MaybeT IO ()
saveMnistData labeledIntData = MaybeT $ do
    putStrLn $ "Saving to " ++ fpath
    let matrix = fromLabeledIntData2Matrix labeledIntData
    let saveAction = HMatrix.saveMatrix fpath "%lf" matrix
    () <- catchIOError saveAction (\_ -> do putStrLn "SaveError"; return ())
    return $ Just ()

fetchMnistData :: MaybeT IO (LabeledIntData, LabeledIntData)
fetchMnistData = do
    (Data_ trD) <- fetchIDX trainImagesPath decodeIDXData
    (Labels_ trL) <- fetchIDX trainLabelsPath decodeIDXLabels
    (Data_ teD) <- fetchIDX testImagesPath decodeIDXData
    (Labels_ teL) <- fetchIDX testLabelsPath decodeIDXLabels
    MaybeT $ return $ do
        tr <- IDX.labeledIntData trL trD
        te <- IDX.labeledIntData teL teD
        return (tr, te)

data IDXW = Data_ IDX.IDXData | Labels_ IDX.IDXLabels

decodeIDXData :: ByteString -> Maybe IDXW
decodeIDXData = fmap Data_ . IDX.decodeIDX
decodeIDXLabels :: ByteString -> Maybe IDXW
decodeIDXLabels = fmap Labels_ . IDX.decodeIDXLabels
fetchIDX :: String -> (ByteString -> Maybe IDXW) -> MaybeT IO IDXW
fetchIDX path decodeF = MaybeT $ do
    let url = toUrl path
    putStrLn $ "Downloading from " ++ url
    idxBytes <- getFileFromUrl url
    let idx = decodeF idxBytes
    return idx

baseUrl = "http://yann.lecun.com/exdb/mnist/"
trainImagesPath = "train-images-idx3-ubyte.gz"
trainLabelsPath = "train-labels-idx1-ubyte.gz"
testImagesPath = "t10k-images-idx3-ubyte.gz"
testLabelsPath = "t10k-labels-idx1-ubyte.gz"
fpath = "mnist.pickle"
toUrl path = baseUrl ++ path

getFileFromUrl :: String -> IO ByteString
getFileFromUrl url = do
    response <- simpleHTTP (getRequest url)
    body <- getResponseBody response
    let bytes = Char8.pack body
    return $ GZip.decompress bytes


module Language.Hakaru.Util.Coda where

import Statistics.Autocorrelation
import qualified Data.Packed.Vector as V
import qualified Data.Vector.Generic as G

effectiveSampleSize :: [Double] -> Double
effectiveSampleSize samples = n / (1 + 2*(G.sum rho))
  where n = fromIntegral (V.dim vec)
        vec = V.fromList samples
        cov = autocovariance vec
        rho = G.map (/ G.head cov) cov


import Criterion.Main
import Numeric.LinearAlgebra

n = 10 :: Int
main = defaultMain [
         bench "eigenvalues" $ w (\_ -> eigenvalues sampleDiag),
         bench "singularValues" $ w (\_ -> singularValues sampleDiag),

         bench "svd" $ w (\_ -> svd sampleDiag),
         bench "thin svd" $ w (\_ -> thinSVD sampleDiag),

         bench "nullspace" $ w (\_ -> nullspace sampleDiag),
         bench "orthogonal" $ w (\_ -> orth sampleDiag),

         bench "determinant" $ w (\_ -> det sampleDiag)
       ]
  where
    w f = f `whnf` 0
    sampleList = replicate n $ replicate n 1
    sampleDiag = diagl [1..fromIntegral n]


{-|
Module: MachineLearning.NeuralNetwork
Description: Neural Network
Copyright: (c) Alexander Ignatyev, 2016
License: BSD-3
Stability: experimental
Portability: POSIX

Simple Neural Networks.
-}

module MachineLearning.NeuralNetwork
(
    Model(..)
  , NeuralNetworkModel(..)
  , MLC.calcAccuracy
  , T.Topology
  , T.initializeTheta
  , T.initializeThetaIO
  , T.initializeThetaM
  , Regularization(..)
)

where

import qualified Numeric.LinearAlgebra as LA
import MachineLearning.Types (R, Vector, Matrix)
import MachineLearning.Utils (reduceByRowsV)
import qualified MachineLearning.Classification.Internal as MLC
import MachineLearning.Model (Model(..))
import qualified MachineLearning.NeuralNetwork.Topology as T
import MachineLearning.Regularization (Regularization(..))

-- | Neural Network Model.
-- Takes neural network topology as a constructor argument.
newtype NeuralNetworkModel = NeuralNetwork T.Topology

instance Model NeuralNetworkModel where
  hypothesis (NeuralNetwork topology) x theta = predictions
    where thetaList = T.unflatten topology theta
          scores = calcScores topology x thetaList
          predictions = reduceByRowsV (fromIntegral . LA.maxIndex) scores

  cost (NeuralNetwork topology) lambda x y theta =
    let (ys, thetaList) = processParams topology y theta
        scores = calcScores topology x thetaList
    in T.loss topology lambda scores thetaList ys

  gradient (NeuralNetwork topology) lambda x y theta =
    let (ys, thetaList) = processParams topology y theta
        (scores, cacheList) = T.propagateForward topology x thetaList
        grad = T.flatten $ T.propagateBackward topology lambda scores cacheList ys
    in grad

-- | Score function. Takes a topology, X and theta list.
calcScores :: T.Topology -> Matrix -> [(Matrix, Matrix)] -> Matrix
calcScores topology x thetaList = fst $ T.propagateForward topology x thetaList

processParams :: T.Topology -> Vector -> Vector -> (Matrix, [(Matrix, Matrix)])
processParams topology y theta =
  let nOutputs = T.numberOutputs topology
      ys = LA.fromColumns $ MLC.processOutputOneVsAll nOutputs y
      thetaList = T.unflatten topology theta
  in (ys, thetaList)


{-# LANGUAGE GADTs, DataKinds, PolyKinds #-}
{-# LANGUAGE FlexibleContexts #-}

module TensorDAG where

import Data.Vinyl
import Data.Vinyl.Functor
import Data.Singletons
import Numeric.LinearAlgebra (Numeric)

import TensorHMatrix
import DAGIO
import VarArgs

makeGrad1 :: (a -> b -> a) -> HList '[a] -> Identity b -> HList '[a]
makeGrad1 g (Identity a :& RNil) (Identity b) = Identity (g a b) :& RNil

makeGrad2 :: (a -> b -> c -> (a, b)) -> HList '[a, b] -> Identity c -> HList '[a, b]
makeGrad2 g (Identity a :& Identity b :& RNil) (Identity c) = Identity a' :& Identity b' :& RNil
  where (a', b') = g a b c

makeDot :: Numeric a => Node (Tensor '[n] a) -> Node (Tensor '[n] a) -> IO (Node a)
makeDot = makeNode (uncurry2 dot, makeGrad2 gradDot)

makeMV :: (SingI n, IntegralN n, Usable a) => Node (Tensor '[n, m] a) -> Node (Tensor '[m] a) -> IO (Node (Tensor '[n] a))
makeMV = makeNode (uncurry2 mv, makeGrad2 gradMV)

makeMM :: (SingI n, IntegralN n, SingI k, Usable a) => Node (Tensor '[n, m] a) -> Node (Tensor '[m, k] a) -> IO (Node (Tensor '[n, k] a))
makeMM = makeNode (uncurry2 mm, makeGrad2 gradMM)

makeSelect i = makeNode (uncurry1 $ select i, makeGrad1 $ gradSelect i)


module STCR2S1Image where

import           Control.Monad           as M
import           Control.Monad.Parallel  as MP
import           Data.Array.Repa         as R
import           Data.Binary             (decodeFile)
import           Data.Complex
import           Data.List               as L
import           Data.Vector.Unboxed     as VU
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.MonteCarlo
import           Image.IO
import           STC.CompletionFieldR2S1
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types
import           Utils.Array

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:numTrailStr:maxTrailStr:initOriStr:initSpeedStr:writeFlagStr:numIterationStr:thresholdStr:histFilePath:cutoffStr:imagePath:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      sigma = read sigmaStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      initOri = read initOriStr :: Double
      initSpeed = read initSpeedStr :: Double
      writeFlag = read writeFlagStr :: Bool
      numIteration = read numIterationStr :: Int
      threshold = read thresholdStr :: Double
      cutoff = read cutoffStr :: Int
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2S1Image"
  imgRepa@(ImageRepa _ img') <- readImageRepa imagePath False
  let (Z :. _ :. cols :. rows) = extent img'
      img =
        computeS . R.traverse img' id $ \f idx@(Z :. k :. i :. j) ->
          if (sqrt . fromIntegral $ (i - div cols 2) ^ 2 + (j - div rows 2) ^ 2) >
             32
            then 0
            else f idx
  createDirectoryIfMissing True folderPath
  plotImageRepa (folderPath </> "input.png") . ImageRepa 8 $ img
  flag <- doesFileExist histFilePath
  arrG <-
    if flag
      then getNormalizedHistogramArr .
           mapHistogram
             (\x ->
                let y = x :: Int
                 in fromIntegral y) <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        computeS . R.map magnitude <$>
          solveMonteCarloR2S1
            numThread
            numTrail
            numPoint
            numPoint
            numOrientation
            sigma
            tao
            cutoff
            histFilePath
            (0, 0, 0, 0, initOri / 180 * pi, initSpeed)
  powerMethod
    emptyPlan
    folderPath
    arrG
    numIteration
    writeFlag
    ""
    threshold
    (R.traverse img (const (Z :. numOrientation :. cols :. rows)) $ \f (Z :. _ :. i :. j) ->
       if f (Z :. (0 :: Int) :. i :. j) > 0
         then 1
         else 0)
    (R.traverse img (const (Z :. numOrientation :. cols :. rows)) $ \f (Z :. _ :. i :. j) ->
       if f (Z :. (0 :: Int) :. i :. j) > 0
         then 1 / fromIntegral numOrientation
         else 0)


{-# LANGUAGE AllowAmbiguousTypes, ConstraintKinds #-}

module GPLVM.GaussianProcess
       ( GaussianProcess (..)
       , GPTrainingData (..)
       , PosteriorSample (..)
       , gpToPosteriorSample
       , kernelGP
       ) where

import Universum hiding (transpose, Vector)

import Control.Lens (makeLenses)

import Data.Array.Repa
import Data.Array.Repa.Repr.Unboxed (Unbox)
import Numeric.LinearAlgebra.Repa hiding (Matrix, Vector)
import System.Random (Random, RandomGen, mkStdGen)

import GPLVM.Types (InputObservations(..), Matrix, Vector, unInputObs)
import GPLVM.Util

data GaussianProcess a = GaussianProcess
    { -- | we define gaussian process by some kernel function
      _kernelGP :: Vector D a -> Vector D a -> Matrix D a
    }

data GPTrainingData a = GPTrainingData
    { _inputTrain  :: Vector D a  -- ^ input training data
    , _outputTrain :: Vector D a  -- ^ output training data
    }

makeLenses ''GaussianProcess
makeLenses ''GPTrainingData

newtype PosteriorSample a = PosteriorSample
    { unSample :: Matrix D a  -- ^ posterior sample
    }

-- | Constraint kind required to getting a posterior sample by a given GP
type GPConstraint a =
  ( Field a
  , Random a
  , Unbox a
  , Floating a
  , Eq a
  )

-- | Main GP function: get a posterior sample by some kernel function, input observations and training data

gpToPosteriorSample
  :: forall a g.
  GPConstraint a
  => InputObservations a        -- ^ input observations
  -> GaussianProcess a          -- ^ kernel function
  -> GPTrainingData a           -- ^ training data
  -> Int                        -- ^ number of samples
  -> Maybe (PosteriorSample a)  -- ^ posterior functional prior
gpToPosteriorSample (InputObservations observe@(ADelayed (Z :. len) _)) gP trainingData sampleNumber = do
  -- | kernel applied to input test points (so-called K_ss)
  let covarianceMatrix = kernel observe observe

  -- | kernel applied to input training points (so-called K)
  let trainingKernel = kernel inputTrain' inputTrain'

  -- | Cholesky decomposition applied to kernel of training points (:), so-called L
  let cholK = cholSH $ trainingKernel +^
                (smap (* 0.00005) . identD . size . extent $ inputTrain')

  -- | covariance between test points and input training points (so-called K_s)
  let testPointMean = kernel observe inputTrain'

  -- | (roots of L * x = K_s)
  cholKSolve <- delay <$> linearSolveS cholK testPointMean

  -- | solve linear system for output training points
  cholKSolveOut <- delay <$> linearSolveS cholK (transposeMatrix $ toMatrix outputTrain' len)

  -- | compute mean
  let mean = (transposeMatrix cholKSolveOut) `mulD` cholKSolve

  -- | posterior
  let postF' = cholSH $
                     covarianceMatrix +^
                     ((smap (* 1.0e-6) (identD len)) -^
                     (transposeMatrix cholKSolve) `mulD` cholKSolve)

  -- | posterior sample
  return $ (PosteriorSample $ mean +^ (functionalPrior postF' (mkStdGen 1) sampleNumber))
    where
       kernel = gP ^. kernelGP
       inputTrain' = trainingData ^. inputTrain
       outputTrain' = trainingData ^. outputTrain
       mulD m n = delay $ m `mulS` n
       functionalPrior matrix@(ADelayed (Z :. rows :. cols) _) gen sampleNumber =
         delay $ matrix `mulS` randomCoeffs
         where
           randomCoeffs = randomMatrixD (mkStdGen (-4)) (rows, sampleNumber)


module Fuml.Optimisation.SGD where

import Numeric.LinearAlgebra
import Data.Random

data SGDOpts = SGDOpts
  { learning_rate :: Double
  , batch_size :: Int
  , epochs :: Int
  }

sgdBatch :: SGDOpts -> (Vector Double -> Vector Double) -> Vector Double -> Vector Double
sgdBatch opts grad p = p - scale (learning_rate opts) (grad p)

sgdEpoch :: SGDOpts -> ([a] -> Vector Double -> Vector Double) -> [a] -> Vector Double -> RVar (Vector Double)
sgdEpoch opts grad pts' p' = do
  pts <- shuffle pts'
  return $ go p' $ inGroupsOf (batch_size opts) pts
    where go p [] = p
          go p (batch:bs) = go (sgdBatch opts (grad batch) p) bs

inGroupsOf :: Int -> [a] -> [[a]]
inGroupsOf _ [] = []
inGroupsOf n xs = let (these,those) = splitAt n xs
                  in these : inGroupsOf n those


module Main where

import Net
import Data.Traversable
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel
import qualified Data.Vector.Storable as V
import System.Random
import Data.List
import Control.Monad

testNet :: DoubleNet -> [Vector Double] -> [Vector Double] -> Double
testNet n ins ans =
  let outs :: [Vector Double]
      outs = map (runNet n) ins
  in sum (zipWith (\x y -> norm_1 $ add x (scale (-1) y)) outs ans)

rankNets :: [DoubleNet] -> IO [(Double,DoubleNet)]
rankNets nets =
  do values <- mapM (\_ -> do x <- randomRIO (0,10)
                              return $ vector [x]) [1..100]
     let answers = map sin values
         scoredNets = map (\x -> (testNet x values answers,x)) nets
     return $ sortOn fst scoredNets

plotSome :: DoubleNet -> IO ()
plotSome n =
  do let pts = map (\x -> (x/1000)*10) $ [1..1000]
     mapM_ (\x -> putStrLn $ (show $ runNet n (vector [x])) ++ " " ++ (show $ (sin x)/2)) pts

sinFunc = DiffFunction sin (\x y -> y*cos x)

errFunc :: Double -> DifferentiableFunction (Vector Double) Double
errFunc c = DiffFunction (\v -> ((atIndex v 0) - c)^2) (\v w -> 2*((atIndex v 0) - c)*(atIndex w 0))

iterateNet :: Int -> DoubleNet -> IO DoubleNet
iterateNet 0 net = return net
iterateNet steps net =
  do let n = 10
     xs <- replicateM n $ randomRIO (0,20)

     --let x = 0.5
     let loop loopSteps currNet =
           if loopSteps ==0 then
             return currNet
           else
             do grads <- forM xs (\x -> do let (val,err,grad) = runAndDiffWithError currNet (vector [x]) (errFunc $ (sin x)/2)
                                         --print (x,val,(sin x)/2,err)
                                           return grad)
                let grad = foldl1 (zipWith (\a b -> (add b a))) grads
                let newNet = updateNet net grad (-0.0005)
                loop (loopSteps - 1) newNet
     loop 5 net >>= iterateNet (steps-1)

main :: IO ()
main = do randNet <- randomSimpleNet 1 1 100
          finalNet <- iterateNet 10000 randNet
          plotSome finalNet


{-# LANGUAGE OverloadedStrings #-}

module Main where

import Numeric.LinearAlgebra

main :: IO ()
main = putStrLn $ show $ vector [1,2,3] * vector [3,0,-2]


-- Kolmogorov-Smirnov tests for distribution
--
-- Note that it's not most powerful test for normality.
module KS (
  tests
  ) where

import qualified Data.Vector.Unboxed as U

import Statistics.Test.KolmogorovSmirnov

import Statistics.Distribution
import Statistics.Distribution.Binomial
import Statistics.Distribution.Exponential
import Statistics.Distribution.Gamma
import Statistics.Distribution.Normal
import Statistics.Distribution.Uniform
import Statistics.Distribution.Beta

import qualified System.Random.MWC               as MWC
import qualified System.Random.MWC.Distributions as MWC

import Test.HUnit hiding (Test)
import Test.Framework
import Test.Framework.Providers.HUnit

tests :: MWC.GenIO -> Test
tests g = testGroup "Kolmogorov-Smirnov"
  [ testCase "standard"           $ testKS standard           MWC.standard    g
  , testCase "normal m=1 s=2"     $ testKS (normalDistr 1 2) (MWC.normal 1 2) g
    -- Gamma distribution
  , testCase "gamma  k=1   θ=1"   $ testKS (gammaDistr  1   1  ) (MWC.gamma  1   1  ) g
  , testCase "gamma  k=0.3 θ=0.4" $ testKS (gammaDistr  0.3 0.4) (MWC.gamma  0.3 0.4) g
  , testCase "gamma  k=0.3 θ=3"   $ testKS (gammaDistr  0.3 3  ) (MWC.gamma  0.3 3  ) g
  , testCase "gamma  k=3   θ=0.4" $ testKS (gammaDistr  3   0.4) (MWC.gamma  3   0.4) g
  , testCase "gamma  k=3   θ=3"   $ testKS (gammaDistr  3   3  ) (MWC.gamma  3   3  ) g
    -- Uniform
  , testCase "uniform -2 .. 3"    $ testKS (uniformDistr (-2) 3) (MWC.uniformR (-2,3)) g
    -- Exponential
  , testCase "exponential l=1"    $ testKS (exponential 1)       (MWC.exponential 1) g
  , testCase "exponential l=3"    $ testKS (exponential 3)       (MWC.exponential 3) g
    -- Beta
  , testCase "beta a=0.3,b=0.5"    $ testKS (betaDistr 0.3 0.5)       (MWC.beta 0.3 0.5) g
  , testCase "beta a=0.1,b=0.8"    $ testKS (betaDistr 0.3 0.5)       (MWC.beta 0.3 0.5) g
  , testCase "beta a=0.8,b=0.1"    $ testKS (betaDistr 0.3 0.5)       (MWC.beta 0.3 0.5) g
  ]

testKS :: (Distribution d) => d -> (MWC.GenIO -> IO Double) -> MWC.GenIO -> IO ()
testKS distr generator g = do
  sample <- U.replicateM 1000 (generator g)
  case kolmogorovSmirnovTest distr 0.01 sample of
    Significant    -> assertFailure "KS test failed"
    NotSignificant -> return ()


-- | Helper for writing tests for numeric code
module Test.QuickCheck.Numeric (
    -- * Approximate equality
    eq
  , eqC
    -- * Function monotonicity
  , Monotonicity(..)
  , monotonicFunction
  , monotonicFunctionIEEE
    -- * Inverse function
  , checkInverse
  , checkInverse2
  ) where

import Data.Complex
import qualified Numeric.IEEE    as IEEE

----------------------------------------------------------------
-- Approximate equality
----------------------------------------------------------------

-- | Approximate equality for 'Double'. Doesn't work well for numbers
--   which are almost zero.
eq :: Double                    -- ^ Relative error
   -> Double -> Double -> Bool
eq eps a b
  | a == 0 && b == 0 = True
  | otherwise        = abs (a - b) <= eps * max (abs a) (abs b)

-- | Approximate equality for 'Complex Double'
eqC :: Double                   -- ^ Relative error
    -> Complex Double
    -> Complex Double
    -> Bool
eqC eps a@(ar :+ ai) b@(br :+ bi)
  | a == 0 && b == 0 = True
  | otherwise        = abs (ar - br) <= eps * d
                    && abs (ai - bi) <= eps * d
  where
    d = max (realPart $ abs a) (realPart $ abs b)

----------------------------------------------------------------
-- Function monotonicity
----------------------------------------------------------------

-- | Function monotonicity type.
data Monotonicity
  = StrictInc   -- ^ Strictly increasing function
  | MonotoneInc -- ^ Monotonically increasing function
  | StrictDec   -- ^ Strictly decreasing function
  | MonotoneDec -- ^ Monotonically decreasing function
  deriving (Show,Eq,Ord)

-- | Check that function is nondecreasing. For floating point number
--   it may give spurious failures so 'monotonicFunction'
--   should be used in this case.
monotonicFunction :: (Ord a, Ord b) => Monotonicity -> (a -> b) -> a -> a -> Bool
monotonicFunction cmp f x1 x2
  = f (min x1 x2) `op` f (max x1 x2)
  where
    op = case cmp of
           StrictInc   -> (< )
           MonotoneInc -> (<=)
           StrictDec   -> (> )
           MonotoneDec -> (>=)

-- | Check that function is nondecreasing taking rounding errors into
--   account. This function makes no distinction between strictly
--   increasing function and monotonically increasing function since
--   distinction is pointless for floating point.
--
--   In fact funstion is allowed to decrease less than one ulp in order
--   to guard againist problems with excess precision. On x86 FPU works
--   with 80-bit numbers but doubles are 64-bit so rounding happens
--   whenever values are moved from registers to memory
monotonicFunctionIEEE :: (Ord a, IEEE.IEEE b) => Monotonicity -> (a -> b) -> a -> a -> Bool
monotonicFunctionIEEE cmp f x1 x2
  =  y1 `op` y2
  || abs (y1 - y2) < abs (y2 * IEEE.epsilon)
  where
    y1 = f (min x1 x2)
    y2 = f (max x1 x2)
    op = case cmp of
           StrictInc   -> (<=)
           MonotoneInc -> (<=)
           StrictDec   -> (>=)
           MonotoneDec -> (>=)

----------------------------------------------------------------
-- Function and its inverse
----------------------------------------------------------------

-- | Check that function is inverse. Breaks down near zero.
checkInverse
  :: (Double -> Double) -- ^ Function @f(x)@
  -> (Double -> Double) -- ^ Inverse function @g@, @g(f(x)) = x@
  -> (Double -> Double) -- ^ Derivative of function @f(x)@
  -> Double             -- ^ Relative error for
                        --   @f(x)@. Usually is machine epsilon.
  -> Double             -- ^ Relative error for inverse function
                        --   @g(x)@. Usually is machine epsilon.
  -> Double -> Bool
checkInverse f invF f' eps eps' x
  = x ~= invF y
  where
    (~=) = eq (eps' + abs (y / f' x * eps))
    y    = f x

-- | Check that function is inverse. Breaks down near zero.
checkInverse2
  :: (Double -> Double) -- ^ Function @f(x)@
  -> (Double -> Double) -- ^ Inverse function @g@, @g(f(x)) = x@
  -> (Double -> Double) -- ^ Derivative of function @g(x)@
  -> Double             -- ^ Relative error for
                        --   @f(x)@. Usually is machine epsilon.
  -> Double             -- ^ Relative error for inverse function
                        --   @g(x)@. Usually is machine epsilon.
  -> Double -> Bool
checkInverse2 f invF invF' eps eps' x
  = x ~= invF y
  where
    (~=) = eq (eps' + abs (y * (invF' y * eps)))
    y    = f x


{-# LANGUAGE OverloadedStrings #-}

module Main where

import Prelude hiding (readFile, drop)
import Control.Monad
import Control.Monad.IO.Class
import Control.Exception
import Foreign.C.Types
import System.Random
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data
import SDL
import Debug.Trace
import qualified Linear
import Data.List.Split (chunksOf)

import Neural
import MNIST

-- Generates a layer with a given number of neurons with inputCount weights.
-- Weights are randomised between 0 and 1.
genLayer :: Int -> Int -> IO Layer
genLayer neurons inputCount = do
  weights <- replicateM (neurons * inputCount) (randomRIO (-0.2, 0.2) :: IO R)
  biases <- replicateM (neurons) (randomRIO (-0.2, 0.2) :: IO R)
  return $ Layer (matrix inputCount weights) (vector biases)

-- Generates a network with the specified number of neurons in each layer,
-- and the number of inputs to the network.
-- Initialised with random weights and biases as in genLayer.
genNetwork :: [Int] -> Int -> IO Network
genNetwork layers inputCount = sequence . map (uncurry genLayer) $ layerDefs
  where
    -- Layer definitions in the form (neuronCount, inputCount)
    layerDefs = zip layers (inputCount:layers)

-- A test network
testNetwork :: IO Network
testNetwork = do
  l1 <- genLayer 5 3
  l2 <- genLayer 5 5
  l3 <- genLayer 1 5
  return [l1, l2, l3]

-- Another test network
testNetwork2 :: Network
testNetwork2 = [hiddenLayer, outputLayer]
  where
    hiddenLayer = Layer (matrix 2 [0.15, 0.2, 0.25, 0.3]) (vector [0.35, 0.35])
    outputLayer = Layer (matrix 2 [0.4, 0.45, 0.5, 0.55]) (vector [0.6, 0.6])

-- Test inputs for testNetwork2
testInputs = vector [0.05, 0.1]

-- Training data for testInputs
trainingData = vector [0.01, 0.99]

-- Train neural network
trainNet = do
  (imageCount, width, height, imageData) <- readMNISTImages "train-images-idx3-ubyte"
  (labelCount, labels) <- readMNISTLabels "train-labels-idx1-ubyte"

  let image = head imageData
  let input = vector (map fromIntegral image)
  let correctNumber = fromIntegral . head $ labels
  let training = vector $ map (\i -> if i == correctNumber then 1.0 else 0.0) [1..10]

  hiddenLayer <- genLayer 20 (width * height)
  outputLayer <- genLayer 10 20

  let initialNetwork = [hiddenLayer, outputLayer]

  let loop net i err =
        case i > 1000 || err < 0.1 of
          True  -> net
          False -> trace (show $ sumElements (last error)) $ loop nextNet (i+1) (sumElements (last error))
            where
              output = evalNetwork net input
              error = evalNetworkError net output input training
              nextNet = gradientDescent net input (map fst output) error 0.1

  let trainedNet = loop initialNetwork 0 1000

  return (input, initialNetwork, trainedNet, training)

-- Train neural network
trainNetM = do
  (imageCount, width, height, imageData) <- readMNISTImages "train-images-idx3-ubyte"
  (labelCount, labels) <- readMNISTLabels "train-labels-idx1-ubyte"

  let image = head imageData
  let input = fromRows [vector (map fromIntegral image)]
  let correctNumber = fromIntegral . head $ labels
  let training = fromRows [vector $ map (\i -> if i == correctNumber then 1.0 else 0.0) [1..10]]

  hiddenLayer <- genLayer 20 (width * height)
  outputLayer <- genLayer 10 20

  let initialNetwork = [hiddenLayer, outputLayer]

  let loop net i err =
        case i > 100000 || err < 0.1 of
          True  -> net
          False -> trace (show $ sumElements (last error)) $ loop nextNet (i+1) (sumElements (last error))
            where
              output = evalNetworkM net input
              error = evalNetworkErrorM net output input training
              nextNet = gradientDescentM net input (map fst output) error 0.1

  let trainedNet = loop initialNetwork 0 1000

  return (input, initialNetwork, trainedNet, training)

-- Make a window and show MNIST image
window image width height = do
  initializeAll

  window <- createWindow "SDL Application" defaultWindow
  renderer <- createRenderer window (-1) defaultRenderer

  -- Convert to sdl texture
  texture <- loadTextureMNIST renderer width height image

  SDL.rendererDrawColor renderer $= Linear.V4 255 255 255 255
  SDL.clear renderer

  copy renderer texture Nothing Nothing

  present renderer

  let loop = do
        pollEvents

        keyState <- getKeyboardState

        case keyState ScancodeEscape of
          True -> (return ()) :: IO ()
          False -> loop

  loop

  destroyTexture texture
  destroyRenderer renderer
  destroyWindow window

-- Train a network using a given set of inputs and training values for the given
-- number of epochs using the given learning rate.
trainNetwork :: Network -> [Vector R] -> [Vector R] -> Int -> Int -> R -> Network
trainNetwork initialNet input trainingValues epochs batchSize learnRate =
  train initialNet 0
    where
      train net epoch = case epoch >= epochs of
        True  -> net
        False -> train nextNet (epoch + 1)
          where
            --inputBatch = fromRows input
            --trainingBatch = fromRows trainingValues
            inputBatches = map fromRows . chunksOf batchSize $ input
            trainingBatches = map fromRows . chunksOf batchSize $ trainingValues
            nextNet = foldl runBatch net (zip inputBatches trainingBatches)
            --nextNet = net
            runBatch network (inputBatch, trainingBatch) =
              gradientDescentM net inputBatch (map fst output) error learnRate
                where
                  output = evalNetworkM net inputBatch
                  error = evalNetworkErrorM net output inputBatch trainingBatch

toOutput = \n -> vector . map (\x -> if n == x then 1.0 else 0.0) $ [0..9]

fromOutput :: Vector R -> Int
fromOutput v = guess
  where
    list = zip [0..] . toList $ v
    possibles = filter ((>=0.5) . snd) list
    guess = case length possibles of
              0 -> -1
              1 -> fst . head $ possibles
              _ -> -2

runTest :: Network -> ([CUChar], Int) -> Bool
runTest network (image, label) = guess == label
  where output = fst . last $ evalNetwork network $ vector . map fromIntegral $ image
        guess = fromOutput output

--runTest :: Network -> ([CUChar], Int) -> IO ()
--runTest network (image, label) = do
--  let output = fst . last $ evalNetwork network $ vector . map fromIntegral $ image
--  let guess = fromOutput output
--  let expected = label
--  putStr "Guess: "
--  putStr . show $ guess
--  putStr ",\tExpected: "
--  putStr . show $ label
--  putStr $ if guess == expected then ".\tCorrect! " else ".\tINCORRECT!"
--  putStrLn ""

main = do

  (trainingImageCount, trainingWidth,
   trainingHeight, trainingImageData)  <- readMNISTImages "train-images-idx3-ubyte"
  (trainingLabelCount, trainingLabels) <- readMNISTLabels "train-labels-idx1-ubyte"

  (testImageCount, testWidth,
   testHeight, testImageData)  <- readMNISTImages "t10k-images-idx3-ubyte"
  (testLabelCount, testLabels) <- readMNISTLabels "t10k-labels-idx1-ubyte"

  assert (trainingWidth == testWidth && trainingHeight == testHeight) (return ())

  let input = take 60000 . map (vector . map fromIntegral) $ trainingImageData
  let trainingData = take 60000 . map toOutput $ trainingLabels

  initialNetwork <- genNetwork [20,10] (trainingWidth * trainingHeight)

  let trainedNetwork = trainNetwork initialNetwork input trainingData 2 10 3.0

  let testImage = head testImageData
  let testInput = vector . map fromIntegral $ testImage
  let testOutput = head testLabels --toOutput $ head testLabels

  --window testImage testWidth testHeight

  --let output = fst . last $ evalNetwork trainedNetwork testInput
  --let filteredOutput = map (>=0.5) . toList $ output

  --print $ fromOutput output
  --print $ testOutput

  let tests = map (runTest trainedNetwork) $ zip testImageData testLabels

  let success = length . filter (==True) $ tests
  let total = length tests

  putStr . show $ success
  putStr "/"
  putStrLn . show $ total

  --flip mapM_ (take 100 $ zip testImageData testLabels) runTest
  --  where
  --    runTest (imageData, label) = output
  --      where
  --        output2 = evalNetwork trainedNetwork
  --        output = putStrLn "test"
  --      --let output = evalNetwork trainedNetwork (vector . map fromIntegral $ imageData)
  --      --in do
  --      --  putStr "Guess: "
  --      --  putStr . show . fromOutput $ output


import Test.HUnit
import FFT.FFTControl
import FFT.Orig
import FFT.Samples
import System.Environment
import Data.Complex

testFFTbflySPipeline = TestCase (assertEqual "FFT-bflyS-Pipline" expected actual)
    where
        expected = (sum $ fftbflySPipeline samples512)
        actual = (sum $ fft samples512)
        samples512 = samples 1 512

testFFTDc = TestCase (assertEqual "DFTDc" expected actual)
    where
        expected = (sum $ fftdc samples512)
        actual = (sum $ fft samples512)
        samples512 = samples 1 512

tests = TestList [TestLabel "testFFTbFlysPipeline" testFFTbflySPipeline,
		TestLabel "testFFTDc" testFFTDc
                    ]

main = runTestTT tests


import qualified Numeric.BMML.RVMR             as RVM
import qualified Numeric.LinearAlgebra.HMatrix as H
import           System.Random.MWC             (asGenIO, uniformVector,
                                                withSystemRandom)

f :: Double -> Double
f x = 2.0 * x

generateX :: IO (H.Vector Double)
generateX =
    withSystemRandom . asGenIO $
    \gen ->
         uniformVector gen 12500

generateT :: H.Vector Double -> H.Vector Double
generateT = H.cmap f

main :: IO ()
main = do
  x <- generateX
  noiseX <- H.rand 12500 2
  let trainT = generateT x
      trainX = (H.fromColumns [x]) H.||| noiseX
  print (f 0.5)
  print (RVM.predict (RVM.fit trainX trainT) (H.fromList [0.5, 10.0, 10000.0]))


module GPLVM.PPCA
       ( PPCA (..)
       , makePPCA
       ) where

import Prelude (isNaN, log)

import Universum hiding (All, Any, Vector, map, toList, transpose, (++))

import qualified Universum as U

import GPLVM.Types
import GPLVM.Util

import Data.Array.Repa
import Data.Array.Repa.Repr.ForeignPtr
import Data.List ((\\))
import Numeric.LinearAlgebra.Repa hiding (Matrix, Vector, diag)
import System.Random

data PPCA = PPCA
  {  _noMissedData        :: Bool
   , _learningData        :: Matrix D Double
   , desiredDimension     :: Int
   , stopParameter        :: Either Int Double
   , _variance            :: Double
   , _W                   :: Matrix D Double
   , _finalExpLikelihood  :: Double
   , _restoredMatrix      :: Maybe (Matrix D Double)
   }

makePPCA
  :: (HasCallStack) => RandomGen gen
  => Matrix D Double   -- ^ learning vectors
  -> Int               -- ^ desired dimention
  -> Either Int Double -- ^ number of iterations or
                       --   the value of stop difference of the likelihood expectation between interations of EM.
  -> gen
  -> PPCA
makePPCA leaningVectors desiredDimension stopParameter generator =
  let _learningData@(ADelayed (Z :. n :. _) _) = leaningVectors
      initMatrix = randomMatrixD generator (n, desiredDimension)
      initDispersion = fromIntegral $ fst . next $ generator
      _noMissedData = not $ isNaN $ runIdentity $ sumAllP _learningData
      (_W, _variance, _finalExpLikelihood, _restoredMatrix) =
        if _noMissedData
        then emStepsFast _learningData initMatrix initDispersion stopParameter
        else emStepsMissed _learningData initMatrix initDispersion stopParameter
  in PPCA{..}

-- TODO: add comments to the most of following local binds
-- Most of them does not have informative names because they are just intermediate calculations.

emStepsFast
  :: (HasCallStack) => Matrix D Double
  -> Matrix D Double
  -> Double
  -> Either Int Double
  -> (Matrix D Double, Double, Double, Maybe (Matrix D Double))
emStepsFast learnMatrix@(ADelayed (Z :. _ :. n) _) initMatrix@(ADelayed (Z :. d :. _) _) initVariance stopParam =
  let learnMatrixCentered = transpose $ substractMean (transpose learnMatrix)

      stepOne :: (HasCallStack) => Array D DIM2 Double -> Double -> Int -> (Matrix D Double, Double, Double, Maybe (Matrix D Double))
      stepOne (!oldW) oldVariance iteration =
        let m = delay $ mapDiagonal (oldVariance +) $ (transpose oldW) `mulS` oldW
            invM = invS m
            expEznZ = delay $ invM `mul` (transpose oldW `mulS` learnMatrixCentered)
            expXnEZnZ = learnMatrixCentered `mulS` (transpose expEznZ)
            expZtrZ = computeS $ (map (*((fromIntegral n)*oldVariance)) $ delay invM) +^ (expEznZ `mulS` (transpose $ expEznZ)) :: Matrix F Double
        in stepTwo (expEznZ, expZtrZ, expXnEZnZ) oldW oldVariance iteration

      stepTwo :: (HasCallStack) => (Matrix D Double, Matrix F Double, Matrix F Double) -> Matrix D Double -> Double -> Int -> (Matrix D Double, Double, Double, Maybe (Matrix D Double))
      stepTwo (expEznZ, expZtrZ, expXnEZnZ) oldW oldVariance iteration =
        let newW = delay $ expXnEZnZ `mul` (inv expZtrZ)
            u = chol $ trustSym expZtrZ
            wr = newW `mulS` (transpose u)
            totalSum = sumAllS $ map (^ (2 :: Int)) learnMatrixCentered
            secondDenominatorInVariance = (*2) $ sumAllS $ expEznZ *^ (delay $ (transpose newW) `mulS` learnMatrixCentered)
            thirdDenominatorInVariance  = sumAllS $ map (^(2 :: Int)) wr
            newVariance = (totalSum - secondDenominatorInVariance + thirdDenominatorInVariance)/(fromIntegral $ n*d)
            maxDiffNewOldW = foldAllS max 0.0 $ map abs $ oldW -^ newW
            diffVariance = abs $ newVariance - oldVariance
            newC = mapDiagonal ((+) newVariance) $ newW `mulS` (transpose newW)
            newInvM = delay $ invS $ mapDiagonal (newVariance +) $ (transpose newW) `mulS` newW
            newInvC = (mapDiagonal (+(1/newVariance))
                        (map (\x -> -x/newVariance) $ (delay $ newW `mulS` newInvM) `mulS` (transpose newW)))
            expLogLikelihood = (-1.0)*((fromIntegral n)/2.0)*((fromIntegral d)*(log $ 2*pi)
              + (log $ detS newC)
              + (trace2S $ computeS $ map (/(fromIntegral $ n-1)) $ newInvC `mulS` (delay $ learnMatrixCentered `mulS` (transpose learnMatrixCentered))))
        in case stopParam of
              Left maxIteration -> if maxIteration > iteration then stepOne newW newVariance (iteration + 1) else (newW,newVariance, expLogLikelihood, Nothing)
              Right stopDiffBetweenIterations -> if (max diffVariance maxDiffNewOldW) > stopDiffBetweenIterations then stepOne newW newVariance (iteration + 1) else (newW,newVariance, expLogLikelihood, Nothing)

  in stepOne initMatrix initVariance 0

emStepsMissed
  :: (HasCallStack) => Matrix D Double
  -> Matrix D Double
  -> Double
  -> Either Int Double
  -> (Matrix D Double, Double, Double, Maybe (Matrix D Double))
emStepsMissed learnMatrix@(ADelayed (Z :. _ :. n) _) initMatrix@(ADelayed (Z :. d :. _) _) initVariance stopParam =
  let initMu = delay $ fromListUnboxed (Z:.d:.1) $ replicate d 0.0 -- extend (Any :. (1 :: Int)) $ meanColumnWithNan learnMatrix

      stepOne :: (HasCallStack) => Array D DIM2 Double -> Array D DIM2 Double -> Double -> Int -> (Matrix D Double, Double, Double, Maybe (Matrix D Double))
      stepOne oldMu (!oldW) oldVariance iteration =
        let yi i = extend (Any :. (1 :: Int)) $ slice learnMatrix (Any :. (i :: Int))
            yList i = toList (yi i)
            unknownIndices i =
              let zipped = zip [0..] (yList i)
              in U.map fst $ filter (isNaN . snd) zipped
            yPi i  = deleteRows (unknownIndices i) (yi i)
            muP i =  deleteRows (unknownIndices i) oldMu
            oldWP i = deleteRows (unknownIndices i) oldW
            mP i = mapDiagonal (+ oldVariance) $ (transpose (oldWP i)) `mulS` (oldWP i)
            invMP i = delay $ invS $ mP i
            expXi i = delay $ (delay $ (delay $ invMP i) `mulS` (transpose (oldWP i))) `mulS` (yPi i -^ muP i)
        in stepTwo ([yPi, expXi, invMP], unknownIndices, oldW) oldVariance iteration

      stepTwo :: (HasCallStack) => ([Int -> Matrix D Double], Int -> [Int], Matrix D Double) -> Double -> Int -> (Matrix D Double, Double, Double, Maybe (Matrix D Double))
      stepTwo ([yP, expXi, invMP], unknownIndices, oldW) oldVariance iteration =
        let expX = foldl1 (\acc i -> acc ++ i) $ U.map expXi [0..(n-1)]
            newMu = extend (Any :. (1 :: Int)) $ meanColumnWithNan $ (learnMatrix -^ (oldW `mulS` expX))

            newW =
              let yj j = extend (Any :. (1 :: Int) :. All) $ slice learnMatrix (Any :. (j :: Int) :. All)
                  yJList j = toList (yj j)
                  unknownColumns j =
                    let zipped = zip [0..] (yJList j)
                    in U.map fst $ filter (isNaN . snd) zipped
                  knownColumns j = [0..(n-1)] \\ (unknownColumns j)
                  expXPj j = deleteColumns (unknownColumns j) expX
                  sumInvMP j = sumListMatrices $ U.map invMP (knownColumns j)
                  gj j = ((expXPj j) `mulS` (transpose $ expXPj j)) +^ (map (*oldVariance) (sumInvMP j))
                  yPj j = deleteColumns (unknownColumns j) (yj j)
                  expXtrXj j = (expXPj j) `mulS` (transpose (map (\x -> x - (newMu ! (Z:.j:.0))) (yPj j)))
                  newWj j = (gj j) `solveS` (delay $ expXtrXj j)
              in transpose $ foldl1 (\acc j -> acc ++ j) $ U.map (\x -> delay $ newWj x) [0..(d-1)]

            newVariance =
              let newWPi i = deleteRows (unknownIndices i) newW
                  newMuP i = deleteRows (unknownIndices i) newMu
                  varianceStep i = sumAllS $
                    (map (^(2 :: Int)) (yP i -^ ((newWPi i) `mulS` (expXi i)) -^ (newMuP i)))
                    +^ (map (*oldVariance) $ diag $ delay $ (delay $ (delay $ newWPi i) `mulS` (invMP i)) `mulS` (transpose $ newWPi i))
                  knownVars = foldAllS (\acc x -> if isNaN x then acc else acc + 1.0) 0.0 learnMatrix
              in (sum $ U.map varianceStep [0..(n-1)])/knownVars

            expLogLikelihood =
              let newMuP i = deleteRows (unknownIndices i) newMu
                  newWPi i = deleteRows (unknownIndices i) newW
                  yCenteredP i = (yP i) -^ (newMuP i)
                  yCenteredProd i = delay $ (yCenteredP i) `mulS` (transpose $ yCenteredP i)
                  invMY i = mapDiagonal (+newVariance) $ (newWPi i) `mulS` (transpose $ newWPi i)
                  knownIndices i = [0..(d-1)] \\ (unknownIndices i)
                  expLogLikelihoodStep i = (fromIntegral $ length $ knownIndices i)*(log $ 2*pi) + (log $ detS $ invMY i) +
                    (trace2S $ computeS $ delay $ (invMY i) `solveS` (yCenteredProd i))

              in (-1.0)*(sum $ U.map expLogLikelihoodStep [0..(n-1)])/2.0

            restoredData =
              let muX = extend (Any :. (1 :: Int)) $ meanColumn (transpose expX)
                  finalMu = slice (newMu +^ (newW `mulS` muX)) (Any :. (0 :: Int))
                  wtrw = (transpose newW) `mulS` newW
                  factor1 = newW `mulS` (delay $ inv wtrw)
                  factor2 = computeS $ mapDiagonal (+newVariance) wtrw
                  restoredCentered =  factor1 `mul` factor2 `mul` (computeS $ expX)
              in (delay restoredCentered) +^ (extend (Any :. (n :: Int)) finalMu)

            diffVariance = abs $ newVariance - oldVariance
            maxDiffNewOldW = foldAllS max 0.0 $ map abs $ oldW -^ newW

        in case stopParam of
              Left maxIteration ->
                if maxIteration > iteration
                then stepOne newMu newW newVariance (iteration + 1)
                else (newW,newVariance, expLogLikelihood, Just restoredData)
              Right stopDiffBetweenIterations ->
                if (max diffVariance maxDiffNewOldW) > stopDiffBetweenIterations
                then stepOne newMu newW newVariance (iteration + 1)
                else (newW,newVariance, expLogLikelihood, Just restoredData)
  in stepOne initMu initMatrix initVariance 0


{-# LANGUAGE LambdaCase #-}
-- |
-- Module: Analysis.PerfModel
--
-- An experimental approach to obtaining strengths, which takes Elo (or
-- Elo-like) ratings as derived from a performance model. With a performance
-- model in hand, we can use https://stats.stackexchange.com/a/44142 to
-- obtain the victory probability of a 1500-rated racer given a field of
-- racers, and use that as a measure of strength.
--
-- According to the model, player performances follow a gamma distribution,
-- with the outcome 0 corresponding to an ideal performance. The shape
-- parameter of the distribution is fixed, while the rate parameter varies
-- across players. A higher rate parameter amounts to results that are both
-- more consistent and closer to the ideal. (Pinning both mean and standard
-- deviation of player performance to a single parameter is the main
-- simplifying assumption of the model.)
--
-- Note that the performance model winning probabilities for pairs of players
-- agree exactly with the conventional Elo formula if the shape parameter of
-- the gamma distribution is taken to be 1.
module Analysis.PerfModel
    ( perfModelStrength
    , perfModelTopStrength
    , example
    ) where

import Orbital (orbitalDistr, initialRating)
import qualified Util.Combinations as Util

import qualified Numeric.Integration.TanhSinh as Integration
import Statistics.Distribution

import Data.Profunctor
import qualified Control.Foldl as L
import qualified Data.IntMap.Strict as IntMap
import Data.IntMap.Strict (IntMap, (!))
import qualified Data.IntSet as IntSet
import Data.IntSet (IntSet)
import qualified Control.Foldl as L

-- | PDF for a race victory against a field of racers.
raceWinPDF
    :: Int       -- ^ Gamma shape.
    -> [Double]  -- ^ Ratings of the field of opponents.
    -> Double    -- ^ Rating of the racer.
    -> Double    -- ^ Performance.
    -> Double
raceWinPDF gsh rs r t = density (orbitalDistr gsh r) t
    * L.fold (lmap winVersus L.product) rs
    where
    winVersus r' = 1 - cumulative (orbitalDistr gsh r') t

-- | Integrating the race win PDF gives the likelihood of victory, whose
-- reciprocal we use as a strength metric.
perfModelStrength
    :: Int       -- ^ Gamma shape.
    -> [Double]  -- ^ Ratings of the field of opponents.
    -> Double
perfModelStrength gsh rs = 1 / integ
    where
    integ = Integration.result . Integration.absolute 1e-6 $
        Integration.nonNegative Integration.trap (raceWinPDF gsh rs initialRating)

-- TODO: positionPDF and positionPDFPre are still the same as in
-- Analysis.PerfModel.Reference.Precompute . Conside redefining them so that
-- topPDF can share code with them.

-- | Like 'raceWinPDF', but for an arbitrary position on the scoreboard.
positionPDF
    :: Int       -- ^ Gamma shape.
    -> Int       -- ^ Target position.
    -> [Double]  -- ^ Ratings of the field of opponents.
    -> Double    -- ^ Rating of the racer.
    -> Double    -- ^ Performance.
    -> Double
positionPDF gsh pos rs r t = density (orbitalDistr gsh r) t
    * L.fold L.sum (L.fold L.product . toFactors <$> combos)
    where
    n = length rs
    combos = Util.combinations n (pos-1) [1..n]
    loseVersus r' = cumulative (orbitalDistr gsh r') t
    -- winVersus r' = 1 - loseVersus r'
    lvs = IntMap.fromList (zip [1..n] (loseVersus <$> rs))
    wvs = (1 -) <$> lvs
    toFactors (above, below) = map (lvs !) above ++ map (wvs !) below

-- | A variant of 'positionPDF' that takes precomputed probabilities.
positionPDFPre :: Int                            -- ^ Target position.
               -> (IntMap Double, IntMap Double) -- ^ Win and loss probabilities
                                                 -- (given a probe laptime) for
                                                 -- for each racer on the
                                                 -- scoreboard.
               -> Double                         -- ^ Probability density of the
                                                 -- probe laptime.
               -> Double
positionPDFPre pos (lvs, wvs) p =
    p * L.fold L.sum (L.fold L.product . toFactors <$> combos)
    where
    n = IntMap.size lvs
    combos = Util.combinations n (pos-1) [1..n]
    toFactors (above, below) = map (lvs !) above ++ map (wvs !) below

-- | Like 'positionPDF', but for a top-N finish.
topPDF
    :: Int       -- ^ Gamma shape.
    -> Int       -- ^ Target position.
    -> [Double]  -- ^ Ratings of the field of opponents.
    -> Double    -- ^ Rating of the racer.
    -> Double    -- ^ Performance.
    -> Double
topPDF gsh pos rs r t = (probeDensity *) . L.fold L.sum $
    L.fold L.sum . partials <$> validPositions
    where
    validPositions = zipWith const [1..pos] (r : rs)
    loseVersus r' = cumulative (orbitalDistr gsh r') t
    lvs = IntMap.fromList (zip [1..] (loseVersus <$> rs))
    wvs = (1 -) <$> lvs
    probeDensity = density (orbitalDistr gsh r) t
    partials = Util.processCombsInt alg (length rs) . subtract 1
    alg = \case
        Util.LeafF Nothing -> 0
        Util.LeafF (Just rest) ->
            L.fold L.product ((wvs !) <$> IntSet.toList rest)
        Util.FlowerF a rest -> (lvs ! a)
            * L.fold L.product ((wvs !) <$> IntSet.toList rest)
        Util.BranchF a bs -> (lvs ! a) * L.fold L.sum bs

-- | Like 'perfModelStrength', but for a top-N finish. Note that the
-- computational cost grows quickly with the length of the list of ratings.
perfModelTopStrength
    :: Int       -- ^ Gamma shape.
    -> Int       -- ^ Target position.
    -> [Double]  -- ^ Ratings of the field of racers.
    -> Double
perfModelTopStrength gsh pos rs = 1 / integ
    where
    integ = Integration.result . Integration.absolute 1e-6 $
        Integration.nonNegative Integration.trap (topPDF gsh pos rs initialRating)

example :: [Double]
example = [2200,2100,1900,1870,1850,1600]
--example = [2200,2100,1900,1870,1850,1600,1590,1570,1510,1420,1370,1350,1225]

-- >$> perfModelStrength [2200,2100,1900,1870,1850]
--
-- >$> :set +s
--

-- >$> perfModelTopStrength 5 example


{-# LANGUAGE GADTs #-}
module Distrib (
    -- * Data Constructors
    Distrib(..)
    -- * Operations
  , compNormalDistr
  , pdf
  ) where

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

import           Classifier
import qualified Data.Vector.Unboxed            as VU
import           Statistics.Distribution
import           Statistics.Distribution.Normal
import qualified Statistics.Sample              as Statistics

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

data Distrib where
  Distrib :: ContDistr d => d -> Distrib

-- | We don't know much about our numeric variable
-- so we will assume normallity to simplify the implementation.
compNormalDistr :: [Double] -> Distrib
compNormalDistr xs =
  Distrib (normalDistr mean' stdDev')
    where
      v = VU.fromList xs
      mean' = Statistics.mean v
      stdDev'  = max 0.1 (Statistics.stdDev v) -- stdDev must be > 0

-- | Probability Density Function
pdf :: Distrib -> Double -> Probability
pdf (Distrib d) = Probability . density d


module Statistics.ANOVA (anova, fisherLSD) where
import Data.List
import Statistics.Distribution hiding (mean, variance)
import Statistics.Distribution.FDistribution
import Statistics.Distribution.StudentT
import Statistics.Sample
import qualified Data.Vector as V

{-# INLINE treatments #-}
treatments :: (Ord a,Eq a) => V.Vector (a, Double) -> V.Vector a
treatments input = V.fromList $! nub $! V.toList $! V.map fst input --not good

treatmentMeans :: (Ord a, Eq a) => V.Vector a -> V.Vector (a, Double) -> V.Vector (a,Double)
treatmentMeans treatments input = V.map (\x -> x input) meanFuncs
  where meanFuncs = V.map treatmentMean $! treatments

treatmentMean :: (Ord a, Eq a) => a -> V.Vector (a,Double) -> (a,Double)
treatmentMean treatment input = (treatment, avg)
  where avg = mean $! V.map snd $! filterTreatment treatment $! input

treatmentSize treatment input = V.length $! filterTreatment treatment $! input

{-# INLINE filterTreatment #-}
filterTreatment :: (Ord a, Eq a) => a -> V.Vector (a, Double) -> V.Vector (a, Double)
filterTreatment treatment = V.filter (\(x,y) -> x == treatment)

sampleMean input = mean $! V.map snd input

ssTreatment :: (Ord a, Eq a) => V.Vector a -> V.Vector (a, Double) -> Double
ssTreatment treatments input = V.sum $! V.zipWith (\size mean -> size*(mean-sMean)*(mean-sMean)) tSizes tMeans
  where sMean = sampleMean input
        tMeans = V.map snd $! treatmentMeans treatments input
        tSizes = V.map (\x -> fromIntegral $! x input) $! V.map treatmentSize $! treatments --input

ssError :: (Ord a, Eq a) => V.Vector (a, Double) -> Double
ssError input = V.sum $! V.map (\x -> x input) $! V.map treatmentError $! treatments input

treatmentError treatment input = V.sum $! V.map (\x -> (x-tMean)*(x-tMean)) filteredInput
  where filteredInput = V.map snd $! filterTreatment treatment input
        tMean = snd $! treatmentMean treatment input

fValue :: (Ord a, Eq a) => V.Vector a -> V.Vector (a, Double) -> Double
fValue treatments input = ((ssTreatment treatments input)/(ssError input)) * ((sampleSize - numberOfTreatments)/(numberOfTreatments -1))
  where numberOfTreatments = fromIntegral $! V.length $! treatments
        sampleSize = fromIntegral $! V.length input

fDist treatments input = fDistribution (numberOfTreatments-1) (sampleSize-numberOfTreatments)
  where numberOfTreatments = fromIntegral $! V.length treatments
        sampleSize = fromIntegral $! V.length input

pValue input = 1 - cumulative (fDist t input) (fValue t input)
  where t = treatments input

anova pvalue input
  | pValue input < pvalue = True
  | otherwise = False

tTest :: (Ord a, Eq a) => Double -> a -> a -> V.Vector (a, Double) -> Bool
tTest pvalue treatment1 treatment2 inputs = abs (tTestT treatment1 treatment2 inputs) > critical
  where critical = quantile (studentT (degreesOfFreedom treatment1 treatment2 inputs)) (1 - pvalue/2)

tTestT :: (Ord a, Eq a) => a -> a -> V.Vector (a, Double) -> Double
tTestT treatment1 treatment2 inputs = (mean1-mean2)/sqrt(var1/size1+var2/size2)
  where mean1 = snd $! treatmentMean treatment1 inputs
        mean2 = snd $! treatmentMean treatment2 inputs
        size1 = fromIntegral $! V.length $! filterTreatment treatment1 inputs
        size2 = fromIntegral $! V.length $! filterTreatment treatment2 inputs
        var1 = variance $! V.map snd $! filterTreatment treatment1 inputs
        var2 = variance $! V.map snd $! filterTreatment treatment2 inputs

degreesOfFreedom
  :: (Ord a, Eq a) => a -> a -> V.Vector (a, Double) -> Double
degreesOfFreedom treatment1 treatment2 inputs = numerator / denominator
  where size1 = fromIntegral $! V.length $! filterTreatment treatment1 inputs
        size2 = fromIntegral $! V.length $! filterTreatment treatment2 inputs
        var1 = variance $! V.map snd $! filterTreatment treatment1 inputs
        var2 = variance $! V.map snd $! filterTreatment treatment2 inputs
        numerator = (var1/size1+var2/size2)^2
        denominator = (((var1/size1)^2)/(size1-1)) + (((var2/size2)^2)/(size2-1))

noDifferenceInMean pvalue treatment input = (:) treatment $!
                                            map fst $!
                                            filter (\(x,y) -> not y) $!
                                            map (\x -> (x, tTest pvalue treatment x input)) othertreatments
  where othertreatments = delete treatment $! V.toList $! treatments input

equivalenceList pvalue input = nub $! V.toList $! V.map (\x -> sort $! noDifferenceInMean pvalue x input) $! treatments input

fisherLSD :: (Eq a, Ord a) => Double -> V.Vector (a, Double) -> [[a]]
fisherLSD pvalue input = case (anova pvalue input) of
  False -> [V.toList $! treatments input]
  True -> equivalenceList pvalue input