{-# LANGUAGE TypeApplications #-}

module SSSFNProps where

import Data.VectorSpace
import qualified GHC.Exts as GHC
import Numeric.LinearAlgebra hiding ((#>), (===), inv)
import SSSFN
import Test.QuickCheck

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

infix 4 ~~~

(~~~) ::
  (GHC.IsList as, Show as, a ~ GHC.Item as, RealFloat a, Show a) =>
  as ->
  as ->
  Property
xs ~~~ ys =
  let xs' = GHC.toList xs
      ys' = GHC.toList ys
      sc = 1 `max` maximum (map abs xs') `max` maximum (map abs ys')
      di = maximum (zipWith absdiff xs' ys')
   in counterexample
        ( show xs
            ++ " ~~~ "
            ++ show ys
            ++ " (difference "
            ++ show di
            ++ ", scale "
            ++ show sc
            ++ ")"
        )
        (length xs' == length ys' && di <= eps * sc)
  where
    absdiff x y = abs (x - y)
    eps = 100 * peps

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

prop_Grid_add_assoc :: Grid Double -> Grid Double -> Grid Double -> Property
prop_Grid_add_assoc x y z = (x ^+^ y) ^+^ z ~~~ x ^+^ (y ^+^ z)

prop_Grid_add_lunit :: Grid Double -> Property
prop_Grid_add_lunit x = zeroV ^+^ x ~~~ x

prop_Grid_add_runit :: Grid Double -> Property
prop_Grid_add_runit x = x ~~~ x ^+^ zeroV

prop_Grid_add_comm :: Grid Double -> Grid Double -> Property
prop_Grid_add_comm x y = x ^+^ y ~~~ y ^+^ x

prop_Grid_add_linv :: Grid Double -> Property
prop_Grid_add_linv x = negateV x ^+^ x ~~~ zeroV

prop_Grid_add_rinv :: Grid Double -> Property
prop_Grid_add_rinv x = x ^+^ negateV x ~~~ zeroV

prop_Grid_add_scale_dist :: Double -> Grid Double -> Grid Double -> Property
prop_Grid_add_scale_dist a x y = a *^ (x ^+^ y) ~~~ a *^ x ^+^ a *^ y

prop_Grid_scale_assoc :: Double -> Double -> Grid Double -> Property
prop_Grid_scale_assoc a b x = a *^ (b *^ x) ~~~ (a * b) *^ x

prop_Grid_scale_lunit :: Grid Double -> Property
prop_Grid_scale_lunit x = 1 *^ x ~~~ x

prop_Grid_scale_szero :: Grid Double -> Property
prop_Grid_scale_szero x = 0 *^ x ~~~ zeroV

prop_Grid_scale_zero :: Double -> Property
prop_Grid_scale_zero a = a *^ (zeroV :: Grid Double) ~~~ zeroV

prop_Grid_scale_comm :: Double -> Double -> Grid Double -> Property
prop_Grid_scale_comm a b x = a *^ (b *^ x) ~~~ b *^ (a *^ x)

prop_Grid_scale_Grid_add_dist :: Double -> Double -> Grid Double -> Property
prop_Grid_scale_Grid_add_dist a b x = (a + b) *^ x ~~~ a *^ x ^+^ b *^ x

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

prop_Op_add_assoc :: Op Double -> Op Double -> Op Double -> Property
prop_Op_add_assoc x y z = (x + y) + z ~~~ x + (y + z)

prop_Op_add_lunit :: Op Double -> Property
prop_Op_add_lunit x = 0 + x ~~~ x

prop_Op_add_runit :: Op Double -> Property
prop_Op_add_runit x = x ~~~ x + 0

prop_Op_add_comm :: Op Double -> Op Double -> Property
prop_Op_add_comm x y = x + y ~~~ y + x

prop_Op_add_linv :: Op Double -> Property
prop_Op_add_linv x = negate x + x ~~~ 0

prop_Op_add_rinv :: Op Double -> Property
prop_Op_add_rinv x = x + negate x ~~~ 0

prop_Op_add_scale_dist :: Double -> Op Double -> Op Double -> Property
prop_Op_add_scale_dist a x y = a *^ (x + y) ~~~ a *^ x + a *^ y

prop_Op_mul_assoc :: Op Double -> Op Double -> Op Double -> Property
prop_Op_mul_assoc x y z = (x * y) * z ~~~ x * (y * z)

prop_Op_mul_lunit :: Op Double -> Property
prop_Op_mul_lunit x = 1 * x ~~~ x

prop_Op_mul_runit :: Op Double -> Property
prop_Op_mul_runit x = x ~~~ x * 1

prop_Op_mul_scale_assoc :: Double -> Op Double -> Op Double -> Property
prop_Op_mul_scale_assoc a x y = a *^ (x * y) ~~~ (a *^ x) * y

prop_Op_scale_assoc :: Double -> Double -> Op Double -> Property
prop_Op_scale_assoc a b x = a *^ (b *^ x) ~~~ (a * b) *^ x

prop_Op_scale_lunit :: Op Double -> Property
prop_Op_scale_lunit x = 1 *^ x ~~~ x

prop_Op_scale_lzero :: Op Double -> Property
prop_Op_scale_lzero x = 0 *^ x ~~~ 0

prop_Op_scale_rzero :: Double -> Property
prop_Op_scale_rzero a = a *^ (0 :: Op Double) ~~~ 0

prop_Op_scale_comm :: Double -> Double -> Op Double -> Property
prop_Op_scale_comm a b x = a *^ (b *^ x) ~~~ b *^ (a *^ x)

prop_Op_scale_add_dist :: Double -> Double -> Op Double -> Property
prop_Op_scale_add_dist a b x = (a + b) *^ x ~~~ (a *^ x) + (b *^ x)


{-# LANGUAGE FlexibleContexts #-}
-- | Pearson's chi squared test.
module Statistics.Test.ChiSquared (
    chi2test
    -- * Data types
  , TestType(..)
  , TestResult(..)
  ) where

import Prelude hiding (sum)
import Statistics.Distribution
import Statistics.Distribution.ChiSquared
import Statistics.Function (square)
import Statistics.Sample.Internal (sum)
import Statistics.Test.Types
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U

-- | Generic form of Pearson chi squared tests for binned data. Data
--   sample is supplied in form of tuples (observed quantity,
--   expected number of events). Both must be positive.
chi2test :: (G.Vector v (Int,Double), G.Vector v Double)
         => Double              -- ^ p-value
         -> Int                 -- ^ Number of additional degrees of
                                --   freedom. One degree of freedom
                                --   is due to the fact that the are
                                --   N observation in total and
                                --   accounted for automatically.
         -> v (Int,Double)      -- ^ Observation and expectation.
         -> TestResult
chi2test p ndf vec
  | ndf < 0        = error $ "Statistics.Test.ChiSquare.chi2test: negative NDF " ++ show ndf
  | n   < 0        = error $ "Statistics.Test.ChiSquare.chi2test: too short data sample"
  | p > 0 && p < 1 = significant $ complCumulative d chi2 < p
  | otherwise      = error $ "Statistics.Test.ChiSquare.chi2test: bad p-value: " ++ show p
  where
    n     = G.length vec - ndf - 1
    chi2  = sum $ G.map (\(o,e) -> square (fromIntegral o - e) / e) vec
    d     = chiSquared n
{-# SPECIALIZE
    chi2test :: Double -> Int -> U.Vector (Int,Double) -> TestResult #-}
{-# SPECIALIZE
    chi2test :: Double -> Int -> V.Vector (Int,Double) -> TestResult #-}


{-# LANGUAGE OverloadedStrings, RecordWildCards, NamedFieldPuns #-}
module Numeric.LP where

import           Control.Monad
import           Data.Array.IArray           (Array, (//))
import qualified Data.Array.IArray           as A
import           Data.Maybe
import qualified Data.Text.Lazy              as T
import qualified Data.Text.Lazy.Builder      as TB
import qualified Formatting                  as F
import           Numeric.LinearAlgebra       (I, Matrix, Vector, atIndex, cols,
                                              flatten, ident, minIndex, rows,
                                              subMatrix, subVector, (!), (><),
                                              (?), (|||), (¿))
import qualified Numeric.LinearAlgebra       as LA
import           Numeric.LinearAlgebra.Devel (atM', modifyMatrix, runSTMatrix,
                                              thawMatrix)
import           Text.UnicodeTable

data LP = LP {
    a        :: Matrix Double,
    lhs, rhs :: Vector I}

lp2Table :: Int -> Int -> LP -> Table
lp2Table nPrec nColW LP{..} = mkTable ta cws hs vs
    where rN = rows a
          cN = cols a
          taElem r c | r>0 && c>0 = F.bprint (F.fixed nPrec) $ a `atIndex` (r-1, c-1)
                     | r==0 && c>1 = TB.fromString $ "x"++show (rhs!(c-1))
                     | r==1 && c==0 =  TB.fromString "z"
                     | r>1 && c==0 = TB.fromString $ "x"++show (lhs!(r-1))
                     | otherwise = TB.fromString ""
          ta = A.array ((0,0), (rN, cN)) [((r,c), taElem r c) | r <- [0..rN], c <- [0..cN]]
          cws = A.listArray (0,cN) (replicate (cN+1) nColW)
          hs = konstArray0 (rN+2) NoLine // [(0, SingleLine), (1, DoubleLine), (rN+1, SingleLine)]
          vs = konstArray0 (cN+2) NoLine // [(0, SingleLine), (1, DoubleLine), (cN+1, SingleLine)]

konstArray0 :: Int -> e -> Array Int e
konstArray0 n x = A.listArray (0,n-1) (replicate n x)

instance Show LP where
    show = T.unpack . TB.toLazyText . fromTable . lp2Table 4 15

transfromToEquationalForm :: Matrix Double -> Matrix Double
transfromToEquationalForm a = a ||| (ident (rows a))

choosePivot :: LP -> Maybe (Int, Int)
choosePivot LP {..} =
    let c = subVector 1 (cols a-1) (flatten $ a ? [0])
        b = subVector 1 (cols a-1) (flatten $ a ¿ [0])
        a' = subMatrix (1,1) (rows a-1, cols a-1) a
    in case LA.find (>0) c of
            s:_ ->
                let as = flatten $ a' ¿ [s] -- column s as a vector
                in Just (1 + minIndex (b / as), 1+s)
            _ -> Nothing

-- see <https://en.wikipedia.org/wiki/Simplex_algorithm>
-- Let's solve the following problem:
--      maximize $c^t x$ subject to
--      Ax ≤ b
--      ∀i, xᵢ ≥ 0
--      x,c ∈ ℝⁿ, b ∈ ℝ^p, A ∈ ℝ^pxn

ex1 :: LP
ex1 = LP {a = (3><3)[0,2,-4, 2,-6,1, 8,3,-4],
          lhs = LA.fromList [0,1,4],
          rhs = LA.fromList [0,2,3]
        }

simplexStep :: LP -> Int -> Int -> LP
simplexStep (LP {..}) sr sc =
    let theta = atM' a sr sc -- save the pivot element
        pc = flatten $ a ¿ [sc] -- save the whole pivot column
        a' = runSTMatrix
            (do
                a'm <- thawMatrix a
                -- replace each row -- except the pivot row -- by that linear combination
                -- of itself and the pivot row which makes its pivot column entry zero.
                forM_ [r | r <- [0..(rows a - 1)], r /= sr] $ \r -> do
                    let m = (-1) * (a `atIndex` (r,sc)) / theta
                    forM_ [0..(cols a - 1)] $ \i -> do
                        modifyMatrix a'm r i (\x -> x + m * (a `atIndex` (sr,i)))
                -- divide the pivot row by the negative of the pivot
                forM_ [0..(cols a - 1)] $ \i -> do
                        modifyMatrix a'm sr i (\x -> x/(-theta))
                -- replace the pivot by the negative of its saved reciprocal
                modifyMatrix a'm sr sc (\_ -> 1/theta)
                -- replace the other elements of the pivot column by its saved values divided
                -- by the saved pivot element
                forM_ [r | r <- [0..(rows a - 1)], r /= sr] $ \j -> do
                        modifyMatrix a'm j sc (\_ -> pc!j/theta)
                return a'm)
    in LP a' lhs rhs

main :: IO ()
main = do
    putStrLn (show ex1)
    let (LP {..}) = ex1
    let (sr, sc) = fromJust $ choosePivot ex1
    putStrLn ("Pivot row/column: "++(show sr)++", "++(show sc))
    let theta = atM' a sr sc
    putStrLn ("Pivot element: "++(show theta))


import           Numeric.LinearAlgebra
import           System.Random ( mkStdGen )
import           Learning ( nrmse )

import           RC.NTC as RC

-- Get training data
takePast :: Element a => Int -> Matrix a -> Matrix a
takePast horizon xs = xs ?? (All, Take (len - horizon))
  where
    len = cols xs

-- Get target data
takeFuture :: Element a => Int -> Matrix a -> Matrix a
takeFuture horizon = (?? (All, Drop horizon))

-- Manipulate the data matrix: horizontal split
splitAtRatio :: Element a => Double -> Matrix a -> (Matrix a, Matrix a)
splitAtRatio ratio m = splitAt' spl m
  where
    total = cols m
    spl = round $ ratio * fromIntegral total

splitAt' :: Element a => Int -> Matrix a -> (Matrix a, Matrix a)
splitAt' i m = (m ?? (All, Take i), m ?? (All, Drop i))

main :: IO ()
main = do
  -- Load and transpose time series to predict
  dta <- tr <$> loadMatrix "examples/data/mg.txt"

  -- Train on 50% of data
  let (train, validate) = splitAtRatio 0.50 dta

  -- Configure a new NTC network
  let p = RC.par0 { RC._inputWeightsRange = (0.1, 0.3) }
      g = mkStdGen 1111
      ntc = RC.new g p (1, 1000, 1)

  -- 20 steps ahead prediction horizon
  let horizon = 20

  let train' = takePast horizon train  -- Past series
      trainTeacher = takeFuture horizon train  -- Predicted series

  let forgetPts = 300  -- Washout

  -- Train
  case RC.learn ntc forgetPts train' trainTeacher of
    Left s -> error s
    Right ntc' -> do
      let target = (takeFuture horizon validate) ?? (All, Drop forgetPts)

      -- Predict
      case RC.predict ntc' forgetPts (takePast horizon validate) of
        Left s -> error s
        Right prediction -> do
          let tgt' = flatten target
              predic' = flatten prediction
              err = nrmse tgt' predic'

          putStrLn $ "Error: " ++ show err

          let result = (tr target) ||| (tr prediction)

          saveMatrix "examples/NTC/result.txt" "%g" result


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE RebindableSyntax #-}
module Models.RSA2d where

import Prelude (Double,Int, ($), (.), fmap, Ord(..), toRational, putStrLn, IO,String,show,unwords,unlines,writeFile,Integer)
import Algebra.Linear.Chebyshev
import Algebra.Linear.Vector
import Algebra.Classes
import Data.Complex
import Data.Foldable
import Models.Gnuplot

type Ivl = (Double,Double)

type F2 a = Double -> Double -> a
type S2 a = Samples (Samples a)
-- type P2 a = Polynomial (Polynomial a)

type C = Complex Double

normal :: Transcendental a => a -> a -> a -> a
normal μ σ x = recip (σ * sqrt (2 * pi)) * exp (-0.5 * ((μ - x) / σ) ^+ 2)

prior :: Transcendental a => a -> p -> a
prior h u = normal 68 3 h

groundFun :: (Ord t, Transcendental t) => t -> t -> t
groundFun h u = if h < u then 0 else prior h u

size__ :: Int
size__ = 256

ivlH :: (Double, Double)
ivlH = (68-10, 68+10)

ivlU :: (Double, Double)
ivlU = (68-20, 68+10)

ground :: S2 Double
ground = sample' groundFun

sample' :: F2 Double -> S2 Double
sample' = sample2d size__ ivlH ivlU

sample2d :: Int -> Ivl -> Ivl -> F2 Double -> S2 Double
sample2d n (lo1,hi1) (lo2,hi2) f =
  sample @Double n lo1 hi1 $ \x ->
  sample @Double n lo2 hi2 $ \y ->
  f x y

normalizeS :: Samples Double -> Samples Double
normalizeS v = (recip (integralUnitIvl @Double $ fmap realPart $ samplesToPoly @C $ fmap (:+ 0) v)) *^ v

-- fit2 :: S2 Double -> P2 Double
-- fit2 = fmap (fmap realPart) . samplesToPoly @C  . fmap (samplesToPoly @C . fmap (:+ 0) )

get :: Samples (Samples a) -> Int -> Int -> a
get m i j = fromSamples (fromSamples m ! i) ! j

pts :: Vec Double
pts = chebPoints size__

normalizeH :: Samples (Samples Double) -> Samples (Samples Double)
normalizeH = xpose . fmap normalizeS . xpose
 where
  xpose :: Samples (Samples Double) -> Samples (Samples Double)
  xpose m = Samples $ generate (size__+1)  $ \i -> Samples $ generate (size__+1) $ \ j -> get m j i

mkSamples :: (Int -> Int -> a) -> Samples (Samples a)
mkSamples f = Samples $ generate (size__+1)  $ \i -> Samples $ generate (size__+1) $ \ j -> f i j

normalizeU :: Samples (Samples Double) -> Samples (Samples Double)
normalizeU = fmap normalizeS

l0 :: Samples (Samples Double)
l0 = normalizeH ground

α :: Natural
α = 40

s1 :: Samples (Samples Double)
s1 =  normalizeU (fmap (fmap (^+ α)) l0)

l1 :: Samples (Samples Double)
l1 = normalizeH $ mkSamples $ \h u -> get s1 h u * get (sample' $ prior) h u

-- >>> toGnuPlot "ground.dat" ground
-- >>> toGnuPlot "l0.dat" l0
-- >>> toGnuPlot "s1.dat" s1
-- >>> toGnuPlot "l1.dat" l1


{-# LANGUAGE DataKinds                     #-}
{-# LANGUAGE FlexibleContexts              #-}
{-# LANGUAGE RankNTypes                    #-}
{-# LANGUAGE ScopedTypeVariables           #-}
{-# LANGUAGE TypeApplications              #-}
{-# LANGUAGE TypeOperators                 #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

module Backprop.Learn.Loss (
  -- * Loss functions
    Loss
  , crossEntropy, crossEntropy1
  , squaredError, absError, totalSquaredError, squaredErrorV
  -- , totalCov
  -- ** Manipulate loss functions
  , scaleLoss
  , sumLoss
  , sumLossDecay
  , lastLoss
  , zipLoss
  , t2Loss
  -- * Regularization
  , Regularizer
  , l2Reg
  , l1Reg
  , noReg
  -- ** Manipulate regularizers
  , addReg
  , scaleReg
  ) where

import           Backprop.Learn.Regularize
import           Control.Applicative
import           Data.Coerce
import           Data.Finite
import           Data.Type.Tuple
import           GHC.TypeNats
import           Lens.Micro
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import qualified Data.Vector.Sized                     as SV
import qualified Prelude.Backprop                      as B

type Loss a = forall s. Reifies s W => a -> BVar s a -> BVar s Double

crossEntropy :: KnownNat n => Loss (R n)
crossEntropy targ res = -(log res <.> auto targ)

crossEntropy1 :: Loss Double
crossEntropy1 targ res = -(log res * auto targ + log (1 - res) * auto (1 - targ))

squaredErrorV :: KnownNat n => Loss (R n)
squaredErrorV targ res = e <.> e
  where
    e = res - auto targ

totalSquaredError
  :: (Backprop (t Double), Num (t Double), Foldable t, Functor t)
    => Loss (t Double)
totalSquaredError targ res = B.sum (e * e)
  where
    e = auto targ - res

squaredError :: Loss Double
squaredError targ res = (res - auto targ) ** 2

absError :: Loss Double
absError targ res = abs (res - auto targ)

-- -- | Sum of covariances between matching components.  Not sure if anyone
-- -- uses this.
-- totalCov :: (Backprop (t Double), Foldable t, Functor t) => Loss (t Double)
-- totalCov targ res = -(xy / fromIntegral n - (x * y) / fromIntegral (n * n))
--   where
--     x = auto $ sum targ
--     y = B.sum res
--     xy = B.sum (auto targ * res)
--     n = length targ

-- klDivergence :: Loss Double
-- klDivergence =

sumLoss
    :: (Traversable t, Applicative t, Backprop a)
    => Loss a
    -> Loss (t a)
sumLoss l targ = sum . liftA2 l targ . sequenceVar

zipLoss
    :: (Traversable t, Applicative t, Backprop a)
    => t Double
    -> Loss a
    -> Loss (t a)
zipLoss xs l targ = sum
                  . liftA3 (\α t -> (* auto α) . l t) xs targ
                  . sequenceVar

sumLossDecay
    :: forall n a. (KnownNat n, Backprop a)
    => Double
    -> Loss a
    -> Loss (SV.Vector n a)
sumLossDecay β = zipLoss $ SV.generate (\i -> β ** (fromIntegral i - n))
  where
    n = fromIntegral $ maxBound @(Finite n)

lastLoss
    :: forall n a. (KnownNat (n + 1), Backprop a)
    => Loss a
    -> Loss (SV.Vector (n + 1) a)
lastLoss l targ = l (SV.last targ)
                . viewVar (coerced . SV.ix @(n + 1) maxBound)
                . B.coerce @_ @(ABP (SV.Vector (n + 1)) a)

coerced :: Coercible a b => Lens' a b
coerced f x = coerce <$> f (coerce x)
{-# INLINE coerced #-}

-- | Scale the result of a loss function.
scaleLoss :: Double -> Loss a -> Loss a
scaleLoss β l x = (* auto β) . l x

-- | Lift and sum a loss function over the components of a ':&'.
t2Loss
    :: (Backprop a, Backprop b)
    => Loss a                   -- ^ loss on first component
    -> Loss b                   -- ^ loss on second component
    -> Loss (a :# b)
t2Loss f g (xT :# yT) (xR :## yR) = f xT xR + g yT yR


import Data.Map.Strict (Map)
import Numeric.LinearAlgebra
import LeastSquares

type VarName = String
data LinExpr = Var VarName Int
             | ...
infix 4 :==:
data LinContr = LinExpr :==: LinExpr
data QuadExpr = SumSquares LinExpr
              | ProdQuad Double QuadExpr
              | SumQuad QuadExpr QuadExpr

minimize :: QuadExpr -> [LinContr]
                -> Map VarName (Vector R)


{- |
Description :  DNA Markov models (substitution models or mutation models)
Copyright   :  (c) Dominik Schrempf 2017
License     :  GPLv3

Maintainer  :  dominik.schrempf@gmail.com
Stability   :  unstable
Portability :  non-portable (not tested)

This module provides rate matrices for the most common DNA substitution models
(they will be added according to need). They are also used as mutation models by
the boundary mutation model.

* Changelog

-}

module DNAModel
  ( DNAModelSpec(..)
  , DNAModel(..)
  , Nuc
  , getDNAModelInfoStr
  , rateMatrix
  , dnaModelSpecGetStateFreqVec
  , dnaModelScale
  ) where

import           Data.List             (intercalate)
import           Numeric.LinearAlgebra
import           RateMatrix

-- | The nucleotides.
data Nuc = A | C | G | T deriving (Eq, Show, Read, Ord, Bounded, Enum)

nNucs :: Int
nNucs = 1 + fromEnum (maxBound :: Nuc)

-- | The DNA model specifications. E.g., for the HKY model, we only have kappa.
data DNAModelSpec = JC |
                    -- | HKY model with transition to transversion ratio kappa.
                    HKY Double StationaryDist |
                    -- | GTR model with five params and state frequency vector.
                    GTR Double Double Double Double Double StationaryDist

instance Show DNAModelSpec where
  show JC        = "JC"
  show (HKY k f) = "HKY[" ++ show k ++ "]" ++ show f
  show (GTR a b c d e f) = "GTR[" ++ paramsStr ++ "]" ++ show f
    where paramsStr = intercalate "," (map show [a,b,c,d,e])

-- | Extract stationary frequency vector of a DNA model specification 'DNAModelSpec'.
dnaModelSpecGetStateFreqVec :: DNAModelSpec -> StationaryDist
dnaModelSpecGetStateFreqVec JC                = vector $ replicate nNucs 0.25
dnaModelSpecGetStateFreqVec (HKY _ f)         = f
dnaModelSpecGetStateFreqVec (GTR _ _ _ _ _ f) = f

-- | A rate matrix of a DNA models is called DNAModel.
data DNAModel = DNAModel { dnaRateMatrix :: RateMatrix
                         , dnaModelSpec  :: DNAModelSpec }

-- | The matrix of exchangeabilities; diagonal entries are not set.
exchangeabilityMatrix :: DNAModelSpec -> Matrix R
exchangeabilityMatrix JC = (4><4)
  [ 0.0, 1.0, 1.0, 1.0
  , 1.0, 0.0, 1.0, 1.0
  , 1.0, 1.0, 0.0, 1.0
  , 1.0, 1.0, 1.0, 0.0 ]
exchangeabilityMatrix (HKY k _) = (4><4)
  [ 0.0, 1.0,   k, 1.0
  , 1.0, 0.0, 1.0,   k
  ,   k, 1.0, 0.0, 1.0
  , 1.0,   k, 1.0, 0.0 ]
exchangeabilityMatrix (GTR a b c d e _) = (4><4)
  [ 0.0,   a,   b,   c
  ,   a, 0.0,   d,   e
  ,   b,   d, 0.0, 1.0
  ,   c,   e, 1.0, 0.0 ]
-- exchangeabilityMatrix _     = error "Model not yet supported."

-- | DNA model rate matrix normalized so that one mutation happens per unit time.
rateMatrix :: DNAModelSpec -> DNAModel
rateMatrix s = DNAModel rm s
  where
    f    = dnaModelSpecGetStateFreqVec s
    exch = exchangeabilityMatrix s
    rm   = normalizeRates f $ setDiagonal $ exch <> diag f
-- rateMatrix _ _       = error "Model not yet supported."

-- | Report a specific model.
getDNAModelInfoStr :: DNAModel -> String
getDNAModelInfoStr (DNAModel m s) =
  case s of
    JC        -> unlines
      [ "JC model with rate matrix"
      , show m ]
    (HKY k f) -> unlines
      [ "HKY model"
      , reportRateMatrix
      , reportExchangeabilityMatrix
      , reportStateFreqVec f
      , "And a kappa value of: " ++ show k ]
    (GTR a b c d e f) -> unlines
      [ "GTR model"
      , reportRateMatrix
      , reportExchangeabilityMatrix
      , reportStateFreqVec f
      , "And parameters: " ++ intercalate ", " (map show [a,b,c,d,e]) ]
      -- _ -> error "Model not yet supported."
  where reportRateMatrix   = "Rate matrix " ++ show m
        reportStateFreqVec f = "This corresponds to state frequencies (A, C, G, T): " ++ show f
        reportExchangeabilityMatrix = "Exchangeability matrix (diagonal set to 0.0) " ++
                                      show (exchangeabilityMatrix s)

-- | Scale a DNA model.
dnaModelScale :: Double -> DNAModel -> DNAModel
dnaModelScale s (DNAModel rm spec) = DNAModel (scale s rm) scaledSpec
  where scaledSpec = case spec of
                       JC              -> JC
                       HKY k f         -> HKY (s*k) f
                       GTR a b c d e f -> GTR (s*a) (s*b) (s*c) (s*d) (s*e) f


module Utils where

import Data.Random.Normal
import Numeric.LinearAlgebra
import System.Random
--import Test.QuickCheck

genNormal :: (Random a, Floating a) => [a]
genNormal = let g = mkStdGen 1001 in
  normals g

testGenNormal :: (Random a, Floating a) => [a]
testGenNormal = let g = mkStdGen 1001 in
  take 10 (normals g)

getAverageMat :: Matrix R -> Double
getAverageMat mat =
  let r_f = fromIntegral $ rows mat
      c_f = fromIntegral $ cols mat
      sumMat = sumElements mat in
  sumMat / (r_f * c_f)


{-# OPTIONS_GHC -fwarn-unused-imports #-}

module Main(main) where

import Control.Monad.Par
import Control.Monad.Par.Internal (runParAsync)
import Control.Monad.Par.Stream as S
import Control.Monad.Par.Logging
import Control.Exception
import Data.Complex
import GHC.Conc as Conc
import Debug.Trace
import Math.FFT (dft)

type Elt = Complex Double

fft_kern :: Window Elt -> Window Elt
fft_kern arr = dft arr

--   -- TEMP, evaluate one element to make sure the fft really gets called:
-- --  trace ("One elt sample: "++ show (arr!10)) $
--   case arr2 ! 10 of _ -> arr2
--  where arr2 = dft arr

--------------------------------------------------------------------------------
-- Main script

-- target  = maxBound
target  = 10 * 1000 * 1000
bufsize = 1024

main = do
  putStrLn$ "numCapabilities: "++ show numCapabilities
  putStrLn$ "  Frequency in measurable ticks:  "++ commaint oneSecond ++ "\n"

  putStrLn$ "Performing FFT of "++ commaint target ++" numbers windowed into buffers of size "++ show bufsize ++"\n"

  results <- evaluate $ runParAsync$ do

       strm1 <- countupWin bufsize target :: Par (WStream Elt)

       print_$ "\n Next, applying FFT filter... "
       strm2 <- streamMap fft_kern strm1

-- Make a pipeline of 10 stages:
--       strm2 <- foldl (\ s _ -> streamMap kern0) strm1 [1..10]

       print_$ "\n Stream graph constructed, returning from Par computation... "
       return strm2

  measureRate results
  putStrLn$ "End of stream reached.  All done."

print_ msg = trace msg $ return ()

-- work pop 1 peek N push 1
-- float->float filter
-- firFilter n coefs =
-- {

--     float sum = 0;
--     for (int i = 0; i < N; i++)
--       sum += peek(i) * COEFF[N-1-i];
--     pop();
--     push(sum);
--   }
-- }

{-
 Notes:

 -}


module Optimization where

import Numeric.LinearAlgebra

class Optimize a where
  --optimizeInit :: a -> Double -> Double -> (Int, Int)  -> a
  paramUpdate :: a -> a

data SGD = SGD {
  sgdLearningRate :: Double,
  sgdDecay :: Double,
  sgdMomentum :: Double,
  sgdParams :: [Matrix R],
  sgdGparams :: [Matrix R],
  sgdDeltaPre :: [Matrix R],
  sgdBatchSize :: Double
                 }

instance Optimize SGD where
  paramUpdate sgd =
    let learningRate = sgdLearningRate sgd in
    let decay = sgdDecay sgd in
    let momentum = sgdMomentum sgd in
    let params = sgdParams sgd in
    let gparams = sgdGparams sgd in
    let deltaPre = sgdDeltaPre sgd in
    let batchSize = sgdBatchSize sgd in
    let _sgd (px:pxs) (gx: gxs) (dx:dxs) paUpdate deUpdate gaZeros  =
          let r = rows gx in
          let c = cols gx in
          let lrMat = (r >< c) [learningRate..]:: (Matrix R) in
          let batchSizeMat = (r >< c) [batchSize..] :: (Matrix R) in
          let momMat = (r >< c) [momentum..] :: (Matrix R) in
          let gUp = gx * lrMat / batchSizeMat in
          let ugd = momMat * dx - gx in
          let pu = px + ugd in _sgd pxs gxs dxs (paUpdate++[pu]) (deUpdate++[ugd]) (gaZeros++[(r >< c) [0..]])
        _sgd [] [] [] paUpdate deUpdate gaZeros = SGD {
          sgdLearningRate=learningRate,
          sgdDecay=decay,
          sgdMomentum=momentum,
          sgdParams=paUpdate,
          sgdGparams=gaZeros,
          sgdDeltaPre=deUpdate,
          sgdBatchSize=0
          } in
    _sgd params gparams deltaPre [] [] []

data Adadelta = Adadelta {
  adaDecay :: Double,
  adaEpsilon :: Double,
  adaBatchSize :: Double,
  adaParams :: [Matrix R],
  adaGparams :: [Matrix R],
  adaAccGrad :: [Matrix R],
  adaAccDelta :: [Matrix R]
                         }

instance Optimize Adadelta where
  paramUpdate ada =
    let decay = adaDecay ada in
    let epsilon = adaEpsilon ada in
    let batchSize = adaBatchSize ada in
    let params = adaParams ada in
    let gparams = adaGparams ada in
    let accGrad = adaAccGrad ada in
    let accDelta = adaAccDelta ada in
    let _ada (px:pxs) (gx:gxs) (accgx:accgxs) (accdx:accdxs) paUpdate gradUpdate deltaUpdate gaZeros =
          let r = rows px in
          let c = cols px in
          let decayMat = (r >< c) [decay ..] :: (Matrix R) in
          let epsilonMat = (r >< c) [epsilon ..] :: (Matrix R) in
          let batchSizeMat = (r >< c) [batchSize ..] :: (Matrix R) in
          let oneMat = (r >< c) [1.0 ..]:: (Matrix R) in
          let gUp = accgx / batchSizeMat in
          let gradUp = decayMat * accgx + (oneMat - decayMat) * gUp * gUp in
          let ugd = (cmap (\x -> -(sqrt x)) (accdx + epsilonMat)) /  (cmap (\x -> sqrt x) (accgx + epsilonMat)) * gUp in
          let deltaUp = decayMat * accdx + (oneMat - decayMat) * ugd * ugd in
          let pu = px + ugd in
          _ada pxs gxs accgxs accdxs (pu:paUpdate) (gradUp:gradUpdate) (deltaUp:deltaUpdate) (((r >< c)[0.0 ..]):gaZeros)
        _ada [] [] [] [] paUpdate gradUpdate deltaUpdate gaZeros = Adadelta {
          adaDecay=decay,
          adaEpsilon=epsilon,
          adaBatchSize=batchSize,
          adaParams=(reverse paUpdate),
          adaGparams=(reverse gaZeros),
          adaAccGrad=(reverse gradUpdate),
          adaAccDelta=(reverse deltaUpdate)
                                                                        }
    in _ada params gparams accGrad accDelta [] [] [] []


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Normal
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The normal distribution.  This is a continuous probability
-- distribution that describes data that cluster around a mean.

module Statistics.Distribution.Normal
    (
      NormalDistribution
    -- * Constructors
    -- , normalDistr
    --, normalDistrE
    , standard
    ) where

import Data.Data             (Data, Typeable)
import GHC.Generics          (Generic)
import Numeric.MathFunctions.Constants (m_sqrt_2, m_sqrt_2_pi)
import Numeric.SpecFunctions (erfc, invErfc)

import qualified Statistics.Distribution as D
import Statistics.Internal

-- | The normal distribution.
data NormalDistribution = ND {
      mean       :: {-# UNPACK #-} !Double
    , stdDev     :: {-# UNPACK #-} !Double
    , ndPdfDenom :: {-# UNPACK #-} !Double
    , ndCdfDenom :: {-# UNPACK #-} !Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show NormalDistribution where
  showsPrec i (ND m s _ _) = defaultShow2 "normalDistr" m s i
instance Read NormalDistribution where
  readPrec = defaultReadPrecM2 "normalDistr" normalDistrE

instance D.Distribution NormalDistribution where
    cumulative      = cumulative
    complCumulative = complCumulative

instance D.ContDistr NormalDistribution where
    logDensity    = logDensity
    quantile      = quantile
    complQuantile = complQuantile

-- | Standard normal distribution with mean equal to 0 and variance equal to 1
standard :: NormalDistribution
standard = ND { mean       = 0.0
              , stdDev     = 1.0
              , ndPdfDenom = log m_sqrt_2_pi
              , ndCdfDenom = m_sqrt_2
              }

-- | Create normal distribution from parameters.
--
-- IMPORTANT: prior to 0.10 release second parameter was variance not
-- standard deviation.
normalDistrE :: Double            -- ^ Mean of distribution
             -> Double            -- ^ Standard deviation of distribution
             -> Maybe NormalDistribution
normalDistrE m sd
  | sd > 0    = Just ND { mean       = m
                        , stdDev     = sd
                        , ndPdfDenom = log $ m_sqrt_2_pi * sd
                        , ndCdfDenom = m_sqrt_2 * sd
                        }
  | otherwise = Nothing

logDensity :: NormalDistribution -> Double -> Double
logDensity d x = (-xm * xm / (2 * sd * sd)) - ndPdfDenom d
    where xm = x - mean d
          sd = stdDev d

cumulative :: NormalDistribution -> Double -> Double
cumulative d x = erfc ((mean d - x) / ndCdfDenom d) / 2

complCumulative :: NormalDistribution -> Double -> Double
complCumulative d x = erfc ((x - mean d) / ndCdfDenom d) / 2

quantile :: NormalDistribution -> Double -> Double
quantile d p
  | p == 0         = -inf
  | p == 1         = inf
  | p == 0.5       = mean d
  | p > 0 && p < 1 = x * ndCdfDenom d + mean d
  | otherwise      =
    error $ "Statistics.Distribution.Normal.quantile: p must be in [0,1] range. Got: "++show p
  where x          = - invErfc (2 * p)
        inf        = 1/0

complQuantile :: NormalDistribution -> Double -> Double
complQuantile d p
  | p == 0         = inf
  | p == 1         = -inf
  | p == 0.5       = mean d
  | p > 0 && p < 1 = x * ndCdfDenom d + mean d
  | otherwise      =
    error $ "Statistics.Distribution.Normal.complQuantile: p must be in [0,1] range. Got: "++show p
  where x          = invErfc (2 * p)
        inf        = 1/0


import HVX
import Numeric.LinearAlgebra

-- declare a symbolic variable
x = EVar "x"
-- give x a value (a 4-element column vector)
vars = [("x", (4><1) [0,-3,1,2])]
-- define the expression to sudifferentiate: max(abs(x))
myexpr = hmax(habs(x))
-- compute the subgradient
mysubgrad = jacobianWrtVar myexpr vars "x"


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}

module LinAlg where

import qualified Numeric.LinearAlgebra as L

import Data.Complex
import Data.List

class Vector v where
    -- vector addition
    (<+>) :: v -> v -> v
    -- scalar multiplication
    (<**>) :: Complex Double -> v -> v
    -- additive inverse
    vneg :: v -> v
    (<~>) :: v -> v -> v
    -- vector subtraction for convenience
    u <~> v = u <+> vneg v

class Vector v => Hilbert v where
    -- Inner product
    (<.>) :: v -> v -> Complex Double
    -- Vector norm, default definition generates norm
    -- from inner product
    norm :: v -> Double
    norm v = sqrt . magnitude $ v <.> v

-- Data structure representing a generic linear operator
-- in some representation
newtype Operator r = Operator {
    repr :: r
}

instance Functor Operator where
    fmap f (Operator r) = Operator $ f r

instance Applicative Operator where
    pure = Operator
    (Operator f) <*> op = f <$> op

-- Type alias for matrix representation of operator
type MatOp = Operator (L.Matrix (Complex Double))

-- Allow operators to be composed
instance Semigroup r => Semigroup (Operator r) where
    (Operator r1) <> (Operator r2) = Operator (r1 <> r2)

-- Operators are vectors
instance Vector MatOp where
    o1 <+> o2 = Operator (repr o1 + repr o2)
    a <**> o = L.cmap (*a) <$> o
    vneg o = negate <$> o

-- Trace of an operator
trace :: MatOp -> Complex Double
trace = L.sumElements . L.takeDiag . repr


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}

{-# LANGUAGE CPP                 #-}
{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeOperators       #-}

module Test.Grenade.Layers.Add where

import           Data.Constraint
import           Data.Maybe                   (fromJust)
import           Data.Proxy
import           Unsafe.Coerce

import           GHC.TypeLits

import           Grenade.Core
import           Grenade.Layers.Add

import           Numeric.LinearAlgebra        hiding (R, konst, randomVector,
                                               uniformSample, (===), reshape)
import qualified Numeric.LinearAlgebra.Data   as D
import           Numeric.LinearAlgebra.Static (L, R)
import qualified Numeric.LinearAlgebra.Static as H

import           Hedgehog

import           Test.Hedgehog.Compat
import           Test.Hedgehog.Hmatrix

prop_add_zero_does_nothing = property $ do
  height   :: Int <- forAll $ choose 2 100
  width    :: Int <- forAll $ choose 2 100
  channels :: Int <- forAll $ choose 2 100

  case (someNatVal (fromIntegral height), someNatVal (fromIntegral width), someNatVal (fromIntegral channels)) of
    (Just (SomeNat (Proxy :: Proxy h)), Just (SomeNat (Proxy :: Proxy w)), Just (SomeNat (Proxy :: Proxy c))) -> do
      input :: S ('D3 h w c) <- forAll genOfShape

      let layer    = initAdd :: Add c 1 1
          S3D out  = snd $ runForwards layer input :: S ('D3 h w c)
          S3D inp' = input :: S ('D3 h w c)

      H.extract inp' === H.extract out

prop_add_per_channel_words_as_expected = property $ do
  height   :: Int <- forAll $ choose 2 100
  width    :: Int <- forAll $ choose 2 100

  case (someNatVal (fromIntegral height), someNatVal (fromIntegral width)) of
    (Just (SomeNat (Proxy :: Proxy h)), Just (SomeNat (Proxy :: Proxy w))) ->
      case ( unsafeCoerce (Dict :: Dict ()) :: Dict ( (h * 3) ~ (h + h + h) ) ) of
        Dict -> do
          channel_one   :: R (h * w) <- forAll randomVector
          channel_two   :: R (h * w) <- forAll randomVector
          channel_three :: R (h * w) <- forAll randomVector
          bias          :: R 3       <- forAll randomVector

          let reshape = fromJust . H.create . (D.reshape width) . H.extract
              c1      = reshape channel_one   :: L h w
              c2      = reshape channel_two   :: L h w
              c3      = reshape channel_three :: L h w

              bias' = H.extract bias

              o1 = H.dmmap (\a -> a + bias' ! 0) c1 :: L h w
              o2 = H.dmmap (\a -> a + bias' ! 1) c2 :: L h w
              o3 = H.dmmap (\a -> a + bias' ! 2) c3 :: L h w

              out = o1 H.=== o2 H.=== o3 :: L (h * 3) w
              inp = c1 H.=== c2 H.=== c3 :: L (h * 3) w

              layer    = Add bias :: Add 3 1 1
              inp'     = S3D inp :: S ('D3 h w 3)
              S3D out' = snd $ runForwards layer inp' :: S ('D3 h w 3)

          H.extract out === H.extract out'

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


{-# LANGUAGE OverloadedStrings, NoImplicitPrelude, BangPatterns #-}
import ClassyPrelude hiding (head, tail)
import Prelude (head, tail)
import Test.Hspec
import Test.QuickCheck
import qualified Criterion.Main as Criterion
import Lens.Micro

import Control.Exception (evaluate)
import qualified Data.Vector.Storable as SVec
import qualified Data.Vector.Generic as Vec
import qualified Data.Binary as Bin

import NLP.Albemarle
import qualified NLP.Albemarle.Tokens as Tokens
import NLP.Albemarle.Dict (Dict, counts, ids)
import qualified NLP.Albemarle.Dict as Dict
import NLP.Albemarle.LSA (termvectors, topicweights)
import qualified NLP.Albemarle.LSA as LSA
import qualified NLP.Albemarle.Bench.Maps
import qualified NLP.Albemarle.Test.Dict
import qualified NLP.Albemarle.Test.GloVe
import qualified NLP.Albemarle.Test.LanguageModelEntropy
import qualified NLP.Albemarle.Examples.Webpage

import System.CPUTime
import Text.Printf
import Data.Text.Encoding.Error (lenientDecode)
import System.Directory (doesFileExist)

import qualified Data.Text as Text
import qualified Data.HashSet as HashSet
import qualified Data.HashMap.Strict as HashMap
import System.IO.Streams (InputStream, OutputStream)
import qualified System.IO.Streams as Streams
import Numeric.LinearAlgebra (Matrix)
import qualified Numeric.LinearAlgebra as HMatrix

main :: IO ()
main = hspec $ do
  NLP.Albemarle.Examples.Webpage.test
  NLP.Albemarle.Test.Dict.test
  NLP.Albemarle.Test.GloVe.test
  NLP.Albemarle.Test.LanguageModelEntropy.test

  let sentences = words <$> [
        "Maybe not today. Maybe not tomorrow. But soon.",
        "Pay no attention to the man behind the curtain.",
        "Life is like a box of chocolates."]:: [[Text]]

  describe "Streaming (semi-Gensim) Style Topic Analysis" $ do
    let mean :: SVec.Vector Double -> Double
        mean x = Vec.sum x / fromIntegral (Vec.length x)
        -- Root mean squared error
        rmse l r = mean (Vec.zipWith (\x y -> (x-y)**2) (flat l) (flat r)) ** 0.5
        -- Mean signal to noise ratio
        -- - we tank the PSNR because most things are really close to 0
        -- - in fact we avoid handling 0's altogether
        msnr l r = mean $ Vec.zipWith (\x y -> abs $ x-y / max 1 (min x y)) (flat l) (flat r)
        flat = HMatrix.flatten
        mult u s vt = mconcat [u, HMatrix.diag s, HMatrix.tr vt] :: HMatrix.Matrix Double

    it "performs sparse stochastic truncated SVD with SVDLIBC" $ do
      let termdoc = "termdoc.small.nonzero.txt"
      [zippedfile, originalfile] <- sequence $ doesFileExist <$> [termdoc++".gz", termdoc]
      when ( zippedfile && not originalfile ) $
        Streams.withFileAsInput "termdoc.small.nonzero.txt.gz" $ \infl ->
          Streams.withFileAsOutput "termdoc.small.nonzero.txt" $ \outfl ->
            Streams.gunzip infl >>= Streams.connectTo outfl
      matrix <- HMatrix.loadMatrix "termdoc.small.nonzero.txt"
      let (u, sigma, vt) = LSA.batchLSA 100 $ LSA.sparsify matrix
      -- typically -> RMSE: 0.06464712163523892
      -- printf "RMSE: %f\n" $ rmse (mconcat [u, HMatrix.diag sigma, HMatrix.tr vt]) matrix
      rmse (mult u sigma vt) matrix < 0.12 `shouldBe` True
      -- typically -> MSNR: 0.010548326765305554
      -- printf "MSNR: %f\n" $ msnr (mconcat [u, HMatrix.diag sigma, HMatrix.tr vt]) matrix
      msnr (mult u sigma vt) matrix < 0.02 `shouldBe` True

    it "performs topic analysis starting from text" $ do
      (final_dict, model) <- Streams.withFileAsInput "kjv.verses.txt.gz" $ \file ->
        Streams.gunzip file
        >>= Streams.lines
        >>= Streams.decodeUtf8With lenientDecode
        >>= Streams.map Tokens.wordTokenize
        >>= Streams.chunkList 373 -- Yep, I'm making this stuff up.
        >>= Streams.map (\chunk -> let
          dict = Dict.dictifyAllWords chunk
          sparsem = Dict.asSparseMatrix dict chunk
          lsa = LSA.lsa 222 sparsem -- Too many topics. Should not cause errors
          in (dict, lsa))
        >>= Streams.fold (\(!d1, !lsa1) (d2, lsa2) -> let
          d3 = Dict.filterDict 2 0.5 100 $ d1 <> d2 -- Not nearly enough words
          lsa3 = LSA.rebase d1 d3 lsa1 <> LSA.rebase d2 d3 lsa2
          in (d3, lsa3)) (mempty, mempty)

      -- It should use all 100 words allowed plus the unknown
      HashMap.size (final_dict^.counts) `shouldBe` 101
      -- It should have a full size LSA as well
      HMatrix.size (model^.termvectors) `shouldBe` (222, 101)
      -- Plus topic weights
      HMatrix.size (model^.topicweights) `shouldBe` 222

  describe "Monoid style Topic Analysis" $
    it "creates LSA Models" $ do
      let dict = Dict.dictifyAllWords sentences
      let lsavecs = LSA.lsa 2
            $ Dict.asSparseMatrix dict sentences
      -- Sadly, I can't think of much to assert in this test so let's use size
      -- There are 21 unique words plus one unknown
      HMatrix.size (lsavecs^.termvectors) `shouldBe` (2,22)
      -- It should use both topics
      HMatrix.size (lsavecs^.topicweights) `shouldBe` 2

-- Thanks to Ying He for the following example text and proper
-- tokenizations.
--weightGain :: Text
--weightGain = unwords [
--  "Independent of current body composition, IGF-I levels at 5 yr were ",
--  "significantly associated with rate of weight gain between 0-2 yr",
--  "(beta = 0.19; P < 0.0005), and children who showed postnatal",
--  "catch-up growth (i.e. those who showed gains in weight or length",
--  "between 0-2 yr by >0.67 SD score) had higher IGF-I levels than other",
--  "children (P = 0.02)."]

--weightGainTokens :: [Text]
--weightGainTokens = words $ unwords [
--  "Independent of current body composition , IGF - I levels at 5 yr were",
--  "significantly associated with rate of weight gain between 0 - 2 yr ( beta",
--  "= 0.19 ; P < 0.0005 ) , and children who showed postnatal catch - up",
--  "growth ( i.e . those who showed gains in weight or length between 0 - 2 yr",
--  "by > 0.67 SD score ) had higher IGF - I levels than other children",
--  "( P = 0.02 ) ."] -- i.e. is wrong, but we're calling this close enough.


module Gradient
    (numericalDiff
    ) where

import Numeric.LinearAlgebra

numericalDiff :: Fractional a => (a -> a) -> a -> a
numericalDiff f x = (f (x+h) - f (x-h)) / (2*h)
    where h = 1e-4


{-# LANGUAGE TypeOperators #-}
module Vect
    ( module Vect
    ) where

import Control.Arrow (first, second)
import Control.Monad (ap)
import Data.Complex
import Data.List (intercalate, sortOn)
import GHC.TypeLits

import QBasis

type C = Complex Double

-- |Free vector space over field
newtype Vect k b = V [(b, k)] deriving (Eq, Ord)

instance (Show k, Show b) => Show (Vect k b) where
    show (V xs) = intercalate " + " $ map (\(e, c) -> show c ++ show e) xs

instance Functor (Vect k) where
    fmap f (V xs) = V $ map (first f) xs

instance Num k => Applicative (Vect k) where
    pure = return
    (<*>) = ap

instance Num k => Monad (Vect k) where
    return a = V [(a, 1)]
    V ts >>= f = V $ concat [ [(b, y*x) | let V us = f a, (b, y) <- us] | (a, x) <- ts]

-- |Convert an element of Vect k b into normal form. Normal form consists in having the basis elements in ascending order,
-- with no duplicates, and all coefficients non-zero
nf :: (Ord b, Eq k, Num k) => Vect k b -> Vect k b
nf (V ts) = V $ nf' $ sortOn fst ts where
    nf' ((b1,x1):(b2,x2):ts) =
        case compare b1 b2 of
        LT -> if x1 == 0 then nf' ((b2,x2):ts) else (b1,x1) : nf' ((b2,x2):ts)
        EQ -> if x1+x2 == 0 then nf' ts else nf' ((b1,x1+x2):ts)
        GT -> error "nf': not pre-sorted"
    nf' [(b,x)] = if x == 0 then [] else [(b,x)]
    nf' [] = []

-- |The zero vector
zero :: Vect k b
zero = V []

-- |Vector addition
add :: (Ord b, Num k, Eq k) => Vect k b -> Vect k b -> Vect k b
add (V ts) (V us) = V $ addmerge ts us
    where addmerge axs@((a,x):xs) bxs@((b,y):ys) =
            case compare a b of
                LT -> (a,x) : addmerge xs bxs
                EQ -> if x+y == 0
                        then addmerge xs ys
                        else (a, x+y) : addmerge xs ys
                GT -> (b,y) : addmerge axs ys
          addmerge xs [] = xs
          addmerge [] ys = ys

-- |(Left) Scalar multiplication
smult :: (Eq k, Num k) => k -> Vect k b -> Vect k b
smult 0 _ = zero
smult k (V xs) = V $ map (second (k*)) xs

-- |Linear application to basis elements
linear :: (Ord b, Eq k, Num k) => (a -> Vect k b) -> Vect k a -> Vect k b
linear f v = nf $ v >>= f

-- |Tensor product of two free vector spaces
te :: (Num k, Eq k) => Vect k (QBasis m) -> Vect k (QBasis n) -> Vect k (QBasis (m+n))
te (V u) (V v) = V [(a `tb` b, x*y) | (a,x) <- u, (b,y) <- v]

-- |Tensor product of two functions between free vector spaces
tf :: (Num k, Eq k, KnownNat m) =>
            (Vect k (QBasis m) -> Vect k (QBasis m)) ->
            (Vect k (QBasis n) -> Vect k (QBasis n)) ->
            Vect k (QBasis (m+n))-> Vect k (QBasis (m+n))
tf f g (V ts) = sum [let (a,b) = bt basis in smult x (f (return a) `te` g (return b)) | (basis,x) <- ts]
        where sum = foldl add zero


import Data.Complex

main = do
  let a = 1.0 :+ 2.0    -- complex number 1+2i
  let b = 4             -- complex number 4+0i
  -- 'b' is inferred to be complex because it's used in
  -- arithmetic with 'a' below.
  putStrLn $ "Add:      " ++ show (a + b)
  putStrLn $ "Subtract: " ++ show (a - b)
  putStrLn $ "Multiply: " ++ show (a * b)
  putStrLn $ "Divide:   " ++ show (a / b)
  putStrLn $ "Negate:   " ++ show (-a)
  putStrLn $ "Inverse:  " ++ show (recip a)
  putStrLn $ "Conjugate:" ++ show (conjugate a)


{-# OPTIONS_GHC -Wall #-}

-- | Ring classes. A distinguishment is made between Rings and Commutative Rings.
module Plankton.Ring
  ( Semiring
  , Ring
  , CRing
  ) where

import Data.Complex (Complex(..))
import Plankton.Additive
import Plankton.Distribution
import Plankton.Multiplicative
import Protolude (Bool(..), Double, Float, Int, Integer)

-- | Semiring
class (MultiplicativeAssociative a, MultiplicativeUnital a, Distribution a) =>
      Semiring a

instance Semiring Double

instance Semiring Float

instance Semiring Int

instance Semiring Integer

instance Semiring Bool

instance (AdditiveGroup a, Semiring a) => Semiring (Complex a)

-- | Ring
-- a summary of the laws inherited from the ring super-classes
--
-- > zero + a == a
-- > a + zero == a
-- > (a + b) + c == a + (b + c)
-- > a + b == b + a
-- > a - a = zero
-- > negate a = zero - a
-- > negate a + a = zero
-- > a + negate a = zero
-- > one `times` a == a
-- > a `times` one == a
-- > (a `times` b) `times` c == a `times` (b `times` c)
-- > a `times` (b + c) == a `times` b + a `times` c
-- > (a + b) `times` c == a `times` c + b `times` c
-- > a `times` zero == zero
-- > zero `times` a == zero
class ( AdditiveGroup a
      , MultiplicativeAssociative a
      , MultiplicativeUnital a
      , Distribution a
      ) =>
      Ring a

instance Ring Double

instance Ring Float

instance Ring Int

instance Ring Integer

instance (Ring a) => Ring (Complex a)

-- | CRing is a Ring with Multiplicative Commutation.  It arises often due to '*' being defined as a multiplicative commutative operation.
class (Multiplicative a, Ring a) =>
      CRing a

instance CRing Double

instance CRing Float

instance CRing Int

instance CRing Integer

instance (CRing a) => CRing (Complex a)


-- Modulo destinado a todas as funcoes relativas a execucao da Rede.

module Execution
(execute,
 feedforward) where

import InputOutput
import Types
import Numeric.LinearAlgebra.HMatrix

type Image = [Double]
type Sample = (Int, Image)

-- execute, por retornar IO, pode interagir chamando outras IO Actions como
-- leitura dos weights, biases e input da Rede, por exemplo.
execute :: IO String -- mero esqueleto da funcao de execucao
execute = do
        network <- readIn
        let image = [1.0 .. 784.0]
            executedNetwork = feedforward image network
        do definitiveAnswer executedNetwork

-- recebe a imagem, a network e computa os calculos,
-- retornando a nova data com os valores de ativacao
-- e zeta do hidden e output alterados.
feedforward :: Image -> Data -> Data
feedforward image network = let zH = add ((wHidden network) #> (fromList image)) (bHidden network)
                                aH = fromList $ sigV (toList zH)
                                zO = add ((wOutput network) #> aH) (bOutput network)
                                aO = fromList $ sigV (toList zO)
                            in Data (wHidden network) (bHidden network) aH zH (wOutput network) (bOutput network) aO zO


module Ch01.Broadcast where

import Numeric.LinearAlgebra

as = matrix 2 [1,2,3,4]
bs = matrix 2 [10,20]
r = as * bs

xs :: Matrix Int
xs = (3><2) [51,55,14,19,0,4]


{-# OPTIONS_GHC -Wno-incomplete-patterns #-}

module VectorSpec where

import Data.Complex

import Test.QuickCheck
import Test.QuickCheck.Poly

import Category
import Comonad
import Functor
import Unboxed
import Vector

prop_Vector_Functor_id :: FnProp (Vector A)
prop_Vector_Functor_id = checkFnEqual law_Functor_id

prop_Vector_Functor_comp :: Fun B C -> Fun A B -> FnProp (Vector A)
prop_Vector_Functor_comp f g = checkFnEqual (law_Functor_comp f g)

prop_Vector_Semicomonad_comm ::
    Fun (Vector B) C -> Fun (Vector A) B -> FnProp (Vector A)
prop_Vector_Semicomonad_comm f g = checkFnEqual (law_Semicomonad_comm f g)

prop_Vector_Semicomonad1_comm ::
    Fun (Vector B) C -> Fun (Vector A) B -> FnProp (Vector A)
prop_Vector_Semicomonad1_comm f g = checkFnEqual (law_Semicomonad1_comm f g)

prop_Vector_Semicomonad1_comm' ::
    Fun (Vector B) C -> Fun (Vector A) B -> FnProp (Vector A)
prop_Vector_Semicomonad1_comm' (Fn f) (Fn g) =
    checkFnEqual (law_Semicomonad1_comm' f g)

type UA = Int
type UB = Double
type UC = Complex Double

prop_UVector_Functor_id :: FnProp (UVector UA)
prop_UVector_Functor_id = checkFnEqual law_Functor_id

prop_UVector_Functor_comp ::
    (UB -#> UC) -> (UA -#> UB) -> FnProp (UVector UA)
prop_UVector_Functor_comp f g = checkFnEqual (law_Functor_comp f g)

prop_UVector_Semicomonad1_comm' ::
    Fun (UVector UB) UC -> Fun (UVector UA) UB -> FnProp (UVector UA)
prop_UVector_Semicomonad1_comm' (Fn f) (Fn g) =
    checkFnEqual (law_Semicomonad1_comm' f g)


{-# LANGUAGE RecordWildCards #-}
{-
The shape of a triangle can be represented by two explicit points (ie vectors)
and a third implicit point at (0,0). A triangle defined by three points can be
converted by subtracting one of the points from the others:
  shape (p1,p2,p3) = (p2-p1, p3-p1)
This keeps the shape intact but drops the position of the triangle.

If P and Q are two such triangles, there exists a matrix A that transforms
P into Q: AP = Q
The 'A' matrix can be calculated by multiplying Q and the inverse of P:
  A = QP_inv

So, if we have a set of triangles, we can drop their positional information
and compute their shape transformation matrices. But can we do it in reverse?
Can we start with a set of transformation matrices (which don't give us the
positions of the triangles) and compute a set of points that represent triangles
with the right transformation?
In other words, if we have a set of triangles (with absolute position) and
a transformation matrix for each triangle, can we compute new absolute positions
such that AP=Q?

Solving this is not straightforward as moving a single point will change the
transformation of multiple triangles. Fortunately, though, this problem can be
expression as linear equations and solved quickly with LAPACK.

To express the problem as a linear system, we need 4 equations for each
triangle and 2 unknowns for each point.

The system has this shape: Mx=B
We're trying to find 'x' which are the positions of the new points.
'M' is a matrix that transforms points into transformation matrices.
'B' is the target transformation matrices.

With a bit of rewriting, we can reduce 'A = QP_inv' to 4 linear equations
with 6 unknowns. The 6 unknowns are the x and y positions for the three corners
in a triangle.
This is just another way of calculating the transformation matrix 'A' and we
can verify it by looking at the output of:
  * computeA src dst
  * applyCoeff (coeffOfB src) dst
The output should be the same because they're both computing 'A'.

'applyCoeff' simply multiplies the coefficients by a 6x1 matrix containing
the x and y positions for 'dst'.

So, this means we have roughly the right shape: Mx=B. M is the coefficients
we found for 'src', 'x' is dst, and 'B' is the transformation matrix that
turns 'src' into 'dst'.
We can get the coefficients for each triangle and place them together in a
single large matrix. Then we can put the desired transformation matrices in 'B'
and ask a linear solver to find 'x' which satisfies the equations.

Phew, step one is done. We now have the machinery for turning arbitrary
transformation matrices into proper, connected triangles. Next step is figuring
out which rotational matrices to use. At the two extremes, we have the identity
matrices (representing our starting shape) and the matrices computed from the
target triangles (represetning out target shape). How can these two matrices be
interpolated for a smooth morph?

There are several ways of interpolating matrices. Linear interpolation is
possible. But, it looks prettier if we separate rotation from shear. Then we
can use spherical linear interpolation on the rotation and linear interpolation
on the shear.

-}
module Reanimate.Morph.Rigid where

import           Data.Foldable (toList)
import           Data.Vector                   (Vector)
import qualified Data.Vector                   as V
import           Linear.Quaternion
import           Linear.Vector
import           Linear.V2
import           Linear.V3
import qualified Numeric.LinearAlgebra         as Matrix
import           Numeric.LinearAlgebra.HMatrix (GMatrix, Matrix,
                                                toLists, (!), (><))

type P = V2 Double
type Trig = (P,P,P)
type RelTrig = (Int,Int,Int)

data Mesh = Mesh
  { meshPointsA :: Vector P
  , meshPointsB :: Vector P
  , meshOutline :: Vector Int
  -- , meshSteiner :: Vector Int
  , meshTriangles :: Vector RelTrig }

-- applyA (computeA a b) a = b + some_constant_translation
-- A = Q P_inv
computeA :: Trig -> Trig -> Matrix Double
computeA p q = matQ <> Matrix.inv matP
  where
    matP = trigToMatrix p
    matQ = trigToMatrix q

-- A = UDV = U(VV_t)DV = (UV)(V_tDV) = RS
-- R = UV
-- S = V_tDV
computeRS :: Matrix Double -> (Matrix Double, Matrix Double)
computeRS a = (r, s)
  where
    (u,d,vt) = Matrix.svd a
    v = Matrix.tr vt
    r = u <> v
    s = vt <> Matrix.diag d <> v

matrixToQuaternion :: Matrix Double -> Quaternion Double
matrixToQuaternion r =  q
  where
    w = 0.5 * sqrt (1 + r!0!0 + r!1!1 + 1)
    z = 1/(4*w) * (r!0!1 - r!1!0)
    q = Quaternion w (V3 0 0 z)

quaternionToMatrix :: Quaternion Double -> Matrix Double
quaternionToMatrix q = (2><2)
    [ 1 - 2*(qj*qj + qk*qk), 2*(qi*qj+qk*qr)
    , 2*(qi*qj - qk*qr), 1 - 2*(qi*qi + qk*qk) ]
  where
    [qr,qi,qj,qk] = toList q

-- A = R((1-t)I + tS)
computeA_RSt :: Matrix Double -> Matrix Double -> Double -> Matrix Double
computeA_RSt r s t = r_t <> (realToFrac (1-t) * Matrix.ident 2 + realToFrac t * s)
  where
    i = Quaternion 1 0
    q = slerp i (matrixToQuaternion r) t
    r_t = quaternionToMatrix q

applyA :: Matrix Double -> Trig -> Trig
applyA a p =
    case toLists (a <> matP) of
      [ [x1, x2], [y1, y2] ] -> (V2 0 0, V2 x1 y1, V2 x2 y2)
      _ -> error "invalid matrix"
  where
    matP = trigToMatrix p

coeffOfB :: Trig -> Matrix Double
coeffOfB p = (4><6)
    [ -a0-a2,      0, a0,   0, a2,  0
    , -a1-a3,      0, a1,   0, a3,  0
    ,      0, -a0-a2,  0,  a0,  0, a2
    ,      0, -a1-a3,  0,  a1,  0, a3 ]
  where
    [[a0,a1],[a2,a3]] = toLists (Matrix.inv (trigToMatrix p))

applyCoeff :: Matrix Double -> Trig -> Matrix Double
applyCoeff b (V2 x1 y1, V2 x2 y2, V2 x3 y3) = b <> matQ
  where
    matQ = (6><1) [ x1, y1, x2, y2, x3, y3]

trigToMatrix :: Trig -> Matrix Double
trigToMatrix (p1,p2,p3) = matP
  where
    V2 p12x p12y = p2-p1
    V2 p13x p13y = p3-p1
    matP = (2><2)
      [p12x, p13x
      ,p12y, p13y]

data Prep = Prep
  { prepPivot    :: (P, P)
  , prepPointsA  :: Vector P
  , prepPointsB  :: Vector P
  , prepRS       :: Vector (Matrix Double, Matrix Double)
  , prepRSRev    :: Vector (Matrix Double, Matrix Double)
  , prepUToB     :: GMatrix
  }

symmetric :: Bool
symmetric = True

prepare :: Mesh -> Prep
prepare Mesh{..} = Prep
    { prepPivot = (aOrigin, bOrigin)
    , prepPointsA =
        V.map (subtract aOrigin) $
        V.take pivotIdx meshPointsA <> V.drop (pivotIdx+1) meshPointsA
    , prepPointsB =
        V.map (subtract bOrigin) $
        V.take pivotIdx meshPointsB <> V.drop (pivotIdx+1) meshPointsB
    , prepRS = rsList
    , prepRSRev = rsRevList
    , prepUToB = Matrix.mkSparse uToB }
  where
    aOrigin = meshPointsA V.! pivotIdx
    bOrigin = meshPointsB V.! pivotIdx
    pivotIdx = case V.head meshTriangles of
        (a,_,_) -> a
    mkAbs p (a,b,c) = (p V.! a,p V.! b,p V.! c)
    absATrigs = V.map (mkAbs meshPointsA) meshTriangles
    absBTrigs = V.map (mkAbs meshPointsB) meshTriangles
    aList = V.zipWith computeA absATrigs absBTrigs
    rsList = V.map computeRS aList
    revList = V.zipWith computeA absBTrigs absATrigs
    rsRevList = V.map computeRS revList
    n = length meshTriangles
    -- nT = if symmetric then n*2 else n
    -- pivotToB = ((nT*4)><2) $ concat
    --   [ [ fromMaybe 0 (lookup (x,pivotIdx*2) bigM)
    --     , fromMaybe 0 (lookup (x,pivotIdx*2+1) bigM)]
    --   | x <- [0..(nT*4)-1] ]
    uToB =
      [ ((x,y-2),key) | ((x,y),key) <- bigM, y /= pivotIdx*2 && y /= pivotIdx*2+1 ]
    bigM =
      concat (zipWith worker [0..] (V.toList meshTriangles)) ++
      if symmetric
        then concat $ zipWith workerRev [n..] (V.toList meshTriangles)
        else []
    worker i src@(a,b,c) = concat $
      let effs = coeffOfB (mkAbs meshPointsA src) in
      [ [((i*4+h, e*2), effs!h!(j*2))
        ,((i*4+h, e*2+1), effs!h!(j*2+1))]
      | h <- [0..3]
      , (e,j) <- zip [a,b,c] [0..]
      ]
    workerRev i dst@(a,b,c) = concat $
      let effs = coeffOfB (mkAbs meshPointsB dst) in
      [ [((i*4+h, e*2), effs!h!(j*2))
        ,((i*4+h, e*2+1), effs!h!(j*2+1))]
      | h <- [0..3]
      , (e,j) <- zip [a,b,c] [0..]
      ]

interpolate :: Prep -> Double -> Vector P
interpolate Prep{..} t = V.fromList $
    pivot : worker (Matrix.toList solution)
  where
    -- solution = Matrix.cgSolve False prepUToB b
    solution = Matrix.cgx $ last solutions
    solutions = Matrix.cgSolve'
      False
      1e-9
      1e-9
      1000
      prepUToB
      b
      (Matrix.fromList $ concat [ [x,y] | V2 x y <- V.toList target ])
      -- 0
    target = if t < 0.5
      then prepPointsA
      else prepPointsB
    worker (x:y:xs) = V2 x y ^+^ pivot : worker xs
    worker _ = []
    pivot = case prepPivot of
      (src, dst) -> lerp t dst src
    n = V.length prepRS
    b = Matrix.vector $
      [ concat (toLists a)!!j
      | i <- [0..n-1]
      , let (r,s) = prepRS V.! i
      , let a = computeA_RSt r s t
      , j <- [0..3]
      ] ++
      [ concat (toLists a)!!j
      | symmetric
      , i <- [0..n-1]
      , let (r,s) = prepRSRev V.! i
      , let a = computeA_RSt r s (1-t)
      , j <- [0..3]
      ]


{-|
Module      : Numeric.MCMC.Diagnostics
Description : Diagnostics for Markov Chain Monte Carlo Sampling
Copyright   : (c) Adam Conner-Sax 2019
License     : BSD
Maintainer  : adam_conner_sax@yahoo.com
Stability   : experimental

Numeric.MCMC.Diagnostics contains functions to analyze the convergence of MCMC sampling and to estimate the distribution of expectations computed
using a set of MCMC chains.
-}
module Numeric.MCMC.Diagnostics
  ( summarize
  , mpsrf
  , ExpectationSummary(..)
  , mannWhitneyUTest
  )
where

import qualified Control.Foldl                 as FL
import qualified Statistics.Types              as S
import qualified Data.List                     as L
--import qualified Data.Vector.Generic           as V
import qualified Data.Vector.Unboxed           as VU
import qualified Statistics.Test.MannWhitneyU  as S
import qualified Numeric.LinearAlgebra         as LA

data ExpectationSummary a = ExpectationSummary { mean :: a, confidence :: (a, a), rHat :: a } deriving (Show)
-- all sub-chains should have same length here
-- do we need to throw away first half of all the chains??
{- |
Compute the exepctation, confidence interval and
PSRF statistic (from <http://www2.stat.duke.edu/~scs/Courses/Stat376/Papers/ConvergeDiagnostics/BrooksGelman.pdf>)
for a given confidence %, expectation function, and set of MCMC chains. In the version in the paper, each chain is
halved before the statistic is computed but it's not clear why.  Here we assume that is "burn-in" and the user may
choose to do that or not before calling this function.

All sub-chains should have the same length here, though that is not checked.
-}
summarize
  :: (RealFrac b, Ord b)
  => S.CL b -- ^ confidence interval to test
  -> (a -> b) -- ^ function to map a sample to a function under consideration
  -> [[a]] -- ^ list of chains
  -> Maybe (ExpectationSummary b) -- ^ @Maybe@ here in case chains aren't long enough
summarize confidence' expectationF chains = do
  let applied    = fmap (fmap expectationF) chains
      allApplied = concat applied
      chainInterval :: Ord b => [b] -> Maybe (b, b)
      chainInterval c =
        let n  = L.length c
            n' = round (realToFrac n * S.confidenceLevel confidence')
            d  = (n - n') `div` 2
            c' = L.take n' $ L.drop d $ L.sort c
        in if length c' >= 2 then (head &&& last) <$> nonEmpty c' else Nothing
--      intervalLength (a, b) = b - a
  intervals <- traverse chainInterval applied
  (lo, hi)  <- chainInterval allApplied
  let (mLo, mHi) = FL.fold
        ((,) <$> FL.premap fst FL.mean <*> FL.premap snd FL.mean)
        intervals
  let rHat' = (hi - lo) / (mHi - mLo)
      mean' = FL.fold FL.mean allApplied
  return $ ExpectationSummary mean' (lo, hi) rHat'

-- if the chain list is empty or chains have different sizes, etc. we return Nothing
mpsrf :: [a -> Double] -> [[a]] -> Maybe Double
mpsrf expectations chains = do
  let --mkVec :: a -> LA.Vector Double
      mkVec x = LA.fromList $ fmap ($ x) expectations
      m = length chains
  n <- length . head <$> nonEmpty chains --if m > 0 then Just (length $ head chains) else Nothing
  traverse_ (\c -> if length c == n then Just () else Nothing) chains -- check that all chains are length n
  let
    sumVecsF = FL.Fold (+) (LA.vector $ L.replicate (length expectations) 0) id
    sumMatsF = FL.Fold (+) (LA.scale 0 $ LA.ident (length expectations)) id
    meanVecsF =
      (\m' l -> LA.scale (1 / realToFrac l) m') <$> sumVecsF <*> FL.length
    avgVec     = FL.fold meanVecsF
    phiChains  = fmap (fmap mkVec) chains
    avgPhiEach = fmap avgVec phiChains
    diffProd :: LA.Vector Double -> LA.Vector Double -> LA.Matrix Double
    diffProd x y = (x - y) `LA.outer` (x - y)
    w1 = (\(ap, ps) -> (`diffProd` ap) <$> ps) <$> zip avgPhiEach phiChains
    w =
      LA.scale (1 / (realToFrac m * realToFrac (n - 1)))
        $ FL.fold sumMatsF
        $ fmap (FL.fold sumMatsF) w1
    avgPhiAll = avgVec avgPhiEach
    b1        = fmap (`diffProd` avgPhiAll) avgPhiEach
    b = LA.scale (realToFrac n / realToFrac (m - 1)) $ FL.fold sumMatsF b1
    mat       = LA.scale (1 / realToFrac n) $ LA.inv w LA.<> b
    eigen1    = LA.eigenvaluesSH (LA.trustSym mat) `LA.atIndex` 0
    rhat =
      (realToFrac (n - 1) / realToFrac n)
        + (realToFrac (m + 1) / realToFrac m)
        * eigen1
  return rhat

mannWhitneyUTest
  :: (Ord b, VU.Unbox b)
  => S.PValue Double
  -> (a -> b)
  -> [a]
  -> [a]
  -> Maybe S.TestResult
mannWhitneyUTest p f c1 c2 =
  let vOfF = VU.fromList . fmap f
  in  S.mannWhitneyUtest S.SamplesDiffer p (vOfF c1) (vOfF c2)


{-|
Module: MachineLearning.Classification.Binary
Description: Binary Classification.
Copyright: (c) Alexander Ignatyev, 2016-2017
License: BSD-3
Stability: experimental
Portability: POSIX

Binary Classification.
-}

module MachineLearning.Classification.Binary
(
    Opt.MinimizeMethod(..)
  , module Log
  , module Model
  , predict
  , learn
  , MLC.calcAccuracy
  , Regularization(..)
)

where

import MachineLearning.Types (R, Vector, Matrix)
import MachineLearning.Regularization (Regularization(..))
import qualified MachineLearning.Optimization as Opt
import qualified MachineLearning.LogisticModel as Log
import qualified MachineLearning.Model as Model
import qualified MachineLearning.Classification.Internal as MLC
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA

-- | Binary Classification prediction function.
-- Takes a matrix of features X and a vector theta.
-- Returns predicted class.
predict :: Matrix -> Vector -> Vector
predict x theta = V.map (\r -> if r >= 0.5 then 1 else 0) h
  where h = Model.hypothesis Log.Logistic x theta

-- | Learns Binary Classification.
learn :: Opt.MinimizeMethod -- ^ (e.g. BFGS2 0.1 0.1)
         -> R                  -- ^ epsilon, desired precision of the solution;
         -> Int                -- ^ maximum number of iterations allowed;
         -> Regularization     -- ^ regularization parameter lambda;
         -> Matrix             -- ^ matrix X;
         -> Vector             -- ^ binary vector y;
         -> Vector             -- ^ initial Theta;
         -> (Vector, Matrix)   -- ^ solution vector and optimization path.
learn mm eps numIters lambda x y initialTheta = Opt.minimize mm Log.Logistic eps numIters lambda x y initialTheta


module LinReg (

  lsr, bulkPDs

  ) where

import Poly
import Data.List
import Data.Monoid
import Control.Arrow
import Data.Packed.Matrix
import Numeric.LinearAlgebra.LAPACK

type Data = [(Double,Double)]

-- Data Acquisition - Very unsafe

getData :: PartialDerivitives -> ([[Double]], [Double])
getData = (map tail &&& map (negate . head)) . fun
  where
    fun                         = map (map (getNums . unSP) . unSS) . unPD
    getNums                     = foldl' getNum 1
    getNum x (SimpleNumber   n) = n * x
    getNum x (SimpleConstant _) = x

-- Concat Partialderivitives over data-points

bulkPDs :: Int -> Data -> PartialDerivitives
bulkPDs ord dat = pard
  where
    poly = polynomial_of_order               ord
    dsqr = delta_squared_term                poly
    ssss = mconcat $ map (point_result dsqr) dat
    pard = partial_derivatives ord           ssss

point_result :: Expression -> (Double,Double) -> SimpleSum
point_result e (x,y) = normal (Env [("x", x), ("y", y)]) e

lsr :: Int -> Data -> [Double]
lsr order points = concat $ toLists $ linearSolveR coefs values
  where
    linear     = bulkPDs (order - 1) points
    (lhs, rhs) = getData linear
    coefs      = (order >< order) (concat lhs)
    values     = (order >< 1    )  rhs


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Qubism.StateVecSpec (genStateVec, spec) where

import GHC.TypeLits
import Data.Singletons

import           Test.Hspec
import           Test.QuickCheck
import           Test.QuickCheck.Monadic
import           Control.Monad.Random
import           Control.Monad.Trans.State
import qualified Numeric.LinearAlgebra as LA
import           Data.Complex

import Qubism.AlgebraTests
import Qubism.StateVec

genC :: Gen (Complex Double)
genC = do
  a <- choose (-1, 1)
  b <- choose (-1, 1)
  pure $ a :+ b

genStateVec :: forall n . KnownNat n => Gen (StateVec n)
genStateVec = normalize . UnsafeMkStateVec . LA.fromList <$> vectorOf l genC
  where l = (2 ^) . fromIntegral . fromSing $ (sing :: Sing n)

-- | I know, it's an orphan instance, but it's not relevant outside of this
-- test suite, and never should be.
instance KnownNat n => Arbitrary (StateVec n) where
  arbitrary = genStateVec

propIdempotent
  :: (Eq a, Show a, Eq b, Show b)
  => StateT a (Rand StdGen) b
  -> a
  -> PropertyM IO Expectation
propIdempotent st a = do
  g <- lift getStdGen
  let one = runStateT st a `evalRand` g
      two = runStateT (st >> st) a `evalRand` g
  pure $ two `shouldBe` one

spec :: Spec
spec = do
  describe "Quantum Registers" $ do
    describe "form a vector space"  $ isVectorSpace  (T :: T (StateVec 1))
    describe "form a hilbert space" $ isHilbertSpace (T :: T (StateVec 1))
    describe "measure" $
      it "is idempotent"
        $ property
        $ monadicIO
        $ forAllM (genStateVec :: Gen (StateVec 1))
        $ propIdempotent measure
    describe "measureQubit" $
      it "is idempotent"
        $ property
        $ monadicIO
        $ forAllM (genStateVec :: Gen (StateVec 1))
        $ propIdempotent (measureQubit 0)


{-# LANGUAGE ParallelListComp #-}
module Main where

import Numeric.LinearAlgebra.Data (fromList, toList, cmap)
import Numeric.GSL.Fitting.Linear (linear_w)
import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Backend.Cairo
import Text.Printf

x = fromList [10, 100, 1000, 10000, 100000, 1000000]
y = fromList [4772, 21148, 134194, 1343814, 13389821, 148509325]
sigma = fromList [2112, 8194, 61719, 804762, 4853113, 42071732]

(intercept, slope, _, _, _, chi_sq) = linear_w x (cmap (\s -> 1 / s^2) sigma) y

errbars = liftEC $ do
    let values = [symErrPoint x' y' 0 dy
                  | x' <- toList x
                  | y' <- toList y
                  | dy <- toList sigma
                 ]
    plot_errbars_values .= values

main = toFile def "benchmark-fit.png" $ do
    setColors (map opaque [blue,red,green])
    layout_title .=
        printf "t = %.2f (ns/element) n + %.2f (ns) [chi^2 = %.2f]" slope intercept chi_sq
    layout_x_axis . laxis_generate .= autoScaledLogAxis def
    layout_y_axis . laxis_generate .= autoScaledLogAxis def
    plot errbars
    plot $ points "Benchmark data"
        (zip (toList x) (toList y))
    plot $ line "Best fit" [[
            (t, slope*t + intercept)
            | t <- map (10**) [0.5, 0.6 .. 6.5 ]
        ]]


#!/usr/bin/env stack
-- stack --install-ghc runghc --resolver lts-10.4 --package hmatrix --package lens -- -Wall

{-# LANGUAGE DataKinds       #-}
{-# LANGUAGE DeriveGeneric   #-}
{-# LANGUAGE TemplateHaskell #-}

-- Imports for network logic

import           Control.Lens hiding          ((<.>))
import           Numeric.LinearAlgebra.Static

-- Imports for utility functions

import           GHC.Generics (Generic)
import           Numeric.OneLiner

data Net = N { _weights1 :: L 250 784
             , _bias1    :: R 250
             , _weights2 :: L 10 250
             , _bias2    :: R 10
             }
  deriving (Generic)
makeLenses ''Net

logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))

softMax
    :: R 10
    -> R 10
softMax x = expx / konst (norm_1 expx)
  where
    expx = exp x

crossEntropy
    :: R 10
    -> R 10
    -> Double
crossEntropy targ res = -(log res <.> targ)

runNet
    :: Net
    -> R 784
    -> R 10
runNet n x = z
  where
    y = logistic $ (n ^. weights1) #> x + (n ^. bias1)
    z = softMax  $ (n ^. weights2) #> y + (n ^. bias2)

netErr
    :: R 784
    -> R 10
    -> Net
    -> Double
netErr x targ n = crossEntropy targ (runNet n x)

main :: IO ()
main = return ()

instance Num Net where
    (+)         = gPlus
    (-)         = gMinus
    (*)         = gTimes
    negate      = gNegate
    abs         = gAbs
    signum      = gSignum
    fromInteger = gFromInteger

instance Fractional Net where
    (/)          = gDivide
    recip        = gRecip
    fromRational = gFromRational


{-# LANGUAGE OverloadedStrings, DeriveGeneric #-}

module Data.Wavelets.Scaling where
import Prelude hiding (length,replicate,map,maximum,minimum)
-- import Data.Wavelets
import Numeric.Statistics
import Statistics.Function
import Statistics.Sample
import GHC.Generics
import Data.Vector.Storable -- hiding (map)
import Data.Packed.Matrix
-- import Linear

{-| Scale factors are designed to make it possible to use a scaled wavelet to reconstruct an approximation
    of the original data series. ordinary least squares is used so that translation is available for approximation,
    in addition to scaling.

|-}

data SeriesFactors = SeriesFactors {
      seriesMin :: Double ,
      seriesMax :: Double ,
      seriesMean :: Double ,
      seriesCount  :: Int } deriving (Show,Generic)

computeSeriesFactors :: Vector Double -> SeriesFactors
computeSeriesFactors v = SeriesFactors mn mx avg tot
    where (mn, mx,avg, tot) = minMaxAndAverage v

newtype OldSeriesFactors = OSF SeriesFactors  deriving (Show)
newtype NewSeriesFactors = NSF SeriesFactors  deriving (Show)
newtype NewSeriesMatrix = NSM (Matrix Double) deriving (Show)
newtype OldSeriesMatrix = OSM (Matrix Double) deriving (Show)

-- | x scaling matrix is B in Y = XB + e matrix formulation of ordinary least squares

newtype ScalingMatrix = SM (Matrix Double) deriving (Show)

type SeriesMatrix = Matrix Double

-- ScaleMatrix is the scaling matrix for the least squares problem
computeSeriesMatrix :: SeriesFactors -> SeriesMatrix
computeSeriesMatrix (SeriesFactors mn mx amean _ ) = (fromColumns vlst)
    where onesList = replicate 3 1
          nsf' = fromList [mn,mx,amean]
          vlst = [onesList,nsf']

computeNewSeriesMatrix :: NewSeriesFactors -> NewSeriesMatrix
computeNewSeriesMatrix (NSF sf) = NSM $ computeSeriesMatrix sf

computeSeriesResult :: SeriesFactors -> SeriesMatrix
computeSeriesResult (SeriesFactors mn mx avg _ ) = (fromColumns vlst)
    where nsf' = fromList [mn,mx,avg]
          vlst = [nsf']

computeOldSeriesMatrix :: OldSeriesFactors -> OldSeriesMatrix
computeOldSeriesMatrix (OSF sf) = OSM $ computeSeriesResult sf

-- | ordinary least squares estimation for the multivariate model
--   Y = X B + e        rows are observations, columns are elements
--   mean e = 0, cov e = kronecker s I

computeScalingMatrix :: NewSeriesFactors -> OldSeriesFactors -> ScalingMatrix
computeScalingMatrix nsf osf = let (OSM _Y) = computeOldSeriesMatrix osf
                                   (NSM _X) = computeNewSeriesMatrix nsf
                                   (_B,_ ,_ ) = ols _X _Y
                               in SM _B

-- | compute the average of the absolute value of a foldable type by setting it up like
-- foldl (\(a,t) b -> absAvgItrFcn a b (t+1) ) (0,0) foldableThing
absAvgItrFcn ::(Num a, Fractional a)=> a -> a -> a -> (a,a)
absAvgItrFcn a b t = avgtuple
   where
     diff = (a - abs(b))
     dv  = abs(diff/t)
     tot  = a + dv
     avgtuple = (tot,t)

-- | compute the maximum absolute value of a foldable type by setting it up like
-- foldl (\mx b -> absMaxItrFcn mx b ) 0 foldableThing
absMaxItrFcn :: (Num a, Ord a) => a -> a -> a
absMaxItrFcn a b = max a (abs b)

-- | compute the max min and average of a Vector in one pass
-- (avg, max, min, tot)
minMaxAndAverage :: Vector Double -> (Double, Double ,Double , Int )
minMaxAndAverage v = (mn,mx,av,l)
    where
      av = mean v
      (mn, mx) = minMax v
      l = length v

applyScalingMatrix :: ScalingMatrix -> Vector Double -> Vector Double
applyScalingMatrix (SM sm) v = map scaleFcn v
    where tCoef = sm @@> (0,0)
          sCoef = sm @@> (1,0)
          scaleFcn x = sCoef * x + tCoef


module SparseLinAlg (solve_tree) where

import Data.Bifunctor (bimap)
import Data.List (sort)
import Data.Maybe (fromJust)
import Data.Sparse.SpMatrix
import Data.Sparse.SpVector
import Numeric.LinearAlgebra.Sparse

import LinEq (Coeff, Var, Equation(..), remove_term, combine_terms, ltree_of_tree, equations_of_ltree)
import Sexp
import Tree
import Util (debug)

-- Here our equations have a slightly different form, where the LHS is
-- a linear combination of the variables and the RHS is a single
-- rational value. E.g., 1x - 1/2y = 1/2.
type MatTerm = (Coeff, Var)
type MatEq = ([MatTerm], Double)

-- mateq_of_equation :: Equation -> MatEq
-- mateq_of_equation (Equation (x, tms)) =
--   case remove_term Nothing $ combine_terms tms of
--     Just (c, tms') -> (second fromJust <$> tms', c)
--     Nothing -> (second fromJust <$> tms, 0)

mateq_of_equation :: Equation -> MatEq
mateq_of_equation (Equation (x, tms)) =
  case remove_term Nothing $ combine_terms tms of
    Just (c, tms') -> ((1, x) : (bimap negate fromJust <$> tms'), c)
    Nothing -> ((1, x) : (bimap negate fromJust <$> tms), 0)

constraint_matrix :: [MatEq] -> SpMatrix Double
constraint_matrix eqs =
  let l = concat $ f <$> zip [0..] eqs in
    debug ("l: " ++ show ((\(x, y, z) -> (x, y, z)) <$> l)) $
    fromListSM (n, n) l
  where
    n = length eqs

    f :: (Int, MatEq) -> [(Int, Int, Double)]
    f (x, (tms, _)) = g x <$> tms

    g :: Int -> MatTerm -> (Int, Int, Double)
    g x (c, y) = (x, y, c)
    -- g x (c, y) = (y, x, c)

-- tree_constraint_matrix :: Tree Bool -> SpMatrix Double
-- tree_constraint_matrix =
--   fmap fromRational . constraint_matrix . map mateq_of_equation .
--   equations_of_ltree . ltree_of_tree

-- tree_constraint_matrix :: Tree Bool -> SpMatrix Double
-- tree_constraint_matrix =
--   fmap fromRational .
--   constraint_matrix .
--   map (\eq -> let mateq = mateq_of_equation eq in
--                 debug ("eq: " ++ show eq) $
--                 debug ("mateq" ++ show mateq) $
--                 mateq) .
--   equations_of_ltree .
--   ltree_of_tree

matrix_of_mateqs :: [MatEq] -> SpMatrix Double
matrix_of_mateqs = constraint_matrix

rhs_of_mateqs :: [MatEq] -> SpVector Double
rhs_of_mateqs eqs = fromListDenseSV n $ snd <$> eqs
  where n = length eqs

solve_system_gmres :: SpMatrix Double -> SpVector Double -> IO (SpVector Double)
solve_system_gmres mat rhs =
  debug ("mat: " ++ show mat) $
  debug ("rhs: " ++ show rhs) $
  debug ("dense:") $
  -- let _ = unsafePerformIO $ prd mat in
  -- mat <\> (fromListDenseSV n )
  mat <\> rhs
  -- where
  --   n = nrows mat -- should also be the length of the rhs vector

solve_tree :: Tree Double -> IO (SpVector Double)
solve_tree (Leaf x) = return $ fromListDenseSV 1 [x]
solve_tree t =
  let lt = ltree_of_tree t
      eqs = sort $ equations_of_ltree lt
      mateqs = mateq_of_equation <$> eqs
      (mat, rhs) = (matrix_of_mateqs mateqs, rhs_of_mateqs mateqs) in
    debug ("eqs: " ++ show eqs) $
    debug ("mateqs: " ++ show mateqs) $
    debug ("ltree: " ++ toSexp lt) $ do
    prd mat
    prd rhs
    sol <- solve_system_gmres mat rhs
    prd sol
    return sol

-- infer_gmres :: Tree Bool -> Maybe Rational
-- infer_gmres t =
--   toRational <$>
--   (lookupSV 0 $ unsafePerformIO $ solve_matrix $ tree_constraint_matrix t)

-- [(0,[(0,1.0),(1,-0.5),(4,-0.5)]),
--  (1,[(0,-0.5),(1,1.0),(3,-0.5)]),
--  (2,[(1,-0.5),(3,1.0)]),
--  (3,[(0,-0.5),(2,-0.5),(4,1.0)]),
--  (4,[(2,1.0),(5,-0.5)]),
--  (5,[(0,-0.5),(2,-0.5),(5,1.0)])]


{-# OPTIONS_HADDOCK hide #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Types
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Basic type classes and enums.
--

module Numeric.LinearAlgebra.Types (

    -- * Matrix views
    Herm(..),
    Tri(..),

    -- * Matrix factorization views
    Chol(..),

    -- * Vector math types
    VNum,
    VFractional,
    VFloating,

    -- * BLAS element types
    BLAS1,
    BLAS2,
    BLAS3,

    -- * LAPACK element types
    LAPACK,

    -- * Enums
    Trans(..),
    Uplo(..),
    Side(..),
    Diag(..),

    -- * Algorithm-specific parameters
    CovMethod(..),

    -- * Re-export of Complex from Data.Complex
    module Data.Complex,

    -- * Re-export of Storable from Foreign.Storable
    module Foreign.Storable,

    ) where

import Foreign.VMath( VNum, VFractional, VFloating )
import Foreign.BLAS( Trans(..), Uplo(..), Side(..), Diag(..), BLAS1, BLAS2, BLAS3 )
import Foreign.LAPACK( LAPACK )
import Data.Complex( Complex(..) )
import Foreign.Storable( Storable() )

-- | A hermitian view of an underlying matrix.  The view can either be
-- of the upper or lower triangular part of the matrix.  The type arguments
-- are as follows:
--
--     * @m@: the underlyting matrix type.
--
--     * @e@: the element type of the matrix.
--
data Herm m e = Herm Uplo (m e) deriving (Show)

-- | A triangular view of an underlying matrix.  The view can either be
-- upper or lower triangular, with a unit or non-unit diagonal.  The type
-- arguments are as follows:
--
--     * @m@: the underlyting matrix type.
--
--     * @e@: the element type of the matrix.
--
data Tri m e = Tri Uplo Diag (m e) deriving (Show)

-- | A Cholesky decomposition view of a matrix.
data Chol m e = Chol Uplo (m e) deriving (Show)

-- | The method of scaling the sample covariance matrix.
data CovMethod =
      UnbiasedCov -- ^ This is the default behavior. Corresponds to a
                  -- scaling of @n/(n-1)@ in the unweighed case, and
                  -- @1/(1 - \\sum w_i^2)@ in the weighted case, where @w_i@
                  -- is the normalized weight. Note the unweighted and
                  -- weighted cases agree when @w_i = 1/n@.

    | MLCov       -- ^ Returns the centered second moment matrix without
                  -- scaling the result.
    deriving (Eq, Show)


{-# LANGUAGE BangPatterns #-}
-- # LANGUAGE UndecidableInstances #
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Data.Network.FullyConnected
  (FeedForwardFC, FullyConnected(FC), Layers,
   weights, biases, activation,
   unGrad, unGradBatch, Grad(GradBatch, Grad))
where

import Numeric.LinearAlgebra as LA hiding (fromList, toList)
import Numeric.LinearAlgebra.Devel
import qualified Data.Vector.Storable as V hiding (foldr, modify)
import System.Random
import Control.DeepSeq
-- import Data.Foldable as F
import Data.Bifunctor
import Data.Network as Network
import Data.Network.Activation
import Algorithm.Training
import Algorithm.GradientDescent
import Data.Network.Utils (squeeze, (.*))

-- Fully connected feedforward layer
data FullyConnected i a = FC {weights :: {-# UNPACK #-} !(Matrix a),
                              biases ::  {-# UNPACK #-} !(Vector a),
                              activation :: ActivationFunction i
                       }

instance (V.Storable a, NFData a) => NFData (FullyConnected i a) where
  rnf (FC w b a) = rnf w `deepseq` rnf b

instance Layer (Vector R) (FullyConnected (Vector R) R) where
  compute (FC weights biases _) input = biases + (input <# weights)
  run     l@(FC weights biases af) input = unAct af (compute l input)
  makeRandomLayer g rf i o af = let rnds = (randomRs (0,1) g) :: [Double]
                                    rnds' = (rf rnds)
                                    rw   = (i><o) rnds'
                                    rb   = o |> (repeat 0)
                                in (FC rw rb af)

instance Layer (Matrix R) (FullyConnected (Matrix R) R) where
  compute (FC weights biases _) input = (fromRows $ replicate (rows input) biases) + (input <> weights)
  run     l@(FC _ _ af) input = unAct af (compute l input)
  makeRandomLayer g rf i o af = let rnds = (randomRs (0,1) g) :: [Double]
                                    rnds' = (rf rnds)
                                    rw   = (i><o) rnds'
                                    rb   = o |> (repeat 0)
                                in (FC rw rb af)

instance (Container Matrix a, Container Vector a, Num a) => Show (FullyConnected i a) where
  show l = let (a,b) = size $ weights l
           in "Layer " ++ (show a) ++ "x" ++ (show b) ++ " " ++ (name $ activation l)

instance Floating a => Num (FullyConnected a Double) where
  (FC w b a) + (FC w' b' _) = FC (w + w') (b + b') a
  (FC w b a) - (FC w' b' _) = FC (w - w') (b - b') a
  (FC w b a) * (FC w' b' _) = FC (w * w') (b * b') a
  abs (FC w b a) = FC (abs w) (abs b) a
  signum (FC w b a) = FC (signum w) (signum b) a
  fromInteger x = FC (fromInteger x) (fromInteger x) (act id "identity")

instance Floating a => Fractional (FullyConnected a Double) where
  (FC w b a) / (FC w' b' _) = FC (w / w') (b / b') a
  recip (FC w b a) = FC (recip w) (recip b) a
  fromRational a = FC (fromRational a) (fromRational a) (act id "identity")

instance (Container Vector a, Container Matrix a) => Linear a (FullyConnected i) where
  scale x (FC w b a) = FC (scale x w) (scale x b) a

-- Container Layer instance?

instance (Layer i (FullyConnected i Double)) => Network i (Layers (FullyConnected i Double)) where
  runNetwork (Output l) i = run l i
  runNetwork (l :=> n)  i = runNetwork n (fst $ run l i)
  runLayers  (Output l) i = [run l i]
  runLayers  (l :=> n)  i = let y = (run l i)
                            in  y:(runLayers n (fst y))

  makeRandomNetwork g t nt i [] (o, af) = let l = makeRandomLayer g t i o af
                                          in  (Output l)
  makeRandomNetwork g t nt i ((h,haf):hs) o = let (g1,g2) = split g
                                                  l = makeRandomLayer g1 t i h haf
                                              in nt $ l :=> (makeRandomNetwork g2 t nt h hs o)

type FeedForwardFC i = Layers (FullyConnected i Double)

instance (Monad m) => Backpropagation Data (NetworkEvaluator Data) (FeedForwardFC Data) m where
  backward lf n (out,tar) das = do
    let δout = (derivate lf (tar, out)) -- dE/dy
        deltas = scanr (\(l, a') δ -> let w = weights l in a' * (w #> δ)) δout (zip (tail $ toList n) das)
    return (deltas) -- deltas for computing bias

instance (Monad m) => Backpropagation Batch (NetworkEvaluator Batch) (FeedForwardFC Batch) m where
  backward lf n (out,tar) das = do
    let δout = tr (derivate lf (tar, out)) -- dE/dy
        deltas = scanr (\(l, a') δ -> let w = weights l in (tr a') * (w <> δ)) δout (zip (tail $ toList n) das)
    return (deltas) -- deltas for computing bias

-- type instance Params lf = forall i. ((FeedForwardFC i), (i,i))

instance (Monad m) => GradientDescent (NetworkEvaluator Batch) ((FeedForwardFC Batch), (Batch,Batch)) m where
  -- parameters of loss function: Network(weights, biases, activation functions for each layer)
  -- and sample((input, output) pair)

  data Grad (NetworkEvaluator Batch) = GradBatch { unGradBatch :: ([Matrix Double], [Vector Double]) } deriving (Show)

  -- Gradients over parameters(weights, biases)
  grad lf (n, (i,t)) = do
    let (as, as') = unzip $ runLayers n i -- forward propagation: compute layers outputs and activation derivatives
        (out) = last as
    (ds) <- backward lf n (out, t) (init as') -- compute deltas with backpropagation
    let r  = fromIntegral $ rows i -- size of minibatch
    let gs = zipWith (\δ a -> tr (δ <> a)) ds (i:init as)
    return $ GradBatch ((recip r .*) <$> gs, ((recip r .*) . squeeze) <$> ds)

  move lr (n, (i,t)) (GradBatch (gs, ds)) = do
    -- update function
    -- divide by batch size
    let update = (\(FC w b af) g δ -> FC (w + (lr).*g) (b + (lr).*δ) af)
        n' = Network.fromList $ zipWith3 update (Network.toList n) gs ds -- new model
    return (n', (i,t))

instance (Monad m) => GradientDescent (NetworkEvaluator Data) ((FeedForwardFC Data), (Data,Data)) m where
  -- parameters of loss function: Network(weights, biases, activation functions for each layer)
  -- and sample((input, output) pair)

  data Grad (NetworkEvaluator Data) = Grad { unGrad :: ([Matrix Double], [Vector Double]) } deriving (Show)

  -- Gradients over parameters(weights, biases)
  grad lf (n, (i,t)) = do
    let aos = runLayers n i -- forward propagation: compute layers outputs and activation derivatives
        (as, as') = (unzip) aos -- juggle data to extract layers inputs and derivatives in right order
        out = last as
    (ds) <- backward lf n (out,t) (init as') -- compute gradients with backpropagation
    let gs = zipWith (\δ i -> tr (δ `outer` i)) ds (i:as)
    return $ Grad (gs, ds)

  move lr (n, (i,t)) (Grad (gs, ds)) = do
    -- update function
    let update = (\(FC w b af) g δ -> FC (w + lr.*g) (b + lr.* δ) af)
        n' = Network.fromList $ zipWith3 update (Network.toList n) gs ds -- new model
    return (n', (i,t))


{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE TypeFamilies          #-}
--------------------------------------------------------------------------------
-- |
-- Robust Principal Component Analysis Outlier Detection a la Netflix/Surus
---------------------------------------------------------------------------------
module RobustPCA
  ( -- * Datatypes
    LR(..)
  , LRPenalty(..)
  , MaxIters(..)
  , FitIters(..)
  , Periodicity(..)
  , RobustPCA(..)
  , Sparse(..)
  , SparsePenalty(..)
  , Tolerance(..)
    -- * Fitting
  , runRobustPCA
  , optimize
  ) where

import           Control.Foldl            (fold)
import           Control.Foldl.Statistics (LMVSK (..), fastLMVSKu, fastStdDev)
import           Control.Lens             (declareWrapped, makeLenses, use,
                                           (+=), (.=), (^.))
import           Control.Monad            (return)
import           Control.Monad.Identity   (runIdentity)
import           Control.Monad.Loops      (whileM_)
import           Control.Monad.State      (State, execStateT)
import           Data.Functor             (void, (<$>))
import           Data.List                (replicate, take)
import           Numeric.LinearAlgebra    (Matrix, R, cmap, cols, diag, flatten,
                                           rows, sumElements, thinSVD, toList,
                                           tr, (<>), (><))
import           Prelude                  (Double, Int, Show, abs, div,
                                           fromIntegral, max, mod, ($), (&&),
                                           (*), (**), (+), (-), (.), (/), (<),
                                           (>))
import           Utils                    (absL1Norm, frobeniusNorm,
                                           softThreshold)

---------------------------------------------------------------------------------
-- Datatypes --------------------------------------------------------------------
---------------------------------------------------------------------------------

declareWrapped [d|
  newtype LRPenalty     = LRPenalty Double     deriving Show
  newtype SparsePenalty = SparsePenalty Double deriving Show
  newtype Periodicity   = Periodicity Int      deriving Show
  newtype MaxIters      = MaxIters Int         deriving Show
  newtype FitIters      = FitIters Int         deriving Show
  newtype LR            = LR [R]               deriving Show
  newtype Sparse        = Sparse [R]           deriving Show
  newtype Tolerance     = Tolerance [R]        deriving Show
  |]

-- |
-- Optimization state
data RobustPCA
  = RobustPCA
  { _maxIters       :: MaxIters
    -- ^ maximum optimization iterations
  , _lrPenalty      :: LRPenalty
    -- ^ low rank fit penalty
  , _sparsePenalty  :: SparsePenalty
    -- ^ sparse fit penalty
  , _normalizedData :: Matrix R
    -- ^ normalized input data matrix
  , _iters          :: Int
    -- ^ iterations count
  , _obj            :: R
    -- ^ objective
  , _tol            :: R
    -- ^ tolerance
  , _diff           :: R
    -- ^ difference
  , _mean           :: R
    -- ^ input data mean
  , _stdDev         :: R
    -- ^ input data standard deviation
  , _mu             :: R
    -- ^ incoherence factor
  , _lr             :: Matrix R
    -- ^ recovered low rank matrix (underlying data 'approximation')
  , _sparse         :: Matrix R
    -- ^ sparse outliers matrix
  , _tolerance      :: Matrix R
    -- ^ tolerated deviation of real data from the 'fit'
  } deriving (Show)
makeLenses ''RobustPCA

-- | Fitting state monad
type RobustPCAM a = State RobustPCA a

---------------------------------------------------------------------------------
-- Fitting ----------------------------------------------------------------------
---------------------------------------------------------------------------------

-- |
-- Run robust PCA optimization
runRobustPCA
  :: [R]
  -- ^ input 'Real' data
  -> LRPenalty
  -- ^ low rank fit penalty
  -> SparsePenalty
  -- ^ sparse fit penalty
  -> Periodicity
  -- ^ input data periodicity
  -> MaxIters
  -- ^ maximum iteration count allowed
  -> (LR, Sparse, Tolerance, FitIters)
runRobustPCA originalData lpen spen p mi =
    (LR lr', Sparse sparse', Tolerance tol', FitIters itrs)
  where rs      = runIdentity $
                   execStateT optimize $
                    initState originalData lpen spen p mi
        mean'   = rs^.mean
        stdDev' = rs^.stdDev
        lr'     = toList $ flatten $ tr $ cmap ((mean' +) . (stdDev' *)) (rs^.lr)
        sparse' = toList $ flatten $ tr $ cmap (stdDev' *) (rs^.sparse)
        tol'    = toList $ flatten $ tr $ cmap (stdDev' *) (rs^.tolerance)
        itrs    = rs^.iters

---------------------------------------------------------------------------------
-- Internal Functions -----------------------------------------------------------
---------------------------------------------------------------------------------

initState
  :: [R]
  -> LRPenalty
  -> SparsePenalty
  -> Periodicity
  -> MaxIters
  -> RobustPCA
initState originalData lp sp (Periodicity p) mi =
    RobustPCA mi lp sp normalizedData' 0 obj' tol' diff' m sdev mu' z z z
  where (LMVSK l m v _ _) = fold fastLMVSKu originalData
        sdev                   = v ** 0.5
        zeroMeanUnitVariance x = (x - m) / sdev
        paddedLength           = l - (l `mod` p)
        normalizedData'        =
          tr $ ((paddedLength `div` p) >< p)
               (zeroMeanUnitVariance <$> take paddedLength originalData)
        r                      = rows normalizedData'
        c                      = cols normalizedData'
        z                      = (r >< c) (replicate (r * c) 0.0)
        mu'                    =
          fromIntegral c * fromIntegral r / (4 * absL1Norm normalizedData')
        obj'                   = 0.5 * ((frobeniusNorm normalizedData') ** 2.0)
        tol'                   = 1e-8 * obj'
        diff'                  = 2.0 * tol'

optimize :: RobustPCAM ()
optimize = whileM_ checkTolerance go
  where go = do
          nuclearNorm <- computeSparse
          l1Norm      <- computeLR
          l2Norm      <- computeTol
          objPrev     <- use obj
          let obj'     = 0.5 * l2Norm + nuclearNorm + l1Norm
          diff        .= abs (objPrev - obj')
          obj         .= obj'
          void updateMu
          iters += 1
        checkTolerance = do
          iters'       <- use iters
          MaxIters mi' <- use maxIters
          tol'         <- use tol
          diff'        <- use diff
          return $ diff' > tol' && iters' < mi'

computeSparse :: RobustPCAM R
computeSparse = do
  normalizedData'          <- use normalizedData
  SparsePenalty spen       <- use sparsePenalty
  mu'                      <- use mu
  lr'                      <- use lr
  let sparsePenalty'        = spen * mu'
  let normalizedDataMinusLr = normalizedData' - lr'
  let penalizedSparse       = softThreshold normalizedDataMinusLr sparsePenalty'
  sparse                   .= penalizedSparse
  return $ (absL1Norm penalizedSparse) * sparsePenalty'

computeLR :: RobustPCAM R
computeLR = do
  normalizedData'              <- use normalizedData
  LRPenalty lpen               <- use lrPenalty
  mu'                          <- use mu
  sparse'                      <- use sparse
  let lrPenalty'                = lpen * mu'
  let normalizedDataMinusSparse = normalizedData' - sparse'
  let (u, svs, v)               = thinSVD normalizedDataMinusSparse
  let penalizedD                = softThreshold svs lrPenalty'
  let d'                        = diag penalizedD
  lr                           .= u <> d' <> (tr v)
  return $ (sumElements penalizedD) * lrPenalty'

computeTol :: RobustPCAM R
computeTol = do
  normalizedData' <- use normalizedData
  lr'             <- use lr
  sparse'         <- use sparse
  let tolerance'   = normalizedData' - lr' - sparse'
  tolerance       .= tolerance'
  return $ (frobeniusNorm tolerance') ** 2

updateMu :: RobustPCAM ()
updateMu = do
  tolerance'  <- use tolerance
  let tdStdDev = fold fastStdDev (toList $ flatten tolerance')
  let m        = rows tolerance'
  let n        = cols tolerance'
  mu .= max 0.01 (tdStdDev * ((2.0 * (fromIntegral $ max m n)) ** 0.5))