{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
-- | Fourier amplitude sensitivity test, a quasi-monte carlo method
--
-- does not yet play well with "ErrorProp.MonteCarlo", though a function
-- like 'ErrorProp.MonteCarlo.sample' should be straightforward to
-- implement. However, 'fast' provides the core of such a function
module ErrorProp.QuasiMC where

import qualified Data.Map as M
import Math.Polynomial
import Math.Polynomial.Bernoulli
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Generic as VG
import Data.List

import Data.Numbers.Primes
import Data.VectorSpace

import ErrorProp.QuasiMC.Cukier (cukierOmegas)

import Statistics.Sample
import Data.Maybe

import ErrorProp.MonteCarlo
import ErrorProp.Common

f1 alpha2 = constPoly 1 ^+^ ((-1)^(1 + alpha2 `mod` 2) * (2*pi)^(2*alpha2)) *^ bx
  where bx = bernoulliPoly !! (2*alpha2)

-- the explicit equations for F2,F4,F6 (equation 4.14,15,16) are correct
prop_f11 x = f11 x ~= f1 1 `evalPoly` x
  where f11 x = 1 + 2 * pi^2 * (x^2 - x + 1/6)

prop_f12 x = f12 x ~= f1 2 `evalPoly` x
  where f12 x = 1 + pi^4/45 * (1 - 30* x^2 *  (1-x)^2)

prop_f13 x = f13 x ~= f1 3 `evalPoly` x
  where f13 x = 1 + (2*pi^6/945) * (1 - 21*x^2 + 105*x^4 -126*x^5 + 42*x^6)

x ~= y = (2 * abs(x-y) / (1 + abs x + abs y)) < (1e-8 :: Double)
infix 4 ~=

-- eqn 4.10. Is the function integrated to generate the "best"
-- lattice points
f2 alpha2 xs = V.product $ V.map (f1 alpha2 `evalPoly`) xs

-- | Korbov's construction for good lattice points. Is an LCG
--
-- 4.32
korbov l n s = V.iterateN s (\p -> (p * l) `mod` n) 1

-- | try all values of L up to n/2
korbovLExhaustive alpha2 n s =
    head $
    sort [ (abs (r1-1) :: Double, l)
            | l <- [1 .. n `div` 2 ],
              let r1 = rank1 (korbov l n s) n (V.replicate s 0) (f2 alpha2) ]

-- | assume L has k prime factors. Provides candidate L values,
-- and the error in the integral used to decide that L is the best.
-- The error is absolute and relative, since the target function
-- integrates to 1.
korbovLCompositeK
  :: Int -- ^ alpha/2, roughness of the function being integrated
  -> Int -- ^ k
  -> Int -- ^ n, number of points in the quadrature
  -> Int -- ^ s, dimension
  -> [(Double, Int)] -- ^ (error, best l)
korbovLCompositeK alpha2 k n s =
    sort [ (abs (rankL l - 1) :: Double, l)
            | l <- go k ps 1 ]
  where
    rankL l = rank1 (korbov l n s) n (V.replicate s 0) (f2 alpha2)
    go 0 _ n = [n]
    go k ps1 n = do
      q : ps1' @ (_:_) <- tails ps1
      go (k-1) ps1' (n*q)

    ps = takeWhile ( < round ( (fromIntegral n / 2) ** (1/fromIntegral k)))
                primes

{- | rank1 shifted lattice rule

Given

> f :: [0,1)^s -> R

calculate

> Q(Z,N,C)f = (1/N) sum_j=0^N f { j z/N + c }

is an approximation to the integral of @f@ (@If@) over that region.

equation 4.42 Sloan&Joe 1994 pg 89

If C is multivariant uniform [0,1)^s, we can do a CI for the actual
integral

P( |Q - I| < sigma v) >= 1 - (1/v^2)

with sigma^2 estimated by sum_k=1^q (Q - Qbar)^2 / ( q (q-1))

is unbiased. Quantiles are biased (Lemieux 2008). Bootstrap/jacknife resample to estimate bias.

trapezoidal / rectangle rule is good for the periodic integral here,
but possibly try "GA Evans, JR Webster 1999"

-}
rank1 :: (Integral i, RealFrac a, VG.Vector v a, VG.Vector vi i,
            VG.Vector vi a)
  => vi i -- ^ z
  -> Int -- ^ N
  -> v a -- ^ c
  -> (v a -> a) -- ^ f
  -> a
rank1 z n c f = go 0 0
  where
    zn = VG.convert $ VG.map (\zi -> fromIntegral zi / fromIntegral n) z

    go !j !accum
      | j == VG.length zn = accum
      | otherwise = go (j+1) $ f $ VG.zipWith (+) c
              $ VG.map (\zni -> fracPart $ fromIntegral j*zni) zn

{- |

>>> map fracPart [12.3, -2.5]
[0.3, 0.5]

-}
fracPart x = x - fromInteger (floor x)

-- | fourier amplitude sensitivity test
fast ::
  (RealFrac b, VG.Vector V.Vector b) =>
    V.Vector Int -- ^ N incommensurate frequencies (see 'korbov', 'cukierOmegas', or others)
  -> (V.Vector b -> Double) -- ^ a function of N variables which will be sampled on @[0, 1)^N@
  -> V.Vector Double -- ^ fraction of the total variance attributed to the i'th variable
fast omegas f =
    let omegaMax = VG.maximum omegas
        q = omegaMax

        xj k = VG.map (\w -> fracPart $ fromIntegral (w*k) / fromIntegral (2*q+1)) omegas
        sks = V.enumFromTo (-q) q
        fks = VG.map (f . xj) sks

        withK k = VG.map (\w -> mean
                              $ VG.zipWith (\s f -> f * k (w*s)) sks fks)
                omegas

        nq = fromIntegral (2*q+1)

        a i = withK (cos . (*(2*i*pi/nq)) . fromIntegral)
        b i = withK (sin . (*(2*i*pi/nq)) . fromIntegral)

        vjs = foldl1 (V.zipWith (+))
              $ map (V.map (^2))
              [  if odd j then b fij else a fij
                  | j <- [1 .. 2],
                    let fij = fromIntegral j ]

        -- normalizeFac = 2 / variance fks -- XXX the
        -- correct normalization should be some analytical
        -- expression... like above, not the value below
        normalizeFac = recip $ V.sum vjs

    in V.map (*normalizeFac) vjs

fastKi s np k =
  let n = primes !! np
      (_, l) = korbovLExhaustive 1 n s
      omegas = korbov l n s

  in fast omegas (\v -> v V.! k)

fastKis np s = [ (k,fastKi s np k)  | k <- [0 .. s-1]]
{- properties.

> ki k = fast omegas (\x -> x V.! k)

> ki k V.! j == 1 if k==j else 0

-}

-- | `pickNP err s` uses 'fastKis' to pick the parameter np, such that the
-- value that should be 1 is within err of 1, for a problem of dimension s
pickNP err s = head [ np
    | np <- [ 2 .. ],
     all ok (fastKis np s) ]
  where ok (i, v) = abs (1 - v V.! i) < err

sampleQMC :: Double -- ^ error tolerance (see 'pickNP')
  -> MV Double
  -> [(String, Double)] -- variable name, fraction of the total variance
sampleQMC err mv =
  let n = wheelSieve 5 !! pickNP err s
      (_, l) = korbovLExhaustive 1 n s
      omegas = korbov l n s
      s = M.size (mv_vars mv)

      f v = mv_sample mv (updateMV v (mv_vars mv))
  in M.keys (mv_vars mv) `zip` VG.toList (fast omegas f)

updateMV :: (Floating b, V.Unbox b) => V.Vector b -> M.Map String (Maybe b) -> M.Map String b
updateMV v m = snd $ M.mapAccum (\ n e -> (n+1, fromMaybe 1 e * toUnitVar (v V.! n))) 0 m
  where
        toUnitVar x = sqrt 12 * (x - 0.5)


{-# LANGUAGE PartialTypeSignatures, ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wno-orphans -funbox-strict-fields #-}
-- | Provides @'Matrix'@ and @'MMatrix'@ instances for
--  @'LA.Matrix'@ type of @hmatrix@ package.
module Algebra.Matrix.HMatrix (STMatrix) where
import           Algebra.Matrix.Generic
import           Algebra.Matrix.Generic.Mutable (MMatrix (..))
import           Algebra.Prelude.Core
import           Control.Monad                  (forM_)
import           Control.Monad.Primitive
import qualified Data.Coerce                    as DC
import qualified Data.Vector.Storable           as SV
import qualified Numeric.LinearAlgebra          as LA
import qualified Numeric.LinearAlgebra.Devel    as LA

type instance Mutable LA.Matrix = STMatrix
type instance Row LA.Matrix = LA.Vector
type instance Column LA.Matrix = LA.Vector
type instance Row STMatrix = LA.Vector
type instance Column STMatrix = LA.Vector

-- | Wrapper type for @hmatrix@'s @'LA.STMatrix'@.
data STMatrix s a = STMat { stmRows :: !Int
                          , stmCols :: !Int
                          , unwrapSTM :: LA.STMatrix s a
                          }

instance (Num a, LA.Container LA.Matrix a) => Matrix LA.Matrix (WrapNum a) where
  basicRowCount = LA.rows
  basicColumnCount = LA.cols
  unsafeFromRows rs = DC.coerce $ LA.fromRows (DC.coerce rs :: [LA.Vector a])
  unsafeFromColumns cs = DC.coerce $ LA.fromColumns (DC.coerce cs :: [LA.Vector a])
  basicUnsafeIndexM m i j =
    return $ WrapNum $ (DC.coerce m :: LA.Matrix a) `LA.atIndex` (i, j)
  unsafeFreeze = stToPrim . LA.unsafeFreezeMatrix . unwrapSTM
  unsafeThaw m = stToPrim $ STMat (LA.rows m) (LA.cols m) <$> LA.unsafeThawMatrix m
  basicUnsafeGetRowM m i = return $ DC.coerce $ getRowLA  i (DC.coerce m :: LA.Matrix a)
  basicUnsafeGetColumnM m i =
    return $ DC.coerce $ head $ LA.toColumns (DC.coerce m LA.¿ [i] :: LA.Matrix a)
  unsafeGenerate w h f =
    DC.coerce $
    LA.fromRows [ SV.generate w $ \j -> unwrapNum $ f i j
                | i <- [0..h - 1]
                ]

instance (Integral a, LA.Container LA.Matrix a) => Matrix LA.Matrix (WrapIntegral a) where
  basicRowCount = LA.rows
  basicColumnCount = LA.cols
  unsafeFromRows rs = DC.coerce $ LA.fromRows (DC.coerce rs :: [LA.Vector a])
  unsafeFromColumns cs = DC.coerce $ LA.fromColumns (DC.coerce cs :: [LA.Vector a])
  basicUnsafeIndexM m i j =
    return $ WrapIntegral $ (DC.coerce m :: LA.Matrix a) `LA.atIndex` (i, j)
  unsafeFreeze = stToPrim . LA.unsafeFreezeMatrix . unwrapSTM
  unsafeThaw m = stToPrim $ STMat (LA.rows m) (LA.cols m) <$> LA.unsafeThawMatrix m
  basicUnsafeGetRowM m i = return $ DC.coerce $ getRowLA  i (DC.coerce m :: LA.Matrix a)
  basicUnsafeGetColumnM m i =
    return $ DC.coerce $ head $ LA.toColumns (DC.coerce m LA.¿ [i] :: LA.Matrix a)
  unsafeGenerate w h f =
    DC.coerce $
    LA.fromRows [ SV.generate w $ \j -> unwrapIntegral $ f i j
                | i <- [0..h - 1]
                ]

instance (Fractional a, LA.Container LA.Matrix a) => Matrix LA.Matrix (WrapFractional a) where
  basicRowCount = LA.rows
  basicColumnCount = LA.cols
  unsafeFromRows rs = DC.coerce $ LA.fromRows (DC.coerce rs :: [LA.Vector a])
  unsafeFromColumns cs = DC.coerce $ LA.fromColumns (DC.coerce cs :: [LA.Vector a])
  basicUnsafeIndexM m i j =
    return $ WrapFractional $ (DC.coerce m :: LA.Matrix a) `LA.atIndex` (i, j)
  unsafeFreeze = stToPrim . LA.unsafeFreezeMatrix . unwrapSTM
  unsafeThaw m = stToPrim $ STMat (LA.rows m) (LA.cols m) <$> LA.unsafeThawMatrix m
  basicUnsafeGetRowM m i = return $ DC.coerce $ getRowLA  i (DC.coerce m :: LA.Matrix a)
  basicUnsafeGetColumnM m i =
    return $ DC.coerce $ head $ LA.toColumns (DC.coerce m LA.¿ [i] :: LA.Matrix a)
  unsafeGenerate w h f =
    DC.coerce $
    LA.fromRows [ SV.generate w $ \j -> unwrapFractional $ f i j
                | i <- [0..h - 1]
                ]

getRowLA :: LA.Element t => Int -> LA.Matrix t -> LA.Vector t
getRowLA i m = head $ LA.toRows $ m LA.? [i]
getColLA :: LA.Element t => Int -> LA.Matrix t -> LA.Vector t
getColLA i m = head $ LA.toColumns $ m LA.¿ [i]

unwrapSTMNum :: STMatrix s (WrapNum a) -> LA.STMatrix s a
unwrapSTMNum = DC.coerce . unwrapSTM
{-# INLINE unwrapSTMNum #-}

castSTMatrixNum :: LA.STMatrix s (WrapNum a) -> LA.STMatrix s a
castSTMatrixNum = DC.coerce
{-# INLINE castSTMatrixNum #-}

instance (Num a, LA.Container LA.Matrix a) => MMatrix STMatrix (WrapNum a) where
  basicUnsafeNew n m = stToPrim $ STMat n m <$> LA.newMatrix 0 n m
  basicInitialise _ = return ()
  basicRowCount = stmRows
  basicColumnCount = stmCols
  unsafeGetRow i m =
    stToPrim $ DC.coerce . getRowLA i <$> LA.unsafeFreezeMatrix (unwrapSTMNum m)
  unsafeGetColumn i m = stToPrim $ DC.coerce . getColLA i <$> LA.unsafeFreezeMatrix (unwrapSTMNum m)
  unsafeFill n m a = stToPrim $ STMat n m <$> LA.newMatrix a n m
  unsafeFromRow r = stToPrim $ STMat 1 (SV.length r) <$> LA.unsafeThawMatrix (LA.asRow r)
  unsafeFromRows rs = stToPrim $ do
    let m = LA.fromRows (DC.coerce rs :: [LA.Vector a])
    STMat (LA.rows m) (LA.cols m) . DC.coerce <$> LA.unsafeThawMatrix m
  unsafeFromColumn c = stToPrim $ STMat (SV.length c) 1 <$> LA.unsafeThawMatrix (LA.asColumn c)
  unsafeFromColumns rs = stToPrim $ do
    let m = LA.fromColumns (DC.coerce rs :: [LA.Vector a])
    STMat (LA.rows m) (LA.cols m) . DC.coerce <$> LA.unsafeThawMatrix m
  unsafeCopy (STMat r c m)  (STMat _ _ m') = stToPrim $ do
    m0 <- LA.unsafeFreezeMatrix m'
    LA.setMatrix (DC.coerce m :: LA.STMatrix w a) r c (DC.coerce m0 :: LA.Matrix a)
  unsafeRead (STMat _ _ m) i j = stToPrim $ LA.readMatrix m i j
  unsafeWrite (STMat _ _ m) i j x = stToPrim $ LA.writeMatrix m i j x
  basicSet (STMat r c m) x = stToPrim $ LA.modifyMatrix m r c $ const x
  basicUnsafeSwapRows (STMat _ _ m) i j =
    stToPrim $ LA.rowOper (LA.SWAP i j LA.AllCols) $ castSTMatrixNum m
  unsafeScaleRow (STMat _ _ m) i c =
    stToPrim $ LA.rowOper (LA.SCAL (unwrapNum c) (LA.Row i) LA.AllCols) $ castSTMatrixNum m
  basicUnsafeIMapRowM (STMat _ c m) i f = forM_ [0..c-1] $ \j ->
    stToPrim . LA.writeMatrix m i j =<< f j =<< stToPrim (LA.readMatrix m i j)
  toRows (STMat _ _ m) =
    stToPrim $ DC.coerce . LA.toRows <$> LA.unsafeFreezeMatrix (castSTMatrixNum m)
  toColumns (STMat _ _ m) =
    stToPrim $ DC.coerce . LA.toColumns <$> LA.unsafeFreezeMatrix (castSTMatrixNum m)

unwrapSTMFractional :: STMatrix s (WrapFractional a) -> LA.STMatrix s a
unwrapSTMFractional = DC.coerce . unwrapSTM
{-# INLINE unwrapSTMFractional #-}

castSTMatrixFractional :: LA.STMatrix s (WrapFractional a) -> LA.STMatrix s a
castSTMatrixFractional = DC.coerce
{-# INLINE castSTMatrixFractional #-}

instance (Fractional a, LA.Container LA.Matrix a) => MMatrix STMatrix (WrapFractional a) where
  basicUnsafeNew n m = stToPrim $ STMat n m <$> LA.newMatrix 0 n m
  basicInitialise _ = return ()
  basicRowCount = stmRows
  basicColumnCount = stmCols
  unsafeGetRow i m =
    stToPrim $ DC.coerce . getRowLA i <$> LA.unsafeFreezeMatrix (unwrapSTMFractional m)
  unsafeGetColumn i m = stToPrim $ DC.coerce . getColLA i <$> LA.unsafeFreezeMatrix (unwrapSTMFractional m)
  unsafeFill n m a = stToPrim $ STMat n m <$> LA.newMatrix a n m
  unsafeFromRow r = stToPrim $ STMat 1 (SV.length r) <$> LA.unsafeThawMatrix (LA.asRow r)
  unsafeFromRows rs = stToPrim $ do
    let m = LA.fromRows (DC.coerce rs :: [LA.Vector a])
    STMat (LA.rows m) (LA.cols m) . DC.coerce <$> LA.unsafeThawMatrix m
  unsafeFromColumn c = stToPrim $ STMat (SV.length c) 1 <$> LA.unsafeThawMatrix (LA.asColumn c)
  unsafeFromColumns rs = stToPrim $ do
    let m = LA.fromColumns (DC.coerce rs :: [LA.Vector a])
    STMat (LA.rows m) (LA.cols m) . DC.coerce <$> LA.unsafeThawMatrix m
  unsafeCopy (STMat r c m)  (STMat _ _ m') = stToPrim $ do
    m0 <- LA.unsafeFreezeMatrix m'
    LA.setMatrix (DC.coerce m :: LA.STMatrix w a) r c (DC.coerce m0 :: LA.Matrix a)
  unsafeRead (STMat _ _ m) i j = stToPrim $ LA.readMatrix m i j
  unsafeWrite (STMat _ _ m) i j x = stToPrim $ LA.writeMatrix m i j x
  basicSet (STMat r c m) x = stToPrim $ LA.modifyMatrix m r c $ const x
  basicUnsafeSwapRows (STMat _ _ m) i j =
    stToPrim $ LA.rowOper (LA.SWAP i j LA.AllCols) $ castSTMatrixFractional m
  unsafeScaleRow (STMat _ _ m) i c =
    stToPrim $ LA.rowOper (LA.SCAL (unwrapFractional c) (LA.Row i) LA.AllCols) $ castSTMatrixFractional m
  basicUnsafeIMapRowM (STMat _ c m) i f = forM_ [0..c-1] $ \j ->
    stToPrim . LA.writeMatrix m i j =<< f j =<< stToPrim (LA.readMatrix m i j)
  toRows (STMat _ _ m) =
    stToPrim $ DC.coerce . LA.toRows <$> LA.unsafeFreezeMatrix (castSTMatrixFractional m)
  toColumns (STMat _ _ m) =
    stToPrim $ DC.coerce . LA.toColumns <$> LA.unsafeFreezeMatrix (castSTMatrixFractional m)

unwrapSTMIntegral :: STMatrix s (WrapIntegral a) -> LA.STMatrix s a
unwrapSTMIntegral = DC.coerce . unwrapSTM
{-# INLINE unwrapSTMIntegral #-}

castSTMatrixIntegral :: LA.STMatrix s (WrapIntegral a) -> LA.STMatrix s a
castSTMatrixIntegral = DC.coerce
{-# INLINE castSTMatrixIntegral #-}

instance (Integral a, LA.Container LA.Matrix a) => MMatrix STMatrix (WrapIntegral a) where
  basicUnsafeNew n m = stToPrim $ STMat n m <$> LA.newMatrix 0 n m
  basicInitialise _ = return ()
  basicRowCount = stmRows
  basicColumnCount = stmCols
  unsafeGetRow i m =
    stToPrim $ DC.coerce . getRowLA i <$> LA.unsafeFreezeMatrix (unwrapSTMIntegral m)
  unsafeGetColumn i m = stToPrim $ DC.coerce . getColLA i <$> LA.unsafeFreezeMatrix (unwrapSTMIntegral m)
  unsafeFill n m a = stToPrim $ STMat n m <$> LA.newMatrix a n m
  unsafeFromRow r = stToPrim $ STMat 1 (SV.length r) <$> LA.unsafeThawMatrix (LA.asRow r)
  unsafeFromRows rs = stToPrim $ do
    let m = LA.fromRows (DC.coerce rs :: [LA.Vector a])
    STMat (LA.rows m) (LA.cols m) . DC.coerce <$> LA.unsafeThawMatrix m
  unsafeFromColumn c = stToPrim $ STMat (SV.length c) 1 <$> LA.unsafeThawMatrix (LA.asColumn c)
  unsafeFromColumns rs = stToPrim $ do
    let m = LA.fromColumns (DC.coerce rs :: [LA.Vector a])
    STMat (LA.rows m) (LA.cols m) . DC.coerce <$> LA.unsafeThawMatrix m
  unsafeCopy (STMat r c m)  (STMat _ _ m') = stToPrim $ do
    m0 <- LA.unsafeFreezeMatrix m'
    LA.setMatrix (DC.coerce m :: LA.STMatrix w a) r c (DC.coerce m0 :: LA.Matrix a)
  unsafeRead (STMat _ _ m) i j = stToPrim $ LA.readMatrix m i j
  unsafeWrite (STMat _ _ m) i j x = stToPrim $ LA.writeMatrix m i j x
  basicSet (STMat r c m) x = stToPrim $ LA.modifyMatrix m r c $ const x
  basicUnsafeSwapRows (STMat _ _ m) i j =
    stToPrim $ LA.rowOper (LA.SWAP i j LA.AllCols) $ castSTMatrixIntegral m
  unsafeScaleRow (STMat _ _ m) i c =
    stToPrim $ LA.rowOper (LA.SCAL (unwrapIntegral c) (LA.Row i) LA.AllCols) $ castSTMatrixIntegral m
  basicUnsafeIMapRowM (STMat _ c m) i f = forM_ [0..c-1] $ \j ->
    stToPrim . LA.writeMatrix m i j =<< f j =<< stToPrim (LA.readMatrix m i j)
  toRows (STMat _ _ m) =
    stToPrim $ DC.coerce . LA.toRows <$> LA.unsafeFreezeMatrix (castSTMatrixIntegral m)
  toColumns (STMat _ _ m) =
    stToPrim $ DC.coerce . LA.toColumns <$> LA.unsafeFreezeMatrix (castSTMatrixIntegral m)

instance {-# OVERLAPPING #-} (LA.Container LA.Matrix a, Show a) => Show (LA.Matrix (WrapNum a)) where
  showsPrec d = showsPrec d . (DC.coerce :: LA.Matrix (WrapNum a) -> LA.Matrix a)

instance {-# OVERLAPPING #-} (LA.Container LA.Matrix a, Show a) => Show (LA.Matrix (WrapIntegral a)) where
  showsPrec d = showsPrec d . (DC.coerce :: LA.Matrix (WrapIntegral a) -> LA.Matrix a)

instance {-# OVERLAPPING #-} (LA.Container LA.Matrix a, Show a) => Show (LA.Matrix (WrapFractional a)) where
  showsPrec d = showsPrec d . (DC.coerce :: LA.Matrix (WrapFractional a) -> LA.Matrix a)


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE BangPatterns #-}

module Numeric.MCMC.Langevin (
          MarkovChain(..), Options(..)
        , runChain
        ) where

import Control.Monad
import Control.Monad.Trans
import Control.Monad.Reader
import Control.Monad.Primitive
import Control.Arrow
import System.Random.MWC
import System.Random.MWC.Distributions
import Data.List
import Statistics.Distribution
import Statistics.Distribution.Normal hiding (standard)
import GHC.Float

-- | State of the Markov chain.  Current parameter values are held in 'theta',
--   while accepts counts the number of proposals accepted.
data MarkovChain = MarkovChain { theta   :: [Double]
                               , accepts :: {-# UNPACK #-} !Int }

-- | Options for the chain.  The target (expected to be a log density), its
--   gradient, and a step size tuning parameter.
data Options = Options { _target  :: [Double] -> Double
                       , _gTarget :: [Double] -> [Double]
                       , _eps     :: {-# UNPACK #-} !Double }

-- | A result with this type has a view of the chain options.
type ViewsOptions = ReaderT Options

-- | Display the current state.
instance Show MarkovChain where
    show config = filter (`notElem` "[]") $ show (map double2Float (theta config))

-- | Density function for an isotropic Gaussian.  The (identity) covariance
--   matrix is multiplied by the scalar 'sig'.
isoGauss :: [Double] -> [Double] -> Double -> Double
isoGauss xs mu sig = foldl1' (*) (zipWith density nds xs)
    where nds = map (`normalDistr` sig) mu
{-# INLINE isoGauss #-}

-- | Mean function for the discretized Langevin diffusion.
localMean :: Monad m
          => [Double]                -- Current state
          -> ViewsOptions m [Double] -- Localized mean
localMean t = do
    Options _ gTarget e <- ask
    return $! zipWith (+) t (map (* (0.5 * e^(2 :: Int))) (gTarget t))
{-# INLINE localMean #-}

-- | Perturb the state, creating a new proposal.
perturb :: PrimMonad m
        => [Double]                   -- Current state
        -> Gen (PrimState m)          -- MWC PRNG
        -> ViewsOptions m [Double]    -- Resulting perturbation.
perturb t g = do
    Options _ _ e <- ask
    zs    <- replicateM (length t) (lift $ standard g)
    t0    <- localMean t
    let perturbedState = zipWith (+) t0 t1
        t1 = map (* e) zs
    return $! perturbedState
{-# INLINE perturb #-}

-- | Perform a Metropolis accept/reject step.
metropolisStep :: PrimMonad m
               => MarkovChain                -- Current state
               -> Gen (PrimState m)          -- MWC PRNG
               -> ViewsOptions m MarkovChain -- New state
metropolisStep state g = do
    Options target _ e <- ask
    let (t0, nacc) = (theta &&& accepts) state
    zc           <- lift $ uniformR (0, 1) g
    proposal     <- perturb t0 g
    t0Mean       <- localMean t0
    proposalMean <- localMean proposal
    let mc = if   zc < acceptProb
             then (proposal, 1)
             else (t0,       0)

        acceptProb = if isNaN val then 0 else val where val = arRatio

        arRatio = exp . min 0 $
            target proposal + log (isoGauss proposal t0Mean (e^(2 :: Int)))
          - target t0       - log (isoGauss t0 proposalMean (e^(2 :: Int)))

    return $! MarkovChain (fst mc) (nacc + snd mc)
{-# INLINE metropolisStep #-}

-- | Diffuse through states.
runChain :: Options         -- Options of the Markov chain.
         -> Int             -- Number of epochs to iterate the chain.
         -> Int             -- Print every nth iteration
         -> MarkovChain     -- Initial state of the Markov chain.
         -> Gen RealWorld   -- MWC PRNG
         -> IO MarkovChain  -- End state of the Markov chain, wrapped in IO.
runChain = go
  where go o n t !c g | n == 0 = return c
                      | n `rem` t /= 0 = do
                            r <- runReaderT (metropolisStep c g) o
                            go o (n - 1) t r g
                      | otherwise = do
                            r <- runReaderT (metropolisStep c g) o
                            print r
                            go o (n - 1) t r g
{-# INLINE runChain #-}


{-# LANGUAGE PatternGuards #-}
-- |
-- Module    : Statistics.Matrix
-- Copyright : 2011 Aleksey Khudyakov, 2014 Bryan O'Sullivan
-- License   : BSD3
--
-- Basic matrix operations.
--
-- There isn't a widely used matrix package for Haskell yet, so
-- we implement the necessary minimum here.

module Statistics.Matrix
    ( -- * Data types
      Matrix(..)
    , Vector
      -- * Conversion from/to lists/vectors
    , fromVector
    , dimension
    -- , center
    , multiplyV
    , transpose
    , norm
    , column
    -- , row
    , for
    , unsafeIndex
    ) where

import Prelude hiding (exponent, map, sum)
import qualified Data.Vector.Unboxed as U

import Statistics.Function (for, square)
import Statistics.Matrix.Types
import Statistics.Sample.Internal (sum)

----------------------------------------------------------------
-- Conversion to/from vectors/lists
----------------------------------------------------------------

-- | Convert from a row-major vector.
fromVector :: Int               -- ^ Number of rows.
           -> Int               -- ^ Number of columns.
           -> U.Vector Double   -- ^ Flat list of values, in row-major order.
           -> Matrix
fromVector r c v
  | r*c /= len = error "input size mismatch"
  | otherwise  = Matrix r c 0 v
  where len    = U.length v

----------------------------------------------------------------
-- Other
----------------------------------------------------------------

-- | Return the dimensions of this matrix, as a (row,column) pair.
dimension :: Matrix -> (Int, Int)
dimension (Matrix r c _ _) = (r, c)

-- | Matrix-vector multiplication.
multiplyV :: Matrix -> Vector -> Vector
multiplyV m v
  | cols m == c = U.generate (rows m) (sum . U.zipWith (*) v . row m)
  | otherwise   = error $ "matrix/vector unconformable " ++ show (cols m,c)
  where c = U.length v

-- | Calculate the Euclidean norm of a vector.
norm :: Vector -> Double
norm = sqrt . sum . U.map square

-- | Return the given column.
column :: Matrix -> Int -> Vector
column (Matrix r c _ v) i = U.backpermute v $ U.enumFromStepN i c r
{-# INLINE column #-}

-- | Return the given row.
row :: Matrix -> Int -> Vector
row (Matrix _ c _ v) i = U.slice (c*i) c v

unsafeIndex :: Matrix
            -> Int              -- ^ Row.
            -> Int              -- ^ Column.
            -> Double
unsafeIndex = unsafeBounds U.unsafeIndex

-- | Given row and column numbers, calculate the offset into the flat
-- row-major vector, without checking.
unsafeBounds :: (Vector -> Int -> r) -> Matrix -> Int -> Int -> r
unsafeBounds k (Matrix _ cs _ v) r c = k v $! r * cs + c
{-# INLINE unsafeBounds #-}

transpose :: Matrix -> Matrix
transpose m@(Matrix r0 c0 e _) = Matrix c0 r0 e . U.generate (r0*c0) $ \i ->
  let (r,c) = i `quotRem` r0
  in unsafeIndex m c r


{-# OPTIONS_GHC -Wall            #-}

{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE OverloadedLists     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

module Heat2D (main, a1, a2) where

import Data.Number.Symbolic
import Data.Proxy

import qualified Naperian as N
import qualified GHC.TypeLits as M
import           Data.Functor
import           System.IO

import           Numeric.Sundials
import           Numeric.LinearAlgebra

kx, ky :: Floating a => a
kx = 0.5
ky = 0.75

-- x mesh spacing
-- y mesh spacing
dx :: Floating a => a
dx = 1 / (fromIntegral nx - 1)

dy :: Floating a => a
dy = 1 / (fromIntegral ny - 1)

c1, c2 :: Floating a => a
c1 = kx/dx/dx
c2 = ky/dy/dy

preA2 :: forall b m n . (M.KnownNat m, M.KnownNat n, Num b) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] b
preA2 = N.Prism $ N.Prism $ N.Prism $ N.Prism $ N.Scalar $
      N.viota @m <&> (\(N.Fin x) ->
      N.viota @n <&> (\(N.Fin w) ->
      N.viota @m <&> (\(N.Fin v) ->
      N.viota @n <&> (\(N.Fin u) ->
      (f m n x w v u)))))
        where
          m = fromIntegral $ M.natVal (undefined :: Proxy m)
          n = fromIntegral $ M.natVal (undefined :: Proxy n)
          f p q i j k l | i == 0               = 0
                        | j == 0               = 0
                        | i == p - 1           = 0
                        | j == q - 1           = 0
                        | k == i - 1 && l == j = 1
                        | k == i     && l == j = -2
                        | k == i + 1 && l == j = 1
                        | otherwise            = 0

a2Num :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] a
a2Num = N.binary (*) (N.Scalar c2) preA2

a2 :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a, Eq a) =>
          N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] (Sym a)
a2 = N.binary (*) (N.Scalar $ var "a") preA2

preA1 :: forall b m n . (M.KnownNat m, M.KnownNat n, Num b) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] b
preA1 = N.Prism $ N.Prism $ N.Prism $ N.Prism $ N.Scalar $
      N.viota @m <&> (\(N.Fin x) ->
      N.viota @n <&> (\(N.Fin w) ->
      N.viota @m <&> (\(N.Fin v) ->
      N.viota @n <&> (\(N.Fin u) ->
      (f m n x w v u)))))
        where
          m = fromIntegral $ M.natVal (undefined :: Proxy m)
          n = fromIntegral $ M.natVal (undefined :: Proxy n)
          f :: Int -> Int -> Int -> Int -> Int -> Int -> b
          f p q i j k l | i == 0                   = 0
                        | j == 0                   = 0
                        | i == p - 1               = 0
                        | j == q - 1               = 0
                        | k == i     && l == j - 1 = 1
                        | k == i     && l == j     = -2
                        | k == i     && l == j + 1 = 1
                        | otherwise                = 0

a1Num :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] a
a1Num = N.binary (*) (N.Scalar c1) preA1

a1 :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a, Eq a) =>
           N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] (Sym a)
a1 = N.binary (*) (N.Scalar $ var "b") preA1

h :: forall m n a . (M.KnownNat m, M.KnownNat n, Floating a) =>
                     N.Hyper '[N.Vector m, N.Vector n] a
h = N.Prism $ N.Prism $ N.Scalar $
     N.viota @n <&> (\(N.Fin x) ->
     N.viota @m <&> (\(N.Fin w) ->
     sin (pi * (fromIntegral w) * dx)
             * sin (2 * pi * (fromIntegral x) * dy)))

-- spatial mesh size
nx, ny :: Int
nx = 30
ny = 60

bigA :: Matrix Double
bigA = fromLists $
       fmap (N.elements . N.Prism . N.Prism . N.Scalar) $
       N.elements $ N.crystal $ N.crystal $
       N.binary (+) (a2Num @Double @60 @30) (a1Num @Double @60 @30)

b :: Vector Double
b = fromList $ N.elements (h @30 @60 @Double)

t0, tf :: Double
t0 = 0.0
tf = 0.3

bigNt :: Int
bigNt = 20

dTout :: Double
dTout = (tf - t0) / (fromIntegral bigNt)

ts :: [Double]
ts = map (dTout *) $ map fromIntegral [1..bigNt]

sol :: Matrix Double
-- sol = odeSolveV SDIRK_5_3_4' Nothing 1.0e-5 1.0e-10 (const bigU') (assoc (nx * ny) 0.0 [] :: Vector Double) (fromList $ ts)
--   where
--     bigU' bigU = bigA #> bigU + b
sol = undefined

main :: IO ()
main = do
  h1 <- openFile "heat2e.000.txt" WriteMode
  mapM_ (hPutStrLn h1) $ map (concatMap (' ':)) $ map (map show) $ toLists sol
  hClose h1
  mapM_ (\i -> putStrLn $ show $ sqrt $ (sol!i) <.> (sol!i) / (fromIntegral nx) / (fromIntegral ny)) ([0 .. length ts - 1] :: [Int])