{-|
Module      : MachineLearning.Model.Regression
Description : TIR expression data structures
Copyright   : (c) Fabricio Olivetti de Franca, 2022
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com
Stability   : experimental
Portability : POSIX

Fitting functions for the coefficients.
-}
module MachineLearning.Model.Regression
  ( fitTask
  , predictTask
  , evalPenalty
  , applyMeasures
  , nonlinearFit
  , tirToMatrix
  ) where

import Data.Bifunctor
import           Data.SRTree                   (evalFun, Function(..), OptIntPow(..), evalFun, inverseFunc)
import MachineLearning.TIR       (TIR(..),  Individual(..), Dataset, Constraint, assembleTree, replaceConsts)
import qualified Data.Vector           as V
import qualified Data.Vector.Storable  as VS
import Data.Vector           ((!))
import           Data.List                     (nub)
import           Data.Vector.Storable          (Vector, splitAt)
import           Numeric.LinearAlgebra         ((<\>), Matrix)
import           Numeric.GSL.Fitting           (nlFitting, FittingMethod(..))
import           Prelude               hiding  (splitAt)

import qualified Numeric.LinearAlgebra                     as LA

import MachineLearning.TIR (Individual(..))
import MachineLearning.Utils.Config (Task(..), Penalty(..))
import MachineLearning.Model.Measure (Measure(..))

-- * IT specific stuff

createQ :: Vector Double -> LA.Matrix Double -> LA.Matrix Double
createQ ys = LA.fromColumns . map (*ys) . LA.toColumns
{-# INLINE createQ #-}

avg :: Vector Double -> Vector Double
avg ys = LA.fromList [LA.sumElements ys / fromIntegral (LA.size ys)]
{-# INLINE avg #-}

-- | transform a data matrix using a TIR expression. This function returns
-- a tuple with the transformed data of the numerator and denominator, respectivelly.
-- Each column of the transformed data represents one term of the TIR expression.
tirToMatrix :: Dataset Double -> TIR -> (LA.Matrix Double, LA.Matrix Double)
tirToMatrix xss (TIR _ p q) = bimap (LA.fromColumns . (bias:)) LA.fromColumns (p', q')
  where
    bias      = V.head xss
    xss'      = V.tail xss
    sigma2mtx = map (\(_, g, ps) -> evalFun g $ evalPi ps)
    evalPi    = foldr (\(ix, k) acc -> acc * (xss' ! ix ^^ k)) 1
    p'        = sigma2mtx p
    q'        = sigma2mtx q
{-# INLINE tirToMatrix #-}

-- | Fits a linear model using l2-penalty
ridge :: Matrix Double -> Vector Double -> Matrix Double
ridge a b = oA <\> oB
  where
   mu = 0.01

   a' = LA.tr a
   b' = LA.tr $ LA.asColumn b
   oA = (a' <> LA.tr a') + (mu * LA.ident (LA.rows a'))
   oB = a' <> LA.tr b'
{-# INLINE ridge #-}

-- | Predicts a linear model
predict :: Matrix Double -> Vector Double -> Vector Double
predict xs w | LA.cols xs == LA.size w = xs LA.#> w
             | otherwise = error $ "predict: " ++ show (LA.size xs) ++ show (LA.size w)
{-# INLINE predict #-}

-- | Solve the OLS *zss*w = ys*
solveOLS :: Matrix Double -> Vector Double -> Vector Double
solveOLS zss ys = zss <\> ys
{-# INLINE solveOLS #-}

-- | Applies OLS and returns a Solution
-- if the expression is invalid, it returns Infinity as a fitness
--regress :: Matrix Double -> Vector Double -> [Vector Double]
regress :: TIR -> Dataset Double -> Vector Double -> [Vector Double]
regress tir xss ys = [ws]
  where
    (zssP, zssQ) = tirToMatrix xss tir
    ys'          = evalFun (inverseFunc $ _funY tir) ys
    zssQy        = createQ ys' zssQ
    zss          = if LA.cols zssQ >= 1
                       then zssP LA.||| negate zssQy
                       else zssP
    ws             = if LA.cols zss == 1
                       then avg ys'
                       else solveOLS zss ys'
-- regress zss ys = [solveOLS zss ys]
{-# INLINE regress #-}

-- | Applies conjugate gradient for binary classification
--classify :: Matrix Double -> Vector Double -> [Vector Double]
classify :: Int -> TIR -> Dataset Double -> Vector Double -> [Vector Double]
classify niter tir xss ys = [ws]
  where
    ws           = nonlinearFit niter zssP zssQ ys sigmoid dsigmoid theta0
    theta0       = LA.konst 0 (LA.cols zssP + LA.cols zssQ)
    (zssP, zssQ) = tirToMatrix xss tir

-- | Applies conjugate gradient for one-vs-all classification
--classifyMult :: Matrix Double -> Vector Double -> [Vector Double]
classifyMult :: Int -> TIR -> Dataset Double -> Vector Double -> [Vector Double]
classifyMult niter tir xss ys = zipWith minimize yss theta0
  where
    numLabels    = length $ nub $ LA.toList ys
    yss          = map f [0 .. numLabels-1]
    minimize y t = nonlinearFit niter zssP zssQ y sigmoid dsigmoid t
    theta0       = replicate numLabels $ LA.konst 0 (LA.cols zssP + LA.cols zssQ)
    (zssP, zssQ) = tirToMatrix xss tir

    f sample = VS.map (\a -> if round a == sample then 1 else 0) ys

-- | chooses the appropriate fitting function
--fitTask :: Task -> Matrix Double -> Vector Double -> [Vector Double]
fitTask :: Task -> TIR -> Dataset Double -> Vector Double -> [Vector Double]
fitTask Regression             = regress
fitTask (RegressionNL niter)   = regressNL niter
fitTask (Classification niter) = classify niter
fitTask (ClassMult niter)      = classifyMult niter
{-# INLINE fitTask #-}

-- | sigmoid function for classification.
sigmoid :: Floating a => a -> a
sigmoid z = 1 / (1+exp(-z))
{-# INLINE sigmoid #-}

-- | derivative sigmoid function for classification.
dsigmoid :: Floating a => a -> a
dsigmoid z = sigmoid z * (1 - sigmoid z)
{-# INLINE dsigmoid #-}

-- | chooses the appropriate prediction function
predictTask :: Task -> [Vector Double] -> Vector Double
predictTask _ []                   = error "predictTask: empty coefficients matrix"
predictTask Regression yss         = head yss
predictTask (RegressionNL _) yss   = head yss
predictTask (Classification _) yss = sigmoid $ head yss
predictTask (ClassMult _) yss      = LA.vector $ map (fromIntegral . LA.maxIndex) $ LA.toRows $ LA.fromColumns $ map sigmoid yss
{-# INLINE predictTask #-}

-- | evals the penalty function
evalPenalty :: Penalty -> Int -> Double -> Double
evalPenalty NoPenalty _   _   = 0.0
evalPenalty (Len c)   len _   = fromIntegral len * c
evalPenalty (Shape c) _   val = val*c
{-# INLINE evalPenalty #-}

-- | applies a list of performance measures
applyMeasures :: [Measure] -> Vector Double -> Vector Double -> [Double]
applyMeasures measures ys ysHat = map ((`uncurry` (ys, ysHat)) . _fun) measures
{-# INLINE applyMeasures #-}

regressNL :: Int -> TIR -> Dataset Double -> Vector Double -> [Vector Double]
regressNL niter tir xss ys = [ws]
  where
    ws           = nonlinearFit niter zssP zssQ ys f f' theta0
    f            = evalFun $ _funY tir
    f'           = derivative $ _funY tir
    --theta0       = LA.konst 1 (LA.cols zssP + LA.cols zssQ)
    theta0       = head $ regress tir xss ys
    (zssP, zssQ) = tirToMatrix xss tir

-- | Non-linear optimization using Levenberg-Marquardt method.
--nonlinearFit :: Monad m => Vector Double -> Matrix Double -> Matrix Double -> Vector Double -> m (Vector Double)
nonlinearFit :: Int
             -> Matrix Double
             -> Matrix Double
             -> Vector Double
             -> (Vector Double -> Vector Double)
             -> (Vector Double -> Vector Double)
             -> Vector Double
             -> Vector Double
nonlinearFit niter zssP zssQ ys f f' theta0 = fst $ nlFitting LevenbergMarquardtScaled 1e-6 1e-6 niter model' jacob' theta0
  where
    model'       = model f ys zssP zssQ
    jacob'       = jacob f' zssP zssQ

-- | calculates the error given the parameter vector beta
model :: (Vector Double -> Vector Double) -> Vector Double -> Matrix Double -> Matrix Double -> Vector Double -> Vector Double
model f ys zssP zssQ beta
  | LA.cols zssQ == 0 = f ysHat_P - ys
  | otherwise         = f ysHat - ys
  where
    (betaP, betaQ) = splitAt (LA.cols zssP) beta
    ysHat_P        = predict zssP betaP
    ysHat_Q        = if LA.cols zssQ == 0 then 0 else predict zssQ betaQ
    ysHat          = ysHat_P / (1 + ysHat_Q)

-- | calculates the Jacobian given the parameter vector beta. Doesn't support
-- the outer transformation function.
jacob :: (Vector Double -> Vector Double) -> Matrix Double -> Matrix Double -> Vector Double -> Matrix Double
jacob f zssP zssQ beta | LA.cols zssQ == 0 = zssP
                       | otherwise         = LA.fromColumns $ pjac <> qjac
  where
    (betaP, betaQ) = splitAt (LA.cols zssP) beta
    ysHat_P        = predict zssP betaP
    ysHat_Q        = predict zssQ betaQ
    ysHat          = f $ ysHat_P / (1 + ysHat_Q)
    pjac           = [ysHat * c / ysHat_Q | c <- LA.toColumns zssP]
    qjac           = [-ysHat * c * (ysHat_P / (1 + ysHat_Q)^2) | c <- LA.toColumns zssQ]

derivative :: (Eq val, Floating val) => Function -> val -> val
derivative Id      = const 1
derivative Abs     = \x -> x / abs x
derivative Sin     = cos
derivative Cos     = negate.sin
derivative Tan     = recip . (**2.0) . cos
derivative Sinh    = cosh
derivative Cosh    = sinh
derivative Tanh    = (1-) . (**2.0) . tanh
derivative ASin    = recip . sqrt . (1-) . (^2)
derivative ACos    = negate . recip . sqrt . (1-) . (^2)
derivative ATan    = recip . (1+) . (^2)
derivative ASinh   = recip . sqrt . (1+) . (^2)
derivative ACosh   = \x -> 1 / (sqrt (x-1) * sqrt (x+1))
derivative ATanh   = recip . (1-) . (^2)
derivative Sqrt    = recip . (2*) . sqrt
derivative Square  = (2*)
derivative Exp     = exp
derivative Log     = recip
{-# INLINE derivative #-}
-- (w1 * p1 + w2 * p2) / (1 + w3 * p3 + w4 * p4)
-- d/dw1 = p1 / (1 + w3 * p3 + w4 * p4)
-- d/dw2 = p2 / (1 + w3 * p3 + w4 * p4)
-- d/dw3 = -p3 * (w1 * p2 + w2 * p2) / (1 + w3 * p3 + w4 * p4)^2
{-
toEv :: SRTree Int Double -> (V.Vector Double -> Double)
toEv (Var !ix) = (`V.unsafeIndex` ix)
toEv (Const !val) = const val
toEv (Add !l !r) = jn (+) (toEv l) (toEv r)
toEv (Mul !l !r) = jn (*) (toEv l) (toEv r)
toEv (Fun Exp !t) = exp . toEv t
{-# INLINE toEv #-}

jn :: (Double -> Double -> Double) -> (V.Vector Double -> Double) -> (V.Vector Double -> Double) -> (V.Vector Double -> Double)
jn op f g = \x -> op (f x) (g x)
{-# INLINE jn #-}
-}


{- |
Description :  Boundary mutation model Markov process
Copyright   :  (c) Dominik Schrempf 2017
License     :  GPLv3

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

This module defines the transition rate matrix of the boundary mutation model.

-}

module BndModel
  ( BMModel(..)
  , MutModel
  , Heterozygosity
  , createBMM
  , createBMMNormalized
  , scaleTreeToBMM
  , getBMMInfoStr
  ) where

import           BndState
import           Data.Maybe
import           DNAModel              hiding (rateMatrix)
import           Numeric.LinearAlgebra
import           RateMatrix            hiding (State)
import           RTree
import           Tools

-- | The boundary mutation models uses an underlying mutation model.
type MutModel = DNAModel

-- | Define a heterozygosity to make function definitions clearer.
type Heterozygosity = Double

-- | A boundary mutation model is defined by its rate matrix. However, it is
-- convenient to keep the underlying mutation model, the population size the
-- heterozygosity and the stationary distribution at hand.
data BMModel = BMModel { bmmRateMatrix     :: RateMatrix
                       , bmmMutModel       :: MutModel
                       , bmmPopSize        :: PopSize
                       , bmmHeterozygosity :: Heterozygosity
                       , bmmStationaryDist :: StationaryDist }

-- Get the mutation rate from one allele to another.
mutRate :: MutModel -> Allele -> Allele -> Double
mutRate (DNAModel m _) a b = m ! i ! j
  where i = fromEnum a
        j = fromEnum b

-- Get state frequency for specific allele.
stationaryFreq :: StationaryDist -> Allele -> Double
stationaryFreq f a = f ! fromEnum a

-- Calculate the rate of frequency shifts.
moranCoef :: PopSize -> Int -> Double
moranCoef n i = iD * (nD - iD) / nD
  where iD = fromIntegral i
        nD = fromIntegral n

-- The transition rate from one boundary state to another.
rate :: MutModel -> State -> State -> Double
rate m s t
  | not $ connected s t = 0.0
  | otherwise           = rate' s t
  where rate' (Ply n i _ _) _ = moranCoef n i
        rate' (Bnd _ a) (Ply _ _ b c)
          | a == b    = mutRate m a c
          | a == c    = mutRate m a b
          | otherwise = error "Cannot compute rate between states."
        rate' _ _ = error "Cannot compute rate between states."

-- The transition rate from one state (index) to another.
rateById :: MutModel -> PopSize -> Int -> Int -> Double
rateById m n i j = rate m s t
  where s = idToState n i
        t = idToState n j

-- The build function (see below) has a weird way of assigning entries to
-- indices. The indices have to be the same data type as the entries. This is
-- just a helper function that changes the indices from Double to Int.
rateByDouble :: MutModel -> PopSize -> Double -> Double -> Double
rateByDouble m n x y = rateById m n (round x) (round y)

-- The rate matrix of the boundary mutation model. The dimension with four
-- alleles will be 4 + 6*(N-1).
rateMatrix :: MutModel -> PopSize -> RateMatrix
rateMatrix m n = setDiagonal $ build (s,s) (rateByDouble m n)
  where s = stateSpaceSize n

-- Normalize the rate matrix such that on average one event happens per unit
-- time.
normalizedRateMatrix :: MutModel -> PopSize -> RateMatrix
normalizedRateMatrix m n = normalizeRates f' m'
  where m' = rateMatrix m n
        -- f  = dnaModelSpecGetStateFreqVec (dnaModelSpec m)
        f' = stationaryDist m n

-- Calculate the heterozygosity at stationarity.
theta :: MutModel -> Heterozygosity
theta (DNAModel m s) = f <.> rDiagZero #> f
  where f         = dnaModelSpecGetStateFreqVec s
        e         = toExchMatrix m f
        (r, _)    = matrixSeparateSymSkew e
        -- The summation excludes the diagonal (a /= b).
        rDiagZero = matrixSetDiagToZero r

-- The normalization constant of the stationary distribution.
norm :: MutModel -> PopSize -> Double
norm m n = 1.0 + harmonic (n-1) * theta m

-- Get entries of the stationary measure (not normalized) for a boundary model state.
stationaryMeasEntry :: MutModel -> State -> Double
stationaryMeasEntry m s =
  let f   = dnaModelSpecGetStateFreqVec (dnaModelSpec m) in
  case s of
    (Bnd _ a)     -> f ! fromEnum a
    (Ply n i a b) -> piA * mAB / (fromIntegral n - fromIntegral i)
                                      + piB * mBA / fromIntegral i
      where piA = stationaryFreq f a
            mAB = mutRate m a b
            piB = stationaryFreq f b
            mBA = mutRate m b a

stationaryMeasEntryById :: MutModel -> PopSize -> Int -> Double
stationaryMeasEntryById m n i = stationaryMeasEntry m s
  where s = idToState n i

stationaryMeasEntryByDouble :: MutModel -> PopSize -> Double -> Double
stationaryMeasEntryByDouble m n x = stationaryMeasEntryById m n (round x)

-- Get the stationary distribution of a boundary mutation model.
stationaryDist :: MutModel -> PopSize -> StationaryDist
stationaryDist m n = scale (1.0/norm m n) $
  build s (stationaryMeasEntryByDouble m n)
  where s = stateSpaceSize n

-- Normalize the mutation coefficients such that the heterozygosity matches a
-- given level (see Eq. 12.14 in my thesis).
normalizeToTheta :: MutModel -> PopSize -> Heterozygosity -> MutModel
normalizeToTheta mo@(DNAModel m _) n h =
  mo { dnaRateMatrix = scale (h / (t * (1.0 - c * h))) m }
  where
    -- The heterozygosity of the boundary mutation model.
    t = theta mo
    -- The branch length multiplicative factor introduced by the coalescent.
    c = harmonic (n-1)

-- | Create a boundary mutation model using the minimal number of necessary
-- ingredients.
createBMM :: DNAModel -> PopSize -> Heterozygosity -> BMModel
createBMM m n h = BMModel rm m' n h f
  where m' = normalizeToTheta m n h
        rm = normalizedRateMatrix m' n
        f  = stationaryDist m n

-- | Create a boundary mutation model without providing a heterozygosity. This
-- means that the mutation model has to be normalized already.
createBMMNormalized :: DNAModel -> PopSize -> BMModel
createBMMNormalized m n = BMModel rm m n h f
  where rm = normalizedRateMatrix m n
        h  = theta m
        f  = stationaryDist m n

-- | The branch lengths of threes in the boundary mutation model are not
-- measured in average number of substitutions per site but in average number of
-- mutations or frequency shifts per site. The conversion factor is just the
-- square of the population size. This function converts the branch lengths of a
-- tree.
scaleTreeToBMM :: BMModel -> RTree a Double -> RTree a Double
scaleTreeToBMM bmm = fmap (* nSq)
  where nSq = fromIntegral (bmmPopSize bmm) ** 2

-- | Report the boundary mutation model specifications.
getBMMInfoStr :: BMModel
              -> Maybe Double
              -> Maybe [Double]
              -> String
getBMMInfoStr bmm ma mrs = unlines $
  [ "Population size: " ++ show (bmmPopSize bmm)
  , "Heterozygosity: " ++ show (bmmHeterozygosity bmm)
  , "Mutation model:"
  , getDNAModelInfoStr (bmmMutModel bmm)
  , "Gamma rate heterogeneity: " ++ show (isJust ma) ] ++
  gammaShape ++ gammaMeans
  where
    gammaShape = maybe [] (\a -> ["Shape parameter: " ++ show a]) ma
    gammaMeans = maybe [] (\rs -> ["This corresponds to uniformly distributed rates: " ++ show rs]) mrs


{-# LANGUAGE FlexibleContexts #-}
-- | Mixed-radix FFT calculation.
--
-- Arbitrary input vector lengths are handled using a mixed-radix
-- Cooley-Tukey decimation in time algorithm with residual prime
-- length vectors being treated using Rader's algorithm or hand-coded
-- codelets for small primes.
module Numeric.FFT
       ( fft, ifft, fftWith, ifftWith
       , plan, planFromFactors, execute
       , Plan (..), Direction (..), BaseTransform (..)
       ) where

import Prelude hiding (length, map, sum, zipWith)
import Data.Vector.Generic
import Data.Complex

import Numeric.FFT.Types
import Numeric.FFT.Plan
import Numeric.FFT.Execute

-- | Forward FFT with embedded plan calculation.  For an input vector
-- /h/ of length /N/, with entries numbered from 0 to /N - 1/,
-- calculates the entries in /H/, the discrete Fourier transform of
-- /h/, as:
--
-- <<doc-formulae/fft-formula.svg>>
fft :: Vector v (Complex Double) =>
       v (Complex Double) -> IO (v (Complex Double))
fft xs = do
  p <- plan $ length xs
  return $ fftWith p xs

-- | Inverse FFT with embedded plan calculation.  For an input vector
-- /H/ of length /N/, with entries numbered from 0 to /N - 1/,
-- representing Fourier amplitudes of a signal, calculates the entries
-- in /h/, the inverse discrete Fourier transform of /H/, as:
--
-- <<doc-formulae/ifft-formula.svg>>
ifft :: Vector v (Complex Double) =>
        v (Complex Double) -> IO (v (Complex Double))
ifft xs = do
  p <- plan $ length xs
  return $ ifftWith p xs

-- | Forward FFT with pre-computed plan.
fftWith :: Vector v (Complex Double) =>
           Plan -> v (Complex Double) -> v (Complex Double)
fftWith p = convert . execute p Forward . convert

-- | Inverse FFT with pre-computed plan.
ifftWith :: Vector v (Complex Double) =>
            Plan -> v (Complex Double) -> v (Complex Double)
ifftWith p = convert . execute p Inverse . convert


module PCA.QModels where

import PCA.Distribution
import Numeric.LinearAlgebra.Data
       (Matrix, Vector, ident, konst, size, toRows, fromRows, diag, diagl,
        takeDiag)
import qualified Numeric.LinearAlgebra.HMatrix as H

initQtau :: Distr Double Double
initQtau = Gamma 1.0e-3 1.0e-3

initQmu :: Int -> Distr (Vector Double) (Matrix Double)
initQmu d = MNormal (konst 0 d) (ident d)

initQalpha :: Int -> Distr (Vector Double) (Vector Double)
initQalpha d = MGamma (konst 1.0e-3 d) (konst 1.0e-3 d)

initQW :: Int -> Distr (Matrix Double) (Matrix Double)
initQW d = MatrixMNormal (ident d) (ident d)

initQX :: Matrix Double -> Distr (Matrix Double) (Matrix Double)
initQX trainT = let (_, d) = size trainT
                    sigmaX = calcSigmaX d initQtau (initQW d)
                    mX = calcMX trainT initQtau sigmaX (initQW d) (initQmu d)
                in MatrixMNormal mX sigmaX

type QModel = (Matrix Double, Vector Double)  -- ^ Mode of W and \mu

calculateQ :: Matrix Double -> QModel
calculateQ trainT =
    go 5 initQtau (initQmu d) (initQalpha d) (initQW d) (initQX trainT) trainT
  where
    (n,d) = size trainT
    go 0 _ mu _ w _ _ = (mean w, mean mu)
    go k (Gamma aTau bTau) mu (MGamma aAlpha bAlpha) w x t =
        go (k - 1) ntau nmu nalpha nw nx t
      where
        ntau = Gamma (calcAtau n d aTau) (calcBtau bTau t mu w x)
        nalpha = MGamma (calcAalpha d aAlpha) (calcBalpha bAlpha w)
        nSigmaMu = calcSigmaMu n ntau mu
        nmu = MNormal (calcMMu t ntau nSigmaMu w x) nSigmaMu
        nSigmaW = calcSigmaW ntau nalpha x
        nw = MatrixMNormal (calcMW ntau nSigmaW x t nmu) nSigmaW
        nSigmaX = calcSigmaX d ntau nw
        nx = MatrixMNormal (calcMX t ntau nSigmaX nw nmu) nSigmaX

calcSigmaX
    :: Int
    -> Distr Double Double
    -> Distr (Matrix Double) (Matrix Double)
    -> Matrix Double
calcSigmaX d tau w = H.inv (ident d + mean tau `H.scale` wSquareMean)
  where
    mW = mean w
    wSquareMean = variance w + (H.tr mW H.<> mW)

calcMX
    :: Matrix Double
    -> Distr Double Double
    -> Matrix Double
    -> Distr (Matrix Double) (Matrix Double)
    -> Distr (Vector Double) (Matrix Double)
    -> Matrix Double
calcMX trainT tau sigmaX w mu =
    H.tr
        (mean tau `H.scale` sigmaX H.<> H.tr (mean w) H.<>
         H.tr
             (fromRows
                  (map
                       (\r ->
                             r - mean mu)
                       (toRows trainT))))

calcSigmaMu
    :: Int
    -> Distr Double Double
    -> Distr (Vector Double) (Matrix Double)
    -> Matrix Double
calcSigmaMu n tau mu =
    vMu + 1 / (fromIntegral n * mean tau) `H.scale` ident d
  where
    vMu = variance mu
    (_,d) = size vMu

calcMMu
    :: Matrix Double
    -> Distr Double Double
    -> Matrix Double
    -> Distr (Matrix Double) (Matrix Double)
    -> Distr (Matrix Double) (Matrix Double)
    -> Vector Double
calcMMu trainT tau sigmaMu w x =
    (mean tau `H.scale` sigmaMu) H.#>
    foldr
         (flip (+))
         (konst 0 d)
         (toRows (trainT - H.tr (mW H.<> H.tr mX)))
  where
    mW = mean w
    mX = mean x
    (_, d) = size trainT

calcSigmaW
    :: Distr Double Double
    -> Distr (Vector Double) (Vector Double)
    -> Distr (Matrix Double) (Matrix Double)
    -> Matrix Double
calcSigmaW tau alpha x =
    H.inv
        (diag (mean alpha) +
         mean tau `H.scale`
         (fromIntegral n `H.scale` variance x + H.tr mX H.<> mX))
  where
    mX = mean x
    (n,_) = size mX

calcMW
    :: Distr Double Double
    -> Matrix Double
    -> Distr (Matrix Double) (Matrix Double)
    -> Matrix Double
    -> Distr (Vector Double) (Matrix Double)
    -> Matrix Double
calcMW tau sigmaW x trainT mu =
    (mean tau `H.scale` sigmaW) H.<>
    (H.tr mX H.<>
     (fromRows
          (map
               (\tn ->
                     tn - mMu)
               (toRows trainT))))
  where
    mX = mean x
    mMu = mean mu

calcAalpha :: Int -> Vector Double -> Vector Double
calcAalpha d = H.cmap (+ fromIntegral d / 2)

calcBalpha :: Vector Double
           -> Distr (Matrix Double) (Matrix Double)
           -> Vector Double
calcBalpha b w = b + 0.5 `H.scale` (takeDiag (H.tr mW H.<> mW))
  where
    mW = mean w

calcAtau :: Int -> Int -> Double -> Double
calcAtau n d a = a + fromIntegral (n * d) / 2

-- | TODO: Formula is not correct.
-- Skip part Tr (\langle W^TW \rangle \langle x_n x^T_n 'rangle')
-- because i think it is not correct in that result must be number
calcBtau
    :: Double
    -> Matrix Double
    -> Distr (Vector Double) (Matrix Double)
    -> Distr (Matrix Double) (Matrix Double)
    -> Distr (Matrix Double) (Matrix Double)
    -> Double
calcBtau b trainT mu w x =
    b +
    0.5 *
    (fromIntegral n * norm mMu +
     H.sumElements (takeDiag (trainT H.<> H.tr trainT)) +
     2 * (H.sumElements (H.tr (mW H.<> H.tr mX) H.#> mMu)) -
     2 * (H.sumElements (takeDiag (trainT H.<> mW H.<> H.tr mX))) -
     2 * (H.sumElements (trainT H.#> mMu)))
  where
    (n,_) = size trainT
    norm v = v `H.dot` v
    mMu = mean mu
    mW = mean w
    mX = mean x


{-# LANGUAGE DeriveDataTypeable, DeriveGeneric, FlexibleContexts #-}
-- |
-- Module    : Statistics.Sample.KernelDensity.Simple
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Kernel density estimation code, providing non-parametric ways to
-- estimate the probability density function of a sample.
--
-- The techniques used by functions in this module are relatively
-- fast, but they generally give inferior results to the KDE function
-- in the main 'Statistics.KernelDensity' module (due to the
-- oversmoothing documented for 'bandwidth' below).

module Statistics.Sample.KernelDensity.Simple
    {-# DEPRECATED "Use Statistics.Sample.KernelDensity instead." #-}
    (
    -- * Simple entry points
      epanechnikovPDF
    , gaussianPDF
    -- * Building blocks
    -- These functions may be useful if you need to construct a kernel
    -- density function estimator other than the ones provided in this
    -- module.

    -- ** Choosing points from a sample
    , Points(..)
    , choosePoints
    -- ** Bandwidth estimation
    , Bandwidth
    , bandwidth
    , epanechnikovBW
    , gaussianBW
    -- ** Kernels
    , Kernel
    , epanechnikovKernel
    , gaussianKernel
    -- ** Low-level estimation
    , estimatePDF
    , simplePDF
    -- * References
    -- $references
    ) where

import Data.Data (Data, Typeable)
import GHC.Generics (Generic)
import Numeric.MathFunctions.Constants (m_1_sqrt_2, m_2_sqrt_pi)
import Prelude hiding (sum)
import Statistics.Function (minMax)
import Statistics.Sample (stdDev)
import Statistics.Sample.Internal (sum)
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U

-- | Points from the range of a 'Sample'.
newtype Points = Points {
      fromPoints :: U.Vector Double
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

-- | Bandwidth estimator for an Epanechnikov kernel.
epanechnikovBW :: Double -> Bandwidth
epanechnikovBW n = (80 / (n * m_2_sqrt_pi)) ** 0.2

-- | Bandwidth estimator for a Gaussian kernel.
gaussianBW :: Double -> Bandwidth
gaussianBW n = (4 / (n * 3)) ** 0.2

-- | The width of the convolution kernel used.
type Bandwidth = Double

-- | Compute the optimal bandwidth from the observed data for the
-- given kernel.
--
-- This function uses an estimate based on the standard deviation of a
-- sample (due to Deheuvels), which performs reasonably well for
-- unimodal distributions but leads to oversmoothing for more complex
-- ones.
bandwidth :: G.Vector v Double =>
             (Double -> Bandwidth)
          -> v Double
          -> Bandwidth
bandwidth kern values = stdDev values * kern (fromIntegral $ G.length values)

-- | Choose a uniform range of points at which to estimate a sample's
-- probability density function.
--
-- If you are using a Gaussian kernel, multiply the sample's bandwidth
-- by 3 before passing it to this function.
--
-- If this function is passed an empty vector, it returns values of
-- positive and negative infinity.
choosePoints :: G.Vector v Double =>
                Int             -- ^ Number of points to select, /n/
             -> Double          -- ^ Sample bandwidth, /h/
             -> v Double        -- ^ Input data
             -> Points
choosePoints n h sample = Points . U.map f $ U.enumFromTo 0 n'
  where lo     = a - h
        hi     = z + h
        (a, z) = minMax sample
        d      = (hi - lo) / fromIntegral n'
        f i    = lo + fromIntegral i * d
        n'     = n - 1

-- | The convolution kernel.  Its parameters are as follows:
--
-- * Scaling factor, 1\//nh/
--
-- * Bandwidth, /h/
--
-- * A point at which to sample the input, /p/
--
-- * One sample value, /v/
type Kernel =  Double
            -> Double
            -> Double
            -> Double
            -> Double

-- | Epanechnikov kernel for probability density function estimation.
epanechnikovKernel :: Kernel
epanechnikovKernel f h p v
    | abs u <= 1 = f * (1 - u * u)
    | otherwise  = 0
    where u = (v - p) / (h * 0.75)

-- | Gaussian kernel for probability density function estimation.
gaussianKernel :: Kernel
gaussianKernel f h p v = exp (-0.5 * u * u) * g
    where u = (v - p) / h
          g = f * 0.5 * m_2_sqrt_pi * m_1_sqrt_2

-- | Kernel density estimator, providing a non-parametric way of
-- estimating the PDF of a random variable.
estimatePDF :: G.Vector v Double =>
               Kernel           -- ^ Kernel function
            -> Bandwidth        -- ^ Bandwidth, /h/
            -> v Double         -- ^ Sample data
            -> Points           -- ^ Points at which to estimate
            -> U.Vector Double
estimatePDF kernel h sample
    | n < 2     = errorShort "estimatePDF"
    | otherwise = U.map k . fromPoints
  where
    k p = sum . G.map (kernel f h p) $ sample
    f   = 1 / (h * fromIntegral n)
    n   = G.length sample
{-# INLINE estimatePDF #-}

-- | A helper for creating a simple kernel density estimation function
-- with automatically chosen bandwidth and estimation points.
simplePDF :: G.Vector v Double =>
             (Double -> Double) -- ^ Bandwidth function
          -> Kernel             -- ^ Kernel function
          -> Double             -- ^ Bandwidth scaling factor (3 for a Gaussian kernel, 1 for all others)
          -> Int                -- ^ Number of points at which to estimate
          -> v Double           -- ^ sample data
          -> (Points, U.Vector Double)
simplePDF fbw fpdf k numPoints sample =
    (points, estimatePDF fpdf bw sample points)
  where points = choosePoints numPoints (bw*k) sample
        bw     = bandwidth fbw sample
{-# INLINE simplePDF #-}

-- | Simple Epanechnikov kernel density estimator.  Returns the
-- uniformly spaced points from the sample range at which the density
-- function was estimated, and the estimates at those points.
epanechnikovPDF :: G.Vector v Double =>
                   Int          -- ^ Number of points at which to estimate
                -> v Double     -- ^ Data sample
                -> (Points, U.Vector Double)
epanechnikovPDF = simplePDF epanechnikovBW epanechnikovKernel 1

-- | Simple Gaussian kernel density estimator.  Returns the uniformly
-- spaced points from the sample range at which the density function
-- was estimated, and the estimates at those points.
gaussianPDF :: G.Vector v Double =>
               Int              -- ^ Number of points at which to estimate
            -> v Double         -- ^ Data sample
            -> (Points, U.Vector Double)
gaussianPDF = simplePDF gaussianBW gaussianKernel 3

errorShort :: String -> a
errorShort func = error ("Statistics.KernelDensity." ++ func ++
                        ": at least two points required")

-- $references
--
-- * Deheuvels, P. (1977) Estimation non paramétrique de la densité
--   par histogrammes
--   généralisés. Mhttp://archive.numdam.org/article/RSA_1977__25_3_5_0.pdf>


{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveDataTypeable #-}
#if __GLASGOW_HASKELL__ >= 786
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}
#endif

module Data.Complex.Polar (
#if __GLASGOW_HASKELL__ >= 800
    Polar((:<)),
#elif __GLASGOW_HASKELL >= 786
    Polar,
    pattern (:<),
#else
    Polar,
#endif
    fromPolar,
    fromComplex,
    realPart,
    imagPart,
    conjugate,
    mkPolar,
    cis,
    polar,
    magnitude,
    phase
) where

import           Data.Complex (Complex(..))
import qualified Data.Complex as C

import Data.Typeable
#ifdef __GLASGOW_HASKELL__
import Data.Data (Data)
#endif

#ifdef __HUGS__
import Hugs.Prelude(Num(fromInt), Fractional(fromDouble))
#endif

import Text.Read (parens)
import Text.ParserCombinators.ReadPrec (prec, step)
import Text.Read.Lex (Lexeme (Ident))

infix 6 :<<

-- -----------------------------------------------------------------------------
-- The Polar type

-- | Complex numbers are an algebraic type.
--
-- For a complex number @z@, @'abs' z@ is a number with the magnitude of @z@,
-- but oriented in the positive real direction, whereas @'signum' z@
-- has the phase of @z@, but unit magnitude.
data (RealFloat a) => Polar a = !a :<< !a -- ^ forms a complex number from its magnitude
                                          --   and its phase in radians.
#if __GLASGOW_HASKELL__
    deriving (Eq, Data, Typeable)
#else
    deriving Eq
#endif

#if __GLASGOW_HASKELL__ >= 786
infix 6 :<

-- | Smart constructor that canonicalizes the magnitude and phase.
pattern r :< theta <-
  ( \(r :<< theta) -> Just (r, theta) ->
      Just (r, theta)
    )
  where
    r :< theta = mkPolar r theta
#endif

instance (RealFloat a, Show a) => Show (Polar a) where
    {-# SPECIALISE instance Show (Polar Float) #-}
    {-# SPECIALISE instance Show (Polar Double) #-}
    showsPrec d (r :<< theta) =
      showParen (d >= 11) (showString "mkPolar "
                           . showsPrec 11 r
                           . showString " "
                           . showsPrec 11 theta)

instance (RealFloat a, Read a) => Read (Polar a) where
    {-# SPECIALISE instance Read (Polar Float) #-}
    {-# SPECIALISE instance Read (Polar Double) #-}
    readsPrec d = readParen (d > 10)
                        (\s -> do ("mkPolar", s2) <- lex s
                                  (r, s3) <- readsPrec 11 s2
                                  (theta, s4) <- readsPrec 11 s3
                                  return (mkPolar r theta, s4))

-- | Wrap phase back in interval @(-'pi', 'pi']@.
wrap :: (RealFloat a) => a -> a
{-# SPECIALISE wrap :: Float  -> Float   #-}
{-# SPECIALISE wrap :: Double -> Double  #-}
wrap phi | phi <= (-pi) = wrap (phi+2*pi)
wrap phi | phi > pi     = wrap (phi-2*pi)
wrap phi                = phi

-- | Convert to rectangular form.
fromPolar :: (RealFloat a) => Polar a -> Complex a
{-# INLINE fromPolar #-}
fromPolar p = realPart p :+ imagPart p

-- | Convert to polar form.
fromComplex :: (RealFloat a) => Complex a -> Polar a
{-# INLINE fromComplex #-}
fromComplex = uncurry mkPolar_ . C.polar

-- | Extracts the real part of a complex number.
realPart :: (RealFloat a) => Polar a -> a
realPart (r :<< theta) = r * cos theta

-- | Extracts the imaginary part of a complex number.
imagPart :: (RealFloat a) => Polar a -> a
imagPart (r :<< theta) = r * sin theta

-- | The conjugate of a complex number.
conjugate :: (RealFloat a) => Polar a -> Polar a
{-# SPECIALISE conjugate :: Polar Float  -> Polar Float #-}
{-# SPECIALISE conjugate :: Polar Double -> Polar Double #-}
conjugate (r :<< theta) = mkPolar r (negate theta)

-- | Form a complex number from polar components of magnitude and phase.
-- The magnitude and phase are expected to be in canonical form.
mkPolar_ :: RealFloat a => a -> a -> Polar a
{-# SPECIALISE mkPolar_ :: Float  -> Float  -> Polar Float  #-}
{-# SPECIALISE mkPolar_ :: Double -> Double -> Polar Double #-}
mkPolar_ r theta = r :<< theta

-- | Form a complex number from polar components of magnitude and phase.
mkPolar :: RealFloat a => a -> a -> Polar a
{-# SPECIALISE mkPolar :: Float  -> Float  -> Polar Float  #-}
{-# SPECIALISE mkPolar :: Double -> Double -> Polar Double #-}
mkPolar r theta | r == 0 = 0 :<< 0
mkPolar r theta | r < 0 = mkPolar (- r) (theta + pi)
mkPolar r theta = mkPolar_ r (wrap theta)

-- | @'cis' t@ is a complex value with magnitude @1@
-- and phase @t@ (modulo @2*'pi'@).
cis :: (RealFloat a) => a -> Polar a
{-# SPECIALISE cis :: Float  -> Polar Float  #-}
{-# SPECIALISE cis :: Double -> Polar Double #-}
cis theta = mkPolar 1 theta

-- | The function 'polar' takes a complex number and
-- returns a (magnitude, phase) pair in canonical form:
-- the magnitude is nonnegative, and the phase in the range @(-'pi', 'pi']@;
-- if the magnitude is zero, then so is the phase.
polar :: (RealFloat a) => Polar a -> (a,a)
{-# SPECIALISE polar :: Polar Double -> (Double,Double) #-}
{-# SPECIALISE polar :: Polar Float  -> (Float,Float)   #-}
polar (r :<< theta) = (r,theta)

-- | The nonnegative magnitude of a complex number.
magnitude :: (RealFloat a) => Polar a -> a
{-# SPECIALISE magnitude :: Polar Float  -> Float  #-}
{-# SPECIALISE magnitude :: Polar Double -> Double #-}
magnitude (r :<< _) = r

-- | The phase of a complex number, in the range @(-'pi', 'pi']@.
-- If the magnitude is zero, then so is the phase.
phase :: (RealFloat a) => Polar a -> a
{-# SPECIALISE phase :: Polar Float  -> Float  #-}
{-# SPECIALISE phase :: Polar Double -> Double #-}
phase (_ :<< theta) = theta

instance (RealFloat a) => Num (Polar a) where
    {-# SPECIALISE instance Num (Polar Float)  #-}
    {-# SPECIALISE instance Num (Polar Double) #-}
    z + z'           = fromComplex (fromPolar z + fromPolar z')
    z - z'           = fromComplex (fromPolar z - fromPolar z')
    z * z'           = mkPolar  (magnitude z * magnitude z') (phase z + phase z')
    negate z         = mkPolar  (negate (magnitude z)) (phase z)
    abs z            = mkPolar_ (magnitude z) 0
    signum (0 :<< _) = 0
    signum z         = mkPolar_ 1 (phase z)
    fromInteger      = flip mkPolar 0 . fromInteger

instance (RealFloat a) => Fractional (Polar a) where
    {-# SPECIALISE instance Fractional (Polar Float)  #-}
    {-# SPECIALISE instance Fractional (Polar Double) #-}
    z / z'          = mkPolar (magnitude z / magnitude z') (phase z - phase z')
    fromRational r  = mkPolar (fromRational r) 0

instance  (RealFloat a) => Floating (Polar a)  where
    {-# SPECIALISE instance Floating (Polar Float) #-}
    {-# SPECIALISE instance Floating (Polar Double) #-}
    pi             = mkPolar_ pi 0
    exp (r :<< theta) = mkPolar (exp (r * cos theta)) (r * sin theta)
    log (r :<< theta) = fromComplex (log r :+ theta)

    -- sqrt (0 :<< _)       =  0
    -- sqrt z@(r :<< theta) =  u :+ (if y < 0 then -v else v)
    --                           where (u,v) = if x < 0 then (v',u') else (u',v')
    --                                 v'    = abs y / (u'*2)
    --                                 u'    = sqrt ((magnitude z + abs x) / 2)
    sqrt = fromComplex.sqrt.fromPolar

    -- sin (x:+y)     =  sin x * cosh y :+ cos x * sinh y
    sin = fromComplex.sin.fromPolar

    -- cos (x:+y)     =  cos x * cosh y :+ (- sin x * sinh y)
    cos = fromComplex.cos.fromPolar

    -- tan (x:+y)     =  (sinx*coshy:+cosx*sinhy)/(cosx*coshy:+(-sinx*sinhy))
    --                   where sinx  = sin x
    --                         cosx  = cos x
    --                         sinhy = sinh y
    --                         coshy = cosh y
    tan = fromComplex.tan.fromPolar

    -- sinh (x:+y)    =  cos y * sinh x :+ sin  y * cosh x
    sinh = fromComplex.sinh.fromPolar

    -- cosh (x:+y)    =  cos y * cosh x :+ sin y * sinh x
    cosh = fromComplex.cosh.fromPolar

    -- tanh (x:+y)    =  (cosy*sinhx:+siny*coshx)/(cosy*coshx:+siny*sinhx)
    --                   where siny  = sin y
    --                         cosy  = cos y
    --                         sinhx = sinh x
    --                         coshx = cosh x
    tanh = fromComplex.tanh.fromPolar

    -- asin z@(x:+y)  =  y':+(-x')
    --                   where  (x':+y') = log (((-y):+x) + sqrt (1 - z*z))
    asin = fromComplex.asin.fromPolar

    -- acos z         =  y'':+(-x'')
    --                   where (x'':+y'') = log (z + ((-y'):+x'))
    --                         (x':+y')   = sqrt (1 - z*z)
    acos = fromComplex.acos.fromPolar

    -- atan z@(x:+y)  =  y':+(-x')
    --                   where (x':+y') = log (((1-y):+x) / sqrt (1+z*z))
    atan = fromComplex.atan.fromPolar

    asinh z        =  log (z + sqrt (1+z*z))
    acosh z        =  log (z + (z+1) * sqrt ((z-1)/(z+1)))
    atanh z        =  log ((1+z) / sqrt (1-z*z))


module Distributions where

import Control.Monad (replicateM)
import qualified Control.Monad.Bayes.Class as B
import Control.Monad.Bayes.Class (MonadSample, MonadCond)
import Numeric.SpecFunctions (logGamma)

import qualified Data.Vector as V
import qualified Data.Map as M
import qualified Data.MultiSet as S

import Numeric.Log (Log (Exp))

import SymbolicArithmetic
import Util.Numeric (logFact)
import Util.Bijection

type Distr = Distr' Double

data Distr' d a = Distr
  { sample :: forall m. MonadSample m => m a
  , score :: a -> d
  }

mapOutputDiscrete :: Bijection a b -> Distr' d a -> Distr' d b
mapOutputDiscrete (Bijection to from) (Distr sample score) =
  Distr (to <$> sample) (score . from)

factor :: MonadCond m => Double -> m ()
factor = B.score . Numeric.Log.Exp

negInf :: Double
negInf = - 1 / 0

observe :: MonadCond m => Distr a -> a -> m ()
observe d x = factor (score d x)

data ExpFam a = ExpFam
  { dim :: Int
  , suffStats :: [Exp () Double]
  , logPartition :: [Double] -> Double
  , logBaseMeas :: Exp () Double
  , efSample :: forall m. MonadSample m => [Double] -> m a
  , efRepr :: a -> Double
  , mle :: [Double] -> [Double]
  , efDescribe :: [Double] -> String }

bernoulliEF :: ExpFam Bool
bernoulliEF = ExpFam
  { dim = 1
  , suffStats = [Var ()]
  , logPartition = undefined
  , logBaseMeas = undefined
  , efSample = \[p] -> B.bernoulli p
  , efRepr = \x -> if x then 1 else 0
  , mle = id
  , efDescribe = \[eta] -> let z = exp eta in "Bernoulli(" ++ show (z / (1 + z)) ++ ")"
  }

betaEF :: ExpFam Double
betaEF = ExpFam
  { dim = 2
  , suffStats = let x = Var () in  [log x, log (1 - x)]
  , logPartition = \[a, b] -> logBeta a b
  , logBaseMeas = let x = Var () in - log x - log (1 - x)
  , efSample = \[a, b] -> B.beta a b
  , efRepr = id
  , mle = \[lx, lox] -> undefined
  , efDescribe = \[a, b] -> "Beta(" ++ show a ++ ", " ++ show b ++ ")"
  }

logBeta :: Double -> Double -> Double
logBeta a b = logGamma a + logGamma b - logGamma (a + b)

logChoose :: Int -> Int -> Double
logChoose n k = logFact n - logFact k - logFact (n - k)

normalEF :: ExpFam Double
normalEF = ExpFam
  { dim = 2
  , suffStats = let x = Var () in [x, x^2]
  , logPartition = \[eta1, eta2] -> - eta1^2 / (4 * eta2) - 1 / 2 * log(-2 * eta2)
  , logBaseMeas = - 1 / 2 * log (2 * pi)
  , efSample = \eta -> let (mu, sigma2) = conv eta in
      B.normal mu (sqrt sigma2)
  , efRepr = id
  , mle = \[mu, ex2] -> let sigma2 = ex2 - mu^2 in normalNatParams mu sigma2
  , efDescribe = \eta -> let (mu, sigma2) = conv eta in
    "Normal(" ++ show mu ++ ", " ++ show sigma2 ++ ")"
  }
  where
  conv [eta1, eta2] = let sigma2 = - 1 / (2 * eta2) in (sigma2 * eta1, sigma2)
  conv _ = error "conv"

-- indep :: ExpFam a -> ExpFam b -> ExpFam (a, b)
-- indep efa efb = ExpFam
--   { dim = dim efa + dim efb
--   , suffStats = \(a, b) -> suffStats efa a ++ suffStats efb b
--   , logPartition = \eta -> let (eta1, eta2) = split eta in
--       logPartition efa eta1 + logPartition efb eta2
--   , logBaseMeas = \(a, b) -> logBaseMeas efa a + logBaseMeas efb b
--   , efSample = \eta -> let (eta1, eta2) = split eta in
--       (,) <$> efSample efa eta1 <*> efSample efb eta2
--   , mle = \ss -> let (ss1, ss2) = split ss in
--       mle efa ss1 ++ mle efb ss2
--   , efDescribe = \eta -> let (eta1, eta2) = split eta in
--     "Indep(" ++ efDescribe efa eta1 ++ ", " ++ efDescribe efb eta2 ++ ")"
--   }
--   where
--   split = splitAt (dim efa)

normalNatParams :: Fractional d => d -> d -> [d]
normalNatParams mu sigma2 = [mu / sigma2, -1 / (2 * sigma2)]

dirac :: Eq a => a -> Distr a
dirac x = Distr (return x) (\y -> if x == y then 0 else negInf)

dot :: Num d => [d] -> [d] -> d
dot xs ys = sum (zipWith (*) xs ys)

efToDistr :: ExpFam a -> [Double] -> Distr a
efToDistr ef naturalParams =
  Distr (efSample ef naturalParams)
  (\x -> let ev = eval (\_ -> efRepr ef x) in
    naturalParams `dot` map ev (suffStats ef) - logPartition ef naturalParams + ev (logBaseMeas ef))

bernoulli :: Double -> Distr Bool
bernoulli p = Distr (B.bernoulli p) (\b -> bernoulli_ll p (if b then 1 else 0))

bernoulli01 :: Double -> Distr Int
bernoulli01 p = Distr sam sco where
  sam :: MonadSample m => m Int
  sam = do
    b <- B.bernoulli p
    return (if b then 1 else 0)
  sco 1 = log p
  sco 0 = log (1 - p)
  sco _ = negInf

categorical :: [Double] -> Distr Int
categorical ps = Distr (B.categorical (V.fromList ps)) (\i -> log (ps !! i))

categoricalM :: Ord a => M.Map a Double -> Distr a
categoricalM ps = Distr (fmap (\i -> fst (M.elemAt i ps)) (B.categorical (V.fromList (M.elems ps))))
  (\k -> log (ps M.! k))

bernoulli_ll :: Floating a => a -> a -> a
bernoulli_ll p b = b * log p + (1 - b) * log (1 - p)

-- WARNING: NOT NORMALIZED!
beta_ll :: Floating a => a -> a -> a -> a
beta_ll a b p = a * log p + b * log (1 - p)

beta :: Double -> Double -> Distr Double
beta a b = Distr (B.beta a b) (\x -> beta_ll a b x + logPartition betaEF [a, b])

gaussian_ll :: Floating a => a -> a -> a -> a
gaussian_ll mu sigma2 x = - x^2 / sigma2 + 2 * x * mu / sigma2 - mu^2 / sigma2 - log (2 * pi) - log sigma2

gaussian_ll' :: Floating a => a -> a -> a -> a
gaussian_ll' mu sigma2 x = - (x - mu)^2 / sigma2 - log (2 * pi) - log sigma2

gaussian :: Double -> Double -> Distr Double
gaussian mu sigma2 = Distr (B.normal mu (sqrt sigma2)) (gaussian_ll' mu sigma2)

normal = gaussian

poisson_ll :: Double -> Int -> Double
poisson_ll lambda k = fromIntegral k * log lambda - lambda - logGamma (fromIntegral k + 1)

poisson :: Double -> Distr Int
poisson lambda = Distr (B.poisson lambda) (poisson_ll lambda)

uniform :: Double -> Double -> Distr Double
uniform a b = Distr (B.uniform a b) $ \x ->
  if a <= x && x <= b
    then - log range
    else negInf
  where
  range = b - a

replicateNIID :: Int -> Distr a -> Distr [a]
replicateNIID n d = Distr (replicateM n (sample d)) (\obs ->
  if length obs == n
    then sum (map (score d) obs)
    else (-1 / 0))

uniformIntRange :: Int -> Distr Int
uniformIntRange max = Distr (B.uniformD [0 .. max - 1])
  (\obs -> if 0 <= obs && obs < max then - log (fromIntegral max) else -1 / 0)

bind :: Distr a -> (a -> Distr b) -> Distr (a, b)
bind d f = Distr (do {x <- sample d; y <- sample (f x); pure (x, y)})
                 (\(x, y) -> score d x + score (f x) y)

replicateIID :: forall a. Distr Int -> Distr a -> Distr [a]
replicateIID howMany d = Distr sam obs where
  sam :: MonadSample m => m [a]
  sam = do
    n <- sample howMany
    replicateM n (sample d)
  obs xs = score howMany (length xs) + sum (map (score d) xs)

geometric' :: MonadSample m => m Int
geometric' = do
  x <- B.bernoulli (0.5 :: Double)
  if x then return 0 else fmap (+1) geometric'

exponential :: MonadSample m => Double -> m Double
exponential lambda = do
  u <- B.random
  pure (- log u / lambda)

passert :: MonadCond m => Bool -> m ()
passert True = pure ()
passert False = B.score 0

data Some f where
  Some :: f a -> Some f

exponentialFamilies :: [Some ExpFam]
exponentialFamilies = [Some betaEF, Some bernoulliEF, Some normalEF]

shuffleWithRepeats' :: forall a. Ord a => M.Map a (Distr Int) -> Distr [a]
shuffleWithRepeats' countDistrs = Distr sam sco where
  sam :: MonadSample m => m [a]
  sam = do
    counts <- mapM sample countDistrs
    randomlyInterleave [ replicate n k | (k, n) <- M.toList counts ]
  sco :: [a] -> Double
  sco idxes = let ns = mconcat [ S.singleton i | i <- idxes ] in
    sum (M.mapWithKey (\k d -> score d (length (filter (== k) idxes))) countDistrs)
    + randomlyInterleaveLogPDF (map snd (S.toOccurList ns))

shuffleWithRepeats :: [Distr Int] -> Distr [Int]
shuffleWithRepeats countDistrs = Distr sam sco where
  sam :: MonadSample m => m [Int]
  sam = do
    counts <- mapM sample countDistrs
    randomlyInterleave [ replicate n i | (i, n) <- zip [0..] counts ]
  sco :: [Int] -> Double
  sco idxes = let ns = [ length (filter (== i) idxes) | i <- take (length countDistrs) [0..] ] in
    sum (zipWith score countDistrs ns) + randomlyInterleaveLogPDF ns

randomlyInterleave :: MonadSample m => [[a]] -> m [a]
randomlyInterleave xs = if ntot == 0
  then pure []
  else do
    i <- sample (categorical [ fromIntegral n / fromIntegral ntot | n <- ns ])
    let (before, (v : vs) : after) = splitAt i xs
    (v :) <$> randomlyInterleave (before ++ vs : after)
  where
    ns = map length xs
    ntot = sum ns

logNumPermutationsWithRepeats :: [Int] -> Double
logNumPermutationsWithRepeats ns = logFact (sum ns) - sum [ logFact n | n <- ns ]

randomlyInterleaveLogPDF :: [Int] -> Double
randomlyInterleaveLogPDF ns = - logNumPermutationsWithRepeats ns

shuffleListLogPDF :: Int -> Double
shuffleListLogPDF n = - logFact n


{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ImpredicativeTypes #-}
{-# LANGUAGE RankNTypes #-}
{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}

module Data.Histogram.Util where

import Control.Foldl (fold)
import qualified Control.Foldl as L
import Data.Histogram (Histogram)
import qualified Data.Histogram as H
import Data.Histogram.Fill
import Data.Monoid
import qualified Data.Vector.Unboxed as V
import Statistics.Distribution
import qualified Statistics.Distribution.Combo as OC
import Statistics.Distribution.Exponential
import Statistics.Distribution.Gamma
import qualified Statistics.Distribution.Laplace as OL
import Statistics.Distribution.Normal
import Statistics.Distribution.Poisson

bin :: BinD
bin = binDn (-10.5) 1 10.5

hist :: BinD -> V.Vector Double -> H.Histogram BinD Double
hist b =
  fillBuilderVec (forceDouble -<< mkSimple b)

-- | make a percentile histogram
perc :: Histogram BinD Double -> Histogram BinD Double
perc h = H.map (\x -> x / H.sum h) h

-- reversing out mean and stdev
-- calculating mean using the trick that the Bin boundaries are constructed to be isomorphic to actual values
meanH :: Histogram BinD Double -> Double
meanH h =
  let t = H.sum h
  in H.bfoldl step begin h /
     t
  where begin = 0.0
        step b bv v = b + bv * v

meanSqH :: Histogram BinD Double -> Double
meanSqH h =
  let t = H.sum h
  in H.bfoldl step begin h /
     t
  where begin = 0.0
        step b bv v =
          b +
          (bv ^
           (2 :: Int)) *
          v

type Pdf = BinD -> [Double] -> V.Vector Double

pdfBin :: (Double -> Double) -> BinD -> V.Vector Double
pdfBin cdf h =
  V.map (\(x,y) -> cdf y - cdf x) binsV
  where binsV =
          binsList h :: V.Vector (Double,Double)

pdfBinNormal :: Pdf
pdfBinNormal b [] =
  pdfBin (cumulative $
          normalDistr 0.0 1.0)
         b
pdfBinNormal b [s] =
  pdfBin (cumulative $
          normalDistr 0.0 s')
         b
  where s' = max 1.0e-5 s
pdfBinNormal b (m:s:_) =
  pdfBin (cumulative $
          normalDistr m s')
         b
  where s' = max 1.0e-5 s

pdfBinLaplace :: Pdf
pdfBinLaplace b (r:l:sl:glue:_) =
  pdfBin (cumulative $
          OL.laplace r' l' sl' glue)
         b
  where l' = max 1.0e-5 l
        r' = max 1.0e-5 r
        sl' = max 1.0e-4 sl
pdfBinLaplace _ _ = mempty

pdfBinExponential :: Pdf
pdfBinExponential b (e:_) =
  pdfBin (cumulative $ exponential e') b
  where e' = max 1.0e-5 e
pdfBinExponential _ _ = mempty

pdfBinPoisson :: Pdf
pdfBinPoisson b (e:_) =
  pdfBin (cumulative $ poisson e') b
  where e' = max 1.0e-5 e
pdfBinPoisson _ _ = mempty

pdfBinGamma :: Pdf
pdfBinGamma b (m:k:v:_) =
  V.map (m *) $
  pdfBin (cumulative $
          gammaDistr k' v')
         b
  where k' = max 1.0e-5 k
        v' = max 1.0e-5 v
pdfBinGamma _ _ = mempty

pdfBinCombo :: Pdf
pdfBinCombo b (split:m:s:r:l:sl:_) =
  pdfBin (cumulative $
          OC.combo split' m s' r' l' sl')
         b
  where split' = max 1.0e-5 split
        s' = max 1.0e-5 s
        l' = max 1.0e-5 l
        r' = max 1.0e-5 r
        sl' = max 1.0e-4 sl
pdfBinCombo _ _ = mempty

err :: Histogram BinD Double -> Pdf -> [Double] -> Double
err h pdf p =
  do let pdfT = pdf (H.bins h) p
         pdfE = H.histData (perc h)
         errV =
           zipWith (-)
                   (V.toList pdfE)
                   (V.toList pdfT)
         err' = fold L.sum (fmap abs errV)
     err'


{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE DeriveGeneric             #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE LambdaCase                #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE StandaloneDeriving        #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeFamilies              #-}
{-# LANGUAGE TypeInType                #-}
{-# LANGUAGE TypeSynonymInstances      #-}
{-# LANGUAGE UndecidableInstances      #-}

module Learn.Neural.Layer.Recurrent.FullyConnected (
    FullyConnectedR
  , FullyConnectedR'
  , CommonMap(..)
  , MapFunc(..)
  ) where

import           Data.Kind
import           Data.Proxy
import           Data.Reflection
import           Data.Singletons.Prelude
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import           Learn.Neural.Layer
import           Learn.Neural.Layer.Mapping
import           Numeric.BLAS
import           Numeric.Backprop
import           Numeric.Backprop.Iso           (iso)
import           Numeric.Backprop.Op
import           Statistics.Distribution
import           Statistics.Distribution.Normal
import qualified Generics.SOP                   as SOP

data FullyConnectedR :: Type

deriving instance Generic (CParam FullyConnectedR b '[i] '[o])
instance SOP.Generic (CParam FullyConnectedR b '[i] '[o])

instance (Num (b '[o,o]), Num (b '[o,i]), Num (b '[o]))
      => Num (CParam FullyConnectedR b '[i] '[o]) where
    FCRP wI1 wS1 b1 + FCRP wI2 wS2 b2 = FCRP (wI1 + wI2) (wS1 + wS2) (b1 + b2)
    FCRP wI1 wS1 b1 - FCRP wI2 wS2 b2 = FCRP (wI1 - wI2) (wS1 - wS2) (b1 - b2)
    FCRP wI1 wS1 b1 * FCRP wI2 wS2 b2 = FCRP (wI1 * wI2) (wS1 * wS2) (b1 * b2)
    negate (FCRP wI wS b) = FCRP (negate wI) (negate wS) (negate b)
    signum (FCRP wI wS b) = FCRP (signum wI) (signum wS) (signum b)
    abs    (FCRP wI wS b) = FCRP (abs    wI) (abs    wS) (abs    b)
    fromInteger x = FCRP (fromInteger x) (fromInteger x) (fromInteger x)

instance (Fractional (b '[o,o]), Fractional (b '[o,i]), Fractional (b '[o]))
      => Fractional (CParam FullyConnectedR b '[i] '[o]) where
    FCRP wI1 wS1 b1 / FCRP wI2 wS2 b2 = FCRP (wI1 / wI2) (wS1 / wS2) (b1 / b2)
    recip (FCRP wI wS b) = FCRP (recip wI) (recip wS) (recip b)
    fromRational x       = FCRP (fromRational x) (fromRational x) (fromRational x)

instance (Floating (b '[o,o]), Floating (b '[o,i]), Floating (b '[o]))
      => Floating (CParam FullyConnectedR b '[i] '[o]) where
    sqrt (FCRP wI wS b) = FCRP (sqrt wI) (sqrt wS) (sqrt b)

instance Num (b '[o]) => Num (CState FullyConnectedR b '[i] '[o]) where
    FCRS s1 + FCRS s2 = FCRS (s1 + s2)
    FCRS s1 - FCRS s2 = FCRS (s1 - s2)
    FCRS s1 * FCRS s2 = FCRS (s1 * s2)
    negate (FCRS s) = FCRS (negate s)
    signum (FCRS s) = FCRS (signum s)
    abs    (FCRS s) = FCRS (abs    s)
    fromInteger x  = FCRS (fromInteger x)

instance Fractional (b '[o]) => Fractional (CState FullyConnectedR b '[i] '[o]) where
    FCRS s1 / FCRS s2 = FCRS (s1 / s2)
    recip (FCRS s)    = FCRS (recip s)
    fromRational x    = FCRS (fromRational x)

instance Floating (b '[o]) => Floating (CState FullyConnectedR b '[i] '[o]) where
    sqrt (FCRS s)    = FCRS (sqrt s)

instance ( BLAS b
         , KnownNat i
         , KnownNat o
         , Floating (b '[o])
         , Floating (b '[o,i])
         , Floating (b '[o,o])
         )
        => Component FullyConnectedR b '[i] '[o] where
    data CParam  FullyConnectedR b '[i] '[o] =
            FCRP { _fcrInpWeights   :: !(b '[o,i])
                 , _fcrStateWeights :: !(b '[o,o])
                 , _fcrBiases       :: !(b '[o])
                 }
    data CState  FullyConnectedR b '[i] '[o] = FCRS { _fcrState :: !(b '[o]) }
    type CConstr FullyConnectedR b '[i] '[o] = (Num (b '[o,i]), Num (b '[o,o]))
    data CConf   FullyConnectedR b '[i] '[o] = forall d. ContGen d => FCRC d

    componentOp = bpOp . withInps $ \(x :< p :< s :< Ø) -> do
        wI :< wS :< b :< Ø <- gTuple #<~ p
        s0 <- opIso (iso _fcrState FCRS) ~$ (s :< Ø)
        y  <- matVecOp ~$ (wI :< x  :< Ø)
        s1 <- matVecOp ~$ (wS :< s0 :< Ø)
        let z = y + s1 + b
        s' <- opIso (iso FCRS _fcrState) ~$ (s1 :< Ø)
        return $ z :< s' :< Ø

    defConf = FCRC (normalDistr 0 0.01)
    initParam = \case
      i `SCons` SNil -> \case
        so@(o `SCons` SNil) -> \(FCRC d) g -> do
          wI <- genA (o `SCons` (i `SCons` SNil)) $ \_ ->
            realToFrac <$> genContVar d g
          wS <- genA (o `SCons` (o `SCons` SNil)) $ \_ ->
            realToFrac <$> genContVar d g
          b <- genA so $ \_ ->
            realToFrac <$> genContVar d g
          return $ FCRP wI wS b
        _ -> error "inaccessible"
      _ -> error "inaccessible"
    initState _ so (FCRC d) g =
        FCRS <$> genA so (\_ -> realToFrac <$> genContVar d g)

instance ( BLAS b
         , KnownNat i
         , KnownNat o
         , Floating (b '[o])
         , Floating (b '[o,i])
         , Floating (b '[o,o])
         )
        => ComponentLayer 'Recurrent FullyConnectedR b '[i] '[o] where
    componentRunMode = RMNotFF

data FullyConnectedR' :: k -> Type

deriving instance Generic (CParam (FullyConnectedR' c) b '[i] '[o])
instance SOP.Generic (CParam (FullyConnectedR' c) b '[i] '[o])

instance (Num (b '[o,o]), Num (b '[o,i]), Num (b '[o]))
      => Num (CParam (FullyConnectedR' s) b '[i] '[o]) where
    FCRP' wI1 wS1 b1 + FCRP' wI2 wS2 b2 = FCRP' (wI1 + wI2) (wS1 + wS2) (b1 + b2)
    FCRP' wI1 wS1 b1 - FCRP' wI2 wS2 b2 = FCRP' (wI1 - wI2) (wS1 - wS2) (b1 - b2)
    FCRP' wI1 wS1 b1 * FCRP' wI2 wS2 b2 = FCRP' (wI1 * wI2) (wS1 * wS2) (b1 * b2)
    negate (FCRP' wI wS b) = FCRP' (negate wI) (negate wS) (negate b)
    signum (FCRP' wI wS b) = FCRP' (signum wI) (signum wS) (signum b)
    abs    (FCRP' wI wS b) = FCRP' (abs    wI) (abs    wS) (abs    b)
    fromInteger x = FCRP' (fromInteger x) (fromInteger x) (fromInteger x)

instance (Fractional (b '[o,o]), Fractional (b '[o,i]), Fractional (b '[o]))
      => Fractional (CParam (FullyConnectedR' s) b '[i] '[o]) where
    FCRP' wI1 wS1 b1 / FCRP' wI2 wS2 b2 = FCRP' (wI1 / wI2) (wS1 / wS2) (b1 / b2)
    recip (FCRP' wI wS b) = FCRP' (recip wI) (recip wS) (recip b)
    fromRational x        = FCRP' (fromRational x) (fromRational x) (fromRational x)

instance (Floating (b '[o,o]), Floating (b '[o,i]), Floating (b '[o]))
      => Floating (CParam (FullyConnectedR' s) b '[i] '[o]) where
    sqrt (FCRP' wI wS b) = FCRP' (sqrt wI) (sqrt wS) (sqrt b)

instance Num (b '[o]) => Num (CState (FullyConnectedR' s) b '[i] '[o]) where
    FCRS' s1 + FCRS' s2 = FCRS' (s1 + s2)
    FCRS' s1 - FCRS' s2 = FCRS' (s1 - s2)
    FCRS' s1 * FCRS' s2 = FCRS' (s1 * s2)
    negate (FCRS' s) = FCRS' (negate s)
    signum (FCRS' s) = FCRS' (signum s)
    abs    (FCRS' s) = FCRS' (abs    s)
    fromInteger x  = FCRS' (fromInteger x)

instance Fractional (b '[o]) => Fractional (CState (FullyConnectedR' s) b '[i] '[o]) where
    FCRS' s1 / FCRS' s2 = FCRS' (s1 / s2)
    recip (FCRS' s)     = FCRS' (recip s)
    fromRational x      = FCRS' (fromRational x)

instance Floating (b '[o]) => Floating (CState (FullyConnectedR' s) b '[i] '[o]) where
    sqrt (FCRS' s)     = FCRS' (sqrt s)

instance ( BLAS b
         , KnownNat i
         , KnownNat o
         , Floating (b '[o])
         , Floating (b '[o,i])
         , Floating (b '[o,o])
         , Reifies s MapFunc
         )
      => Component (FullyConnectedR' s) b '[i] '[o] where
    data CParam  (FullyConnectedR' c) b '[i] '[o] =
            FCRP' { _fcrInpWeights'   :: !(b '[o,i])
                  , _fcrStateWeights' :: !(b '[o,o])
                  , _fcrBiases'       :: !(b '[o])
                  }
    data CState  (FullyConnectedR' c) b '[i] '[o] = FCRS' { _fcrState' :: !(b '[o]) }
    type CConstr (FullyConnectedR' c) b '[i] '[o] =
      ( Num (b '[o,i])
      , Num (b '[o,o])
      )
    data CConf   (FullyConnectedR' c) b '[i] '[o] = forall d. ContGen d => FCRC' d

    componentOp = bpOp . withInps $ \(x :< p :< s :< Ø) -> do
        wI :< wS :< b :< Ø <- gTuple #<~ p
        s0 <- opIso (iso _fcrState' FCRS') ~$ (s :< Ø)
        y  <- matVecOp ~$ (wI :< x  :< Ø)
        s1 <- matVecOp ~$ (wS :< s0 :< Ø)
        s2 <- tmapOp (runMapFunc mf) ~$ (s1 :< Ø)
        let z = y + s2 + b
        s' <- opIso (iso FCRS' _fcrState') ~$ (s2 :< Ø)
        return $ z :< s' :< Ø
      where
        mf :: MapFunc
        mf = reflect (Proxy @s)

    defConf = FCRC' (normalDistr 0 0.01)

    initParam = \case
      i `SCons` SNil -> \case
        so@(o `SCons` SNil) -> \(FCRC' d) g -> do
          wI <- genA (o `SCons` (i `SCons` SNil)) $ \_ ->
            realToFrac <$> genContVar d g
          wS <- genA (o `SCons` (o `SCons` SNil)) $ \_ ->
            realToFrac <$> genContVar d g
          b <- genA so $ \_ ->
            realToFrac <$> genContVar d g
          return $ FCRP' wI wS b
        _ -> error "inaccessible"
      _ -> error "inaccessible"

    initState _ so (FCRC' d) g =
        FCRS' <$> genA so (\_ -> realToFrac <$> genContVar d g)

instance ( BLAS b
         , KnownNat i
         , KnownNat o
         , Floating (b '[o])
         , Floating (b '[o,i])
         , Floating (b '[o,o])
         , Reifies s MapFunc
         )
      => ComponentLayer 'Recurrent (FullyConnectedR' s) b '[i] '[o] where
    componentRunMode = RMNotFF


{-# LANGUAGE RankNTypes, BangPatterns, ScopedTypeVariables, FlexibleInstances #-}

module Math.Probably.MALA where

import Math.Probably.MCMC
import qualified Math.Probably.PDF as PDF
import Math.Probably.RandIO
import Math.Probably.FoldingStats
import Math.Probably.Sampler
import qualified Math.Probably.PDF as PDF
import Control.Applicative

import Numeric.LinearAlgebra
import Numeric.AD
import Text.Printf
import System.IO
import Data.Maybe

import Statistics.Test.KolmogorovSmirnov
import Statistics.Test.MannWhitneyU
import qualified Data.Vector.Unboxed as U
import qualified Control.Monad.State.Strict as S
import qualified Data.Vector.Storable as VS

import Debug.Trace
import Data.IORef
import Control.Spoon

class IsNaN a where
  isItNaN :: a -> Bool

instance IsNaN Double where
  isItNaN = nanOrInf

instance IsNaN (Vector Double) where
  isItNaN = any nanOrInf . toList

instance IsNaN (Matrix Double) where
  isItNaN = or . (map (any nanOrInf . toList)) . toRows

checkNaN :: IsNaN a => String -> a -> String
checkNaN s x | isItNaN x = " "++s++" "
             | otherwise = ""

class Covariance a where
  mvnSampler :: Vector Double ->  Double -> a -> Sampler (Vector Double)
  mvnPDF :: Vector Double -> Double -> a -> PDF.PDF (Vector Double)
  covMul :: a -> Vector Double -> Vector Double
  restrict :: Vector Int -> a -> a

instance Covariance (Matrix Double,Matrix Double,Matrix Double) where
  mvnSampler mu sigma (cov, covInv, covChol)
     = multiNormalByChol mu (scale (sqrt sigma) covChol)
  mvnPDF mu sigma (cov, covInv, covChol)
     = PDF.multiNormalByInvFixCov (scale (recip sigma) covInv) mu
  covMul (cov, covInv, covChol) v = cov <> v

instance Covariance (Vector Double) where
  mvnSampler mu sigma varv
    = multiNormalIndep (scale sigma varv) mu
  mvnPDF mu sigma varv
    = PDF.multiNormalIndep (scale sigma varv) mu
  covMul varv v = VS.zipWith (*) varv  v
  restrict vixs v = VS.backpermute v vixs

data MalaPar = MalaPar { mpXi :: !(Vector Double),
                         mpPi :: !Double,
                         mpGradLast :: !(Vector Double),
                         mpSigma :: !Double,
                         mpCount :: !Double,
                         mpAccept :: !Double,
                         mpFreezeSigma :: !Bool,
                         mpThin :: !Int,
                         mpMaxGradLen :: !Double } --,
--                         mpLastRatio :: !Double }
  deriving Show

data Pair a b = Pair !a !b

data List a = Cons !a !(List a)
            | Nil

unPair (Pair x y) = (x,y)
unList (Cons x xs) = x: unList xs
unList Nil =[]

trunc t v = scale (t/(max t (norm2 v))) v

mala1 :: Covariance a => a -> (Vector Double -> (Double,Vector Double))
         -> Bool
         -> MalaPar
         -> Sampler MalaPar
mala1 cov postgrad useCache
      (MalaPar xi piCached gradientiCached sigma tr tracc freeze thinn maxGrad) = do
  let (pi,gradienti) = if useCache
                          then (piCached, gradientiCached)
                          else let (p, gr_untr) = postgrad xi
                               in (p, trunc maxGrad gr_untr)
  let xstarMean = xi + scale (sigma/2) (cov `covMul` gradienti)
  xstar <- mvnSampler xstarMean sigma cov
  u <- unitSample
  let (!pstar, gradientStarUntrunc) = postgrad xstar
      gradientStar = trunc maxGrad gradientStarUntrunc
  let !revJumpMean = xstar + scale (sigma/2) (cov `covMul` gradientStar)
      ptop = mvnPDF revJumpMean sigma cov xi
      pbot = mvnPDF xstarMean sigma cov xstar
      ratio = exp $   pstar -pi + ptop - pbot
      tr' = max 1 tr
      sigmaNext = case () of
         _ | freeze -> sigma
      accept = tracc / tr
      freezeNext = freeze {-| freeze = True
                 | not freeze = tr > 100 && accept > 0.5 && accept < 0.6  -}
  if trace (show $ (tr, pstar ,pi , ratio, sigma)) $ u < ratio
     then return $ MalaPar xstar pstar (gradientStar)
                           (if freezeNext then sigma else (min 1.4 $ 1+kmala/tr')*sigma)
                           (tr+1) (tracc+1) freezeNext  thinn maxGrad
--                           sigma (tr+1) (tracc+1)
     else return $ MalaPar xi pi ( gradienti)
                           (if freezeNext then sigma else (max 0.7143 $ 1-kmala/tr')**1.3*sigma)
                           (tr+1) tracc freezeNext thinn maxGrad
--                           sigma (tr+1) tracc

blockMala1 :: Covariance a =>
  (Vector Double -> (Double,Vector Double)) ->
  Vector Double ->
  [(a, MalaPar, Vector Int)] ->
  Sampler (Vector Double, [(a, MalaPar, Vector Int)])
blockMala1 postgrad xi blocks = do
  let go totalV [] accblocks = return (totalV, reverse accblocks)
      go totalV ((cov, mp0, vixs) : blocks) accblocks = do
          let myPostGrad vpars
               = let (p, grad) = postgrad $ VS.update_ totalV vixs vpars
                 in (p, VS.backpermute grad vixs)
          mp1 <- mala1 cov myPostGrad False mp0
          let newTotalV = VS.update_ totalV vixs (mpXi mp1)
          go newTotalV blocks $ (cov, mp1, vixs) : accblocks
  go xi blocks []

runMalaBlocks :: Covariance a => a -> (Vector Double -> (Double,Vector Double))
         ->  Int -> Double -> Int -> Vector Double -> [Vector Int]
         -> RIO [Vector Double]
runMalaBlocks cov  postgrad nsam truncN thinN init  vixs = go nsam initBlocks init []  where
  initBlocks = flip map vixs $ \vix->
     (restrict vix cov,
      MalaPar (VS.backpermute init vix) 0 (fromList [0]) 0.001 0 0 False thinN truncN,
      vix)
  go 0 _ _ vs = return vs
  go n blocks0 v0 vs = do
        (!v1, blocks1) <- sample $ blockMala1 postgrad v0 blocks0
        let (_,mpLast,_) = last blocks1
        io $ print $ (n, mpPi mpLast, map (\(_,mp,_) -> mpSigma mp) blocks1)
        let newChainRes = if thinN == 0 || n `mod` thinN ==0
                             then v1 : vs
                             else vs
        go (n-1) blocks1 v1 newChainRes

{-runMala :: Matrix Double -> (Vector Double -> (Double,Vector Double))
         ->  Int -> Vector Double -> RIO [(Double,Vector Double)]
runMala cov postgrad nsam init = go nsam mp1 [] where
  go 0 mpar xs = do io $ putStrLn $ "MALA accept = "++show (mpAccept mpar/mpCount mpar)
                    io $ putStrLn $ "MALA sigma = "++show (mpSigma mpar)
                    return xs
  go n y xs = do y1 <- sample $ mala1 cov postgrad y
                 go (n-1) y1 $ (mpPi y1, mpXi y1):xs
  (pi, gradi) = postgrad init
  mp1 =  MalaPar init pi (gradi)  1 0 0 False -}

runMalaMP :: Covariance a => a -> (Vector Double -> (Double,Vector Double))
         ->  Int -> MalaPar -> [(Double,Vector Double)] -> RIO (MalaPar, [(Double,Vector Double)])
runMalaMP cov  pdf nsam init xs0 = go nsam init xs0 where
  go 0 mpar xs = do io $ putStrLn $ "MALA accept = "++show (mpAccept mpar/mpCount mpar)
                    io $ putStrLn $ "MALA sigma = "++show (mpSigma mpar)
                    return (mpar, xs)
  go n y xs = do y1 <- sample $ mala1 cov pdf True y
--                 io $ do putStrLn $ show (mpCount y1, mpPi y1, mpSigma y1)
--                         hFlush stdout
                 let newChainRes = if mpThin y1 == 0 || round (mpCount y1) `mod` mpThin y1 ==0
                                      then let !xi = mpXi y1
                                               !pi = mpPi y1
                                               !more = (pi,xi)
                                           in more:xs
                                      else xs
                 go (n-1) y1 newChainRes

runMalaUntilBetter :: Covariance a => a ->(Vector Double -> (Double,Vector Double))
         ->  Int -> MalaPar -> RIO (MalaPar, [(Double,Vector Double)])
runMalaUntilBetter cov  pdf nsam init = go nsam init [] where
  pTarget = mpPi init
  go 0 mpar xs = do io $ putStrLn $ "MALA accept = "++show (mpAccept mpar/mpCount mpar)
                    io $ putStrLn $ "MALA sigma = "++show (mpSigma mpar)
                    return (mpar, xs)
  go n y xs = do y1 <- sample $ mala1 cov  pdf True y
                 io $ do putStrLn $ show (mpCount y1, mpPi y1, mpSigma y1)
                         hFlush stdout
                 let newChainRes = if mpThin y1 == 0 || round (mpCount y1) `mod` mpThin y1 ==0
                                      then (mpPi y1, mpXi y1):xs
                                      else xs
                 if mpPi y1 > pTarget
                    then return (y1, newChainRes)
                    else go (n-1) y1 $ newChainRes

{-runMalaMPaccept' :: Matrix Double -> (Vector Double -> (Double,Vector Double))
         ->  Int -> MalaPar -> RIO (MalaPar, [(Double,Vector Double)])
runMalaMPaccept'  cov pdf nsam init  = do
  stseed <- S.get
  res <- io $ do
   seedref <- newIORef stseed
   resRef <- newIORef Nil
   let go !mpar
        | (round $ mpAccept mpar) >= nsam
          = return mpar
        | otherwise
          = do seed <- readIORef seedref
               putStr ((show (round $ mpAccept mpar)) ++".") >> hFlush stdout
               let (!mpar1, !seed1) = unSam (mala1 cov pdf mpar) seed
               writeIORef seedref seed1
               modifyIORef resRef (Cons (Pair (mpPi mpar1) ( mpXi mpar1)))
               go mpar1
   mpar <- go init
   seed <- readIORef seedref
   reslist <- readIORef resRef
   return (seed, (mpar,reslist))
  S.put $ fst res
  let mp = fst $ snd res
  return $ (mp, map unPair $ unList $ snd $ snd res)

--  pi = pdf $ toList $ mpXi init -}

runMalaMPaccept :: Covariance a => a -> (Vector Double -> (Double,Vector Double))
         -> Int -> MalaPar ->  RIO (MalaPar, [(Double,Vector Double)])
runMalaMPaccept cov pdf nsam init = go init [] where
  go !mpar !xs
    | (round $ mpAccept mpar) >= nsam = do
         io $ putStrLn $ "MALA accept = "++show (mpAccept mpar/mpCount mpar)
         io $ putStrLn $ "MALA sigma = "++show (mpSigma mpar)
         return (mpar, xs)
    | mpAccept mpar < 0.5 && mpCount mpar > 20 =
         return (mpar, [])
    | otherwise = do
         !mpar1 <- sample $ mala1 cov  pdf True mpar
         io $ putStr ((show (round $ mpAccept mpar)) ++".") >> hFlush stdout
         go mpar1 $ (mpPi mpar1, mpXi mpar1):xs

kmala = 5

{-runMalaRioESS ::  Matrix Double -> (Vector Double -> (Double,Vector Double))
                  ->  Int -> Vector Double -> RIO [Vector Double]
runMalaRioESS cov pdf want_ess xi = do
    let (p0,grad0) = pdf xi
    let sigma0 = 1.5 / (realToFrac $ dim xi) -- determined empirically
    let mp0 = MalaPar xi p0 grad0 sigma0 0 0 False
        nsam0 = want_ess*20
    io $ putStrLn $ "initial sigma = "++show (mpSigma mp0)
    (mp1, xs1) <- runMalaMP cov pdf nsam0 mp0 []
    let have_ess = min (mpAccept mp1) $ calcESSprim $ map snd xs1
    if have_ess > realToFrac want_ess
       then return $ map snd xs1
       else do let need_ess =  max 1 $ realToFrac want_ess - have_ess
                   samples_per_es = realToFrac nsam0/have_ess
                   to_do = round $ samples_per_es * need_ess
               (mp2, xs2) <- runMalaMP cov pdf to_do mp1 xs1
               return  $ map snd  xs2 -}

runMalaRioCodaESS ::  Covariance a => a-> (Vector Double -> (Double,Vector Double))
                  ->  Int -> Double -> Vector Double -> RIO [Vector Double]
runMalaRioCodaESS cov pdf want_ess truncN xi = do
    let (p0,grad0) = pdf xi
    let sigma0 = 1.5 / (realToFrac $ dim xi) -- determined empirically
    let mp0 = MalaPar xi p0 (trunc truncN grad0) sigma0 0 0 False 0 truncN
        nsam0 = want_ess*1
    let converged mp  xs = do
         let have_ess = min (mpAccept mp) $ calcESSprim $ map snd xs
         io $ putStrLn $ "ESS=" ++show have_ess
         if have_ess > realToFrac want_ess
            then return $ map snd xs
            else do let need_ess =  max 1 $ realToFrac want_ess - have_ess
                        samples_per_es = realToFrac nsam0/have_ess
                        to_do = round $ samples_per_es * need_ess
                    io $ putStrLn $ "running converged for "++show to_do
                    (mp2, xs2) <- runMalaMP cov pdf to_do (mp {mpFreezeSigma = True}) xs
                    io $ putStrLn $ "All done"
                    return $ map snd xs2
    let go mp  n xs = do
            (mp2, xs2) <- runMalaMP cov  pdf n mp []
            let testres = mannWhitneyUtest TwoTailed 0.05 (U.fromList $ map fst xs2)
                                            (U.fromList $ map fst xs)

            if testres/= Just NotSignificant
               then do io$ putStrLn $ "not converged: "++show testres++" at "++show (mpPi mp2)
                       go mp2  (round $ realToFrac n*2) xs2
               else do io$ putStrLn "converged!"
                       converged mp2  (xs2++xs)

    io $ putStrLn $ "initial sigma = "++show (mpSigma mp0)
    (mp1, xs1) <- runMalaMPaccept cov  pdf nsam0 mp0
    case xs1 of
            [] -> return []
            _ -> go mp1  (round $ mpCount mp1) xs1

{-    case (spoon (invlndet cov), mbCholSH cov) of
       (Just (covInv, (lndt,_)), Just covChol) -> go_rest (covInv, covChol)
       _ -> let cov' = PDF.posdefify cov in
            case (spoon (invlndet cov'), mbCholSH cov') of
              (Just (covInv, (lndt,_)), Just covChol) -> go_rest (covInv, covChol)
              _ ->  do io $ putStrLn "non-invertible covariance matrix"
                       return [] -}

runMalaRioSimple ::  Covariance a => a -> (Vector Double -> (Double,Vector Double))
                  ->  Int -> Double -> Int -> Vector Double -> RIO [Vector Double]
runMalaRioSimple cov pdf samples maxGrad thinN xi = do
    let (p0,grad0) = pdf xi
    let sigma0 = 2.7e-4 --1.5 / (realToFrac $ dim xi) -- determined empirically
    let mp0 = MalaPar xi p0 (trunc maxGrad grad0) sigma0 0 0 False thinN maxGrad

    (mp2, xs2) <- runMalaMP cov  pdf samples mp0 []
    io $ putStrLn $ "All done"
    return $ map snd xs2

calcCovariance :: Vector Double ->
                  Vector Double ->
                  (Vector Double -> (Double,Vector Double)) ->
                  (Vector Double -> Double) ->
                  Either (Vector Double) (Matrix Double, Matrix Double, Matrix Double)
calcCovariance vinit vnear postgrad posterior = finalcov where
   ndim = dim vinit
   finalcov
     | ndim > 0 -- > 20000 --FIXME
        = Left $ calcFDindepVars vinit vnear posterior --Left $ iCov vinit --
     | otherwise
        = hessToCov (calcFDhess vinit vnear postgrad) Nothing

iCov v = VS.replicate (VS.length v) 1 -- in (m,m,m)

calcFDindepVars v v' post = trace ("FDVars = "++show (VS.take 10 vars))  $ vars where
   hv =  mapVector (*1e-4) v --mapVector (max 1e-9 .  abs) $ v - v'
   n = dim v
   postv = post v
   postPlus i = post $ v VS.// [(i,v @>i + hv @> i)]
   postMinus i = post $ v VS.// [(i,v @>i - hv@> i)]
   vars = buildVector n fvar
   fvar i = negate $ recip $ (postPlus i - 2*postv + postMinus i)/((hv @> i)*(hv @> i))

calcFDhess v v' postgrad = hess2 where
   grad =  snd . postgrad
   gradi i = (@>i) . grad
   hv = mapVector (max 1e-9 . abs) $ v - v'
   n = dim v
   gradv = grad v
   grads =  fromRows $ flip map [0..(n-1)] $ \i ->
               grad (v VS.// [(i,v @>i + hv @> i)])
   fhess (i,j) | i<j = 0
               | otherwise = (grads @@>(j,i) - (gradv @> i))
                                       /(2*(hv @> j)) +
                             (grads @@>(i,j)- (gradv @> j))
                                       /(2*(hv @> i))

--   hess3 = scale (recip $ realToFrac n) $ sum $ map outerSelf $ grads

   hess1 = buildMatrix n n fhess
   hess2 = buildMatrix n n $ \(i,j) -> if i>=j then hess1 @@> (i,j)
                                               else hess1 @@> (j,i)

outerSelf v= v `outer` v

hessToCov hess mOriginal =
   let mTryPosdefify = case mOriginal of
            Nothing -> hessToCov (PDF.posdefify hess) (Just hess)
            Just horiginal -> let ds =  mapVector (negate . recip)
                                         $ takeDiag horiginal
                              in trace ("USING DIAGS"++show (VS.take 10 ds)) $ Left $ ds
   in
   case spoon $ inv $ negate $ hess of
     Just cov -> case mbCholSH cov of
                   Just cholm ->  trace ("invert success:"++show (VS.take 10 $ takeDiag cov) ++ "det="++show (det cov)) $ Right (cov, negate hess, cholm)
                   Nothing -> trace ("chol fail") mTryPosdefify
     Nothing -> trace ("inv fail") mTryPosdefify
                  {- Just cov -> case mbCholSH cov of
                                Just cholm ->  Right (cov, negate hessToCov,cholm)
                                Nothing -> case mbCholSH $ PDF.posdefify cov of
                                            Just cholm -> Right (cov, negate hess,cholm)
                                            Nothing -> vars
                  Nothing -> vars -}

acceptSM ampar  | mpCount ampar == 0 = "0/0"
               | otherwise = printf "%.3g" (rate::Double) ++ " ("++show yes++"/"++show total++")" where
   rate = realToFrac (yes) / realToFrac (total)
   yes = mpAccept ampar
   total = mpCount ampar