-----------------------------------------------------------------------------
-- |
-- Module      :  Finance.Hqfl.Pricer.BlackScholes
-- Copyright   :  (C) 2016 Mika'il Khan
-- License     :  (see the file LICENSE)
-- Maintainer  :  Mika'il Khan <co.kleisli@gmail.com>
-- Stability   :  stable
-- Portability :  portable
--
----------------------------------------------------------------------------
{-# LANGUAGE FlexibleInstances #-}

module Finance.Hqfl.Pricer.BlackScholes where

import Finance.Hqfl.Instrument
import Statistics.Distribution.Normal
import Data.Random

class BS a where
  price :: a -> Double -> Double -> Double

instance BS (Option Equity) where
  price (Option (Equity s q) m European k t) r v =
    case m of
       Call -> s * exp (-q * t) * cdf normal d1 - k * exp (-r * t) * cdf normal d2
       Put  -> k * exp (-r * t) * cdf normal (-d2) - s * exp (-q * t) * cdf normal (-d1)
    where d1 = (log (s / k) + (r - q + (v * v) / 2) * t) / (v * sqrt t)
          d2 = d1 - v * sqrt t
          normal = Normal (0 :: Double) 1

instance BS (Option StockIndex) where
  price (Option (StockIndex s q) m European k t) r v =
    case m of
       Call -> s * exp (-q * t) * cdf normal d1 - k * exp (-r * t) * cdf normal d2
       Put  -> k * exp (-r * t) * cdf normal (-d2) - s * exp (-q * t) * cdf normal (-d1)
    where d1 = (log (s / k) + (r - q + (v * v) / 2) * t) / (v * sqrt t)
          d2 = d1 - v * sqrt t
          normal = Normal (0 :: Double) 1

-- Garman & Kohlhagen (1983)
--instance BS (Option Currency) where
--  price (Option (Currency s rf) m European k t) r v =
--    case m of
--       Call -> s * exp (-rf * t) * cdf normal d1 - k * exp (-r * t) * cdf normal d2
--       Put  -> k * exp (-r * t) * cdf normal (-d2) - s * exp (-rf * t) * cdf normal (-d1)
--    where d1 = (log (s / k) + (r - rf + (v * v) / 2) * t) / (v * sqrt t)
--          d2 = d1 - v * sqrt t
--          normal = Normal (0 :: Double) 1

-- cdi :: Double -> Double -> Double -> Double -> Double -> Double -> Double -> Double
-- cdi s k r v t q h = s * exp(-q * t) * (h / s)**(2 * lambda) * (cdf normal y)
--                     - k * exp(-r * t) * (h / s)**(2 * lambda - 2) * (cdf normal (y - (v * sqrt t)))
--   where lambda = ((r - q + (v * v) / 2)) / (v * v)
--         y = (log ( (h*h) / (s * k)) / v * sqrt t) + (lambda * v * sqrt t)
--         normal = Normal (0 :: Double) 1

-- cdo :: Double -> Double -> Double -> Double -> Double -> Double -> Double -> Double
-- cdo s k r v t q h = s * (cdf normal x1) * exp (-q * t) - k * exp(-r * t) * (cdf normal (x1 - v * sqrt t)) -
--                     s * exp(-q * t) * (h / s)**(2 * lambda) * (cdf normal y1) + k * exp(-r * t) * (h/ s)**(2 * lambda - 2) * (cdf normal (y1 - v * sqrt t))
--   where lambda = ((r - q + (v * v) / 2)) / (v * v)
--         y1 = (log (h / s) / v * sqrt t) + (lambda * v * sqrt t)
--         x1 = (log ( s / h ) / (v * sqrt t)) + (lambda * v * sqrt t)
--         normal = Normal (0 :: Double) 1

-- cui :: Double -> Double -> Double -> Double -> Double -> Double -> Double -> Double
-- cui s k r v t q h = s * (cdf normal x1) * exp(-q * t) - k * exp(-r * t) * (cdf normal (x1 - v * sqrt t)) -
--       s * exp(-q *t) * (h / s)**(2 * lambda) * ((cdf normal (-y)) - (cdf normal (-y1))) +
--       k * exp(-r * t) * (h / s)**(2 * lambda - 2) * ((cdf normal (-y + v * sqrt t)) - (cdf normal (-y1 + v * sqrt t)))
--   where lambda = ((r - q + (v * v) / 2)) / (v * v)
--         y1 = (log (h / s) / v * sqrt t) + (lambda * v * sqrt t)
--         x1 = (log ( s / h ) / (v * sqrt t)) + (lambda * v * sqrt t)
--         y = (log ( (h*h) / (s * k)) / v * sqrt t) + (lambda * v * sqrt t)
--         normal = Normal (0 :: Double) 1

-- pui :: Double -> Double -> Double -> Double -> Double -> Double -> Double -> Double
-- pui s k r v t q h = -s * exp(-q * t) * (h / s)**(2 * lambda) * (cdf normal (-y)) + k * exp(-r * t) * (h / s)**(2 * lambda - 2) * (cdf normal ((-y) + (v * sqrt t)))
--   where lambda = ((r - q + (v * v) / 2)) / (v * v)
--         y = (log ( (h*h) / (s * k)) / v * sqrt t) + (lambda * v * sqrt t)
--         normal = Normal (0 :: Double) 1

-- pdi :: Double -> Double -> Double -> Double -> Double -> Double -> Double -> Double
-- pdi s k r v t q h = -s * (cdf normal (-x1)) * exp(-q * t)
--                     + k * exp(-r * t) * (cdf normal (-x1 + v * sqrt t))
--                     + s * exp(-q * t) * (h / s)**(2 * lambda) * (cdf normal y - cdf normal y1)
--                     - k * exp(-r * t) * (h / s)**(2 * lambda -2) * (cdf normal (y - v * sqrt t) - cdf normal (y1 - v * sqrt t))
--   where lambda = ((r - q + (v * v) / 2)) / (v * v)
--         y1 = (log (h / s) / v * sqrt t) + (lambda * v * sqrt t)
--         x1 = (log ( s / h ) / (v * sqrt t)) + (lambda * v * sqrt t)
--         y = (log ( (h*h) / (s * k)) / v * sqrt t) + (lambda * v * sqrt t)
--         normal = Normal (0 :: Double) 1

-- puo s k r v t q h = -s * (cdf normal (-x1)) * exp(-q * t) + k * exp(-r * t) * (cdf normal (-x1 + v * sqrt t)) +
--                     s * exp(-q * t) * (h / s)**(2 * lambda) * (cdf normal (-y1)) - k * exp(-r * t) * (h / s)**(2 * lambda -2) * (cdf normal ((-y1) + v * sqrt t))
--   where lambda = ((r - q + (v * v) / 2)) / (v * v)
--         y1 = (log (h / s) / v * sqrt t) + (lambda * v * sqrt t)
--         x1 = (log ( s / h ) / (v * sqrt t)) + (lambda * v * sqrt t)
--         y = (log ( (h*h) / (s * k)) / v * sqrt t) + (lambda * v * sqrt t)
--         normal = Normal (0 :: Double) 1

-- instance BS (Option StockIndex) where
--   blackscholesprice (Option (StockIndex s) Call European k t) r v = call s k r v t q
--   blackscholesprice (Option (StockIndex s) Put European k t) r v = put s k r v t q

-- instance BS (Option Future) where
--   blackscholesprice (Option (Future f) Call European k t) r v = call f k r v t q
--   blackscholesprice (Option (Future f) Put European k t) r v = put f k r v t q

-- TODO : factor in rebates
-- instance BS (BarrierOption Equity) where
--   blackscholesprice (BarrierOption (Equity s q) x European y h k t) r v =
--     case (x, y) of
--       (DownAndIn, Call)
-- --        | s <= h -> Left "Contract is converted to a vanilla call"
--         | h <= k ->  cdi'
--         | h >= k ->   call' - cdo'
--       (DownAndOut, Call)
-- --        | s <= h -> Left "Contract is void"
--         | h <= k ->   call' - cdi'
--         | h >= k ->  cdo'
--       (UpAndIn, Call)
-- --        | h <= s -> Left "Contract is converted to a vanilla call"
--         | h <= k ->  call'
--         | h >= k ->   cui'
--       (UpAndOut, Call)
-- --        | s <= h -> Left "Contract is void"
--         | h <= k ->   call' - cdi'
--         | h >= k ->  cdo'
--       (DownAndIn, Put)
-- --        | s <= h -> Left "Contract is converted to a vanilla call"
--         | h <  k ->  pdi'
--         | h >= k ->  put'
--       (DownAndOut, Put)
-- --        | s <= h -> Left "Contract is void"
--         | h <  k ->   put' - pdi'
--         | h >=  k ->  0
--       (UpAndIn, Put)
-- --        | h <= s -> Left "Contract is converted to a vanilla call"
--         | h <= k ->   put' - puo'
--         | h >= k ->  pui'
--       (UpAndOut, Put)
-- --        | h <= s -> Left "Contract is void"
--         | h <= k ->   puo'
--         | h >= k ->   put' - pui'
--     where
--       call' = call s k r v t q
--       cdi'  = cdi  s k r v t q h
--       cdo'  = cdo  s k r v t q h
--       cui'  = cui  s k r v t q h
--       pdi'  = pdi  s k r v t q h
--       put'  = put  s k r v t q
--       pui'  = pui  s k r v t q h
--       puo'  = puo  s k r v t q h


{-# LANGUAGE BangPatterns, FlexibleContexts #-}
-- |
-- Module    : Statistics.Transform
-- Copyright : (c) 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Fourier-related transformations of mathematical functions.
--
-- These functions are written for simplicity and correctness, not
-- speed.  If you need a fast FFT implementation for your application,
-- you should strongly consider using a library of FFTW bindings
-- instead.

module Statistics.Transform
    (
    -- * Type synonyms
      CD
    -- * Discrete cosine transform
    , dct
    , dct_
    , idct
    , idct_
    -- * Fast Fourier transform
    , fft
    , ifft
    ) where

import Control.Monad (when)
import Control.Monad.ST (ST)
import Data.Bits (shiftL, shiftR)
import Data.Complex (Complex(..), conjugate, realPart)
import Numeric.SpecFunctions (log2)
import qualified Data.Vector.Generic         as G
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Unboxed         as U
import qualified Data.Vector                 as V

type CD = Complex Double

-- | Discrete cosine transform (DCT-II).
dct :: (G.Vector v CD, G.Vector v Double, G.Vector v Int) => v Double -> v Double
dct = dctWorker . G.map (:+0)
{-# INLINABLE  dct #-}
{-# SPECIAlIZE dct :: U.Vector Double -> U.Vector Double #-}
{-# SPECIAlIZE dct :: V.Vector Double -> V.Vector Double #-}

-- | Discrete cosine transform (DCT-II). Only real part of vector is
--   transformed, imaginary part is ignored.
dct_ :: (G.Vector v CD, G.Vector v Double, G.Vector v Int) => v CD -> v Double
dct_ = dctWorker . G.map (\(i :+ _) -> i :+ 0)
{-# INLINABLE  dct_ #-}
{-# SPECIAlIZE dct_ :: U.Vector CD -> U.Vector Double #-}
{-# SPECIAlIZE dct_ :: V.Vector CD -> V.Vector Double#-}

dctWorker :: (G.Vector v CD, G.Vector v Double, G.Vector v Int) => v CD -> v Double
{-# INLINE dctWorker #-}
dctWorker xs
  -- length 1 is special cased because shuffle algorithms fail for it.
  | G.length xs == 1 = G.map ((2*) . realPart) xs
  | vectorOK xs      = G.map realPart $ G.zipWith (*) weights (fft interleaved)
  | otherwise        = error "Statistics.Transform.dct: bad vector length"
  where
    interleaved = G.backpermute xs $ G.enumFromThenTo 0 2 (len-2) G.++
                                     G.enumFromThenTo (len-1) (len-3) 1
    weights = G.cons 2 . G.generate (len-1) $ \x ->
              2 * exp ((0:+(-1))*fi (x+1)*pi/(2*n))
      where n = fi len
    len = G.length xs

-- | Inverse discrete cosine transform (DCT-III). It's inverse of
-- 'dct' only up to scale parameter:
--
-- > (idct . dct) x = (* length x)
idct :: (G.Vector v CD, G.Vector v Double) => v Double -> v Double
idct = idctWorker . G.map (:+0)
{-# INLINABLE  idct #-}
{-# SPECIAlIZE idct :: U.Vector Double -> U.Vector Double #-}
{-# SPECIAlIZE idct :: V.Vector Double -> V.Vector Double #-}

-- | Inverse discrete cosine transform (DCT-III). Only real part of vector is
--   transformed, imaginary part is ignored.
idct_ :: (G.Vector v CD, G.Vector v Double) => v CD -> v Double
idct_ = idctWorker . G.map (\(i :+ _) -> i :+ 0)
{-# INLINABLE  idct_ #-}
{-# SPECIAlIZE idct_ :: U.Vector CD -> U.Vector Double #-}
{-# SPECIAlIZE idct_ :: V.Vector CD -> V.Vector Double #-}

idctWorker :: (G.Vector v CD, G.Vector v Double) => v CD -> v Double
{-# INLINE idctWorker #-}
idctWorker xs
  | vectorOK xs = G.generate len interleave
  | otherwise   = error "Statistics.Transform.dct: bad vector length"
  where
    interleave z | even z    = vals `G.unsafeIndex` halve z
                 | otherwise = vals `G.unsafeIndex` (len - halve z - 1)
    vals = G.map realPart . ifft $ G.zipWith (*) weights xs
    weights
      = G.cons n
      $ G.generate (len - 1) $ \x -> 2 * n * exp ((0:+1) * fi (x+1) * pi/(2*n))
      where n = fi len
    len = G.length xs

-- | Inverse fast Fourier transform.
ifft :: G.Vector v CD => v CD -> v CD
ifft xs
  | vectorOK xs = G.map ((/fi (G.length xs)) . conjugate) . fft . G.map conjugate $ xs
  | otherwise   = error "Statistics.Transform.ifft: bad vector length"
{-# INLINABLE  ifft #-}
{-# SPECIAlIZE ifft :: U.Vector CD -> U.Vector CD #-}
{-# SPECIAlIZE ifft :: V.Vector CD -> V.Vector CD #-}

-- | Radix-2 decimation-in-time fast Fourier transform.
fft :: G.Vector v CD => v CD -> v CD
fft v | vectorOK v  = G.create $ do mv <- G.thaw v
                                    mfft mv
                                    return mv
      | otherwise   = error "Statistics.Transform.fft: bad vector length"
{-# INLINABLE  fft #-}
{-# SPECIAlIZE fft :: U.Vector CD -> U.Vector CD #-}
{-# SPECIAlIZE fft :: V.Vector CD -> V.Vector CD #-}

-- Vector length must be power of two. It's not checked
mfft :: (M.MVector v CD) => v s CD -> ST s ()
{-# INLINE mfft #-}
mfft vec = bitReverse 0 0
 where
  bitReverse i j | i == len-1 = stage 0 1
                 | otherwise  = do
    when (i < j) $ M.swap vec i j
    let inner k l | k <= l    = inner (k `shiftR` 1) (l-k)
                  | otherwise = bitReverse (i+1) (l+k)
    inner (len `shiftR` 1) j
  stage l !l1 | l == m    = return ()
              | otherwise = do
    let !l2 = l1 `shiftL` 1
        !e  = -6.283185307179586/fromIntegral l2
        flight j !a | j == l1   = stage (l+1) l2
                    | otherwise = do
          let butterfly i | i >= len  = flight (j+1) (a+e)
                          | otherwise = do
                let i1 = i + l1
                xi1 :+ yi1 <- M.read vec i1
                let !c = cos a
                    !s = sin a
                    d  = (c*xi1 - s*yi1) :+ (s*xi1 + c*yi1)
                ci <- M.read vec i
                M.write vec i1 (ci - d)
                M.write vec i (ci + d)
                butterfly (i+l2)
          butterfly j
    flight 0 0
  len = M.length vec
  m   = log2 len

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

fi :: Int -> CD
fi = fromIntegral

halve :: Int -> Int
halve = (`shiftR` 1)

vectorOK :: G.Vector v a => v a -> Bool
{-# INLINE vectorOK #-}
vectorOK v = (1 `shiftL` log2 n) == n where n = G.length v


{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-----------------------------------------------------------------------------
{-|
Module      : Math.Tensor.LinearAlgebra.Equations
Description : Linear tensor equations.
Copyright   : (c) Nils Alex, 2020
License     : MIT
Maintainer  : nils.alex@fau.de

Linear tensor equations.
-}
-----------------------------------------------------------------------------
module Math.Tensor.LinearAlgebra.Equations
  ( Equation
  , tensorToEquations
  , equationFromRational
  , equationsToSparseMat
  , equationsToMat
  , tensorsToSparseMat
  , tensorsToMat
  , systemRank
  , Solution
  , fromRref
  , fromRrefRev
  , fromRow
  , fromRowRev
  , applySolution
  , solveTensor
  , solveSystem
  , redefineIndets
  ) where

import Math.Tensor
  ( T
  , removeZerosT
  , toListT
  )
import Math.Tensor.LinearAlgebra.Scalar
  ( Poly (Const, Affine, NotSupported)
  , Lin (Lin)
  , polyMap
  , normalize
  )
import Math.Tensor.LinearAlgebra.Matrix
  ( rref
  , isrref
  , verify
  )

import qualified Numeric.LinearAlgebra.Data as HM
  ( Matrix
  , R
  , Z
  , fromLists
  , toLists
  )
import Numeric.LinearAlgebra (rank)

import Data.Maybe (mapMaybe)
import qualified Data.IntMap.Strict as IM
  ( IntMap
  , foldl'
  , map
  , assocs
  , null
  , findWithDefault
  , lookupMax
  , keys
  , fromList
  , (!)
  , difference
  , intersectionWith
  , mapKeys
  , empty
  )
import Data.List (nub, sort)
import Data.Ratio (numerator, denominator)

-- |A linear equation is a mapping from variable
-- indices to coefficients
type Equation a = IM.IntMap a

-- |Extract linear equations from tensor components.
-- The equations are normalized, sorted, and made unique.
tensorToEquations :: Integral a => T (Poly Rational) -> [Equation a]
tensorToEquations = nub . sort . filter (not . IM.null) . fmap (equationFromRational . normalize . snd) . toListT

-- |Extract linear equation with integral coefficients from polynomial
-- tensor component with rational coefficients.
-- Made made integral by multiplying with the @lcm@ of all denominators.
equationFromRational :: forall a.Integral a => Poly Rational -> Equation a
equationFromRational (Affine x (Lin lin))
    | x == 0 = lin'
    | otherwise = error "affine equation not supported for the moment!"
  where
    fac :: a
    fac = IM.foldl' (\acc v -> lcm (fromIntegral (denominator v)) acc) 1 lin
    lin' = IM.map (\v -> fromIntegral (numerator v) * (fac `div` fromIntegral (denominator v))) lin
equationFromRational (Const c)
    | c == 0    = IM.empty
equationFromRational _ = error "equation can only be extracted from linear scalar!"

-- |Convert list of equations to sparse matrix representation of the
-- linear system.
equationsToSparseMat :: [Equation a] -> [((Int,Int), a)]
equationsToSparseMat xs = concat $ zipWith (\i m -> fmap (\(j,v) -> ((i,j),v)) (IM.assocs m)) [1..] xs

-- |Convert list of equations to dense matrix representation of the
-- linear system.
equationsToMat :: Integral a => [Equation a] -> [[a]]
equationsToMat eqns = mapMaybe (\m -> if IM.null m
                                      then Nothing
                                      else Just $ fmap (\j -> IM.findWithDefault 0 j m) [1..maxVar]) eqns
  where
    maxVar = maximum $ mapMaybe (fmap fst . IM.lookupMax) eqns

-- |Extract sparse matrix representation for the linear system given
-- by a list of existentially quantified tensors with polynomial values.
tensorsToSparseMat :: Integral a => [T (Poly Rational)] -> [((Int,Int), a)]
tensorsToSparseMat = equationsToSparseMat . concatMap tensorToEquations

-- |Extract dense matrix representation for the linear system given
-- by a list of existentially quantified tensors with polynomial values.
tensorsToMat :: Integral a => [T (Poly Rational)] -> [[a]]
tensorsToMat = equationsToMat . concatMap tensorToEquations

-- |Rank of a linear system. Uses dense svd provided by hmatrix.
matRank :: forall a.Integral a => [[a]] -> Int
matRank []  = 0
matRank mat = rank (hmat :: HM.Matrix HM.R)
  where
    hmat = HM.fromLists $ fmap (fmap (fromIntegral @a @HM.R)) mat

-- |Rank of the linear system given by a list of existentially
-- quantified tensors with polynomial values.
systemRank :: [T (Poly Rational)] -> Int
systemRank = matRank . tensorsToMat @Int

-- |The solution to a linear system is represented as a list of
-- substitution rules, stored as @'IM.IntMap' ('Poly' 'Rational')@.
type Solution = IM.IntMap (Poly Rational)

-- |Read substitution rules from reduced row echelon form
-- of a linear system.
fromRref :: HM.Matrix HM.Z -> Solution
fromRref ref = IM.fromList assocs
  where
    rows   = HM.toLists ref
    assocs = mapMaybe fromRow rows

fromRrefRev :: HM.Matrix HM.Z -> Solution
fromRrefRev ref = IM.fromList assocs
  where
    rows   = fmap reverse $ HM.toLists ref
    assocs = mapMaybe fromRowRev rows

-- |Read single substitution rule from single
-- row of reduced row echelon form.
fromRow :: forall a.Integral a => [a] -> Maybe (Int, Poly Rational)
fromRow xs = case assocs of
               []             -> Nothing
               [(i,_)]        -> Just (i, Const 0)
               (i, v):assocs' -> let assocs'' = fmap (\(i',v') -> (i', - fromIntegral @a @Rational v' / fromIntegral @a @Rational v)) assocs'
                                 in Just (i, Affine 0 (Lin (IM.fromList assocs'')))
  where
    assocs = filter ((/=0). snd) $ zip [(1::Int)..] xs

fromRowRev :: forall a.Integral a => [a] -> Maybe (Int, Poly Rational)
fromRowRev xs = case assocs of
                  []             -> Nothing
                  [(i,_)]        -> Just (i, Const 0)
                  (i, v):assocs' -> let assocs'' = fmap (\(i',v') -> (i', - fromIntegral @a @Rational v' / fromIntegral @a @Rational v)) assocs'
                                    in Just (i, Affine 0 (Lin (IM.fromList assocs'')))
  where
    assocs = reverse $ filter ((/=0). snd) $ zip [(1::Int)..] xs

-- |Apply substitution rules to tensor component.
applySolution :: Solution -> Poly Rational -> Poly Rational
applySolution s (Affine x (Lin lin))
    | x == 0 = case p of
                 Affine xFin (Lin linFin) -> if IM.null linFin
                                             then Const xFin
                                             else p
                 _ -> p
    | otherwise = error "affine equations not yet supported"
  where
    s' = IM.intersectionWith (\row v -> if v == 0
                                        then error "value 0 encountered in linear scalar"
                                        else polyMap (v*) row) s lin
    lin' = IM.difference lin s
    p0 = if IM.null lin'
         then Const 0
         else Affine 0 (Lin lin')

    p = IM.foldl' (+) p0 s'

    {-
    lin' = IM.foldlWithKey'
             (\lin' i sub ->
                 case Affine 0 (Lin lin') + sub of
                   Affine 0 (Lin lin'') -> IM.delete i lin''
                   Const 0              -> IM.empty
                   _                    -> error "affine equations not yet supported")
             lin s'
    -}
applySolution _ _ = error "only linear equations supported"

-- |Apply substitution rules to all components of a tensor.
solveTensor :: Solution -> T (Poly Rational) -> T (Poly Rational)
solveTensor sol = removeZerosT . fmap (applySolution sol)

-- |Solve a linear system and apply solution to the tensorial
-- indeterminants.
solveSystem ::
     [T (Poly Rational)] -- ^ Tensorial linear system
  -> [T (Poly Rational)] -- ^ List of indeterminant tensors
  -> [T (Poly Rational)] -- ^ Solved indeterminant tensors
solveSystem system indets
    | wrongSolution = error "Wrong solution found. May be an Int64 overflow."
    | otherwise     = indets'
  where
    mat = HM.fromLists $ tensorsToMat @HM.Z system
    ref = rref mat
    wrongSolution = not (isrref ref && verify mat ref)
    sol = fromRref ref
    indets' = fmap (solveTensor sol) indets

-- |Relabelling of the indeterminants present in a list of tensors.
-- Redefines the labels of @n@ indeterminants as @[1..n]@, preserving
-- the previous order.
redefineIndets :: [T (Poly v)] -> [T (Poly v)]
redefineIndets indets = fmap (fmap (\case
                                       Const c -> Const c
                                       NotSupported -> NotSupported
                                       Affine a (Lin lin) ->
                                         Affine a (Lin (IM.mapKeys (varMap IM.!) lin)))) indets
  where
    comps = snd <$> concatMap toListT indets
    vars = nub $ concat $ mapMaybe (\case
                                       Affine _ (Lin lin) -> Just $ IM.keys lin
                                       _                  -> Nothing) comps
    varMap = IM.fromList $ zip vars [(1::Int)..]


{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE PolyKinds                 #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# OPTIONS_GHC -fwarn-incomplete-patterns #-}

module Math.Regression.Regression where

import qualified Math.HMatrixUtils             as HU

import qualified Colonnade                     as C

import qualified Data.List                     as List
import qualified Data.Text                     as T
import qualified Data.Vector.Storable          as V
import qualified Lucid                         as LH
import qualified Lucid.Colonnade               as LC
import qualified Statistics.Distribution       as S
import qualified Statistics.Distribution.FDistribution
                                               as S
import qualified Statistics.Distribution.StudentT
                                               as S
import qualified Statistics.Types              as S
import qualified Text.Blaze.Colonnade          as BC
import qualified Text.Blaze.Html5              as BH
import qualified Text.Blaze.Html5.Attributes   as BHA
import qualified Text.Printf                   as TP

import           Numeric.LinearAlgebra          ( (#>)
                                                , (<.>)
                                                )
import qualified Numeric.LinearAlgebra         as LA
import           Numeric.LinearAlgebra.Data     ( Matrix
                                                , R
                                                , Vector
                                                )
import qualified Say
-- NB: for normal errors, we will request the ci and get it.  For interval errors, we may request it but we can't produce
-- what we request, we can only produce what we have.  So we make an interface for both.  We take CL as input to predict
-- *and* return it as output, in the prediction.

class Predictor e a b p where
  predict :: p -> b -> S.Estimate e a -- do we need a predict :: (b, b) -> (a, a) for uncertain inputs?

predictFromEstimateAtConfidence
  :: (RealFrac a, Predictor S.NormalErr a b p)
  => Double
  -> p
  -> S.CL Double
  -> b
  -> (a, a)
predictFromEstimateAtConfidence dof p cl b =
  let
    S.Estimate pt (S.NormalErr sigma) = predict p b
    prob =
      S.confidenceLevel $ S.mkCLFromSignificance (S.significanceLevel cl / 2)
    predCI =
      2 * S.quantile (S.studentTUnstandardized dof 0 (realToFrac sigma)) prob
  in
    (pt, realToFrac predCI)

data RegressionResult a = RegressionResult
                          {
                            parameterEstimates :: [S.Estimate S.NormalErr a]
                          , degreesOfFreedom   :: Double -- since N may be an effective N
                          , meanSquaredError   :: a
                          , rSquared           :: Double
                          , adjRSquared        :: Double
                          , fStatistic         :: Maybe Double -- we can't compute this for fits with one parameter
                          , covariances        :: Matrix Double
                          } deriving (Show)

instance (LA.Element a, LA.Numeric a, RealFloat a) => Predictor S.NormalErr a (Vector a) (RegressionResult a) where
  predict rr va =
    let mse = meanSquaredError rr
        sigmaPred = sqrt (mse + va <.> (LA.cmap realToFrac (covariances rr) #> va))
        prediction = va LA.<.> V.fromList (S.estPoint <$> parameterEstimates rr)
    in S.estimateNormErr prediction sigmaPred

data NamedEstimate a = NamedEstimate { regressorName :: T.Text, regressorEstimate :: a, regressorCI :: a, regressorPValue :: a }

namedEstimates :: [T.Text] -> RegressionResult R -> S.CL R -> [NamedEstimate R]
namedEstimates pNames res cl =
  let prob = S.confidenceLevel
        $ S.mkCLFromSignificance (S.significanceLevel cl / 2)
      dof = degreesOfFreedom res
      sigma e = S.normalError $ S.estError e
      dist e = S.studentTUnstandardized dof 0 (sigma e)
      ci e = 2 * S.quantile (dist e) prob
      pValue e = S.complCumulative (dist e) (abs $ S.estPoint e)
  in  List.zipWith (\n e -> NamedEstimate n (S.estPoint e) (ci e) (pValue e))
                   pNames
                   (parameterEstimates res)

namedEstimatesColonnade
  :: S.CL R -> C.Colonnade C.Headed (NamedEstimate R) T.Text
namedEstimatesColonnade cl =
  C.headed "parameter" regressorName
    <> C.headed "estimate" (T.pack . TP.printf "%4.3f" . regressorEstimate)
    <> C.headed
         (  T.pack (TP.printf "%.0f" (100 * S.confidenceLevel cl))
         <> "% confidence"
         )
         (T.pack . TP.printf "%4.3f" . regressorCI)
    <> C.headed "p-value" (T.pack . TP.printf "%4.3f" . regressorPValue)

namedSummaryStats :: RegressionResult R -> [(T.Text, T.Text)]
namedSummaryStats r =
  let
    p             = List.length (parameterEstimates r)
    effN          = degreesOfFreedom r + realToFrac p
    d1            = realToFrac $ p - 1
    d2            = effN - realToFrac p
    pValM = fmap (S.complCumulative (S.fDistributionReal d1 d2)) (fStatistic r)
    printNum      = T.pack . TP.printf "%4.3f"
    printMaybeNum = maybe "N/A" printNum
  in
    [ ("R-Squared"         , printNum (rSquared r))
    , ("Adj. R-squared"    , printNum (adjRSquared r))
    , ("F-stat"            , printMaybeNum (fStatistic r))
    , ("p-value"           , printMaybeNum pValM)
    , ("Mean Squared Error", printNum (meanSquaredError r))
    ]

namedSummaryStatsColonnade :: C.Colonnade C.Headed (T.Text, T.Text) T.Text
namedSummaryStatsColonnade =
  C.headed "Summary Stat." fst <> C.headed "Value" snd

prettyPrintRegressionResult
  :: T.Text -> [T.Text] -> RegressionResult R -> S.CL R -> T.Text
prettyPrintRegressionResult header xNames r cl =
  let nEsts = namedEstimates xNames r cl
      nSS   = namedSummaryStats r
  in  header
      <> "\n"
      <> T.pack (C.ascii (T.unpack <$> namedEstimatesColonnade cl) nEsts)
      <> T.pack (C.ascii (fmap T.unpack namedSummaryStatsColonnade) nSS)

prettyPrintRegressionResultLucid
  :: T.Text -> [T.Text] -> RegressionResult R -> S.CL R -> LH.Html ()
prettyPrintRegressionResultLucid header xNames r cl = do
  let nEsts = namedEstimates xNames r cl
      nSS   = namedSummaryStats r
      toCell t = LC.Cell [LH.style_ "border: 1px solid black"] (LH.toHtmlRaw t)
  LH.div_
      [ LH.style_
          "display: inline-block; padding: 7px; border-collapse: collapse"
      ]
    $ do
        LH.span_ (LH.toHtmlRaw header)
        LC.encodeCellTable
          [LH.style_ "border: 1px solid black; border-collapse: collapse"]
          (toCell <$> namedEstimatesColonnade cl)
          nEsts
        LC.encodeCellTable
          [LH.style_ "border: 1px solid black; border-collapse: collapse"]
          (fmap toCell namedSummaryStatsColonnade)
          nSS

prettyPrintRegressionResultBlaze
  :: T.Text -> [T.Text] -> RegressionResult R -> S.CL R -> BH.Html
prettyPrintRegressionResultBlaze header xNames r cl = do
  let nEsts = namedEstimates xNames r cl
      nSS   = namedSummaryStats r
      toCell t = BC.Cell (BHA.style "border: 1px solid black") (BH.toHtml t)
  BH.div
    BH.! BHA.style
           "display: inline-block; padding: 7px; border-collapse: collapse"
    $    do
           BH.span (BH.toHtml header)
           BC.encodeCellTable
             (BHA.style "border: 1px solid black; border-collapse: collapse")
             (toCell <$> namedEstimatesColonnade cl)
             nEsts
           BC.encodeCellTable
             (BHA.style "border: 1px solid black; border-collapse: collapse")
             (fmap toCell namedSummaryStatsColonnade)
             nSS

data FitStatistics a = FitStatistics { fsRSquared :: a, fsAdjRSquared :: a, fsFStatistic :: Maybe a}

goodnessOfFit
  :: MonadIO m
  => Int
  -> Vector R
  -> Maybe (Vector R)
  -> Vector R
  -> m (FitStatistics R)
goodnessOfFit pInt vB vWM vU = do
  let
    n  = LA.size vB
    p  = realToFrac pInt
    vW = maybe (V.fromList $ List.replicate n (1.0 / realToFrac n))
               (\v -> LA.scale (1 / realToFrac (LA.sumElements v)) v)
               vWM
    mW     = LA.diag vW
    vWB    = mW #> vB
--      vWU = mW #> vU
    meanWB = LA.sumElements vWB
    effN   = 1 / (vW <.> vW)
    ssTot  = let x = LA.cmap (\y -> y - meanWB) vB in x <.> (mW #> x)
    ssRes  = vU <.> (mW #> vU) -- weighted ssq of residuals
    rSq    = 1 - (ssRes / ssTot)
    arSq =
      1 - (1 - rSq) * realToFrac (effN - 1) / realToFrac (effN - p - 1)
    fStatM = if pInt > 1
      then Just (((ssTot - ssRes) / (p - 1.0)) / (ssRes / (effN - p)))
      else Nothing
  Say.say $ "n=" <> show n
  Say.say $ "p=" <> show p
--  Log.log Log.Diagnostic $ "vW=" <> show vW
  Say.say $ "effN=" <> show effN
  Say.say $ "ssTot=" <> show ssTot
  Say.say $ "ssRes=" <> show ssRes
  return $ FitStatistics rSq arSq fStatM

estimates :: Matrix R -> Vector R -> [S.Estimate S.NormalErr R]
estimates cov means =
  let sigmas = LA.cmap sqrt (LA.takeDiag cov)
  in  List.zipWith S.estimateNormErr
                   (LA.toList means)
                   (LA.toList sigmas)

eickerHeteroscedasticityEstimator
  :: MonadIO m
  => Matrix R
  -> Vector R
  -> Vector R
  -> m (Matrix R)
eickerHeteroscedasticityEstimator mA vB vB' = do
  HU.checkVectorMatrix "b" "A" vB mA
  HU.checkVectorMatrix "b'" "A" vB' mA
  let u     = vB - vB'
      diagU = LA.diag $ V.zipWith (*) u u
      mA'   = LA.tr mA
      mC    = LA.inv (mA' LA.<> mA)
  return $ mC LA.<> (mA' LA.<> diagU LA.<> mA) LA.<> mC


module Model where

import Control.Monad.State.Strict
--import Control.Monad.State
import System.Random (StdGen, mkStdGen, randomRs)
import Numeric.LinearAlgebra (Matrix(..), (><), toLists)

type Ising = State IsingState
type Spin = Double
type Model = Matrix Spin

data IsingState = IsingState
     { dim :: Int
     , j :: Double
     , h :: Double
     , step :: Int
     , nAccept :: Int
     , propFreq :: Int
     , model :: Model
     , rng :: StdGen
     }

instance Show IsingState where
    show state =
        "IsingState { d = " ++ show (dim state)
            ++ ", j = " ++ show (j state)
            ++ ", step = " ++ show (step state)
            ++ ", accept = " ++ show (nAccept state)
            ++ ", rng = " ++ show (rng state)
            ++ " }"
            ++ showModel (model state)

showModel :: Model -> String
showModel m =
    let rows = toLists m
        showRow r = "|" ++ foldl (\acc s -> acc ++ showSpin s) "" r ++ "|"
        showSpin s = if s == -1 then " " else "*"
        stringRows = map showRow rows
    in foldl (\acc r -> acc ++ "\n" ++ r) "" stringRows

newModel :: Int -> Int -> Double -> [Int] -> IsingState
newModel seed n j spins =
    let rng = mkStdGen seed
        model = (n><n) $ map fromIntegral spins
    in IsingState
        { dim = n
        , j = j
        , h = 0.0
        , step = 0
        , nAccept = 0
        , propFreq = 1000
        , model = model
        , rng = rng
        }

randModel seed n j =
    let r = mkStdGen seed
        spins = take (n*n) $ randomRs (0,1) r :: [Int]
    in newModel seed n j [if s == 0 then -1 else s | s <- spins]

upSpins   n = [ 1 | _ <- [1..n*n]]
downSpins n = [-1 | _ <- [1..n*n]]

setPropertyFreq :: Int -> IsingState -> IsingState
setPropertyFreq f model = model { propFreq = f }

setH :: Double -> IsingState -> IsingState
setH h model = model { h = h }


module Statistics.Covariance.Shrinkage.Strimmer (
    psmallSVD
  , nsmallSVD
  , posSVD
  , covShrink
  , corShrink
  , corCoef
  , varShrink
  , varCoef
  ) where

import Numeric.LinearAlgebra hiding ((<>))
import Numeric.Sum (kbn, sumVector)
import Statistics.Quantile (quantile, s)
import qualified Data.Vector.Storable as VS

covShrink :: Matrix Double -> Herm Double
covShrink x = trustSym $ (asColumn sc * unSym c) * asRow sc
  where
    sc = sqrt $ varShrink x
    c = corShrink x

psmallSVD :: Matrix Double -> (Matrix Double, Vector Double, Matrix Double)
psmallSVD x = (u', d', v')
  where
    b = unSym $ mTm x
    (_, d, v) = svd b
    tol = 1e-6
    positive = find (> tol) d
    u' = (x <> v') * asRow (recip d')
    d' = sqrt . vector $ map (atIndex d) positive
    v' = v ¿ positive

nsmallSVD :: Matrix Double -> (Matrix Double, Vector Double, Matrix Double)
nsmallSVD x = (u', d', v')
  where
    b = unSym $ mTm (tr x)
    (u, d, _) = svd b
    tol = 1e-6
    positive = find (> tol) d
    u' = u ¿ positive
    d' = sqrt . vector $ map (atIndex d) positive
    v' = (tr x <> u') * asRow (recip d')

posSVD :: Matrix Double -> (Matrix Double, Vector Double, Matrix Double)
posSVD x = (u', d', v')
  where
    (u, d, v) = svd x
    tol = 1e-6
    positive = find (> tol) d
    u' = u ¿ positive
    d' = vector $ map (atIndex d) positive
    v' = v ¿ positive

fastSVD :: Matrix Double -> (Matrix Double, Vector Double, Matrix Double)
fastSVD x
  | n > edgeRatio * p = psmallSVD x
  | edgeRatio * n < p = nsmallSVD x
  | otherwise = posSVD x
  where
    edgeRatio = 2
    n = rows x
    p = cols x

mpower :: Herm Double -> Double -> Herm Double
mpower x p = trustSym $ (evec * asRow (cmap exp $ scale p (cmap log eval))) <> tr evec
  where (eval, evec) = eigSH x

corShrink :: Matrix Double -> Herm Double
corShrink x
  | lambda == 1 || alpha == 0 = trustSym $ ident p
  | alpha == 1 = unitDiag $ scale (1 - lambda) r0
  | lambda == 0 = unitDiag . trustSym $ svdxsV <> unSym (mpower c2 alpha) <> tr svdxsV
  | otherwise = unitDiag . trustSym . scale (exp (alpha * log lambda)) $ ident p - (svdxsV <> f <> tr svdxsV)
  where
    unitDiag :: Herm Double -> Herm Double
    unitDiag powr = trustSym (diag (1 - takeDiag (unSym powr))) `add` powr
    n = rows x
    p = cols x
    alpha = 1 :: Double
    w2 = recip $ fromIntegral n -- w = 1 / n => sum(w^2) =  1 / n
    h1 = 1 / (1 - w2)
    xs = let (u, v) = meanCov x in (x - asRow u) / asRow (sqrt (takeDiag (unSym v)))
    lambda = corCoef x
    r0 = scale h1 $ mTm (xs / sqrt (fromIntegral n))
    (svdxsU, svdxsS, svdxsV) = fastSVD xs
    m = size svdxsS
    utwu = tr svdxsU <> svdxsU / fromIntegral n
    c0 = utwu * asRow svdxsS * asColumn svdxsS
    c1 = scale ((1 - lambda) * h1) c0
    c2 = sym c1
    f = ident m - unSym (mpower (trustSym $ unSym (scale (recip lambda) c2) + ident m) alpha)

corCoef :: Matrix Double -> Double
corCoef x
  | denominator == 0 = 1
  | otherwise = numerator / denominator * h1w2
  where
    n = rows x
    p = cols x
    xs = let (u, v) = meanCov x in (x - asRow u) / asRow (sqrt (takeDiag (unSym v)))
    w2 = recip $ fromIntegral n
    h1w2 = w2 / (1 - w2)
    sw = sqrt . recip $ fromIntegral n
    xsw = scale sw xs
    (xswsvdU, xswsvdS, xswsvdV) = fastSVD xsw
    colSums = vector . map (sumVector kbn) . toColumns
    -- colCumSums = fromColumns . map (VS.postscanl (+) 0) . toColumns
    rowCumSums = fromRows . map (VS.postscanl (+) 0) . toRows
    sE2R = sumVector kbn (flatten (xsw * ((xswsvdU * asRow (cb xswsvdS)) <> tr xswsvdV)))
      - sumVector kbn (sq (colSums (sq xsw)))
    xs2w = scale sw $ sq xs
    sER2 = 2 * (sumVector kbn . flatten
      $ fliprl (takeColumns (p - 1) xs2w) * (rowCumSums . fliprl $ dropColumns 1 xs2w))
    denominator = sE2R
    numerator = sER2 - sE2R

varShrink :: Matrix Double -> Vector Double
varShrink x = scalar (lambda * target) + scale (1 - lambda) v
  where
    lambda = varCoef x
    v = takeDiag . unSym . snd $ meanCov x
    target = quantile s 1 2 v

varCoef :: Matrix Double -> Double
varCoef x
  | denominator == 0 = 1
  | otherwise = max 0 . min 1 $ numerator / denominator * h1w2
  where
    n = rows x
    w2 = recip (fromIntegral n)
    h1 = 1 / (1 - w2)
    h1w2 = w2 / (1 - w2)
    xc = x - asRow (fst (meanCov x))
    v = scale h1 q1
    target = quantile s 1 2 v
    zz = xc * xc
    q1 = fst $ meanCov zz
    q2 = fst (meanCov (zz * zz) ) - q1 * q1
    numerator = sumVector kbn q2
    denominator = sumVector kbn . sq $ q1 - scalar (target / h1)

sq, cb :: Num a => a -> a
sq a = a * a
cb a = a * sq a


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

module Numeric.Optimization.TwoPhase.Tableau.Pivot
    ( index
    , allCells
    , pivot
    ) where

import           Data.Finite                        (Finite)
import           GHC.TypeLits                       (KnownNat)
import           Numeric.LinearAlgebra              ((!))
import qualified Numeric.LinearAlgebra as LA        (atIndex)
import           Numeric.LinearAlgebra.Static       (L)
import qualified Numeric.LinearAlgebra.Static as LS

type Index rows cols = (Finite rows, Finite cols)

index :: (KnownNat rows, KnownNat cols)
    => Index rows cols
    -> L rows cols
    -> Double
index (i, j) x =
    LA.atIndex (LS.unwrap x) ix
  where
    ix = (fromIntegral i, fromIntegral j)

allCells :: (KnownNat rows, KnownNat cols)
    => (Double -> Bool)
    -> [Index rows cols]
    -> L rows cols
    -> Bool
allCells f is x =
    all f $ fmap (`index` x) is

pivot :: (KnownNat rows, KnownNat cols)
    => (Finite rows, Finite cols) -> L rows cols -> L rows cols
pivot (x', y') table =
    LS.build (newValue $ LS.unwrap table)
  where
    (x, y)          = (fromIntegral x', fromIntegral y')

    newValue mat i' j'
        | i == x    = mat!i!j / mat!x!y
        | j == y    = 0
        | otherwise = mat!x!y * mat!i!j - mat!x!j * mat!i!y
      where
        (i, j) = (round i', round j')


{-# OPTIONS_GHC -Wall                     #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing  #-}
{-# OPTIONS_GHC -fno-warn-type-defaults   #-}
{-# OPTIONS_GHC -fno-warn-unused-do-bind  #-}
{-# OPTIONS_GHC -fno-warn-missing-methods #-}
{-# OPTIONS_GHC -fno-warn-orphans         #-}

{-# LANGUAGE DataKinds                    #-}
{-# LANGUAGE ScopedTypeVariables          #-}
{-# LANGUAGE TypeOperators                #-}
{-# LANGUAGE TypeFamilies                 #-}

-- | An extended Kalman filter. Note that this could be generalized
-- further. If you need a Kalman filter and have a state update matrix
-- \(A\) then you can just use @const bigA@ as the function which
-- returns the Jacobian.

module Kalman (
  extKalman
  )where

import GHC.TypeLits
import Numeric.LinearAlgebra.Static hiding ( create )
import Data.Maybe ( fromJust )

-- | The model for the extended Kalman filter is given by
--
-- \[
-- \begin{aligned}
-- \boldsymbol{x}_i &= \boldsymbol{a}(\boldsymbol{x}_{i-1}) + \boldsymbol{\psi}_{i-1} \\
-- \boldsymbol{y}_i &= {H}\boldsymbol{x}_i + \boldsymbol{\upsilon}_i
-- \end{aligned}
-- \]
--
-- where
--
-- * \(\boldsymbol{a}\)\ is some non-linear vector-valued state update
-- function.
--
-- * \(\boldsymbol{\psi}_{i}\) are independent identically normally
-- distributed random variables with mean 0 representing the fact that
-- the state update is noisy: \(\boldsymbol{\psi}_{i} \sim {\cal{N}}(0,Q)\).
--
-- * \({H}\) is a matrix which represents how we observe the
-- hidden state process.
--
-- * \(\boldsymbol{\upsilon}_i\) are independent identically normally
-- distributed random variables with mean 0 represent the fact that
-- the observations are noisy: \(\boldsymbol{\upsilon}_{i} \sim {\cal{N}}(0,R)\).
--
-- We assume the whole process starts at 0 with our belief of the state (aka the
-- prior state) being given by
-- \(\boldsymbol{x}_0 \sim {\cal{N}}(\boldsymbol{\mu}_0, \Sigma_0)\)
extKalman ::  forall m n .
              (KnownNat m, KnownNat n, (1 <=? n) ~ 'True, (1 <=? m) ~ 'True) =>
              R n              -- ^ Prior mean \(\boldsymbol{\mu}_0\)
              -> Sq n          -- ^ Prior variance \(\Sigma_0\)
              -> L m n         -- ^ Observation map \(H\)
              -> Sq m          -- ^ Observation noise \(R\)
              -> (R n -> R n)  -- ^ State update function \(\boldsymbol{a}\)
              -> (R n -> Sq n) -- ^ A function which returns the
                               -- Jacobian of the state update
                               -- function at a given point
                               -- \(\frac{\partial \boldsymbol{a}}{\partial \boldsymbol{x}}\)
              -> Sq n          -- ^ State noise \(Q\)
              -> [R m]         -- ^ List of observations \(\boldsymbol{y}_i\)
              -> [(R n, Sq n)] -- ^ List of posterior means and variances \((\hat{\boldsymbol{x}}_i, \hat{\Sigma}_i)\)
extKalman muPrior sigmaPrior bigH bigSigmaY
  littleA bigABuilder bigSigmaX ys = result
  where
    result = scanl update (muPrior, sigmaPrior) ys

    update :: (R n, Sq n) -> R m -> (R n, Sq n)
    update (xHatFlat, bigSigmaHatFlat) y =
      (xHatFlatNew, bigSigmaHatFlatNew)
      where

        v :: R m
        v = y - (bigH #> xHatFlat)

        bigS :: Sq m
        bigS = bigH <> bigSigmaHatFlat <> (tr bigH) + bigSigmaY

        bigK :: L n m
        bigK = bigSigmaHatFlat <> (tr bigH) <> (inv bigS)

        xHat :: R n
        xHat = xHatFlat + bigK #> v

        bigSigmaHat :: Sq n
        bigSigmaHat = bigSigmaHatFlat - bigK <> bigS <> (tr bigK)

        bigA :: Sq n
        bigA = bigABuilder xHat

        xHatFlatNew :: R n
        xHatFlatNew = littleA xHat

        bigSigmaHatFlatNew :: Sq n
        bigSigmaHatFlatNew = bigA <> bigSigmaHat <> (tr bigA) + bigSigmaX

inv :: (KnownNat n, (1 <=? n) ~ 'True) => Sq n -> Sq n
inv m = fromJust $ linSolve m eye


{-# LANGUAGE DataKinds #-}

module Rotation
  ( Rotation
  , rotationMatrix
  , rotate
  , mkRotation
  , mkXYRotation
  , mkXYRotationCamera
  , rotateV
  ) where

import qualified Numeric.LinearAlgebra        as L
import           Numeric.LinearAlgebra.Static

import           Matrix

type Quaternion = R 4

type Rotation = Quaternion

rotate :: Quaternion -> Quaternion -> Quaternion
q1 `rotate` q2 =
  vec4
    (r1 * r2 - i1 * i2 - j1 * j2 - k1 * k2)
    (r1 * i2 + i1 * r2 + j1 * k2 - k1 * j2)
    (r1 * j2 + j1 * r2 + k1 * i2 - i1 * k2)
    (r1 * k2 + k1 * r2 + i1 * j2 - j1 * i2)
  where
    [r1, i1, j1, k1] = L.toList . unwrap $ q1
    [r2, i2, j2, k2] = L.toList . unwrap $ q2

mkRotation :: Double -> Vector3 -> Quaternion
mkRotation angle axis =
  normalize $
  vec4
    (cos (angle / 2))
    (x * sin (angle / 2))
    (y * sin (angle / 2))
    (z * sin (angle / 2))
  where
    [x, y, z] = L.toList . unwrap $ axis

rotationMatrix :: Quaternion -> Matrix4
rotationMatrix q = eye + 2 * m
  where
    [r, i, j, k] = L.toList . unwrap $ q
    m =
      row (vec4 (-j * j - k * k) (i * j - k * r) (i * k + j * r) 0) ===
      row (vec4 (i * j + k * r) (-i * i - k * k) (j * k - i * r) 0) ===
      row (vec4 (i * k - j * r) (j * k + i * r) (-i * i - j * j) 0) ===
      row (vec4 0 0 0 0)

mkXYRotationCamera x y = q2 `rotate` q1
  where
    q1 = mkRotation x yHat
    q2 = mkRotation y xHat

mkXYRotation x y = q2 `rotate` q1
  where
    q1 = mkRotation y xHat
    q2 = mkRotation x yHat

conj :: Quaternion -> Quaternion
conj q = q * vec4 1 (-1) (-1) (-1)

qmult = rotate

rotateV :: Vector3 -> Quaternion -> Vector3
rotateV v q = v'
  where
    (_, v') = split $ q `qmult` p `qmult` conj q
    p = vector [0] # v


module Newton.Render where

-- this module comes from JuicyPixels package
import qualified Codec.Picture            as J
import           Codec.Picture.ColorQuant (defaultPaletteOptions)
import           Data.Complex
import           Newton.Utils             (root)

-- | Add brightness to an RGB8 Image
brightnessRGB8 :: Int -> J.Palette -> J.Palette
brightnessRGB8 add = J.pixelMap brightFunction
  where
    up v = fromIntegral (fromIntegral v + add)
    brightFunction (J.PixelRGB8 r g b) = J.PixelRGB8 (up r) (up g) (up b)

-- | Save an PNG image as a file
saveBasinsPic ::
     FilePath -- ^ Where to write the image file.
  -> J.Palette
  -> IO ()
saveBasinsPic = J.writePng

saveBasinsAnim ::
     FilePath -- ^ Filename of an image
  -> J.GifDelay -- ^ Delay between frames
  -> [J.Palette] -- ^ List of pictures to store within the GIF
  -> IO ()
saveBasinsAnim path delay images =
  case J.writeGifImages path J.LoopingForever finalImgs of
    Left msg  -> error msg
    Right ans -> ans
  where
    palettized = map (J.palettize defaultPaletteOptions) images
    finalImgs = map (\(img, plt) -> (plt, delay, img)) palettized

-- | Render attraction basins for zeros (roots) of a complex function
-- as a PNG image.
renderBasins ::
     (Fractional a, RealFloat a)
  => Int -- ^ Size of the square image in pixels
  -> (Complex a -> Complex a) -- ^ A function f to find zero for.
  -> (Complex a -> Complex a) -- ^ A derivative of function f.
  -> (Complex a -> J.PixelRGB8) -- ^ How to color roots.
  -> J.Palette
renderBasins commonSize f f' color = image
  where
    image = J.generateImage render width height
    width = commonSize -- width in pixels
    height = commonSize -- height in pixels
    (xfrom, xto) = (-2, 2) -- range for the real axis
    (yfrom, yto) = (-2, 2) -- range for the imaginary axis
    -- | Convert pixel coordinates to a complex number.
    pixelToComplex (i, j) = re :+ im
      where
        re = xfrom + (xto - xfrom) * fromIntegral i / fromIntegral width
        im = yfrom + (yto - yfrom) * fromIntegral j / fromIntegral height
    n = 100 -- maximum number of iterations
    p e = magnitude e < 1e-7 -- precision tolerance
    -- | Convert actual number of interations into
    -- into a grey pixel using logarithmic scale
    -- (this helps bring out more details in the image).
    logscale :: Integral a => Int -> a -> a
    logscale k c = floor (fromIntegral c * (1 - logBase n' k'))
      where
        n' = fromIntegral n + 1
        k' = fromIntegral k + 1
    -- | Render one pixel.
    render :: Int -> Int -> J.PixelRGB8
    render i j =
      case root n p f f' (pixelToComplex (i, j)) of
        Nothing -> J.PixelRGB8 0 0 0 -- black
        Just (k, xr) ->
          let J.PixelRGB8 r g b = color xr
           in J.PixelRGB8 (logscale k r) (logscale k g) (logscale k b)

-- | Simple (unstructured) complex number coloring.
simpleColoring :: Complex Double -> J.PixelRGB8
simpleColoring z = J.PixelRGB8 r g b
  where
    r = floor ((2 ^ 8 - 1) * r')
    g = floor ((2 ^ 8 - 1) * g')
    b = floor ((2 ^ 8 - 1) * b')
    (r', g', b') = hsl (phase z) 1 (1 - 0.5 ** magnitude z)

-- | Convert HSL color to RGB.
-- This function is copied with slight changes
-- from http://hackage.haskell.org/package/colour-2.3.4
hsl :: (RealFrac a, Floating a, Ord a) => a -> a -> a -> (a, a, a)
hsl h s l = (component tr, component tg, component tb)
  where
    hk = h / (2 * pi)
    (tr, tg, tb) = (mod1 (hk + 1 / 3), mod1 hk, mod1 (hk - 1 / 3))
    q
      | l < 0.5 = l * (1 + s)
      | otherwise = l + s - l * s
    p = 2 * l - q
    component t
      | t < 1 / 6 = p + ((q - p) * 6 * t)
      | t < 1 / 2 = q
      | t < 2 / 3 = p + ((q - p) * 6 * (2 / 3 - t))
      | otherwise = p
    mod1 x =
      if pf < 0
        then pf + 1
        else pf
      where
        (_, pf) = properFraction x


{-# LANGUAGE BangPatterns #-}
{- |
   Description :  Functions to work with transition probability matrices on rooted trees
   Copyright   :  (c) Dominik Schrempf 2017
   License     :  GPLv3

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

Calculate transition probability matrices, map rate matrices on trees, populate
a tree with states according to a stationary distribution, etc.

The implementation of the Markov process is more than basic and can be improved in a lot of ways.

TODO: Try https://hackage.haskell.org/package/Dist-0.4.1.0/docs/Numeric-Probability-Distribution.html

* Changelog

-}

module Transition where
import           Control.Monad.Random.Strict hiding (fromList)
import           Data.Distribution           hiding (toList)
import           Numeric.LinearAlgebra       hiding (fromList)
import           RateMatrix
import           RTree

-- | This may be moved to a different module ProbMatrix or alike.
type ProbMatrix   = Matrix R

-- | The important matrix that gives the probabilities to move from one state to
-- another in a specific time (branch length).
probMatrix :: RateMatrix -> BranchLn -> ProbMatrix
probMatrix m t = exponential
  where !exponential = expm $ scale t m

-- | Convert a tree with branch lengths into a tree that has transition
-- probability matrices assigned to each of its branches.
branchLengthsToTransitionProbs :: RateMatrix -> RTree a Double -> RTree a ProbMatrix
branchLengthsToTransitionProbs m = fmap (probMatrix m)

-- | Move from a given state to a new one according to a transition probability
-- matrix (for performance reasons this probability matrix needs to be given as
-- a list of generators, see
-- https://hackage.haskell.org/package/distribution-1.1.0.0/docs/Data-Distribution-Sample.html).
-- This function is the bottleneck of the simulator and takes up most of the
-- computation time. However, I was not able to find a faster implementation
-- than the one from Data.Distribution.
jump :: (MonadRandom m) => State -> [Generator State] -> m State
jump (State s) p = target
  where !target = getSample $ p !! s

-- | Perform N jumps from a given state and according to a transition
-- probability matrix transformed to a list of generators. This implementation
-- uses 'foldM' and I am not sure how to access or store the actual chain. This
-- could be done by an equivalent of 'scanl' for general monads, which I was
-- unable to find. This function is neat, but will most likely not be needed.
-- However, it is instructive and is left in place.
jumpN :: (MonadRandom m) => State -> [Generator State] -> Int -> m State
jumpN s p n = foldM jump s (replicate n p)

-- | This is the heart of the simulation. Take a tree (most probably with node
-- labels of type a, but this could be anything) and a root state. If there is a
-- split (i.e., if we have a node and not a leaf), jump down the left branch and
-- populate and flatten the tree and jump down the right branch and populate and
-- flatten the tree and append the two results. This is what 'liftM2' is doing,
-- it lifts the append function of lists (++) to the 'Rand g' monad. If we
-- encounter a leaf, just return the node label (or whatever the type a is) and
-- the state that we ended up at.
populateAndFlattenTree :: (MonadRandom m) => RTree a [Generator State] -> State -> m [(a, State)]
populateAndFlattenTree (Leaf a) s = return [(a, s)]
populateAndFlattenTree (Node _ lp lc rp rc) s = liftM2 (++) (jumpDownBranch lp lc) (jumpDownBranch rp rc)
  where jumpDownBranch p t = jump s p >>= populateAndFlattenTree t

-- | Convert a stationary distribution to a generator which speeds up random
-- picks. Cf. 'jump'.
stationaryDistToGenerator :: StationaryDist -> Generator State
stationaryDistToGenerator f = fG
  -- This is a little complicated. I need to convert the vector to a list to be
  -- able to create a distribution.
  where !fL = toList f
        !fD = fromList $ zip (map State [0..]) fL
        !fG = fromDistribution fD

-- | See 'jump'.
treeProbMatrixToTreeGenerator :: RTree a ProbMatrix -> RTree a [Generator State]
treeProbMatrixToTreeGenerator t = tG
  where
    -- Create a tree with the probability matrices as list of row vectors.
    !tL = fmap toLists t
    -- A complicated double map. We need to create generators for each branch on
    -- the tree (fmap) and for each target state on each branch (map).
    !tG = (fmap . map) (fromDistribution . fromList . zip (map State [0..])) tL

-- | Simulate data (states at the leaves) for a tree with transition
-- probabilities on its branches and with the stationary distribution of states
-- at the root. The state at the root is randomly chosen from the stationary
-- distribution and the states at the nodes and leaves are randomly chosen
-- according to the transition probabilities.
simulateSite :: (MonadRandom m) =>
                Generator State
             -> RTree a [Generator State]
             -> m [(a, State)]
simulateSite f t = do
  !rootState <- getSample f
  populateAndFlattenTree t rootState

-- | Randomly draw an index according to a given generator. Use the stationary
-- distribution and rooted tree at the drawn index to simulate a site. This is
-- useful for simulation, e.g., Gamma rate heterogeneity models.
simulateSiteGen :: (MonadRandom m) =>
                   Generator Int
                -> [Generator State]
                -> [RTree a [Generator State]]
                -> m [(a, State)]
simulateSiteGen gen fs trs = do
  !i <- getSample gen
  let !f = fs  !! i
      !t = trs !! i
  simulateSite f t


{-
 - *Main Data.Complex> h = concat [[1..5], replicate 11 0]
 - *Main Data.Complex> x = concat [[1..10], replicate 6 0]
 - *Main Data.Complex> convFft h x
 - [1.0,4.000000000000007,10.0,20.0,35.0,49.99999999999999,65.0,80.0,95.0,110.0,114.0,106.0,85.0,50.00000000000001,7.105427357601002e-15,7.105427357601002e-15]
 - *Main Data.Complex>
 -}
import Data.Complex

convOne :: (Num a) => [a] -> [a] -> Int -> a
convOne h x n    = sum (zipWith (*) kernel input)
    where kernel = drop (-begin) (reverse h)
          input  = drop begin x
          begin  = n - length h + 1

conv :: (Num a) => [a] -> [a] -> [a]
conv h x = [convOne h x n | n <- [0..(length x + length h - 2)]]

isPowerOfTwo :: Int -> Bool
isPowerOfTwo 1 = True
isPowerOfTwo x
    | mod x 2 == 0 = isPowerOfTwo (div x 2)
    | otherwise = False

fftRaw :: (RealFloat a) => [a] -> Int -> Int -> [Complex a]
fftRaw _ 0 _ = []
fftRaw [] _  _ = []
fftRaw (x0:_) 1 _ = [x0 :+ 0]
fftRaw x n s = zipWith (+) x1 x2 ++ (zipWith (-) x1 x2)
    where x1 = fftRaw x (div n 2) (2 * s)
          x2 = zipWith (*) [exp (0 :+ (-2 * pi * fromIntegral k / (fromIntegral n))) | k <- [0..((div n 2) - 1)]] (fftRaw (drop s x) (div n 2) (2 * s))

fft :: (RealFloat a) => [a] -> [Complex a]
fft x
    | isPowerOfTwo n = fftRaw x n 1
    | otherwise = error "FFT works only for powers of two"
    where n = length x

ifftRaw :: (RealFloat a) => [Complex a] -> Int -> Int -> [Complex a]
ifftRaw _ 0 _ = []
ifftRaw [] _  _ = []
ifftRaw (x0:_) 1 _ = [x0]
ifftRaw x n s = zipWith (+) x1 x2 ++ (zipWith (-) x1 x2)
    where x1 = ifftRaw x (div n 2) (2 * s)
          x2 = zipWith (*) [exp (0 :+ (2 * pi * fromIntegral k / (fromIntegral n))) | k <- [0..((div n 2) - 1)]] (ifftRaw (drop s x) (div n 2) (2 * s))

ifft :: (RealFloat a) => [Complex a] -> [a]
ifft x
    | isPowerOfTwo n = [realPart v / (fromIntegral n) | v <- ifftRaw x n 1]
    | otherwise = error "IFFT works only for powers of two"
    where n = length x

convFft :: (RealFloat a) => [a] -> [a] -> [a]
convFft h x
    | length x == length h = ifft (zipWith (*) (fft h) (fft x))
    | otherwise = error "Kernel and input data must have the same length"


{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE Trustworthy #-}
---------------------------------------------------------------------------
-- |
-- Copyright   :  (C) 2012-2015 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Simple matrix operation for low-dimensional primitives.
---------------------------------------------------------------------------
module Linear.Trace
  ( Trace(..)
  , frobenius
  ) where

import Control.Monad as Monad
import Linear.V0
import Linear.V1
import Linear.V2
import Linear.V3
import Linear.V4
import Linear.Plucker
import Linear.Quaternion
import Linear.V
import Linear.Vector
import Data.Complex
import Data.Distributive
import Data.Foldable as Foldable
import Data.Functor.Bind as Bind
import Data.Functor.Compose
import Data.Functor.Product
import Data.Hashable
import Data.HashMap.Lazy
import Data.IntMap (IntMap)
import Data.Map (Map)

-- $setup
-- >>> import Data.Complex
-- >>> import Debug.SimpleReflect.Vars
-- >>> import Linear.V2

class Functor m => Trace m where
  -- | Compute the trace of a matrix
  --
  -- >>> trace (V2 (V2 a b) (V2 c d))
  -- a + d
  trace :: Num a => m (m a) -> a
#ifndef HLINT
  default trace :: (Foldable m, Num a) => m (m a) -> a
  trace = Foldable.sum . diagonal
  {-# INLINE trace #-}
#endif

  -- | Compute the diagonal of a matrix
  --
  -- >>> diagonal (V2 (V2 a b) (V2 c d))
  -- V2 a d
  diagonal :: m (m a) -> m a
#ifndef HLINT
  default diagonal :: Monad m => m (m a) -> m a
  diagonal = Monad.join
  {-# INLINE diagonal #-}
#endif

instance Trace IntMap where
  diagonal = Bind.join
  {-# INLINE diagonal #-}

instance Ord k => Trace (Map k) where
  diagonal = Bind.join
  {-# INLINE diagonal #-}

instance (Eq k, Hashable k) => Trace (HashMap k) where
  diagonal = Bind.join
  {-# INLINE diagonal #-}

instance Dim n => Trace (V n)
instance Trace V0
instance Trace V1
instance Trace V2
instance Trace V3
instance Trace V4
instance Trace Plucker
instance Trace Quaternion

instance Trace Complex where
  trace ((a :+ _) :+ (_ :+ b)) = a + b
  {-# INLINE trace #-}
  diagonal ((a :+ _) :+ (_ :+ b)) = a :+ b
  {-# INLINE diagonal #-}

instance (Trace f, Trace g) => Trace (Product f g) where
  trace (Pair xx yy) = trace (pfst <$> xx) + trace (psnd <$> yy) where
    pfst (Pair x _) = x
    psnd (Pair _ y) = y
  {-# INLINE trace #-}
  diagonal (Pair xx yy) = diagonal (pfst <$> xx) `Pair` diagonal (psnd <$> yy) where
    pfst (Pair x _) = x
    psnd (Pair _ y) = y
  {-# INLINE diagonal #-}

instance (Distributive g, Trace g, Trace f) => Trace (Compose g f) where
  trace = trace . fmap (fmap trace . distribute) . getCompose . fmap getCompose
  {-# INLINE trace #-}
  diagonal = Compose . fmap diagonal . diagonal . fmap distribute . getCompose . fmap getCompose
  {-# INLINE diagonal #-}

-- | Compute the <http://mathworld.wolfram.com/FrobeniusNorm.html Frobenius norm> of a matrix.
frobenius :: (Num a, Foldable f, Additive f, Additive g, Distributive g, Trace g) => f (g a) -> a
frobenius m = trace $ fmap (\ f' -> Foldable.foldl' (^+^) zero $ liftI2 (*^) f' m) (distribute m)


{-|
Module      : Qubism.StateVec
Description : Types and functions for quantum state vectors
Copyright   : (c) Keith Pearson, 2018
License     : MIT
Maintainer  : keith@qubitrot.org
-}

{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}

module Qubism.StateVec
  ( StateVec (UnsafeMkStateVec)
  , mkStateVec
  , mkStateVec'
  , mkQubit
  , normalize
  , tensor
  , collapse
  , measureQubit
  , measure
  , dimension
  ) where

-- For dependent typing
import GHC.TypeLits
import Data.Singletons
import Data.Singletons.TypeLits
import Data.Finite

import           Control.Monad.Random
import           Control.Monad.Trans.State.Strict
import           Data.Complex
import           Numeric.LinearAlgebra ((|>))
import qualified Numeric.LinearAlgebra as LA
import           Text.Printf

import Qubism.Algebra
import Qubism.CReg

newtype StateVec (n :: Nat) =
  UnsafeMkStateVec (LA.Vector C)

-- | To provide "close enough" equality testing.
instance Eq (StateVec n) where
  UnsafeMkStateVec zs == UnsafeMkStateVec ws =
     LA.norm_2 (zs - ws) < 0.000001

instance KnownNat n => VectorSpace (StateVec n) where
  zero = UnsafeMkStateVec $ internalLen (sing :: Sing n) |> repeat 0
  z .: (UnsafeMkStateVec a) = UnsafeMkStateVec $ LA.scalar z * a
  (UnsafeMkStateVec a) +: (UnsafeMkStateVec b) = UnsafeMkStateVec $ a + b
  neg (UnsafeMkStateVec a) = UnsafeMkStateVec $ -a

instance KnownNat n => HilbertSpace (StateVec n) where
  (UnsafeMkStateVec zs) <.> (UnsafeMkStateVec ws) = zs LA.<.> ws

instance KnownNat n => Show (StateVec n) where
  show (UnsafeMkStateVec zs) = foldl show' "" $ zip [0..] (LA.toList zs)
    where show' str (i,z) = str ++ pf z ++ "  " ++ basis i ++ "\n"
          pf z    =  printf "% 6.4f" (realPart z) ++ "  + "
                  ++ printf "% 6.4f" (imagPart z) ++ "i"
          basis i = let bit j = if i `quot` 2^(n-j-1) `mod` 2 == 0
                                then '0' else '1'
                    in  "|" ++ fmap bit (take n [0..]) ++ ">"
          n = fromIntegral $ fromSing (sing :: Sing n)

internalLen :: (KnownNat n, Num a) => Sing n -> a
internalLen = (2 ^) . fromIntegral . fromSing

-- | Number of qubits in the StateVec
dimension :: forall (n :: Nat) a. (KnownNat n, Num a) => StateVec n -> a
dimension _ = fromIntegral $ fromSing (sing :: SNat n)

-- | StateVec's are intialized to |0>
mkStateVec :: forall n . KnownNat n => StateVec n
mkStateVec = UnsafeMkStateVec $ internalLen (sing :: Sing n) |> (1 : repeat 0)

-- | Make a statevector with an explicity specified size.
-- StateVec's are intialized to |0>
mkStateVec' :: Sing n -> StateVec n
mkStateVec' sn = UnsafeMkStateVec $ l |> (1 : repeat 0)
  where l = (2^) . fromIntegral $ fromSing sn

-- | A qubit is just a StateVec 1, initalized to |0>
mkQubit :: StateVec 1
mkQubit = UnsafeMkStateVec $ LA.fromList [1, 0]

normalize :: StateVec n -> StateVec n
normalize (UnsafeMkStateVec zs) = UnsafeMkStateVec $ LA.normalize zs

adjoint :: KnownNat n => StateVec n -> (StateVec n -> C)
adjoint qr = (qr <.>)

-- | The tensor product on elements of our Hilbert space
tensor :: StateVec n -> StateVec m -> StateVec (n + m)
tensor (UnsafeMkStateVec zs) (UnsafeMkStateVec ws) =
  UnsafeMkStateVec $ LA.flatten (zs `LA.outer` ws)

-- | Collapse a StateVec to a state compatable with qubit <n> being measured <bit>.
-- Only occurs physically with measurement, so internal use only.
collapse
  :: forall n . KnownNat n
  => Finite n -> Bit -> StateVec n -> StateVec n
collapse i b (UnsafeMkStateVec zs) = normalize . UnsafeMkStateVec $ zs * mask
  where
    mask   = l |> altseq
    altseq = replicate m ifZero ++ replicate m ifOne ++ altseq
    ifZero = if b == Zero then 1 else 0
    ifOne  = if b == One  then 1 else 0
    l      = internalLen (sing :: Sing n)
    m      = l `quot` (2 ^ (getFinite i + 1))

-- | Preforms a measurement of an induvidual qubit (indexed Finite n)
-- in a quantum register.
measureQubit
  :: (MonadRandom m, KnownNat n)
  => Finite n -> StateT (StateVec n) m Bit
measureQubit i = do
  qr <- get
  r  <- getRandomR (0, 1)
  let qrZero = collapse i Zero qr
      qrOne  = collapse i One qr
      pOne   = realPart $ qrOne <.> qr -- guaranteed to be real, so this
  if r < pOne                          -- is just a type-cast.
    then put qrOne  >> pure One
    else put qrZero >> pure Zero

-- | Preform a measurement on all qubits in a quantum register, returning a
-- computational basis state.
measure
  :: forall m n . (MonadRandom m, KnownNat n)
  => StateT (StateVec n) m CReg
measure = mkCReg <$> traverse measureQubit (take n [0 ..])
  where n = fromIntegral $ fromSing (sing :: Sing n)


{-# LANGUAGE OverloadedStrings #-}

module Main where

import Control.Monad (forever)
import Control.Monad.IO.Class (MonadIO (liftIO))
import Control.Monad.State.Strict
  ( MonadState (get, put),
    execStateT,
    forever,
    lift,
    runStateT,
  )
import Data.Complex (Complex (..), magnitude)
import Data.List (foldl')
import qualified Data.Text as T
import System.Random (Random (randomR), newStdGen)
import UI.NCurses
  ( Color (..),
    CursorMode (..),
    defaultWindow,
    drawText,
    moveCursor,
    render,
    runCurses,
    setColor,
    setCursorMode,
    setEcho,
    updateWindow,
  )
import UI.NCurses.HalfHeight
  ( Buffer,
    colorId,
    drawBuffer,
    initHexColors,
    mkBuffer,
    setXY,
  )

-- Set up a 15-color greyscale pallette
greyScale :: [String]
greyScale =
  [ "000000",
    "111111",
    "222222",
    "333333",
    "444444",
    "666666",
    "777777",
    "888888",
    "999999",
    "aaaaaa",
    "bbbbbb",
    "cccccc",
    "dddddd",
    "eeeeee",
    "ffffff"
  ]

-- Calculate the Mandelbrot steps-to-divergence and plot as Color 1 to Color 15
mandelbrot :: (Fractional a, RealFloat a) => a -> a -> Color
mandelbrot x y = Color (16 - ((+) 1 $ steps (0 :+ 0) (x' :+ y') 0))
  where
    x' = (x - 100) * 0.025
    y' = (y - 50) * 0.025
    steps z c i
      | i > 13 || magnitude z' > 2 = i
      | otherwise = steps z' c (i + 1)
      where
        z' = z * z + c

-- Draw the Mandelbrot set to the buffer point by point
drawMandelbrot :: Buffer -> Buffer
drawMandelbrot buffer =
  foldl'
    (\b (x, y) -> setXY x y (mandelbrot (fromIntegral x) (fromIntegral y)) b)
    buffer
    [(x, y) | x <- [0 .. 149], y <- [0 .. 99]]

mandelbrotMain :: IO ()
mandelbrotMain = do
  -- Register the 15 greyscale hex colors as Color 1 through Color 15.
  -- Every combination of these colors (fg and bg) is registered with Curses
  -- with a unique ID, which we return in the colorMap.
  colorMap <- runCurses $ initHexColors nord

  -- Creates a new buffer and draws the Mandelbrot set in each cell using the
  -- greyscale colors 1 through 15
  let buffer = drawMandelbrot $ mkBuffer 150 100 (Color 1)

  runCurses $ do
    -- Hide the cursor
    setEcho False
    setCursorMode CursorInvisible

    -- Enter an infinite drawing loop displaying the buffer.
    forever $ do
      let drawOp = do
            -- First draw the buffer to the screen.
            drawBuffer colorMap 0 0 buffer
            -- Now draw some text in the centre.
            -- We set the color to curses ID corresponding to fg 15 (white) and bg 1 (black)
            let (Just c) = colorId colorMap (Color 15) (Color 1)
            setColor c
            moveCursor 25 50
            drawText "| Text demonstrating that each cell is half the row height |"
      w <- defaultWindow
      updateWindow w drawOp
      render

nord :: [String]
nord =
  [ "3b4252",
    "bf616a",
    "a3be8c",
    --"ebcb8b",
    "81a1c1",
    "b48ead",
    "88c0d0",
    "e5e9f0",
    "4c566a",
    "bf616a",
    "a3be8c",
    "ebcb8b",
    "81a1c1",
    "b48ead",
    "8fbcbb",
    "eceff4"
  ]

performanceMain :: IO ()
performanceMain = do
  colorMap <- runCurses $ initHexColors greyScale
  -- Precompute the buffers
  let buffer = mkBuffer 100 100 (Color 1)
      col x y t =
        let x' = fromIntegral (x - 50)
            y' = fromIntegral (y - 50)
            t' = fromIntegral t
         in abs $ round $ (14 * ((sin $ 0.1 * x') * (cos $ 0.1 * y')) / (sin $ 0.1 * t')) + 1
      nextBuf b t = foldl' (\b (x, y) -> setXY x y (Color $ col x y t) b) b [(x, y) | x <- [0 .. 99], y <- [0 .. 99]]
      bufs = cycle $ scanl nextBuf buffer [0 .. 63]
  runCurses $ do
    _ <- (flip runStateT) (bufs, 0) $ do
      lift $ setEcho False
      lift $ setCursorMode CursorInvisible
      forever $ do
        (bufs, t) <- get
        let b = head bufs
            drawOp = do
              drawBuffer colorMap 0 0 b
              let (Just c) = colorId colorMap (Color 15) (Color 1)
              setColor c
              moveCursor 0 0
              drawText $ T.pack $ show t
        w <- lift defaultWindow
        lift $ updateWindow w drawOp
        lift render
        put (tail bufs, t + 1)
    return ()

main :: IO ()
main = do
  --mandelbrotMain
  performanceMain


module Problem215 where
{-
  Task description:
  Consider the problem of building a wall out of 2×1 and 3×1 bricks (horizontal×vertical dimensions) such that,
  for extra strength, the gaps between horizontally-adjacent bricks never line up in consecutive layers, i.e. never form a "running crack".

  For example, the following 9×3 wall is not acceptable due to the running crack shown in red:

  There are eight ways of forming a crack-free 9×3 wall, written W(9,3) = 8.

  Calculate W(32,10).
-}

import Control.Monad
import Data.Function (on)
import Data.Graph (Graph, Vertex)
import Data.Maybe (fromJust)
import Numeric.LinearAlgebra (Matrix, (<>))
import qualified GHC.Arr as Arr
import qualified Data.Graph as Graph
import qualified Data.HashSet as HashSet
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Tree as Tree
import qualified Numeric.LinearAlgebra as LA

{-
A row may only ever depend on the preceding row.
We may have cycles of rows that are allowed to follow each other.
An interesting question will be wether collapsing such cycles will be easier than simply 'exploring them'.
In both ways memoization may come in handy here.
-}

type Brick = Int
type Line = [Brick]
type Width = Int
type Height = Int
type Wall = [Line]

bricks :: [Brick]
bricks = [2, 3]

buildLines :: Width -> [Line]
buildLines 0 = [[]]
buildLines w = do
  brick <- bricks
  guard $ brick <= w
  lineRest <- buildLines $ w - brick
  return $ brick:lineRest

matchingLines :: Line -> Line -> Bool
matchingLines l1 l2 = not $ (matchingCracks `on` crackPositions) l1 l2
  where
    crackPositions :: Line -> [Int]
    crackPositions = fmap sum . init . tail . List.inits

    matchingCracks :: [Int] -> [Int] -> Bool
    matchingCracks [] _ = False
    matchingCracks _ [] = False
    matchingCracks xss@(x:xs) yss@(y:ys)
      | x == y = True
      | x < y = matchingCracks xs yss
      | x > y = matchingCracks xss ys

{-
With the matchingLines predicate I can generate a graph out of the possible lines.
We're searching for all routes of a given length that traverse this graph
and that start in any of it's nodes.
-}

mkLineEdges :: [Line] -> [(Line, Line)]
mkLineEdges [] = []
mkLineEdges (l:ls) = findMatchingLines l ls `mappend` mkLineEdges ls
  where
    findMatchingLines :: Line -> [Line] -> [(Line, Line)]
    findMatchingLines l lines = concat $ do
      l' <- lines
      guard $ matchingLines l l'
      return [(l,l'), (l',l)]

mkLineToVertex :: [Line] -> (Line -> Vertex)
mkLineToVertex lines = let lvMap = Map.fromList $ zip lines [0..]
                       in fromJust . (`Map.lookup` lvMap)

graphForWidth :: Width -> Graph
graphForWidth w = let lines = buildLines w
                      lEdges = mkLineEdges lines
                      bounds = (0, length lines - 1)
                      toVertex = mkLineToVertex lines
                      edges = [(toVertex a, toVertex b)|(a,b) <- lEdges]
                  in Graph.buildG bounds edges

graphToDenseMatrix :: Graph -> Matrix Double
graphToDenseMatrix = LA.toDense . graphToAssocMatrix
  where
    graphToAssocMatrix :: Graph -> LA.AssocMatrix
    graphToAssocMatrix g = do
      (v, targets) <- Arr.assocs g
      [((v, t), 1)| t <- targets]

dissect :: Graph -> [Graph]
dissect g = do
  tree <- Graph.components g
  let vertices = Tree.flatten tree
      (renumber, bounds) = mkRenumber vertices
  return . Graph.buildG bounds $ do
    v <- vertices
    w <-  Arr.unsafeAt g v
    return (renumber v, renumber w)
  where
    mkRenumber ls = let mapping = Map.fromList $ zip ls [0..]
                        λ = fromJust . (`Map.lookup` mapping)
                    in (λ, (0, Map.size mapping))

mPow :: Matrix Double -> Int -> Matrix Double
mPow m 0 = LA.ident $ LA.rows m
mPow m 1 = m
mPow m n
  | even n = mPow (m <> m) (n `div` 2)
  | otherwise = m <> mPow m (n - 1)

countWalls :: Width -> Height -> Int
countWalls w h
  | h < 1 = 0
  | h == 1 = length $ buildLines w
-- Using dissect reduced computation time from ~140s to ~9s.
  | otherwise = let g = graphForWidth w
                    ms = graphToDenseMatrix <$> dissect g
                    ms' = fmap (mPow `flip` (h - 1)) ms
                    calcSum = sum . concat . LA.toLists
                in round . sum $ fmap calcSum ms'

example = countWalls 9 3

-- 806844323190414
solution = countWalls 32 10

main = print solution


module Main where

import Control.Monad.Par
import Data.Complex
import qualified Data.Vector as V

import Graphics.Gloss.Data.Display
import Graphics.Gloss.Data.Color
import Graphics.Gloss.Data.Picture
import Graphics.Gloss.Interface.Pure.Display

windowSpec :: Display
windowSpec = InWindow "hello" (800,800) (0,0)

blackColor :: Color
blackColor = makeColor 0.0 0.0 0.0 1.0

whiteColor :: Color
whiteColor = makeColor 1.0 1.0 1.0 1.0

pseudoPixel :: Picture
pseudoPixel = rectangleSolid 1.0 1.0

blockf :: ((Float, Float), Int) -> Picture
blockf ((x,y),t) = let b = ((fromIntegral t) / 25) :: Float
                       c = makeColor b b b 1.0
                   in color c (translate x y pseudoPixel)

fractalRadius :: Float
fractalRadius = 2.0

fractalCalc :: (Complex Float, Complex Float) -> ((Float,Float),Int)
fractalCalc n = let a1 = V.iterateN 50 (\(start,c) -> (start*start + c,c)) (fmap (*0.005) n) :: V.Vector (Complex Float, Complex Float)
                    b =  V.length . V.filter (> fractalRadius) . V.map (realPart . abs . fst)  $ a1 :: Int
                    c = snd n
                    d = (realPart c, imagPart c) :: (Float, Float)
                in (d, b)

complexBlock :: V.Vector (Complex Float, Complex Float)
complexBlock = let x = [-400.0..400.0]
                   xv = V.fromList x
                   c = 0 :+ 0
               in V.concatMap (\x -> V.map (\z -> (c, x :+ z)) xv) xv

fractalP :: V.Vector (Par ((Float, Float), Int))
fractalP = V.map (\x -> (spawnP . fractalCalc $ x) >>= get) complexBlock

fractal :: Picture
fractal =  mconcat . V.toList . V.map blockf . runPar . sequence $ fractalP

main :: IO ()
main = display windowSpec whiteColor fractal


-- |
-- Module    : Statistics.Test.MannWhitneyU
-- Copyright : (c) 2010 Neil Brown
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Mann-Whitney U test (also know as Mann-Whitney-Wilcoxon and
-- Wilcoxon rank sum test) is a non-parametric test for assessing
-- whether two samples of independent observations have different
-- mean.
module Statistics.Test.MannWhitneyU (
    -- * Mann-Whitney U test
    mannWhitneyUtest
  , mannWhitneyU
  , mannWhitneyUCriticalValue
  , mannWhitneyUSignificant
    -- ** Wilcoxon rank sum test
  , wilcoxonRankSums
  , module Statistics.Test.Types
    -- * References
    -- $references
  ) where

import Control.Applicative ((<$>))
import Data.List (findIndex)
import Data.Ord (comparing)
import Numeric.SpecFunctions (choose)
import Prelude hiding (sum)
import Statistics.Distribution (quantile)
import Statistics.Distribution.Normal (standard)
import Statistics.Function (sortBy)
import Statistics.Sample.Internal (sum)
import Statistics.Test.Internal (rank, splitByTags)
import Statistics.Test.Types (TestResult(..), PositionTest(..), significant)
import Statistics.Types (PValue,pValue)
import qualified Data.Vector.Unboxed as U

-- | The Wilcoxon Rank Sums Test.
--
-- This test calculates the sum of ranks for the given two samples.
-- The samples are ordered, and assigned ranks (ties are given their
-- average rank), then these ranks are summed for each sample.
--
-- The return value is (W₁, W₂) where W₁ is the sum of ranks of the first sample
-- and W₂ is the sum of ranks of the second sample.  This test is trivially transformed
-- into the Mann-Whitney U test.  You will probably want to use 'mannWhitneyU'
-- and the related functions for testing significance, but this function is exposed
-- for completeness.
wilcoxonRankSums :: (Ord a, U.Unbox a) => U.Vector a -> U.Vector a -> (Double, Double)
wilcoxonRankSums xs1 xs2 = (sum ranks1, sum ranks2)
  where
    -- Ranks for each sample
    (ranks1,ranks2) = splitByTags $ U.zip tags (rank (==) joinSample)
    -- Sorted and tagged sample
    (tags,joinSample) = U.unzip
                      $ sortBy (comparing snd)
                      $ tagSample True xs1 U.++ tagSample False xs2
    -- Add tag to a sample
    tagSample t = U.map (\x -> (t,x))

-- | The Mann-Whitney U Test.
--
-- This is sometimes known as the Mann-Whitney-Wilcoxon U test, and
-- confusingly many sources state that the Mann-Whitney U test is the same as
-- the Wilcoxon's rank sum test (which is provided as 'wilcoxonRankSums').
-- The Mann-Whitney U is a simple transform of Wilcoxon's rank sum test.
--
-- Again confusingly, different sources state reversed definitions for U₁
-- and U₂, so it is worth being explicit about what this function returns.
-- Given two samples, the first, xs₁, of size n₁ and the second, xs₂,
-- of size n₂, this function returns (U₁, U₂)
-- where U₁ = W₁ - (n₁(n₁+1))\/2
-- and U₂ = W₂ - (n₂(n₂+1))\/2,
-- where (W₁, W₂) is the return value of @wilcoxonRankSums xs1 xs2@.
--
-- Some sources instead state that U₁ and U₂ should be the other way round, often
-- expressing this using U₁' = n₁n₂ - U₁ (since U₁ + U₂ = n₁n₂).
--
-- All of which you probably don't care about if you just feed this into 'mannWhitneyUSignificant'.
mannWhitneyU :: (Ord a, U.Unbox a) => U.Vector a -> U.Vector a -> (Double, Double)
mannWhitneyU xs1 xs2
  = (fst summedRanks - (n1*(n1 + 1))/2
    ,snd summedRanks - (n2*(n2 + 1))/2)
  where
    n1 = fromIntegral $ U.length xs1
    n2 = fromIntegral $ U.length xs2

    summedRanks = wilcoxonRankSums xs1 xs2

-- | Calculates the critical value of Mann-Whitney U for the given sample
-- sizes and significance level.
--
-- This function returns the exact calculated value of U for all sample sizes;
-- it does not use the normal approximation at all.  Above sample size 20 it is
-- generally recommended to use the normal approximation instead, but this function
-- will calculate the higher critical values if you need them.
--
-- The algorithm to generate these values is a faster, memoised version of the
-- simple unoptimised generating function given in section 2 of \"The Mann Whitney
-- Wilcoxon Distribution Using Linked Lists\"
mannWhitneyUCriticalValue
  :: (Int, Int)     -- ^ The sample size
  -> PValue Double  -- ^ The p-value (e.g. 0.05) for which you want the critical value.
  -> Maybe Int      -- ^ The critical value (of U).
mannWhitneyUCriticalValue (m, n) p
  | m < 1 || n < 1 = Nothing    -- Sample must be nonempty
  | p' <= 1        = Nothing    -- p-value is too small. Null hypothesis couldn't be disproved
  | otherwise      = findIndex (>= p')
                   $ take (m*n)
                   $ tail
                   $ alookup !! (m+n-2) !! (min m n - 1)
  where
    mnCn = (m+n) `choose` n
    p'   = mnCn * pValue p

{-
-- Original function, without memoisation, from Cheung and Klotz:
-- Double is needed to avoid integer overflows.
a :: Int -> Int -> Int -> Double
a u bigN m
  | u < 0            = 0
  | u >= m * n       = bigN `choose` m
  | m == 1 || n == 1 = fromIntegral (u + 1)
  | otherwise        = a  u      (bigN - 1)  m
                     + a (u - n) (bigN - 1) (m-1)
  where
    n = bigN - m
-}

-- Memoised version of the original a function, above.
--
-- Doubles are stored to avoid integer overflow. 32-bit Ints begin to
-- overflow for bigN as small as 33 (64-bit one at 66) while Double to
-- go to infinity till bigN=1029
--
--
-- outer list is indexed by big N - 2
-- inner list by (m-1) (we know m < bigN)
-- innermost list by u
--
-- So: (alookup !! (bigN - 2) !! (m - 1) ! u) == a u bigN m
alookup :: [[[Double]]]
alookup = gen 2 [1 : repeat 2]
  where
    gen bigN predBigNList
       = let bigNlist = [ [ amemoed u m
                          | u <- [0 .. m*(bigN-m)]
                          ] ++ repeat (bigN `choose` m)
                        | m <- [1 .. (bigN-1)]] -- has bigN-1 elements
         in bigNlist : gen (bigN+1) bigNlist
      where
        amemoed :: Int -> Int -> Double
        amemoed u m
          | m == 1 || n == 1 = fromIntegral (u + 1)
          | otherwise        = mList !! u
                             + if u < n then 0 else predmList !! (u-n)
          where
            n = bigN - m
            (predmList : mList : _) = drop (m-2) predBigNList
            -- Lists for m-1 and m respectively. i-th list correspond to m=i+1
            --
            -- We know that predBigNList has bigN - 2 elements
            -- (and we know that n > 1 therefore bigN > m + 1)
            -- So bigN - 2 >= m, i.e. predBigNList must have at least m elements
            -- elements, so dropping (m-2) must leave at least 2

-- | Calculates whether the Mann Whitney U test is significant.
--
-- If both sample sizes are less than or equal to 20, the exact U critical value
-- (as calculated by 'mannWhitneyUCriticalValue') is used.  If either sample is
-- larger than 20, the normal approximation is used instead.
--
-- If you use a one-tailed test, the test indicates whether the first sample is
-- significantly larger than the second.  If you want the opposite, simply reverse
-- the order in both the sample size and the (U₁, U₂) pairs.
mannWhitneyUSignificant
  :: PositionTest     -- ^ Perform one-tailed test (see description above).
  -> (Int, Int)       -- ^ The samples' size from which the (U₁,U₂) values were derived.
  -> PValue Double    -- ^ The p-value at which to test (e.g. 0.05)
  -> (Double, Double) -- ^ The (U₁, U₂) values from 'mannWhitneyU'.
  -> Maybe TestResult -- ^ Return 'Nothing' if the sample was too
                      --   small to make a decision.
mannWhitneyUSignificant test (in1, in2) pVal (u1, u2)
  -- Use normal approximation
  | in1 > 20 || in2 > 20 =
    let mean  = n1 * n2 / 2     -- (u1+u2) / 2
        sigma = sqrt $ n1*n2*(n1 + n2 + 1) / 12
        z     = (mean - u1) / sigma
    in Just $ case test of
                AGreater      -> significant $ z     < quantile standard p
                BGreater      -> significant $ (-z)  < quantile standard p
                SamplesDiffer -> significant $ abs z > abs (quantile standard (p/2))
  -- Use exact critical value
  | otherwise = do crit <- fromIntegral <$> mannWhitneyUCriticalValue (in1, in2) pVal
                   return $ case test of
                              AGreater      -> significant $ u2        <= crit
                              BGreater      -> significant $ u1        <= crit
                              SamplesDiffer -> significant $ min u1 u2 <= crit
  where
    n1 = fromIntegral in1
    n2 = fromIntegral in2
    p  = pValue pVal

-- | Perform Mann-Whitney U Test for two samples and required
-- significance. For additional information check documentation of
-- 'mannWhitneyU' and 'mannWhitneyUSignificant'. This is just a helper
-- function.
--
-- One-tailed test checks whether first sample is significantly larger
-- than second. Two-tailed whether they are significantly different.
mannWhitneyUtest
  :: (Ord a, U.Unbox a)
  => PositionTest     -- ^ Perform one-tailed test (see description above).
  -> PValue Double    -- ^ The p-value at which to test (e.g. 0.05)
  -> U.Vector a       -- ^ First sample
  -> U.Vector a       -- ^ Second sample
  -> Maybe TestResult -- ^ Return 'Nothing' if the sample was too small to
                      --   make a decision.
mannWhitneyUtest ontTail p smp1 smp2 =
  mannWhitneyUSignificant ontTail (n1,n2) p $ mannWhitneyU smp1 smp2
    where
      n1 = U.length smp1
      n2 = U.length smp2

-- $references
--
-- * Cheung, Y.K.; Klotz, J.H. (1997) The Mann Whitney Wilcoxon
--   distribution using linked lists. /Statistica Sinica/
--   7:805–813.
-- <http://www3.stat.sinica.edu.tw/statistica/oldpdf/A7n316.pdf>.


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

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

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

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

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

type Count = CountG Int

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

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

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

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

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

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

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

asMean :: Mean -> Mean
asMean = id

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

asMeanKBN :: MeanKBN -> MeanKBN
asMeanKBN = id

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

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

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

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

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

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

asVarWelfordKBN :: VarWelfordKBN -> VarWelfordKBN
asVarWelfordKBN = id

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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