{-# LANGUAGE BangPatterns #-}
module FourierSeries where

import           Control.Monad              as M
import           Data.Array.Repa            as R
import           Data.Complex
import           Data.List                  as L
import           Data.Vector.Storable       as VS
import           Data.Vector.Unboxed        as VU
import           FokkerPlanck.FourierSeries
import           Foreign.CUDA.BLAS          as BLAS
import           Foreign.CUDA.Driver        as CUDA
import           Image.IO
import           Numeric.LinearAlgebra      as NL hiding (All)
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Utils.BLAS
import           Utils.Distribution
import           Utils.Parallel
import           Utils.SimpsonRule
import           Utils.Time

main = do
  args@(gpuIDStr:sizeStr:deltaStr:periodStr:freqStr:angularFreqStr:radialFreqStr:stdStr:_) <-
    getArgs
  printCurrentTime . show $ args
  let gpuID = read gpuIDStr :: Int
      numPoints = read sizeStr :: Int
      delta = read deltaStr :: Double
      period = read periodStr :: Double
      freq = read freqStr :: Int
      numFreq = 2 * freq + 1
      angularFreq = read angularFreqStr :: Int
      radialFreq = read radialFreqStr :: Int
      std = read stdStr :: Double
      folderPath = "output/test/FourierSeries"
      len' = round $ (fromIntegral numPoints / delta :: Double)
      len =
        if odd len'
          then len'
          else len' + 1
      center = div len 2
      !numAngularFreq = 2 * angularFreq + 1
      !numRadialFreq = 2 * radialFreq + 1
      !numPinwheels = numAngularFreq * numRadialFreq
      pinwheel =
        fromFunction (Z :. len :. len :. numAngularFreq :. numRadialFreq) $ \(Z :. i :. j :. af :. rf) ->
          let x = delta * (fromIntegral $ i - center)
              y = delta * (fromIntegral $ j - center)
          in fourierMellin 0.5 (af - angularFreq) (rf - radialFreq) (x, y)
      rectangularHarmonics' =
        fromFunction (Z :. numFreq :. numFreq :. len :. len) $ \(Z :. xFreq :. yFreq :. i :. j) ->
          let !x = fromIntegral $ i - center
              !y = fromIntegral $ j - center
          in cis
               ((-2) * pi * delta *
                ((fromIntegral $ xFreq - freq) * x +
                 (fromIntegral $ yFreq - freq) * y) /
                period)
      simpsonWeights =
        computeUnboxedS $ computeWeightArrFromListOfShape [len, len] :: R.Array U DIM2 (Complex Double)
      inverseR2Harmonics' =
        fromFunction (Z :. numPoints :. numPoints :. numFreq :. numFreq) $ \(Z :. i :. j :. xFreq :. yFreq) ->
          cis
            (2 * pi *
             fromIntegral
               ((xFreq - freq) * (i - (div numPoints 2)) +
                (yFreq - freq) * (j - (div numPoints 2))) /
             period)
      !weight = (delta / 3) ^ 2 :+ 0
  createDirectoryIfMissing True folderPath
  createDirectoryIfMissing True (folderPath </> "Recons")
  createDirectoryIfMissing True (folderPath </> "Coefficients")
  print len
  !weightedPinwheels <-
    computeUnboxedP . R.traverse2 pinwheel simpsonWeights const $ \fP fS idx@(Z :. a :. b :. _ :. _) ->
      fP idx * fS (Z :. a :. b)
  !rectangularHarmonics <- computeUnboxedP rectangularHarmonics'
  !inverseR2Harmonics <- computeUnboxedP inverseR2Harmonics'
  -- printCurrentTime "CPU starts .."
  -- let !gaussianCoefficients =
  --       computeUnboxedS . R.fromFunction (Z :. numFreq :. numFreq) $ \(Z :. xFreq :. yFreq) ->
  --         weight *
  --         (gaussian2D
  --            (fromIntegral $ xFreq - freq)
  --            (fromIntegral $ yFreq - freq)
  --            std :+
  --          0)
  -- M.mapM_
  --   (\(af, rf) -> do
  --      let coefficients =
  --            sumS .
  --            sumS .
  --            R.traverse2
  --              rectangularHarmonics
  --              (R.slice weightedPinwheels (Z :. All :. All :. af :. rf))
  --              const $ \fR fP idx@(Z :. _ :. _ :. i :. j) ->
  --              fR idx * fP (Z :. i :. j)
  --          recon =
  --            sumS .
  --            sumS .
  --            R.traverse2
  --              inverseR2Harmonics
  --              (coefficients *^ gaussianCoefficients)
  --              const $ \fR fC idx@(Z :. _ :. _ :. i :. j) ->
  --              fR idx * fC (Z :. i :. j)
  --      plotImageRepaComplex
  --        (folderPath </> "Recon" </> printf "Recon%03d.png" (af * numRadialFreq + rf + 1)) .
  --        ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
  --        recon)
  --   [(af, rf) | af <- [0 .. numAngularFreq - 1], rf <- [0 .. numRadialFreq - 1]]
  -- printCurrentTime "Done"
  -- coefficients <-
  --   (sumP . R.traverse2 rectangularHarmonics weightedPinwheel const $ \fR fP idx@(Z :. _ :. _ :. i :. j) ->
  --      fR idx * fP (Z :. i :. j)) >>=
  --   sumP
  --     gaussianCoefficients =
  --       computeUnboxedS . R.traverse coefficients id $ \f idx@(Z :. xFreq :. yFreq) ->
  --         weight * (f idx) *
  --         (gaussian2D
  --            (fromIntegral $ xFreq - freq)
  --            (fromIntegral $ yFreq - freq)
  --            std :+
  --          0)
  -- recon <-
  --   (sumP . R.traverse2 inverseR2Harmonics gaussianCoefficients const $ \fR fC idx@(Z :. _ :. _ :. i :. j) ->
  --      fR idx * fC (Z :. i :. j)) >>=
  --   sumP
  -- -- Repa
  -- let coefficients =
  --       sumS . sumS . R.traverse2 rectangularHarmonics weightedPinwheel const $ \fR fP idx@(Z :. _ :. _ :. i :. j) ->
  --         fR idx * fP (Z :. i :. j)
  --     gaussianCoefficients =
  --       computeUnboxedS . R.traverse coefficients id $ \f idx@(Z :. xFreq :. yFreq) ->
  --         weight * (f idx) *
  --         (gaussian2D
  --            (fromIntegral $ xFreq - freq)
  --            (fromIntegral $ yFreq - freq)
  --            std :+
  --          0)
  --     recon =
  --       sumS . sumS . R.traverse2 inverseR2Harmonics gaussianCoefficients const $ \fR fC idx@(Z :. _ :. _ :. i :. j) ->
  --         fR idx * fC (Z :. i :. j)
  -- -- HMatrix
  -- let rectangularHarmonicsMat =
  --       ((numFreq ^ 2) >< (len ^ 2)) . R.toList $ rectangularHarmonics
  --     inverseR2HarmonicsMat =
  --       ((numPoints ^ 2) >< (numFreq ^ 2)) . R.toList $ inverseR2Harmonics
  --     weightedPinwheelVec = NL.fromList . R.toList $ weightedPinwheel
  --     gaussianCoefficientsVec =
  --       NL.fromList . R.toList . R.fromFunction (Z :. numFreq :. numFreq) $ \(Z :. xFreq :. yFreq) ->
  --         weight *
  --         (gaussian2D
  --            (fromIntegral $ xFreq - freq)
  --            (fromIntegral $ yFreq - freq)
  --            std :+
  --          0)
  --     coefficientsVec = rectangularHarmonicsMat #> weightedPinwheelVec
  --     recon =
  --       fromListUnboxed (Z :. numPoints :. numPoints) . NL.toList $
  --       inverseR2HarmonicsMat #> (coefficientsVec * gaussianCoefficientsVec)
  -- --BLAS mxv
  -- let rectangularHarmonicsMat = VU.convert . toUnboxed $ rectangularHarmonics
  --     inverseR2HarmonicsMat = VU.convert . toUnboxed $ inverseR2Harmonics
  --     weightedPinwheelVec = VU.convert . toUnboxed $ weightedPinwheel
  --     gaussianCoefficientsVec =
  --       VU.convert .
  --       toUnboxed . computeS . R.fromFunction (Z :. numFreq :. numFreq) $ \(Z :. xFreq :. yFreq) ->
  --         weight *
  --         (gaussian2D
  --            (fromIntegral $ xFreq - freq)
  --            (fromIntegral $ yFreq - freq)
  --            std :+
  --          0)
  -- coefficientsVec <-
  --   mxv (numFreq ^ 2) (len ^ 2) rectangularHarmonicsMat weightedPinwheelVec
  -- reconVec <-
  --   mxv
  --     (numPoints ^ 2)
  --     (numFreq ^ 2)
  --     inverseR2HarmonicsMat
  --     (VS.zipWith (*) coefficientsVec gaussianCoefficientsVec)
  -- let recon = fromUnboxed (Z :. numPoints :. numPoints) . VS.convert $ reconVec
  --BLAS mxm
  -- let !rectangularHarmonicsMat = VU.convert . toUnboxed $ rectangularHarmonics
  --     !inverseR2HarmonicsMat = VU.convert . toUnboxed $ inverseR2Harmonics
  --     !weightedPinwheelVec = VU.convert . toUnboxed $ weightedPinwheels
  --     !gaussianCoefficients =
  --       computeUnboxedS . R.fromFunction (Z :. numFreq :. numFreq) $ \(Z :. xFreq :. yFreq) ->
  --         weight *
  --         (gaussian2D
  --            (fromIntegral $ xFreq - freq)
  --            (fromIntegral $ yFreq - freq)
  --            std :+
  --          0)
  -- printCurrentTime "Start.."
  -- coefficientsVec <-
  --   gemmBLAS
  --     (numFreq ^ 2)
  --     numPinwheels
  --     (len ^ 2)
  --     rectangularHarmonicsMat
  --     weightedPinwheelVec
  -- printCurrentTime "coefficients done"
  -- let coefficientsArr =
  --       fromUnboxed (Z :. numFreq :. numFreq :. numAngularFreq :. numRadialFreq) .
  --       VS.convert $
  --       coefficientsVec
  --     weightedCoefficientsArr =
  --       computeS .
  --       R.traverse2
  --         coefficientsArr
  --         gaussianCoefficients
  --         const $ \f1 f2 idx@(Z :. i :. j :. _ :. _) ->
  --         f1 idx * f2 (Z :. i :. j)
  --     weightedCoefficients = VS.convert . toUnboxed $ weightedCoefficientsArr
  -- M.mapM_
  --   (\(af, rf) ->
  --      plotImageRepaComplex
  --        (folderPath </> "Coefficients" </>
  --         printf
  --           "Coef(%d,%d).png"
  --           (af - div numAngularFreq 2)
  --           (rf - div numRadialFreq 2)) .
  --      ImageRepa 8 .
  --      computeS .
  --      extend (Z :. (1 :: Int) :. All :. All) . R.slice coefficientsArr $
  --      (Z :. All :. All :. af :. rf))
  --   [(af, rf) | af <- [0 .. numAngularFreq - 1], rf <- [0 .. numRadialFreq - 1]]
  -- printCurrentTime "Done"
  -- reconVec <-
  --   gemmBLAS
  --     (numPoints ^ 2)
  --     numPinwheels
  --     (numFreq ^ 2)
  --     inverseR2HarmonicsMat
  --     weightedCoefficients
  -- let !recon =
  --       fromUnboxed
  --         (Z :. numPoints :. numPoints :. numAngularFreq :. numRadialFreq) .
  --       VS.convert $
  --       reconVec
  -- M.mapM_
  --   (\(af, rf) ->
  --      plotImageRepaComplex
  --        (folderPath </> "Recons" </>
  --         printf
  --           "Recon(%d,%d).png"
  --           (af - div numAngularFreq 2)
  --           (rf - div numRadialFreq 2)) .
  --      ImageRepa 8 .
  --      computeS . extend (Z :. (1 :: Int) :. All :. All) . R.slice recon $
  --      (Z :. All :. All :. af :. rf))
  --   [(af, rf) | af <- [0 .. numAngularFreq - 1], rf <- [0 .. numRadialFreq - 1]]
  -- printCurrentTime "Done"
  --BLAS mxmGPU
  let !rectangularHarmonicsMat = VU.convert . toUnboxed $ rectangularHarmonics
      !inverseR2HarmonicsMat = VU.convert . toUnboxed $ inverseR2Harmonics
      !weightedPinwheelVec = VU.convert . toUnboxed $ weightedPinwheels
      !gaussianCoefficients =
        computeUnboxedS . R.fromFunction (Z :. numFreq :. numFreq) $ \(Z :. xFreq :. yFreq) ->
          weight *
          (gaussian2D
             (fromIntegral $ xFreq - freq)
             (fromIntegral $ yFreq - freq)
             std :+
           0)
  initialise []
  dev <- device gpuID
  ctx <- CUDA.create dev []
  handle <- BLAS.create
  printCurrentTime "Start.."
  coefficientsVec <-
    gemmCuBLAS11
      handle
      (numFreq ^ 2)
      numPinwheels
      (len ^ 2)
      rectangularHarmonicsMat
      weightedPinwheelVec
  let coefficientsArr =
        fromUnboxed (Z :. numFreq :. numFreq :. numAngularFreq :. numRadialFreq) .
        VS.convert $
        coefficientsVec
      weightedCoefficients =
        VU.convert .
        toUnboxed .
        computeS . R.traverse2 coefficientsArr gaussianCoefficients const $ \f1 f2 idx@(Z :. i :. j :. _ :. _) ->
          f1 idx * f2 (Z :. i :. j)
  printCurrentTime "coefficients done"
  reconVec <-
    gemmCuBLAS11
      handle
      (numPoints ^ 2)
      numPinwheels
      (numFreq ^ 2)
      inverseR2HarmonicsMat
      weightedCoefficients
  let !recon =
        fromUnboxed (Z :. numPoints :. numPoints :. numPinwheels) . VS.convert $
        reconVec
  M.mapM_
    (\i ->
       plotImageRepaComplex
         (folderPath </> "Recons" </> printf "Recon%03d.png" (i + 1)) .
       ImageRepa 8 .
       computeS . extend (Z :. (1 :: Int) :. All :. All) . R.slice recon $
       (Z :. All :. All :. i))
    [0 .. numPinwheels - 1]
  printCurrentTime "Done"
  M.mapM_
    (\(af, rf) ->
       plotImageRepaComplex
         (folderPath </> "Coefficients" </>
          printf
            "Coef(%d,%d).png"
            (af - div numAngularFreq 2)
            (rf - div numRadialFreq 2)) .
       ImageRepa 8 .
       computeS .
       extend (Z :. (1 :: Int) :. All :. All) . R.slice coefficientsArr $
       (Z :. All :. All :. af :. rf))
    [(af, rf) | af <- [0 .. numAngularFreq - 1], rf <- [0 .. numRadialFreq - 1]]
  -- plotImageRepaComplex (folderPath </> "Grating.png") .
  --   ImageRepa 8 .
  --   computeS .
  --   extend (Z :. (1 :: Int) :. All :. All) . R.slice rectangularHarmonics $
  --   (Z :. (0 :: Int) :. freq :. All :. All)
  -- plotImageRepaComplex (folderPath </> "Pinwheel.png") .
  --   ImageRepa 8 .
  --   computeS . fromFunction (Z :. (1 :: Int) :. numPoints :. numPoints) $ \(Z :. _ :. i :. j) ->
  --   fourierMellin
  --     angularFreq
  --     radialFreq
  --     ( (fromIntegral $ i - (div numPoints 2))
  --     , (fromIntegral $ j - (div numPoints 2)))


import Data.Complex
import Seq

avg xs = (*rec) $ sum xs
  where rec = recip $ fromIntegral $ length xs

harmonic xs = (*rec) . recip . sum . map recip $ xs
  where rec = fromIntegral $ length xs

--TODO use unfold
--
--util function to iteratively generate a list from results
--nothing about this implies a particular "zero"
--scan2 :: Eq a => a -> (a -> (b,a)) -> a -> [b]
--scan2 z f x | x == z    = []
--            | otherwise = fist:scan2 z f sound
--  where (fist, sound) = f x

scan2' = scan2 []
--semi-stable geometric mean that doesn't use too many numbers at once
stableGeometric xs = product $ map ((**rec) . product) $ scan2' (splitAt 10) xs
  where rec = recip $ fromIntegral $ length xs

agm a b = map fst $ iterate (\(x,y) -> (arith x y, geo x y)) (a,b)
  where arith a b = (/2) $ a + b
        geo a b   = sqrt $ a * b

agm1i = realPart $ flip (!!) 10 $ agm (1 :+ 0) (0 :+ 1)


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MagicHash             #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE UnboxedTuples         #-}
{-# LANGUAGE ViewPatterns          #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Numeric.Matrix.Internal.Float () where

import qualified Control.Monad.ST        as ST
import qualified Numeric.DataFrame.ST    as ST
import           Numeric.Dimensions      (KnownDim)
import           Numeric.Matrix.Internal
import           Numeric.Matrix.LU
import           Numeric.Scalar.Internal
import           Numeric.Vector.Internal

{-# INLINE mkMat #-}
mkMat ::
  Float -> Float -> Float -> Float ->
  Float -> Float -> Float -> Float ->
  Float -> Float -> Float -> Float ->
  Float -> Float -> Float -> Float ->
  Mat44f
mkMat
  _11 _12 _13 _14
  _21 _22 _23 _24
  _31 _32 _33 _34
  _41 _42 _43 _44
  = ST.runST $ do
    df <- ST.newDataFrame
    ST.writeDataFrameOff df 0  $ scalar _11
    ST.writeDataFrameOff df 1  $ scalar _12
    ST.writeDataFrameOff df 2  $ scalar _13
    ST.writeDataFrameOff df 3  $ scalar _14
    ST.writeDataFrameOff df 4  $ scalar _21
    ST.writeDataFrameOff df 5  $ scalar _22
    ST.writeDataFrameOff df 6  $ scalar _23
    ST.writeDataFrameOff df 7  $ scalar _24
    ST.writeDataFrameOff df 8  $ scalar _31
    ST.writeDataFrameOff df 9  $ scalar _32
    ST.writeDataFrameOff df 10 $ scalar _33
    ST.writeDataFrameOff df 11 $ scalar _34
    ST.writeDataFrameOff df 12 $ scalar _41
    ST.writeDataFrameOff df 13 $ scalar _42
    ST.writeDataFrameOff df 14 $ scalar _43
    ST.writeDataFrameOff df 15 $ scalar _44
    ST.unsafeFreezeDataFrame df

instance HomTransform4 Float where
  {-# INLINE translate4 #-}
  translate4 (unpackV4# -> (# x, y, z, _ #)) = mkMat
    1 0 0 0
    0 1 0 0
    0 0 1 0
    x y z 1

  {-# INLINE translate3 #-}
  translate3 (unpackV3# -> (# x, y, z #)) = mkMat
    1 0 0 0
    0 1 0 0
    0 0 1 0
    x y z 1

  {-# INLINE rotateX #-}
  rotateX a = mkMat
    1 0 0 0
    0 c s 0
    0 n c 0
    0 0 0 1
    where
      c = cos a
      s = sin a
      n = -s

  {-# INLINE rotateY #-}
  rotateY a = mkMat
    c 0 n 0
    0 1 0 0
    s 0 c 0
    0 0 0 1
    where
      c = cos a
      s = sin a
      n = -s

  {-# INLINE rotateZ #-}
  rotateZ a = mkMat
    c s 0 0
    n c 0 0
    0 0 1 0
    0 0 0 1
    where
      c = cos a
      s = sin a
      n = -s

  {-# INLINE rotate #-}
  rotate (unpackV3# -> (# x, y, z #)) a = mkMat
    (c+xxv)  (yxv+zs) (zxv-ys) 0
    (xyv-zs) (c+yyv)  (zyv+xs) 0
    (xzv+ys) (yzv-xs) (c+zzv)  0
     0        0        0       1
    where
      c = cos a
      v = 1 - c -- v for versine
      s = sin a
      xxv = x * x * v
      xyv = x * y * v
      xzv = x * z * v
      yxv = xyv
      yyv = y * y * v
      yzv = y * z * v
      zxv = xzv
      zyv = yzv
      zzv = z * z * v
      xs = x * s
      ys = y * s
      zs = z * s

  {-# INLINE rotateEuler #-}
  rotateEuler x y z = mkMat
    (cy*cz)  (cx*sz+sx*sy*cz) (sx*sz-cx*sy*cz) 0
    (-cy*sz) (cx*cz-sx*sy*sz) (sx*cz+cx*sy*sz) 0
     sy      (-sx*cy)         (cx*cy)          0
     0        0                0               1
    where
      cx = cos x
      sx = sin x
      cy = cos y
      sy = sin y
      cz = cos z
      sz = sin z

  {-# INLINE lookAt #-}
  lookAt up cam foc = mkMat
    xb1 yb1 zb1 0
    xb2 yb2 zb2 0
    xb3 yb3 zb3 0
    tx  ty  tz  1
    where
      (# xb1, xb2, xb3 #) = unpackV3# xb
      (# yb1, yb2, yb3 #) = unpackV3# yb
      (# zb1, zb2, zb3 #) = unpackV3# zb
      zb = normalized $ cam - foc -- Basis vector for "backward", since +Z is behind the camera
      xb = normalized $ up `cross` zb -- Basis vector for "right"
      yb = zb `cross` xb -- Basis vector for "up"
      ncam = -cam
      tx = unScalar $ xb `dot` ncam
      ty = unScalar $ yb `dot` ncam
      tz = unScalar $ zb `dot` ncam

  {-# INLINE perspective #-}
  perspective n f fovy aspect = mkMat
    dpw 0   0   0
    0   dph 0   0
    0   0   a (-1)
    0   0   b   0
    where
      hpd = tan (fovy * 0.5) -- height/distance
      wpd = aspect * hpd; -- width/distance
      dph = recip hpd -- distance/height
      dpw = recip wpd -- distance/width
      nmf = n - f
      a = (n + f) / nmf
      b = 2 * n * f / nmf

  {-# INLINE orthogonal #-}
  orthogonal n f w h = mkMat
    iw 0  0 0
    0  ih 0 0
    0  0  a 0
    0  0  b 1
    where
      ih = 2 / h
      iw = 2 / w
      nmf = n - f
      a = 2 / nmf
      b = (n + f) / nmf

  {-# INLINE toHomPoint #-}
  toHomPoint (unpackV3# -> (# x, y, z #)) = vec4 x y z 1

  {-# INLINE toHomVector #-}
  toHomVector (unpackV3# -> (# x, y, z #)) = vec4 x y z 0

  {-# INLINE fromHom #-}
  fromHom (unpackV4# -> (# x, y, z, w #))
    | w == 0    = vec3 x y z
    | otherwise = vec3 (x/w) (y/w) (z/w)

instance KnownDim n => MatrixInverse Float n where
    inverse = inverseViaLU

instance KnownDim n => MatrixDeterminant Float n where
    det = detViaLU


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.StudentT
-- Copyright : (c) 2011 Aleksey Khudyakov
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Student-T distribution
module Statistics.Distribution.StudentT (
    StudentT
    -- * Constructors
  , studentT
  , studentTE
  , studentTUnstandardized
    -- * Accessors
  , studentTndf
  ) where

import Control.Applicative
import Data.Aeson          (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary         (Binary(..))
import Data.Data           (Data, Typeable)
import GHC.Generics        (Generic)
import Numeric.SpecFunctions (
  logBeta, incompleteBeta, invIncompleteBeta, digamma)

import qualified Statistics.Distribution as D
import Statistics.Distribution.Transform (LinearTransform (..))
import Statistics.Internal

-- | Student-T distribution
newtype StudentT = StudentT { studentTndf :: Double }
                   deriving (Eq, Typeable, Data, Generic)

instance Show StudentT where
  showsPrec i (StudentT ndf) = defaultShow1 "studentT" ndf i
instance Read StudentT where
  readPrec = defaultReadPrecM1 "studentT" studentTE

instance ToJSON StudentT
instance FromJSON StudentT where
  parseJSON (Object v) = do
    ndf <- v .: "studentTndf"
    maybe (fail $ errMsg ndf) return $ studentTE ndf
  parseJSON _ = empty

instance Binary StudentT where
  put = put . studentTndf
  get = do
    ndf <- get
    maybe (fail $ errMsg ndf) return $ studentTE ndf

-- | Create Student-T distribution. Number of parameters must be positive.
studentT :: Double -> StudentT
studentT ndf = maybe (error $ errMsg ndf) id $ studentTE ndf

-- | Create Student-T distribution. Number of parameters must be positive.
studentTE :: Double -> Maybe StudentT
studentTE ndf
  | ndf > 0   = Just (StudentT ndf)
  | otherwise = Nothing

errMsg :: Double -> String
errMsg _ = modErr "studentT" "non-positive number of degrees of freedom"

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

instance D.ContDistr StudentT where
  density    d@(StudentT ndf) x = exp (logDensityUnscaled d x) / sqrt ndf
  logDensity d@(StudentT ndf) x = logDensityUnscaled d x - log (sqrt ndf)
  quantile = quantile

cumulative :: StudentT -> Double -> Double
cumulative (StudentT ndf) x
  | x > 0     = 1 - 0.5 * ibeta
  | otherwise = 0.5 * ibeta
  where
    ibeta = incompleteBeta (0.5 * ndf) 0.5 (ndf / (ndf + x*x))

complCumulative :: StudentT -> Double -> Double
complCumulative (StudentT ndf) x
  | x > 0     = 0.5 * ibeta
  | otherwise = 1 - 0.5 * ibeta
  where
    ibeta = incompleteBeta (0.5 * ndf) 0.5 (ndf / (ndf + x*x))

logDensityUnscaled :: StudentT -> Double -> Double
logDensityUnscaled (StudentT ndf) x =
    log (ndf / (ndf + x*x)) * (0.5 * (1 + ndf)) - logBeta 0.5 (0.5 * ndf)

quantile :: StudentT -> Double -> Double
quantile (StudentT ndf) p
  | p >= 0 && p <= 1 =
    let x = invIncompleteBeta (0.5 * ndf) 0.5 (2 * min p (1 - p))
    in case sqrt $ ndf * (1 - x) / x of
         r | p < 0.5   -> -r
           | otherwise -> r
  | otherwise = modErr "quantile" $ "p must be in [0,1] range. Got: "++show p

instance D.MaybeMean StudentT where
  maybeMean (StudentT ndf) | ndf > 1   = Just 0
                           | otherwise = Nothing

instance D.MaybeVariance StudentT where
  maybeVariance (StudentT ndf) | ndf > 2   = Just $! ndf / (ndf - 2)
                               | otherwise = Nothing

instance D.Entropy StudentT where
  entropy (StudentT ndf) =
    0.5 * (ndf+1) * (digamma ((1+ndf)/2) - digamma(ndf/2))
    + log (sqrt ndf)
    + logBeta (ndf/2) 0.5

instance D.MaybeEntropy StudentT where
  maybeEntropy = Just . D.entropy

instance D.ContGen StudentT where
  genContVar = D.genContinuous

-- | Create an unstandardized Student-t distribution.
studentTUnstandardized :: Double -- ^ Number of degrees of freedom
                       -> Double -- ^ Central value (0 for standard Student T distribution)
                       -> Double -- ^ Scale parameter
                       -> LinearTransform StudentT
studentTUnstandardized ndf mu sigma
  | sigma > 0 = LinearTransform mu sigma $ studentT ndf
  | otherwise = modErr "studentTUnstandardized" $ "sigma must be > 0. Got: " ++ show sigma

modErr :: String -> String -> a
modErr fun msg = error $ "Statistics.Distribution.StudentT." ++ fun ++ ": " ++ msg


{-# LANGUAGE ScopedTypeVariables #-}

{-|
Module      : Numeric.LinearAlgebra.Static.Util
Description : Functions to manipulate vectors and matrices
Copyright   : (c) Ryan Orendorff, 2020
License     : BSD3
Stability   : experimental
-}
module Numeric.LinearAlgebra.Static.Util
    ( atM
    , atV
    , foldR
    , foldR1
    )
where

import           GHC.TypeLits
import           Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra         as LA
import           Data.Fin
import qualified Data.Vector.Storable          as V

---------------------------------------------------------------
-- Get data from vector/matrix with compile time bound check --
---------------------------------------------------------------

-- | Get a value from a vector at a given coordinate/index with guarantee
-- that the element index is within the length of the vector at compile
-- time.
atV
    :: forall n
     . (KnownNat n)
    => R n
    -> Fin n
    -> Double
atV v p = extract v LA.! (fromFin p)

-- | Get element from matrix at a given coordinate/index with guarantee that
-- the element index is within the shape of the matrix at compile time.
atM
    :: forall m n
     . (KnownNat m, KnownNat n)
    => L m n
    -> Fin m
    -> Fin n
    -> Double
atM mat m n = extract mat `LA.atIndex` (fromFin m, fromFin n)

-- | Fold a vector.
foldR :: KnownNat n => (b -> Double -> b) -> b -> R n -> b
foldR f initial x = V.foldl' f initial (extract x)

-- | Fold a vector without an initial element.
foldR1 :: KnownNat n => (Double -> Double -> Double) -> R n -> Double
foldR1 f x = V.foldl1' f (extract x)


{-# LANGUAGE FlexibleContexts #-}
-- |
-- Module    : Statistics.Sample
-- Copyright : (c) 2008 Don Stewart, 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Commonly used sample statistics, also known as descriptive
-- statistics.

module Statistics.Sample
    (
    -- * Types
      Sample
    , WeightedSample
    -- * Descriptive functions
    , range

    -- * Statistics of location
    , mean
    , welfordMean
    , meanWeighted
    , harmonicMean
    , geometricMean

    -- * Statistics of dispersion
    -- $variance

    -- ** Functions over central moments
    , centralMoment
    , centralMoments
    , skewness
    , kurtosis

    -- ** Two-pass functions (numerically robust)
    -- $robust
    , variance
    , varianceUnbiased
    , meanVariance
    , meanVarianceUnb
    , stdDev
    , varianceWeighted
    , stdErrMean

    -- ** Single-pass functions (faster, less safe)
    -- $cancellation
    , fastVariance
    , fastVarianceUnbiased
    , fastStdDev

    -- * Joint distributions
    , covariance
    , correlation
    , pair
    -- * References
    -- $references
    ) where

import Statistics.Function (minMax)
import Statistics.Sample.Internal (robustSumVar, sum)
import Statistics.Types.Internal  (Sample,WeightedSample)
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U

-- Operator ^ will be overridden
import Prelude hiding ((^), sum)

-- | /O(n)/ Range. The difference between the largest and smallest
-- elements of a sample.
range :: (G.Vector v Double) => v Double -> Double
range s = hi - lo
    where (lo , hi) = minMax s
{-# INLINE range #-}

-- | /O(n)/ Arithmetic mean.  This uses Kahan-Babuška-Neumaier
-- summation, so is more accurate than 'welfordMean' unless the input
-- values are very large.
mean :: (G.Vector v Double) => v Double -> Double
mean xs = sum xs / fromIntegral (G.length xs)
{-# SPECIALIZE mean :: U.Vector Double -> Double #-}
{-# SPECIALIZE mean :: V.Vector Double -> Double #-}

-- | /O(n)/ Arithmetic mean.  This uses Welford's algorithm to provide
-- numerical stability, using a single pass over the sample data.
--
-- Compared to 'mean', this loses a surprising amount of precision
-- unless the inputs are very large.
welfordMean :: (G.Vector v Double) => v Double -> Double
welfordMean = fini . G.foldl' go (T 0 0)
  where
    fini (T a _) = a
    go (T m n) x = T m' n'
        where m' = m + (x - m) / fromIntegral n'
              n' = n + 1
{-# SPECIALIZE welfordMean :: U.Vector Double -> Double #-}
{-# SPECIALIZE welfordMean :: V.Vector Double -> Double #-}

-- | /O(n)/ Arithmetic mean for weighted sample. It uses a single-pass
-- algorithm analogous to the one used by 'welfordMean'.
meanWeighted :: (G.Vector v (Double,Double)) => v (Double,Double) -> Double
meanWeighted = fini . G.foldl' go (V 0 0)
    where
      fini (V a _) = a
      go (V m w) (x,xw) = V m' w'
          where m' | w' == 0   = 0
                   | otherwise = m + xw * (x - m) / w'
                w' = w + xw
{-# INLINE meanWeighted #-}

-- | /O(n)/ Harmonic mean.  This algorithm performs a single pass over
-- the sample.
harmonicMean :: (G.Vector v Double) => v Double -> Double
harmonicMean = fini . G.foldl' go (T 0 0)
  where
    fini (T b a) = fromIntegral a / b
    go (T x y) n = T (x + (1/n)) (y+1)
{-# INLINE harmonicMean #-}

-- | /O(n)/ Geometric mean of a sample containing no negative values.
geometricMean :: (G.Vector v Double) => v Double -> Double
geometricMean = exp . mean . G.map log
{-# INLINE geometricMean #-}

-- | Compute the /k/th central moment of a sample.  The central moment
-- is also known as the moment about the mean.
--
-- This function performs two passes over the sample, so is not subject
-- to stream fusion.
--
-- For samples containing many values very close to the mean, this
-- function is subject to inaccuracy due to catastrophic cancellation.
centralMoment :: (G.Vector v Double) => Int -> v Double -> Double
centralMoment a xs
    | a < 0  = error "Statistics.Sample.centralMoment: negative input"
    | a == 0 = 1
    | a == 1 = 0
    | otherwise = sum (G.map go xs) / fromIntegral (G.length xs)
  where
    go x = (x-m) ^ a
    m    = mean xs
{-# SPECIALIZE centralMoment :: Int -> U.Vector Double -> Double #-}
{-# SPECIALIZE centralMoment :: Int -> V.Vector Double -> Double #-}

-- | Compute the /k/th and /j/th central moments of a sample.
--
-- This function performs two passes over the sample, so is not subject
-- to stream fusion.
--
-- For samples containing many values very close to the mean, this
-- function is subject to inaccuracy due to catastrophic cancellation.
centralMoments :: (G.Vector v Double) => Int -> Int -> v Double -> (Double, Double)
centralMoments a b xs
    | a < 2 || b < 2 = (centralMoment a xs , centralMoment b xs)
    | otherwise      = fini . G.foldl' go (V 0 0) $ xs
  where go (V i j) x = V (i + d^a) (j + d^b)
            where d  = x - m
        fini (V i j) = (i / n , j / n)
        m            = mean xs
        n            = fromIntegral (G.length xs)
{-# SPECIALIZE
    centralMoments :: Int -> Int -> U.Vector Double -> (Double, Double) #-}
{-# SPECIALIZE
    centralMoments :: Int -> Int -> V.Vector Double -> (Double, Double) #-}

-- | Compute the skewness of a sample. This is a measure of the
-- asymmetry of its distribution.
--
-- A sample with negative skew is said to be /left-skewed/.  Most of
-- its mass is on the right of the distribution, with the tail on the
-- left.
--
-- > skewness $ U.to [1,100,101,102,103]
-- > ==> -1.497681449918257
--
-- A sample with positive skew is said to be /right-skewed/.
--
-- > skewness $ U.to [1,2,3,4,100]
-- > ==> 1.4975367033335198
--
-- A sample's skewness is not defined if its 'variance' is zero.
--
-- This function performs two passes over the sample, so is not subject
-- to stream fusion.
--
-- For samples containing many values very close to the mean, this
-- function is subject to inaccuracy due to catastrophic cancellation.
skewness :: (G.Vector v Double) => v Double -> Double
skewness xs = c3 * c2 ** (-1.5)
    where (c3 , c2) = centralMoments 3 2 xs
{-# SPECIALIZE skewness :: U.Vector Double -> Double #-}
{-# SPECIALIZE skewness :: V.Vector Double -> Double #-}

-- | Compute the excess kurtosis of a sample.  This is a measure of
-- the \"peakedness\" of its distribution.  A high kurtosis indicates
-- that more of the sample's variance is due to infrequent severe
-- deviations, rather than more frequent modest deviations.
--
-- A sample's excess kurtosis is not defined if its 'variance' is
-- zero.
--
-- This function performs two passes over the sample, so is not subject
-- to stream fusion.
--
-- For samples containing many values very close to the mean, this
-- function is subject to inaccuracy due to catastrophic cancellation.
kurtosis :: (G.Vector v Double) => v Double -> Double
kurtosis xs = c4 / (c2 * c2) - 3
    where (c4 , c2) = centralMoments 4 2 xs
{-# SPECIALIZE kurtosis :: U.Vector Double -> Double #-}
{-# SPECIALIZE kurtosis :: V.Vector Double -> Double #-}

-- $variance
--
-- The variance — and hence the standard deviation — of a
-- sample of fewer than two elements are both defined to be zero.

-- $robust
--
-- These functions use the compensated summation algorithm of Chan et
-- al. for numerical robustness, but require two passes over the
-- sample data as a result.
--
-- Because of the need for two passes, these functions are /not/
-- subject to stream fusion.

data V = V {-# UNPACK #-} !Double {-# UNPACK #-} !Double

-- | Maximum likelihood estimate of a sample's variance.  Also known
-- as the population variance, where the denominator is /n/.
variance :: (G.Vector v Double) => v Double -> Double
variance samp
    | n > 1     = robustSumVar (mean samp) samp / fromIntegral n
    | otherwise = 0
    where
      n = G.length samp
{-# SPECIALIZE variance :: U.Vector Double -> Double #-}
{-# SPECIALIZE variance :: V.Vector Double -> Double #-}

-- | Unbiased estimate of a sample's variance.  Also known as the
-- sample variance, where the denominator is /n/-1.
varianceUnbiased :: (G.Vector v Double) => v Double -> Double
varianceUnbiased samp
    | n > 1     = robustSumVar (mean samp) samp / fromIntegral (n-1)
    | otherwise = 0
    where
      n = G.length samp
{-# SPECIALIZE varianceUnbiased :: U.Vector Double -> Double #-}
{-# SPECIALIZE varianceUnbiased :: V.Vector Double -> Double #-}

-- | Calculate mean and maximum likelihood estimate of variance. This
-- function should be used if both mean and variance are required
-- since it will calculate mean only once.
meanVariance ::  (G.Vector v Double) => v Double -> (Double,Double)
meanVariance samp
  | n > 1     = (m, robustSumVar m samp / fromIntegral n)
  | otherwise = (m, 0)
    where
      n = G.length samp
      m = mean samp
{-# SPECIALIZE meanVariance :: U.Vector Double -> (Double,Double) #-}
{-# SPECIALIZE meanVariance :: V.Vector Double -> (Double,Double) #-}

-- | Calculate mean and unbiased estimate of variance. This
-- function should be used if both mean and variance are required
-- since it will calculate mean only once.
meanVarianceUnb :: (G.Vector v Double) => v Double -> (Double,Double)
meanVarianceUnb samp
  | n > 1     = (m, robustSumVar m samp / fromIntegral (n-1))
  | otherwise = (m, 0)
    where
      n = G.length samp
      m = mean samp
{-# SPECIALIZE meanVarianceUnb :: U.Vector Double -> (Double,Double) #-}
{-# SPECIALIZE meanVarianceUnb :: V.Vector Double -> (Double,Double) #-}

-- | Standard deviation.  This is simply the square root of the
-- unbiased estimate of the variance.
stdDev :: (G.Vector v Double) => v Double -> Double
stdDev = sqrt . varianceUnbiased
{-# SPECIALIZE stdDev :: U.Vector Double -> Double #-}
{-# SPECIALIZE stdDev :: V.Vector Double -> Double #-}

-- | Standard error of the mean. This is the standard deviation
-- divided by the square root of the sample size.
stdErrMean :: (G.Vector v Double) => v Double -> Double
stdErrMean samp = stdDev samp / (sqrt . fromIntegral . G.length) samp
{-# SPECIALIZE stdErrMean :: U.Vector Double -> Double #-}
{-# SPECIALIZE stdErrMean :: V.Vector Double -> Double #-}

robustSumVarWeighted :: (G.Vector v (Double,Double)) => v (Double,Double) -> V
robustSumVarWeighted samp = G.foldl' go (V 0 0) samp
    where
      go (V s w) (x,xw) = V (s + xw*d*d) (w + xw)
          where d = x - m
      m = meanWeighted samp
{-# INLINE robustSumVarWeighted #-}

-- | Weighted variance. This is biased estimation.
varianceWeighted :: (G.Vector v (Double,Double)) => v (Double,Double) -> Double
varianceWeighted samp
    | G.length samp > 1 = fini $ robustSumVarWeighted samp
    | otherwise         = 0
    where
      fini (V s w) = s / w
{-# SPECIALIZE varianceWeighted :: U.Vector (Double,Double) -> Double #-}
{-# SPECIALIZE varianceWeighted :: V.Vector (Double,Double) -> Double #-}

-- $cancellation
--
-- The functions prefixed with the name @fast@ below perform a single
-- pass over the sample data using Knuth's algorithm. They usually
-- work well, but see below for caveats. These functions are subject
-- to array fusion.
--
-- /Note/: in cases where most sample data is close to the sample's
-- mean, Knuth's algorithm gives inaccurate results due to
-- catastrophic cancellation.

fastVar :: (G.Vector v Double) => v Double -> T1
fastVar = G.foldl' go (T1 0 0 0)
  where
    go (T1 n m s) x = T1 n' m' s'
      where n' = n + 1
            m' = m + d / fromIntegral n'
            s' = s + d * (x - m')
            d  = x - m

-- | Maximum likelihood estimate of a sample's variance.
fastVariance :: (G.Vector v Double) => v Double -> Double
fastVariance = fini . fastVar
  where fini (T1 n _m s)
          | n > 1     = s / fromIntegral n
          | otherwise = 0
{-# INLINE fastVariance #-}

-- | Unbiased estimate of a sample's variance.
fastVarianceUnbiased :: (G.Vector v Double) => v Double -> Double
fastVarianceUnbiased = fini . fastVar
  where fini (T1 n _m s)
          | n > 1     = s / fromIntegral (n - 1)
          | otherwise = 0
{-# INLINE fastVarianceUnbiased #-}

-- | Standard deviation.  This is simply the square root of the
-- maximum likelihood estimate of the variance.
fastStdDev :: (G.Vector v Double) => v Double -> Double
fastStdDev = sqrt . fastVariance
{-# INLINE fastStdDev #-}

-- | Covariance of sample of pairs. For empty sample it's set to
--   zero
covariance :: (G.Vector v (Double,Double), G.Vector v Double)
           => v (Double,Double)
           -> Double
covariance xy
  | n == 0    = 0
  | otherwise = mean $ G.zipWith (*)
                         (G.map (\x -> x - muX) xs)
                         (G.map (\y -> y - muY) ys)
  where
    n       = G.length xy
    (xs,ys) = G.unzip xy
    muX     = mean xs
    muY     = mean ys
{-# SPECIALIZE covariance :: U.Vector (Double,Double) -> Double #-}
{-# SPECIALIZE covariance :: V.Vector (Double,Double) -> Double #-}

-- | Correlation coefficient for sample of pairs. Also known as
--   Pearson's correlation. For empty sample it's set to zero.
correlation :: (G.Vector v (Double,Double), G.Vector v Double)
           => v (Double,Double)
           -> Double
correlation xy
  | n == 0    = 0
  | otherwise = cov / sqrt (varX * varY)
  where
    n       = G.length xy
    (xs,ys) = G.unzip xy
    (muX,varX) = meanVariance xs
    (muY,varY) = meanVariance ys
    cov = mean $ G.zipWith (*)
            (G.map (\x -> x - muX) xs)
            (G.map (\y -> y - muY) ys)
{-# SPECIALIZE correlation :: U.Vector (Double,Double) -> Double #-}
{-# SPECIALIZE correlation :: V.Vector (Double,Double) -> Double #-}

-- | Pair two samples. It's like 'G.zip' but requires that both
--   samples have equal size.
pair :: (G.Vector v a, G.Vector v b, G.Vector v (a,b)) => v a -> v b -> v (a,b)
pair va vb
  | G.length va == G.length vb = G.zip va vb
  | otherwise = error "Statistics.Sample.pair: vector must have same length"
{-# INLINE pair #-}

------------------------------------------------------------------------
-- Helper code. Monomorphic unpacked accumulators.

-- (^) operator from Prelude is just slow.
(^) :: Double -> Int -> Double
x ^ 1 = x
x ^ n = x * (x ^ (n-1))
{-# INLINE (^) #-}

-- don't support polymorphism, as we can't get unboxed returns if we use it.
data T = T {-# UNPACK #-}!Double {-# UNPACK #-}!Int

data T1 = T1 {-# UNPACK #-}!Int {-# UNPACK #-}!Double {-# UNPACK #-}!Double

{-

Consider this core:

with data T a = T !a !Int

$wfold :: Double#
               -> Int#
               -> Int#
               -> (# Double, Int# #)

and without,

$wfold :: Double#
               -> Int#
               -> Int#
               -> (# Double#, Int# #)

yielding to boxed returns and heap checks.

-}

-- $references
--
-- * Chan, T. F.; Golub, G.H.; LeVeque, R.J. (1979) Updating formulae
--   and a pairwise algorithm for computing sample
--   variances. Technical Report STAN-CS-79-773, Department of
--   Computer Science, Stanford
--   University. <ftp://reports.stanford.edu/pub/cstr/reports/cs/tr/79/773/CS-TR-79-773.pdf>
--
-- * Knuth, D.E. (1998) The art of computer programming, volume 2:
--   seminumerical algorithms, 3rd ed., p. 232.
--
-- * 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>
--
-- * West, D.H.D. (1979) Updating mean and variance estimates: an
--   improved method. /Communications of the ACM/
--   22(9):532–535. <http://doi.acm.org/10.1145/359146.359153>


{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Statistics
-- Copyright   :  (c) A. V. H. McPhail 2010, 2012, 2014
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  portable
--
-- Useful statistical functions
--
-----------------------------------------------------------------------------

module Numeric.Statistics (
                           Sample,Samples
                          , covarianceMatrix, correlationCoefficientMatrix
                          , meanList, meanArray, meanMatrix
                          , varianceList, varianceArray, varianceMatrix
                          --
                          , centre, cloglog, corcoeff, cut
                          , ranks, kendall, logit
                          , mahalanobis
                          , mode, moment
                          , ols, percentile, range
                          , run_count
                          , spearman, studentize
                          ) where

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

import Numeric.LinearAlgebra hiding(range)
--import Numeric.LinearAlgebra.Data hiding(range)
--import Numeric.LinearAlgebra.Devel
import Prelude hiding ((<>))
import qualified Data.Array.IArray as I
import qualified Data.List as DL
import qualified Data.Vector.Generic as GV
--import qualified Data.Vector.Storable as SV
import Foreign.Storable
import Numeric.GSL.Statistics
import Numeric.GSL.Sort(sort)

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

type Sample a = Vector a
type Samples a = I.Array Int (Vector a)

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

-- | the covariance matrix
covarianceMatrix :: Samples Double              -- ^ the dimensions of data (each vector being one dimension)
                 -> Matrix Double               -- ^ the symmetric covariance matrix
covarianceMatrix d = let (s,f) = I.bounds d
                      in fromArray2D $ I.array ((s,s),(f,f)) $ concat $ map (\(x,y) -> let c = covariance (d I.! x) (d I.! y) in if x == y then [((x,y),c)] else [((x,y),c),((y,x),c)]) $ filter (\(x,y) -> x <= y) $ I.range ((s,s),(f,f))

-- | the correlation coefficient: (cov x y) / (std x) (std y)
correlationCoefficientMatrix :: Samples Double -> Matrix Double
correlationCoefficientMatrix d = let (s,f) = I.bounds d
                           in fromArray2D $ I.array ((s,s),(f,f)) $ concat $ map (\(x,y) -> let { x' = d I.! x ; y' = d I.! y ; c = (covariance x' y') / ((stddev x') * (stddev y')) } in if x == y then [((x,y),c)] else [((x,y),c),((y,x),c)]) $ filter (\(x,y) -> x <= y) $ I.range ((s,s),(f,f))

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

-- | the mean of a list of vectors
meanList :: (Container Vector a, Num (Vector a), Fractional a) => [Sample a] -> Sample a
meanList []     = error "meanVectors: empty list"
meanList [s]    = s
meanList (s:ss) = let ln = fromIntegral $ length ss + 1
                  in scale (recip ln) $ foldl (+) s ss

-- | the mean of an array of vectors
meanArray :: (Container Vector a, Num (Vector a), Fractional a) => Samples a -> Sample a
meanArray a = meanList $ I.elems a

-- | the mean of a matrix with data series in rows
meanMatrix :: (Container Vector a, Num (Vector a), Fractional a) => Matrix a -> Sample a
meanMatrix a = meanList $ toRows a

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

-- | the variance of a list of vectors
varianceList :: (Container Vector a, Floating (Vector a), Num a, Fractional a) => [Sample a] -> Sample a
varianceList []  = error "varianceList: empty list"
varianceList [s] = konst 0 (size s)
varianceList l   = let mxs = meanList (map (** 2) l)
                       msx = (meanList l) ** 2
                   in mxs - msx

-- | the variance of an array of vectors
varianceArray :: (Container Vector a, Floating (Vector a), Fractional a) => Samples a -> Sample a
varianceArray a = varianceList $ I.elems a

-- | the variance of a matrix with data series in rows
varianceMatrix :: (Container Vector a, Floating (Vector a), Fractional a) => Matrix a -> Sample a
varianceMatrix a = varianceList $ toRows a

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

-- | centre the data to 0: (x - (mean x))
centre :: Vector Double -> Vector Double
centre v = v - (realToFrac (mean v))

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

-- | complementary log-log function
--cloglog :: Vector Double -> Vector Double
cloglog :: Floating a => a -> a
cloglog v = - log (- (log v))

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

-- | corcoeff = covariance x / (std dev x * std dev y)
corcoeff :: Vector Double -> Vector Double -> Double
corcoeff x y = (covariance x y)/((stddev x)*(stddev y))

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

-- | cut numerical data into intervals, data must fall inside the bounds
cut :: Vector Double
    -> Vector Double -- ^ intervals
    -> Vector Int    -- ^ data indexed by bin
cut v c  = let c' = sort c
           in GV.map (\x -> cut_helper 0 x c') v
    where
      cut_helper j x d
          | j >= size d                       = error "Numeric.Statistics: cut: data point not within interval"
          | x >= (d `atIndex` j) && x <= (d `atIndex` (j+1)) = j
          | otherwise                       = cut_helper (j + 1) x d

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

-- | return the rank of each element of the vector
--     multiple identical entries result in the average rank of those entries
--ranks :: Vector Double -> Vector Double
ranks :: (Fractional b, Storable b) => Vector Double -> Vector b
ranks v = let v' = sort v
          in GV.map (\x -> 1 + rank_helper x v') v
              where rank_helper x v' = let is = GV.elemIndices x v'
                                       in (realToFrac (GV.foldl (+) 0 is)) / (fromIntegral $ GV.length is)

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

-- | kendall's rank correlation τ
kendall :: Vector Double -> Vector Double -> Matrix Double
kendall x y = let ln = size x
                  rx = ranks x
                  ry = ranks y
                  r = fromColumns [rx,ry]
                  m = signum $ (kronecker r (asColumn $ konst 1.0 ln)) - (kronecker (asRow $ konst 1.0 ln) r)
                  c = rows m - 1
              in correlationCoefficientMatrix $ I.listArray (0,c) (toColumns m)

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

-- | (logit p) = log(p/(1-p))
--logit :: Vector Double -> Vector Double
logit :: (Floating b, Storable b)
        => Vector b -> Vector b
logit v =  GV.map (\x -> - (log ((1 / x) - 1))) v

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

-- | the Mahalanobis D-square distance between samples
--     columns are components and rows are observations (uses pseudoinverse)
mahalanobis :: Samples Double        -- ^ the data set
            -> Maybe (Sample Double) -- ^ (Just sample) to be measured or use mean when Nothing
            -> Double                -- ^ D^2
mahalanobis x u = let (_,xr) = I.bounds x
                      xl     = I.elems x
                      s'     = pinv $ covarianceMatrix x
                      xu     = case u of
                                 Nothing -> fromList $ map mean xl
                                 Just m  -> m
                      xm     = fromRows $ map ((-) xu) $ toRows $ fromColumns xl
                      --um     = asColumn xu
                      --w      = ((tr' xm) <> xm + (tr' um) <> um)/(fromIntegral $ xr - 1)
                      --w'     = inv w
                  in ((xm <> s' <> (tr' xm)) `atIndex` (0,0))

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

-- | a list of element frequencies
mode :: Vector Double -> [(Double,Integer)]
mode v = let w = sort v
         in DL.sortBy (\(_,n) (_,n') -> compare n' n) $ GV.foldr freqs [] w
            where freqs x []          = [(x,1)]
                  freqs x ((f,n):fns)
                      | f == x         = ((f,n+1):fns)
                      | otherwise     = ((x,1):(f,n):fns)

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

-- | the p'th moment of a vector
moment :: Integral a
       => a             -- ^ moment
       -> Bool          -- ^ calculate central moment
       -> Bool          -- ^ calculate absolute moment
       -> Vector Double -- ^ data
       -> Double
moment p c a v
    | p <= 0     = error "Numeric.Statistics.moment: negative moment requested"
--    | p == 1     = mean v
--    | p == 2     = variance v -- gives sample variance
    | otherwise = let u = if c then centre v else v
                      w = if a then abs u else u
                      x = GV.map (** (fromIntegral p)) w
                  in mean x

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

-- | ordinary least squares estimation for the multivariate model
--   Y = X B + e        rows are observations, columns are elements
--   mean e = 0, cov e = kronecker s I
ols :: (Num (Vector t), Field t)
      => Matrix t         -- ^ X
    -> Matrix t           -- ^ Y
    -> (Matrix t, Matrix t, Matrix t) -- ^ (OLS estimator for B, OLS estimator for s, OLS residuals)
ols x y
    | rows x /= rows y = error "Numeric.Statistics: ols: incorrect matrix dimensions"
    | otherwise       = let (xr,xc) = (rows x,cols x)
                            (yr,yc) = (rows y,cols y)
                            z = (tr' x) <> x
                            r = rank z
                            beta = if r == xc
                                      then (inv z) <> (tr' x) <> y
                                      else (pinv x) <> y
                            rr = y - x <> beta
                            sigma = ((tr' rr) <> rr) / (fromIntegral $ xr - r)
                        in (beta,rr,sigma)

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

-- | compute quantiles in percent
percentile :: Double        -- ^ percentile (0 - 100)
           -> Vector Double -- ^ data
           -> Double        -- ^ result
percentile p d = quantile (0.01*p) d

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

-- | the difference between the maximum and minimum of the input
range :: (Container c e, Num e) => c e -> e
range v = maxElement v - minElement v

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

-- | count the number of runs greater than or equal to @n@ in the data
run_count :: (Num a, Num t, Ord b, Ord a, Container Vector b)
            => a             -- ^ longest run to count
          -> Vector b        -- ^ data
          -> [(a, t)]        -- ^ [(run length,count)]
run_count n v = let w = subVector 1 (size v - 1) v
                    x = GV.foldr run_count' [(1,v `atIndex` 0)] w
                    y = map fst x
                    z = takeWhile (<= n) $  DL.sort y
                in foldr count [] z
    where run_count' m ((c,g):cs)
              | m < g             = ((c+1,m):cs)
              | otherwise         = ((1,m):(c,g):cs)
          run_count' _ [] = error "Lib.Numeric.Statistics.run_count.run_count'"
          count x []           = [(x,1)]
          count x ((yv,yc):ys)
              | x == yv         = ((yv,yc+1):ys)
              | otherwise      = ((x,1):(yv,yc):ys)

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

-- | Spearman's rank correlation coefficient
spearman :: Vector Double -> Vector Double -> Double
spearman x y = corcoeff (ranks x) (ranks y)

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

-- | centre and normalise a vector
studentize :: Vector Double -> Vector Double
studentize x = (centre x)/(fromList $ [stddev x])

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

--table

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

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


import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.LAPACK

m :: Matrix Double
m = (3><3)
  [7.589183,1.703609,-4.477162,
    -4.597851,9.434889,-6.543450,
    0.4588202,-6.115153,1.331191]

v :: Matrix Double
v = (3><1)
  [1.745005,-4.448092,-4.160842]


{-# LANGUAGE BangPatterns #-}

-- |
-- Module    : Statistics.Matrix.Fast.Algorithms
-- Copyright : 2019 Magalame
-- License   : BSD3
--
-- Useful matrix functions.

module Statistics.Matrix.Fast.Algorithms
    (
      qr
    ) where

import Control.Applicative ((<$>), (<*>))
import Control.Monad.ST (ST, runST)
import Prelude hiding (replicate)
import Statistics.Matrix (Matrix (..),dimension, for)
import qualified Statistics.Matrix.Mutable as M
import qualified Data.Vector.Unboxed as U

-- | /O(r*c)/ Compute the QR decomposition of a matrix.
-- The result returned is the matrices (/q/,/r/).
qr :: Matrix -> (Matrix, Matrix)
qr mat = runST $ do
  let (m,n) = dimension mat

  r <- M.replicate n n 0
  a <- M.thaw mat
  for 0 n $ \j -> do
    cn <- M.immutably a $ \aa -> sqrt $ normCol j aa
    M.unsafeWrite r j j cn
    for 0 m $ \i -> M.unsafeModify a i j (/ cn)
    for (j+1) n $ \jj -> do
      p <- innerProduct a j jj
      M.unsafeWrite r j jj p
      for 0 m $ \i -> do
        aij <- M.unsafeRead a i j
        M.unsafeModify a i jj $ subtract (p * aij)
  (,) <$> M.unsafeFreeze a <*> M.unsafeFreeze r

normCol :: Int -> Matrix -> Double
normCol jthcol (Matrix r c v) = sub 0 0
  where sub !acc !ij | ij == r = acc
                     | otherwise = sub ( valCol*valCol + acc ) (ij+1)
                                   where
                                    valCol = U.unsafeIndex v (ij*c+jthcol)

innerProduct :: M.MMatrix s -> Int -> Int -> ST s Double
innerProduct mmat j k = M.immutably mmat $ \mat ->
  dotCol j mat k mat

dotCol :: Int -> Matrix -> Int -> Matrix -> Double
dotCol jthcol (Matrix r1 c1 v1) kthcol (Matrix _ c2 v2) = sub 0 0
  where sub !acc !ij | ij == r1 = acc
                     | otherwise = sub ( valColj*valColk + acc ) (ij+1)
                                   where
                                    valColk = U.unsafeIndex v2 (ij*c2+kthcol)
                                    valColj = U.unsafeIndex v1 (ij*c1+jthcol)


module LogisticRegression where

import Numeric.LinearAlgebra

sigmoid :: Matrix Double -> Matrix Double
sigmoid x = fromRows $ map (cmap sigmoid') (toRows x)

-- sigmoid applied to a single real number.
sigmoid' :: Double -> Double
sigmoid' x = 1 / (1 + (exp (-1 * x)))


{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DataKinds #-}
module Lib
    (
    ) where
import Control.Lens
import Numeric.LinearAlgebra.Static as NLS
import Numeric.Backprop
import Debug.SimpleReflect as SR
data Network = Net { _weight1 :: NLS.L 20 100
                   , _bias1   :: NLS.R 20
                   , _weight2 :: NLS.L  5  20
                   , _bias2   :: NLS.R  5
                   }
Control.Lens.makeLenses ''Network
runNet :: Network -> NLS.R 100 -> NLS.R 5
runNet net x = z
  where
    y = logistic $ (net ^. weight1) #> x + (net ^. bias1)
    z = logistic $ (net ^. weight2) #> y + (net ^. bias2)
logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))
myFun :: Floating a => a -> a -> a
myFun x y = x + y
myFunc :: Floating a => a -> a
myFunc x = sqrt (x*4)
{-
data Point = Point { _x :: Double, _y :: Double }deriving (Show)
data Atom = Atom { _element :: String, _point :: Point } deriving (Show)
data Molecule = Molecule { _atoms :: [Atom] } deriving (Show)
Control.Lens.makeLenses ''Atom
Control.Lens.makeLenses ''Point
Control.Lens.makeLenses ''Molecule
shiftAtomX :: Atom -> Atom
shiftAtomX = over (point . x) (+ 1)
shiftMoleculeX :: Molecule -> Molecule
shiftMoleculeX = over (atoms . traverse . point . x) (+ 1)
atom1 = Atom { _element = "C", _point = Point { _x = 1.0, _y = 2.0 } }
atom2 = Atom { _element = "O", _point = Point { _x = 3.0, _y = 4.0 } }
molecule = Molecule { _atoms = [atom1, atom2] }
newMolecule = shiftMoleculeX molecule
t = atom1 ^. (point . x)
-}


module FinancialTimeseries.Algorithm.MovingAverage where

import qualified Data.Vector as Vec

import qualified Statistics.Sample as Sample

import FinancialTimeseries.Type.Labeled (Labeled(..))
import FinancialTimeseries.Type.Segment (Segment(..), HalfSegment(..))
import FinancialTimeseries.Type.Strategy (Strategy(..))
import FinancialTimeseries.Type.Timeseries (TimeseriesRaw(..), Timeseries(..))
import FinancialTimeseries.Type.Types (StripPrice, stripPrice)

import FinancialTimeseries.Util.Pretty (Pretty, pretty)
import FinancialTimeseries.Util.ToFileString (ToFileString, toFileString)

newtype Window = Window {
  unWindow :: Int
  } deriving (Show)

instance ToFileString Window where
  toFileString (Window w) = "Window-" ++ show w

instance Pretty Window where
  pretty = show

data UpDown =
  Up Int
  | Down Int
  deriving (Show)

movingAverage ::
  (StripPrice price, Real a, Fractional a) =>
  Window -> TimeseriesRaw price a -> Timeseries price a
movingAverage (Window m) ts@(TimeseriesRaw _ pvs) =
  let vs = stripPrice pvs
      idx = Vec.fromList [0 .. Vec.length vs - m]

      slc i =
        let ws = Vec.slice i m vs
            (tn, xn) = Vec.last ws
        in (tn, realToFrac (Sample.mean (Vec.map (realToFrac . snd) ws)), xn)
      us = Vec.map slc idx

      f i ((_, m0, x0), (_, m1, x1)) acc
        | m0 <= x0 && m1 > x1 = Up (i+m) : acc
        | m0 > x0 && m1 <= x1 = Down (i+m) : acc
        | otherwise = acc

      xs = Vec.ifoldr' f [] (Vec.zip us (Vec.tail us))

      g [] = (Nothing, [])
      g [Down i] = (Just (HalfSegment i), [])
      g (Down i:Up j:zs) = fmap (Segment i j :) (g zs)
      g (Up _:zs) = g zs
      g (Down _:zs) = g zs

      h (t, s, _) = (t, s)

      (hs, isegs) = g xs

      res = Timeseries {
        timeseriesRaw = ts
        , investedSegments = isegs
        , lastSegment = hs
        , additionalSeries = [Labeled "Moving Average" (Vec.map h us)]
        }
  in res

movingAverageStrategy ::
  (StripPrice price, Real a, Fractional a) =>
  Window -> Strategy Window price a
movingAverageStrategy w = Strategy "MovingAverage" w movingAverage


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE CPP #-}

{- |
Module      :  Physics.Learn.BlochSphere
Copyright   :  (c) Scott N. Walck 2016
License     :  BSD3 (see LICENSE)
Maintainer  :  Scott N. Walck <walck@lvc.edu>
Stability   :  experimental

This module contains functions for displaying the
state of a spin-1/2 particle or other quantum two-level
system as a point on the Bloch Sphere.
-}

module Physics.Learn.BlochSphere
    ( VisObj
    , toPos
    , ketToPos
    , staticBlochSphere
    , displayStaticState
    , animatedBlochSphere
    , simulateBlochSphere
    , simulateBlochSphereK
    , stateProp
    , statePropK
    , evolutionBlochSphere
    , evolutionBlochSphereK
    , hamRabi
    )
    where

import qualified Physics.Learn.QuantumMat as M
import qualified Physics.Learn.Ket as K
import Physics.Learn.Ket
    ( Ket
    , Operator
    , (<>)
    , dagger
    )
import Numeric.LinearAlgebra
    ( Vector
    , Matrix
    , C
    , iC
--    , (<>)  -- matrix multiplication
--    , (|>)  -- vector definition
    , (!)   -- vector element access
    , (><)  -- matrix definition
    , scale
    , size
    )
import Data.Complex
    ( Complex(..)
    , conjugate
    , realPart
    , imagPart
    )
import Physics.Learn
    ( Position
    , v3FromPos
    , cart
    )
import SpatialMath
    ( Euler(..)
    )
import Vis
    ( VisObject(..)
    , Flavour(..)
    , Options(..)
    , Camera0(..)
    , defaultOpts
    , display
    , simulate
    , blue
    , red
    )
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif

{-
3 ways to specify the state of a spin-1/2 particle:
Vector C
Ket
Position  (Bloch vector)

2 ways to specify a Hamiltonian:
Matrix C
Operator

3 choices for Vis' world:
(Float, Vector C)
(Float, Ket)
(Float, Position)
-}

-- | A Vis object.
type VisObj = VisObject Double

-- | Convert a 2x1 complex state vector for a qubit
--   into Bloch (x,y,z) coordinates.
toPos :: Vector C -> Position
toPos v
    = if size v /= 2
      then error "toPos only for size 2 vectors"
      else let z1 = v ! 0
               z2 = v ! 1
           in cart (2 * realPart (conjugate z1 * z2))
                   (2 * imagPart (conjugate z1 * z2))
                   (realPart (conjugate z1 * z1 - conjugate z2 * z2))

-- | Convert a qubit ket
--   into Bloch (x,y,z) coordinates.
ketToPos :: Ket -> Position
ketToPos psi
    = if K.dim psi /= 2
      then error "ketToPos only for qubit kets"
      else let z1 = dagger K.zp <> psi
               z2 = dagger K.zm <> psi
           in cart (2 * realPart (conjugate z1 * z2))
                   (2 * imagPart (conjugate z1 * z2))
                   (realPart (conjugate z1 * z1 - conjugate z2 * z2))

-- | A static 'VisObj' for the state of a qubit.
staticBlochSphere :: Position -> VisObj
staticBlochSphere r
    = RotEulerDeg (Euler 270 0 0) $ RotEulerDeg (Euler 0 180 0) $
      VisObjects [ Sphere 1 Wireframe blue
                 , Trans (v3FromPos r) (Sphere 0.05 Solid red)
                 ]

displayStaticBlochSphere :: Position -> IO ()
displayStaticBlochSphere r
    = display myOptions (staticBlochSphere r)

-- | Display a qubit state vector as a point on the Bloch Sphere.
displayStaticState :: Vector C -> IO ()
displayStaticState = displayStaticBlochSphere . toPos

-- | Given a Bloch vector as a function of time,
--   return a 'VisObj' as a function of time.
animatedBlochSphere :: (Double -> Position) -> (Float -> VisObj)
animatedBlochSphere f
    = staticBlochSphere . f . realToFrac

-- | Given a sample rate, initial qubit state vector, and
--   state propagation function, produce a simulation.
--   The 'Float' in the state propagation function is the time
--   since the beginning of the simulation.
simulateBlochSphere :: Double -> Vector C -> (Float -> (Float,Vector C) -> (Float,Vector C)) -> IO ()
simulateBlochSphere sampleRate initial statePropFunc
    = simulate myOptions sampleRate (0,initial) (staticBlochSphere . toPos . snd) statePropFunc

-- | Given a sample rate, initial qubit state ket, and
--   state propagation function, produce a simulation.
--   The 'Float' in the state propagation function is the time
--   since the beginning of the simulation.
simulateBlochSphereK :: Double -> Ket -> (Float -> (Float,Ket) -> (Float,Ket)) -> IO ()
simulateBlochSphereK sampleRate initial statePropFuncK
    = simulate myOptions sampleRate (0,initial) (staticBlochSphere . ketToPos . snd) statePropFuncK

{-
-- | Given a sample rate, initial qubit state vector, and
--   state propagation function, produce a simulation.
--   The 'Float' in the state propagation function is the time
--   since the beginning of the simulation.
playBlochSphere :: Double -> Vector C -> (Float -> (Float,Vector C) -> (Float,Vector C)) -> IO ()
playBlochSphere sampleRate initial statePropFunc
    = play myOptions sampleRate (0,initial) (staticBlochSphere . toPos . snd) statePropFunc
-}

-- | Produce a state propagation function from a time-dependent Hamiltonian.
stateProp :: (Double -> Matrix C) -> Float -> (Float,Vector C) -> (Float,Vector C)
stateProp ham tNew (tOld,v)
    = (tNew, M.timeEv (realToFrac dt) (ham tMid) v)
      where
        dt = tNew - tOld
        tMid = realToFrac $ (tNew + tOld) / 2

-- | Produce a state propagation function from a time-dependent Hamiltonian.
statePropK :: (Double -> Operator) -> Float -> (Float,Ket) -> (Float,Ket)
statePropK ham tNew (tOld,psi)
    = (tNew, K.timeEv (realToFrac dt) (ham tMid) psi)
      where
        dt = tNew - tOld
        tMid = realToFrac $ (tNew + tOld) / 2

-- | Given an initial qubit state and a time-dependent Hamiltonian,
--   produce a visualization.
evolutionBlochSphere :: Vector C -> (Double -> Matrix C) -> IO ()
evolutionBlochSphere psi0 ham
    = simulateBlochSphere 0.01 psi0 (stateProp ham)

-- | Given an initial qubit ket and a time-dependent Hamiltonian,
--   produce a visualization.
evolutionBlochSphereK :: Ket -> (Double -> Operator) -> IO ()
evolutionBlochSphereK psi0 ham
    = simulateBlochSphereK 0.01 psi0 (statePropK ham)

myOptions :: Options
myOptions = defaultOpts {optWindowName = "Bloch Sphere"
                        ,optInitialCamera = Just (Camera0 75 20 4)}

{-
staticBz1 :: IO ()
staticBz1 = evolutionBlochSphere M.xp (const (scale 0.9 M.sz))

staticBz2 :: IO ()
staticBz2 = evolutionBlochSphere ((2|>) [(cos (pi / 8)), (sin (pi / 8))]) (const M.sz)

staticBy1 :: IO ()
staticBy1 = evolutionBlochSphere M.xp (const M.sy)
-}

-- | Hamiltonian for nuclear magnetic resonance.
--   Explain omega0, omegaR, omega.
hamRabi :: Double ->  Double ->  Double ->  Double -> Matrix C
hamRabi omega0 omegaR omega t
    = let h11 = omega0 :+ 0
          h12 = (omegaR :+ 0) * exp (-iC * ((omega * t) :+ 0))
      in scale (1/2) $ (2><2) [h11, h12, (conjugate h12), (-h11)]

-- need to scale time

-- a pi pulse


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE Safe #-}
{-# LANGUAGE MultiParamTypeClasses, FunctionalDependencies #-}
{-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-}
{-# LANGUAGE CPP #-}

{- |
Module      :  Physics.Learn.Ket
Copyright   :  (c) Scott N. Walck 2016-2018
License     :  BSD3 (see LICENSE)
Maintainer  :  Scott N. Walck <walck@lvc.edu>
Stability   :  experimental

This module contains ket vectors, bra vectors,
and operators for quantum mechanics.
-}

-- a Ket layer on top of QuantumMat

module Physics.Learn.Ket
    (
    -- * Basic data types
      C
    , i
    , magnitude
    , Ket
    , Bra
    , Operator
    -- * Kets for spin-1/2 particles
    , xp
    , xm
    , yp
    , ym
    , zp
    , zm
    , np
    , nm
    -- * Operators for spin-1/2 particles
    , sx
    , sy
    , sz
    , sn
    , sn'
    -- * Quantum Dynamics
    , timeEvOp
    , timeEv
    -- * Composition
    , Kron(..)
    -- * Measurement
    , possibleOutcomes
    , outcomesProjectors
    , outcomesProbabilities
--    , prob
--    , probs
    -- * Generic multiplication
    , Mult(..)
    -- * Adjoint operation
    , Dagger(..)
    -- * Normalization
    , HasNorm(..)
    -- * Representation
    , Representable(..)
    -- * Orthonormal bases
    , OrthonormalBasis
    , makeOB
    , listBasis
    , size
    -- * Orthonormal bases for spin-1/2 particles
    , xBasis
    , yBasis
    , zBasis
    , nBasis
    -- , angularMomentumXMatrix
    -- , angularMomentumYMatrix
    -- , angularMomentumZMatrix
    -- , angularMomentumPlusMatrix
    -- , angularMomentumMinusMatrix
    -- , jXMatrix
    -- , jYMatrix
    -- , jZMatrix
    -- , matrixCommutator
    -- , rotationMatrix
    -- , jmColumn
    )
    where

-- We try to import only from QuantumMat
-- and not from Numeric.LinearAlgebra

import qualified Data.Complex as C
import Data.Complex
    ( Complex(..)
    , conjugate
    )
import qualified Physics.Learn.QuantumMat as M
import Physics.Learn.QuantumMat
    ( C
    , Vector
    , Matrix
    , (#>)
    , (<#)
    , conjugateTranspose
    , scaleV
    , scaleM
    , conjV
    , fromList
    , toList
    , fromLists
    )
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif

infixl 7 <>

-- | A ket vector describes the state of a quantum system.
data Ket = Ket (Vector C)

instance Show Ket where
    show k =
        let message = "Use 'rep <basis name> <ket name>'."
        in if dim k == 2
           then "Representation in zBasis:\n" ++
                show (rep zBasis k) ++ "\n" ++ message
           else message

-- | An operator describes an observable (a Hermitian operator)
--   or an action (a unitary operator).
data Operator = Operator (Matrix C)

instance Show Operator where
    show op =
        let message = "Use 'rep <basis name> <operator name>'."
        in if dim op == 2
           then "Representation in zBasis:\n" ++
                show (rep zBasis op) ++ "\n" ++ message
           else message

-- | A bra vector describes the state of a quantum system.
data Bra = Bra (Vector C)

instance Show Bra where
    show _ = "<bra>\nTry 'rep zBasis <bra name>'"

magnitude :: C -> Double
magnitude = C.magnitude

i :: C
i = 0 :+ 1

-- | Generic multiplication including inner product,
--   outer product, operator product, and whatever else makes sense.
--   No conjugation takes place in this operation.
class Mult a b c | a b -> c where
    (<>) :: a -> b -> c

instance Mult C C C where
    z1 <> z2 = z1 * z2

instance Mult C Ket Ket where
    c <> Ket matrixKet = Ket (scaleV c matrixKet)

instance Mult C Bra Bra where
    c <> Bra matrixBra = Bra (scaleV c matrixBra)

instance Mult C Operator Operator where
    c <> Operator m = Operator (scaleM c m)

instance Mult Ket C Ket where
    Ket matrixKet <> c = Ket (scaleV c matrixKet)

instance Mult Bra C Bra where
    Bra matrixBra <> c = Bra (scaleV c matrixBra)

instance Mult Operator C Operator where
    Operator m <> c = Operator (scaleM c m)

instance Mult Bra Ket C where
    Bra matrixBra <> Ket matrixKet
        = sum $ zipWith (*) (toList matrixBra) (toList matrixKet)

instance Mult Bra Operator Bra where
    Bra matrixBra <> Operator matrixOp
        = Bra (matrixBra <# matrixOp)

instance Mult Operator Ket Ket where
    Operator matrixOp <> Ket matrixKet
        = Ket (matrixOp #> matrixKet)

instance Mult Ket Bra Operator where
    Ket k <> Bra b =
        Operator
        (fromLists [[ x*y | y <- toList b] | x <- toList k])

instance Mult Operator Operator Operator where
    Operator m1 <> Operator m2 = Operator (m1 M.<> m2)

instance Num Ket where
    Ket v1 + Ket v2 = Ket (v1 + v2)
    Ket v1 - Ket v2 = Ket (v1 - v2)
    (*)         = error "Multiplication is not defined for kets"
    negate (Ket v) = Ket (negate v)
    abs         = error "abs is not defined for kets"
    signum      = error "signum is not defined for kets"
    fromInteger = error "fromInteger is not defined for kets"

instance Num Bra where
    Bra v1 + Bra v2 = Bra (v1 + v2)
    Bra v1 - Bra v2 = Bra (v1 - v2)
    (*)         = error "Multiplication is not defined for bra vectors"
    negate (Bra v) = Bra (negate v)
    abs         = error "abs is not defined for bra vectors"
    signum      = error "signum is not defined for bra vectors"
    fromInteger = error "fromInteger is not defined for bra vectors"

instance Num Operator where
    Operator v1 + Operator v2 = Operator (v1 + v2)
    Operator v1 - Operator v2 = Operator (v1 - v2)
    Operator v1 * Operator v2 = Operator (v1 M.<> v2)
    negate (Operator v) = Operator (negate v)
    abs         = error "abs is not defined for operators"
    signum      = error "signum is not defined for operators"
    fromInteger = error "fromInteger is not defined for operators"

-- | The adjoint operation on complex numbers, kets,
--   bras, and operators.
class Dagger a b | a -> b where
    dagger :: a -> b

instance Dagger Ket Bra where
    dagger (Ket v) = Bra (conjV v)

instance Dagger Bra Ket where
    dagger (Bra v) = Ket (conjV v)

instance Dagger Operator Operator where
    dagger (Operator m) = Operator (conjugateTranspose m)

instance Dagger C C where
    dagger c = conjugate c

class HasNorm a where
    norm      :: a -> Double
    normalize :: a -> a

instance HasNorm Ket where
    norm (Ket v) = M.norm v
    normalize k  = (1 / norm k :+ 0) <> k

instance HasNorm Bra where
    norm (Bra v) = M.norm v
    normalize b  = (1 / norm b :+ 0) <> b

-- | An orthonormal basis of kets.
newtype OrthonormalBasis = OB [Ket]
    deriving (Show)

-- | Make an orthonormal basis from a list of linearly independent kets.
makeOB :: [Ket] -> OrthonormalBasis
makeOB = OB . gramSchmidt

size :: OrthonormalBasis -> Int
size (OB ks) = length ks

listBasis :: OrthonormalBasis -> [Ket]
listBasis (OB ks) = ks

{-
newOrthonormalBasis :: Int -> OrthonormalBasis
newOrthonormalBasis = undefined
-}

class Representable a b | a -> b where
    rep :: OrthonormalBasis -> a -> b
    dim :: a -> Int

instance Representable Ket (Vector C) where
    rep (OB ks) k = fromList (map (\bk -> dagger bk <> k) ks)
    dim (Ket v) = M.dim v

instance Representable Bra (Vector C) where
    rep (OB ks) b = fromList (map (\bk -> b <> bk) ks)
    dim (Bra v) = M.dim v

instance Representable Operator (Matrix C) where
    rep (OB ks) op = fromLists [[ dagger k1 <> op <> k2 | k2 <- ks ] | k1 <- ks ]
    dim (Operator m) = let (p,q) = M.size m
                       in if p == q then p else error "dim: non-square operator"

--------------
-- Spin 1/2 --
--------------

-- | State of a spin-1/2 particle if measurement
--   in the x-direction would give angular momentum +hbar/2.
xp :: Ket
xp = Ket M.xp

-- | State of a spin-1/2 particle if measurement
--   in the x-direction would give angular momentum -hbar/2.
xm :: Ket
xm = Ket M.xm

-- | State of a spin-1/2 particle if measurement
--   in the y-direction would give angular momentum +hbar/2.
yp :: Ket
yp = Ket M.yp

-- | State of a spin-1/2 particle if measurement
--   in the y-direction would give angular momentum -hbar/2.
ym :: Ket
ym = Ket M.ym

-- | State of a spin-1/2 particle if measurement
--   in the z-direction would give angular momentum +hbar/2.
zp :: Ket
zp = Ket M.zp

-- | State of a spin-1/2 particle if measurement
--   in the z-direction would give angular momentum -hbar/2.
zm :: Ket
zm = Ket M.zm

-- | State of a spin-1/2 particle if measurement
--   in the n-direction, described by spherical polar angle theta
--   and azimuthal angle phi, would give angular momentum +hbar/2.
np :: Double -> Double -> Ket
np theta phi
    = (cos (theta / 2) :+ 0) <> zp
      + (sin (theta / 2) :+ 0) * (cos phi :+ sin phi) <> zm

-- | State of a spin-1/2 particle if measurement
--   in the n-direction, described by spherical polar angle theta
--   and azimuthal angle phi, would give angular momentum -hbar/2.
nm :: Double -> Double -> Ket
nm theta phi
    = (sin (theta / 2) :+ 0) <> zp
      - (cos (theta / 2) :+ 0) * (cos phi :+ sin phi) <> zm

-- | The orthonormal basis composed of 'xp' and 'xm'.
xBasis :: OrthonormalBasis
xBasis = makeOB [xp,xm]

-- | The orthonormal basis composed of 'yp' and 'ym'.
yBasis :: OrthonormalBasis
yBasis = makeOB [yp,ym]

-- | The orthonormal basis composed of 'zp' and 'zm'.
zBasis :: OrthonormalBasis
zBasis = makeOB [zp,zm]

-- | Given spherical polar angle theta and azimuthal angle phi,
--   the orthonormal basis composed of 'np' theta phi and 'nm' theta phi.
nBasis :: Double -> Double -> OrthonormalBasis
nBasis theta phi = makeOB [np theta phi,nm theta phi]

-- | The Pauli X operator.
sx :: Operator
sx = xp <> dagger xp - xm <> dagger xm

-- | The Pauli Y operator.
sy :: Operator
sy = yp <> dagger yp - ym <> dagger ym

-- | The Pauli Z operator.
sz :: Operator
sz = zp <> dagger zp - zm <> dagger zm

-- | Pauli operator for an arbitrary direction given
--   by spherical coordinates theta and phi.
sn :: Double -> Double -> Operator
sn theta phi
    = (sin theta * cos phi :+ 0) <> sx +
      (sin theta * sin phi :+ 0) <> sy +
      (cos theta           :+ 0) <> sz

-- | Alternative definition
--   of Pauli operator for an arbitrary direction.
sn' :: Double -> Double -> Operator
sn' theta phi
    = np theta phi <> dagger (np theta phi) -
      nm theta phi <> dagger (nm theta phi)

----------------------
-- Quantum Dynamics --
----------------------

-- | Given a time step and a Hamiltonian operator,
--   produce a unitary time evolution operator.
--   Unless you really need the time evolution operator,
--   it is better to use 'timeEv', which gives the
--   same numerical results without doing an explicit
--   matrix inversion.  The function assumes hbar = 1.
timeEvOp :: Double -> Operator -> Operator
timeEvOp dt (Operator m) = Operator (M.timeEvMat dt m)

-- | Given a time step and a Hamiltonian operator,
--   advance the state ket using the Schrodinger equation.
--   This method should be faster than using 'timeEvOp'
--   since it solves a linear system rather than calculating
--   an inverse matrix.  The function assumes hbar = 1.
timeEv :: Double -> Operator -> Ket -> Ket
timeEv dt (Operator m) (Ket k) = Ket $ M.timeEv dt m k

-----------------
-- Composition --
-----------------

class Kron a where
    kron :: a -> a -> a

instance Kron Ket where
    kron (Ket v1) (Ket v2) = Ket (M.kron v1 v2)

instance Kron Bra where
    kron (Bra v1) (Bra v2) = Bra (M.kron v1 v2)

instance Kron Operator where
    kron (Operator m1) (Operator m2) = Operator (M.kron m1 m2)

-----------------
-- Measurement --
-----------------

-- | The possible outcomes of a measurement
--   of an observable.
--   These are the eigenvalues of the operator
--   of the observable.
possibleOutcomes :: Operator -> [Double]
possibleOutcomes (Operator observable) = M.possibleOutcomes observable

-- | Given an obervable, return a list of pairs
--   of possible outcomes and projectors
--   for each outcome.
outcomesProjectors :: Operator -> [(Double,Operator)]
outcomesProjectors (Operator m)
    = [(val,Operator p) | (val,p) <- M.outcomesProjectors m]

-- | Given an observable and a state ket, return a list of pairs
--   of possible outcomes and probabilites
--   for each outcome.
outcomesProbabilities :: Operator -> Ket -> [(Double,Double)]
outcomesProbabilities (Operator m) (Ket v)
    = M.outcomesProbabilities m v

{-
prob :: Ket -> Ket -> Double
prob k1 k2 = magnitude c ** 2
    where
      c = dagger k1 <> k2

probs :: OrthonormalBasis -> Ket -> [Double]
probs (OB ks) k = map (\bk -> let c = dagger bk <> k in magnitude c ** 2) ks
-}

{-
----------------------------------------
-- Angular Momentum of arbitrary size --
----------------------------------------

angularMomentumZMatrix :: Rational -> Matrix Cyclotomic
angularMomentumZMatrix j
    = case twoJPlusOne j of
        Nothing -> error "j must be a nonnegative integer or half-integer"
        Just d  -> matrix d d (\(r,c) -> if r == c then fromRational (j + 1 - fromIntegral r) else 0)

twoJPlusOne :: Rational -> Maybe Int
twoJPlusOne j
    | j >= 0 && (denominator j == 1 || denominator j == 2)  = Just $ fromIntegral $ numerator (2 * j + 1)
    | otherwise                                             = Nothing

angularMomentumPlusMatrix :: Rational -> Matrix Cyclotomic
angularMomentumPlusMatrix j
    = case twoJPlusOne j of
        Nothing -> error "j must be a nonnegative integer or half-integer"
        Just d  -> matrix d d (\(r,c) -> let mr = j + 1 - fromIntegral r
                                             mc = j + 1 - fromIntegral c
                                         in if mr == mc + 1
                                            then sqrtRat (j*(j+1) - mc*mr)
                                            else 0
                              )

angularMomentumMinusMatrix :: Rational -> Matrix Cyclotomic
angularMomentumMinusMatrix j
    = case twoJPlusOne j of
        Nothing -> error "j must be a nonnegative integer or half-integer"
        Just d  -> matrix d d (\(r,c) -> let mr = j + 1 - fromIntegral r
                                             mc = j + 1 - fromIntegral c
                                         in if mr + 1 == mc
                                            then sqrtRat (j*(j+1) - mc*mr)
                                            else 0
                              )

angularMomentumXMatrix :: Rational -> Matrix Cyclotomic
angularMomentumXMatrix j
    = scaleMatrix (1/2) (angularMomentumPlusMatrix j + angularMomentumMinusMatrix j)

angularMomentumYMatrix :: Rational -> Matrix Cyclotomic
angularMomentumYMatrix j
    = scaleMatrix (1/(2*i)) (angularMomentumPlusMatrix j - angularMomentumMinusMatrix j)

jXMatrix :: Rational -> Matrix Cyclotomic
jXMatrix = angularMomentumXMatrix

jYMatrix :: Rational -> Matrix Cyclotomic
jYMatrix = angularMomentumYMatrix

jZMatrix :: Rational -> Matrix Cyclotomic
jZMatrix = angularMomentumZMatrix

matrixCommutator :: Matrix Cyclotomic -> Matrix Cyclotomic -> Matrix Cyclotomic
matrixCommutator m1 m2 = m1 * m2 - m2 * m1

-----------------------
-- Rotation matrices --
-----------------------

r2i :: Rational -> Integer
r2i r
    | denominator r == 1  = numerator r
    | otherwise           = error "r2i:  not an integer"

-- from Sakurai, revised, (3.8.33)
-- beta in degrees
smallDRotElement :: Rational -> Rational -> Rational -> Rational -> Cyclotomic
smallDRotElement j m' m beta
    = sum [parity(k-m+m') * sqrtRat (fact(j+m) * fact(j-m) * fact(j+m') * fact(j-m'))
                   / fromRational (fact(j+m-k) * fact(k) * fact(j-k-m') * fact(k-m+m'))
                         * cosDeg (beta/2) ^ r2i(2*j-2*k+m-m')
                         * sinDeg (beta/2) ^ r2i(2*k-m+m') | k <- [max 0 (m-m') .. min (j+m) (j-m')]]

parity :: Rational -> Cyclotomic
parity = fromIntegral . parityInteger . r2i

-- | (-1)^n, where n might be negative
parityInteger :: Integer -> Integer
parityInteger n
    | odd n      = -1
    | otherwise  =  1

factInteger :: Integer -> Integer
factInteger n
    | n == 0     = 1
    | n >  0     = n * factInteger (n-1)
    | otherwise  = error "factInteger called on negative argument"

fact :: Rational -> Rational
fact = fromIntegral . factInteger . r2i

-- | Rotation matrix elements.
--   From Sakurai, Revised Edition, (3.5.50).
--   The matrix desribes a rotation by gamma about the z axis,
--   followed by a rotation by beta about the y axis,
--   followed by a rotation by alpha about the z axis.
bigDRotElement :: Rational  -- ^ j, a nonnegative integer or half-integer
               -> Rational  -- ^ m', an integer or half-integer index for the row
               -> Rational  -- ^ m, an integer or half-integer index for the column
               -> Rational  -- ^ alpha, in degrees
               -> Rational  -- ^ beta, in degrees
               -> Rational  -- ^ gamma, in degrees
               -> Cyclotomic  -- ^ rotation matrix element
bigDRotElement j m' m alpha beta gamma
    = polarRat 1 (-(m' * alpha + m * gamma) / 360) * smallDRotElement j m' m beta

-- | Rotation matrix for a spin-j particle.
--   The matrix desribes a rotation by gamma about the z axis,
--   followed by a rotation by beta about the y axis,
--   followed by a rotation by alpha about the z axis.
rotationMatrix :: Rational  -- ^ j, a nonnegative integer or half-integer
               -> Rational  -- ^ alpha, in degrees
               -> Rational  -- ^ beta, in degrees
               -> Rational  -- ^ gamma, in degrees
               -> Matrix Cyclotomic  -- ^ rotation matrix
rotationMatrix j alpha beta gamma
    = case twoJPlusOne j of
        Nothing -> error "bigDRotMatrix: j must be a nonnegative integer or half-integer"
        Just d  -> matrix d d (\(r,c) -> let m' = j + 1 - fromIntegral r
                                             m  = j + 1 - fromIntegral c
                                         in bigDRotElement j m' m alpha beta gamma
                              )

----------------------------------
-- Angular Momentum eigenstates --
----------------------------------

jmColumn :: Rational -> Rational -> Matrix Cyclotomic
jmColumn j m
    = case twoJPlusOne j of
        Nothing -> error "bigDRotMatrix: j must be a nonnegative integer or half-integer"
        Just d  -> matrix d 1 (\(r,_) -> let m' = j + 1 - fromIntegral r
                                         in if m == m'
                                            then 1
                                            else 0
                              )
-}

------------------
-- Gram-Schmidt --
------------------

-- | Form an orthonormal list of kets from
--   a list of linearly independent kets.
gramSchmidt :: [Ket] -> [Ket]
gramSchmidt [] = []
gramSchmidt [k] = [normalize k]
gramSchmidt (k:ks) = let nks = gramSchmidt ks
                         nk = normalize (foldl (-) k [w <> dagger w <> k | w <- nks])
                     in nk:nks

-- todo:  Clebsch-Gordon coeffs


{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

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

import           Control.Applicative
import           Data.Finite
import           Data.Type.Tuple hiding                (T2(..), T3(..))
import           GHC.TypeNats
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import           Numeric.Opto.Update hiding            ((<.>))
import qualified Data.Type.Tuple                       as T
import qualified Data.Vector.Sized                     as SV
import qualified Prelude.Backprop                      as B

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

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

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

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

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

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

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

-- klDivergence :: Loss Double
-- klDivergence =

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

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

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

lastLoss
    :: (KnownNat (n + 1), Backprop a)
    => Loss a
    -> Loss (SV.Vector (n + 1) a)
lastLoss l targ = l (SV.last targ) . viewVar (SV.ix maxBound)

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

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

-- | Lift and sum a loss function over the components of a 'T.T3'.
t3Loss
    :: (Backprop a, Backprop b, Backprop c)
    => Loss a                   -- ^ loss on first component
    -> Loss b                   -- ^ loss on second component
    -> Loss c                   -- ^ loss on third component
    -> Loss (T.T3 a b c)
t3Loss f g h (T.T3 xT yT zT) xRyRzR
        = f xT (xRyRzR ^^. t3_1)
        + g yT (xRyRzR ^^. t3_2)
        + h zT (xRyRzR ^^. t3_3)

-- | A regularizer on parameters
type Regularizer p = forall s. Reifies s W => BVar s p -> BVar s Double

-- | L2 regularization
--
-- \[
-- \sum_w \frac{1}{2} w^2
-- \]
l2Reg
    :: (Metric Double p, Backprop p)
    => Double                   -- ^ scaling factor (often 0.5)
    -> Regularizer p
l2Reg λ = liftOp1 . op1 $ \x ->
            ( λ * quadrance x / 2, (.* x) . (* λ))

-- | L1 regularization
--
-- \[
-- \sum_w \lvert w \rvert
-- \]
l1Reg
    :: (Num p, Metric Double p, Backprop p)
    => Double                   -- ^ scaling factor (often 0.5)
    -> Regularizer p
l1Reg λ = liftOp1 . op1 $ \x ->
            ( λ * norm_1 x, (.* signum x) . (* λ)
            )

-- | No regularization
noReg :: Regularizer p
noReg _ = constVar 0

-- | Add together two regularizers
addReg :: Regularizer p -> Regularizer p -> Regularizer p
addReg = liftA2 (+)

-- | Scale a regularizer's influence
scaleReg :: Double -> Regularizer p -> Regularizer p
scaleReg λ reg = (* constVar λ) . reg


-- |
-- Module      : Control.Monad.Bayes.Class
-- Description : Types for probabilistic modelling
-- Copyright   : (c) Adam Scibior, 2015-2020
-- License     : MIT
-- Maintainer  : leonhard.markert@tweag.io
-- Stability   : experimental
-- Portability : GHC
--
-- This module defines 'MonadInfer', which can be used to represent a simple model
-- like the following:
--
-- @
-- import Control.Monad (when)
-- import Control.Monad.Bayes.Class
--
-- model :: MonadInfer m => m Bool
-- model = do
--   rain <- bernoulli 0.3
--   sprinkler <-
--     bernoulli $
--     if rain
--       then 0.1
--       else 0.4
--   let wetProb =
--     case (rain, sprinkler) of
--       (True,  True)  -> 0.98
--       (True,  False) -> 0.80
--       (False, True)  -> 0.90
--       (False, False) -> 0.00
--   score wetProb
--   return rain
-- @
module Control.Monad.Bayes.Class
  ( MonadSample,
    random,
    uniform,
    normal,
    gamma,
    beta,
    bernoulli,
    categorical,
    logCategorical,
    uniformD,
    geometric,
    poisson,
    dirichlet,
    MonadCond,
    score,
    factor,
    condition,
    MonadInfer,
    normalPdf,
  )
where

import Control.Monad (when)
import Control.Monad.Trans.Class
import Control.Monad.Trans.Cont
import Control.Monad.Trans.Identity
import Control.Monad.Trans.List
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.RWS hiding (tell)
import Control.Monad.Trans.Reader
import Control.Monad.Trans.State
import Control.Monad.Trans.Writer
import qualified Data.Vector as V
import Data.Vector.Generic as VG
import Numeric.Log
import Statistics.Distribution
import Statistics.Distribution.Beta (betaDistr)
import Statistics.Distribution.Gamma (gammaDistr)
import Statistics.Distribution.Geometric (geometric0)
import Statistics.Distribution.Normal (normalDistr)
import qualified Statistics.Distribution.Poisson as Poisson
import Statistics.Distribution.Uniform (uniformDistr)

-- | Monads that can draw random variables.
class Monad m => MonadSample m where
  -- | Draw from a uniform distribution.
  random ::
    -- | \(\sim \mathcal{U}(0, 1)\)
    m Double

  -- | Draw from a uniform distribution.
  uniform ::
    -- | lower bound a
    Double ->
    -- | upper bound b
    Double ->
    -- | \(\sim \mathcal{U}(a, b)\).
    m Double
  uniform a b = draw (uniformDistr a b)

  -- | Draw from a normal distribution.
  normal ::
    -- | mean μ
    Double ->
    -- | standard deviation σ
    Double ->
    -- | \(\sim \mathcal{N}(\mu, \sigma^2)\)
    m Double
  normal m s = draw (normalDistr m s)

  -- | Draw from a gamma distribution.
  gamma ::
    -- | shape k
    Double ->
    -- | scale θ
    Double ->
    -- | \(\sim \Gamma(k, \theta)\)
    m Double
  gamma shape scale = draw (gammaDistr shape scale)

  -- | Draw from a beta distribution.
  beta ::
    -- | shape α
    Double ->
    -- | shape β
    Double ->
    -- | \(\sim \mathrm{Beta}(\alpha, \beta)\)
    m Double
  beta a b = draw (betaDistr a b)

  -- | Draw from a Bernoulli distribution.
  bernoulli ::
    -- | probability p
    Double ->
    -- | \(\sim \mathrm{B}(1, p)\)
    m Bool
  bernoulli p = fmap (< p) random

  -- | Draw from a categorical distribution.
  categorical ::
    Vector v Double =>
    -- | event probabilities
    v Double ->
    -- | outcome category
    m Int
  categorical ps = fromPMF (ps !)

  -- | Draw from a categorical distribution in the log domain.
  logCategorical ::
    (Vector v (Log Double), Vector v Double) =>
    -- | event probabilities
    v (Log Double) ->
    -- | outcome category
    m Int
  logCategorical = categorical . VG.map (exp . ln)

  -- | Draw from a discrete uniform distribution.
  uniformD ::
    -- | observable outcomes @xs@
    [a] ->
    -- | \(\sim \mathcal{U}\{\mathrm{xs}\}\)
    m a
  uniformD xs = do
    let n = Prelude.length xs
    i <- categorical $ V.replicate n (1 / fromIntegral n)
    return (xs !! i)

  -- | Draw from a geometric distribution.
  geometric ::
    -- | success rate p
    Double ->
    -- | \(\sim\) number of failed Bernoulli trials with success probability p before first success
    m Int
  geometric = discrete . geometric0

  -- | Draw from a Poisson distribution.
  poisson ::
    -- | parameter λ
    Double ->
    -- | \(\sim \mathrm{Pois}(\lambda)\)
    m Int
  poisson = discrete . Poisson.poisson

  -- | Draw from a Dirichlet distribution.
  dirichlet ::
    Vector v Double =>
    -- | concentration parameters @as@
    v Double ->
    -- | \(\sim \mathrm{Dir}(\mathrm{as})\)
    m (v Double)
  dirichlet as = do
    xs <- VG.mapM (`gamma` 1) as
    let s = VG.sum xs
    let ys = VG.map (/ s) xs
    return ys

-- | Draw from a continuous distribution using the inverse cumulative density
-- function.
draw :: (ContDistr d, MonadSample m) => d -> m Double
draw d = fmap (quantile d) random

-- | Draw from a discrete distribution using a sequence of draws from
-- Bernoulli.
fromPMF :: MonadSample m => (Int -> Double) -> m Int
fromPMF p = f 0 1
  where
    f i r = do
      when (r < 0) $ error "fromPMF: total PMF above 1"
      let q = p i
      when (q < 0 || q > 1) $ error "fromPMF: invalid probability value"
      b <- bernoulli (q / r)
      if b then pure i else f (i + 1) (r - q)

-- | Draw from a discrete distributions using the probability mass function.
discrete :: (DiscreteDistr d, MonadSample m) => d -> m Int
discrete = fromPMF . probability

-- | Monads that can score different execution paths.
class Monad m => MonadCond m where
  -- | Record a likelihood.
  score ::
    -- | likelihood of the execution path
    Log Double ->
    m ()

-- | Synonym for 'score'.
factor ::
  MonadCond m =>
  -- | likelihood of the execution path
  Log Double ->
  m ()
factor = score

-- | Hard conditioning.
condition :: MonadCond m => Bool -> m ()
condition b = score $ if b then 1 else 0

-- | Monads that support both sampling and scoring.
class (MonadSample m, MonadCond m) => MonadInfer m

-- | Probability density function of the normal distribution.
normalPdf ::
  -- | mean μ
  Double ->
  -- | standard deviation σ
  Double ->
  -- | sample x
  Double ->
  -- | relative likelihood of observing sample x in \(\mathcal{N}(\mu, \sigma^2)\)
  Log Double
normalPdf mu sigma x = Exp $ logDensity (normalDistr mu sigma) x

----------------------------------------------------------------------------
-- Instances that lift probabilistic effects to standard tranformers.

instance MonadSample m => MonadSample (IdentityT m) where
  random = lift random
  bernoulli = lift . bernoulli

instance MonadCond m => MonadCond (IdentityT m) where
  score = lift . score

instance MonadInfer m => MonadInfer (IdentityT m)

instance MonadSample m => MonadSample (MaybeT m) where
  random = lift random

instance MonadCond m => MonadCond (MaybeT m) where
  score = lift . score

instance MonadInfer m => MonadInfer (MaybeT m)

instance MonadSample m => MonadSample (ReaderT r m) where
  random = lift random
  bernoulli = lift . bernoulli

instance MonadCond m => MonadCond (ReaderT r m) where
  score = lift . score

instance MonadInfer m => MonadInfer (ReaderT r m)

instance (Monoid w, MonadSample m) => MonadSample (WriterT w m) where
  random = lift random
  bernoulli = lift . bernoulli
  categorical = lift . categorical

instance (Monoid w, MonadCond m) => MonadCond (WriterT w m) where
  score = lift . score

instance (Monoid w, MonadInfer m) => MonadInfer (WriterT w m)

instance MonadSample m => MonadSample (StateT s m) where
  random = lift random
  bernoulli = lift . bernoulli
  categorical = lift . categorical

instance MonadCond m => MonadCond (StateT s m) where
  score = lift . score

instance MonadInfer m => MonadInfer (StateT s m)

instance (MonadSample m, Monoid w) => MonadSample (RWST r w s m) where
  random = lift random

instance (MonadCond m, Monoid w) => MonadCond (RWST r w s m) where
  score = lift . score

instance (MonadInfer m, Monoid w) => MonadInfer (RWST r w s m)

instance MonadSample m => MonadSample (ListT m) where
  random = lift random
  bernoulli = lift . bernoulli
  categorical = lift . categorical

instance MonadCond m => MonadCond (ListT m) where
  score = lift . score

instance MonadInfer m => MonadInfer (ListT m)

instance MonadSample m => MonadSample (ContT r m) where
  random = lift random

instance MonadCond m => MonadCond (ContT r m) where
  score = lift . score

instance MonadInfer m => MonadInfer (ContT r m)