{-# LANGUAGE BangPatterns, FlexibleContexts, UnboxedTuples #-}
-- |
-- Module    : Statistics.Sample.KernelDensity
-- Copyright : (c) 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Kernel density estimation.  This module provides a fast, robust,
-- non-parametric way to estimate the probability density function of
-- a sample.
--
-- This estimator does not use the commonly employed \"Gaussian rule
-- of thumb\".  As a result, it outperforms many plug-in methods on
-- multimodal samples with widely separated modes.

module Statistics.Sample.KernelDensity
    (
    -- * Estimation functions
      kde
    -- , kde_
    -- * References
    -- $references
    ) where

import Numeric.MathFunctions.Constants (m_sqrt_2_pi)
import Prelude hiding (const, min, max, sum)
import Statistics.Function (minMax, nextHighestPowerOfTwo)
import Statistics.Math.RootFinding (fromRoot, ridders)
import Statistics.Sample.Histogram (histogram_)
import Statistics.Sample.Internal (sum)
import Statistics.Transform (CD, dct, idct)
import qualified Data.Vector.Generic  as G
import qualified Data.Vector.Unboxed  as U
import qualified Data.Vector          as V

-- | Gaussian kernel density estimator for one-dimensional data, using
-- the method of Botev et al.
--
-- The result is a pair of vectors, containing:
--
-- * The coordinates of each mesh point.  The mesh interval is chosen
--   to be 20% larger than the range of the sample.  (To specify the
--   mesh interval, use 'kde_'.)
--
-- * Density estimates at each mesh point.
kde :: (G.Vector v CD, G.Vector v Double, G.Vector v Int)
    => Int
    -- ^ The number of mesh points to use in the uniform discretization
    -- of the interval @(min,max)@.  If this value is not a power of
    -- two, then it is rounded up to the next power of two.
    -> v Double -> (v Double, v Double)
kde n0 xs = kde_ n0 (lo - range / 10) (hi + range / 10) xs
  where
    (lo,hi) = minMax xs
    range   | G.length xs <= 1 = 1       -- Unreasonable guess
            | lo == hi         = 1       -- All elements are equal
            | otherwise        = hi - lo
{-# INLINABLE  kde #-}
{-# SPECIAlIZE kde :: Int -> U.Vector Double -> (U.Vector Double, U.Vector Double) #-}
{-# SPECIAlIZE kde :: Int -> V.Vector Double -> (V.Vector Double, V.Vector Double) #-}

-- | Gaussian kernel density estimator for one-dimensional data, using
-- the method of Botev et al.
--
-- The result is a pair of vectors, containing:
--
-- * The coordinates of each mesh point.
--
-- * Density estimates at each mesh point.
kde_ :: (G.Vector v CD, G.Vector v Double, G.Vector v Int)
     => Int
     -- ^ The number of mesh points to use in the uniform discretization
     -- of the interval @(min,max)@.  If this value is not a power of
     -- two, then it is rounded up to the next power of two.
     -> Double
     -- ^ Lower bound (@min@) of the mesh range.
     -> Double
     -- ^ Upper bound (@max@) of the mesh range.
     -> v Double
     -> (v Double, v Double)
kde_ n0 min max xs
  | G.null xs = error "Statistics.KernelDensity.kde: empty sample"
  | n0 <= 1   = error "Statistics.KernelDensity.kde: invalid number of points"
  | otherwise = (mesh, density)
  where
    mesh = G.generate ni $ \z -> min + (d * fromIntegral z)
        where d = r / (n-1)
    density = G.map (/(2 * r)) . idct $ G.zipWith f a (G.enumFromTo 0 (n-1))
      where f b z = b * exp (sqr z * sqr pi * t_star * (-0.5))
    !n  = fromIntegral ni
    !ni = nextHighestPowerOfTwo n0
    !r  = max - min
    a   = dct . G.map (/ sum h) $ h
        where h = G.map (/ len) $ histogram_ ni min max xs
    !len    = fromIntegral (G.length xs)
    !t_star = fromRoot (0.28 * len ** (-0.4)) . ridders 1e-14 (0,0.1) $ \x ->
              x - (len * (2 * sqrt pi) * go 6 (f 7 x)) ** (-0.4)
      where
        f q t = 2 * pi ** (q*2) * sum (G.zipWith g iv a2v)
          where g i a2 = i ** q * a2 * exp ((-i) * sqr pi * t)
                a2v = G.map (sqr . (*0.5)) $ G.tail a
                iv = G.map sqr $ G.enumFromTo 1 (n-1)
        go s !h | s == 1    = h
                | otherwise = go (s-1) (f s time)
          where time  = (2 * const * k0 / len / h) ** (2 / (3 + 2 * s))
                const = (1 + 0.5 ** (s+0.5)) / 3
                k0    = U.product (G.enumFromThenTo 1 3 (2*s-1)) / m_sqrt_2_pi
    sqr x = x * x
{-# INLINABLE  kde_ #-}
{-# SPECIAlIZE kde_ :: Int -> Double -> Double -> U.Vector Double -> (U.Vector Double, U.Vector Double) #-}
{-# SPECIAlIZE kde_ :: Int -> Double -> Double -> V.Vector Double -> (V.Vector Double, V.Vector Double) #-}

-- $references
--
-- Botev. Z.I., Grotowski J.F., Kroese D.P. (2010). Kernel density
-- estimation via diffusion. /Annals of Statistics/
-- 38(5):2916–2957. <http://arxiv.org/pdf/1011.2602>


{-# LANGUAGE FlexibleContexts #-}

module Main (main) where

import Control.Monad (mzero)
import Control.Monad.IO.Class (liftIO)
import Data.Complex (Complex)
import Data.Foldable (toList)
import Data.Typeable (Typeable)
import System.Console.GetOpt
import Text.PrettyPrint.Mainland
import Text.PrettyPrint.Mainland.Class

import Spiral
import Spiral.Backend.C
import Spiral.Config
import Spiral.Driver
import Spiral.Exp
import Spiral.FFT.CooleyTukey
import Spiral.Monad
import Spiral.OpCount
import Spiral.Program
import Spiral.SPL
import Spiral.SPL.Run
import Spiral.Search
import Spiral.Search.FFTBreakdowns
import Spiral.Util.Uniq

main :: IO ()
main = defaultMainWith' options mempty $ \fs args -> do
    useComplexType <- asksConfig $ testDynFlag UseComplex
    n <- case args of
           [s] -> return (read s)
           _   -> return 4
    f <- formula fs n
    pprint f
    if useComplexType
      then toProgram ("hspiral_dft_" ++ show n) f >>= go
      else toProgram ("hspiral_dft_" ++ show n) (Re f) >>= go
  where
    go :: (Typed a, Num (Exp a)) => Program a -> Spiral ()
    go prog = do
      pprint prog
      ops <- countProgramOps prog
      resetUnique
      defs <- evalCg $ cgProgram prog
      outp <- asksConfig output
      case outp of
        Nothing -> return ()
        Just{}  -> writeOutput (toList defs)
      liftIO $ putDocLn $
          text "Multiplications:" <+> ppr (mulOps ops) </>
          text "      Additions:" <+> ppr (addOps ops) </>
          text "          Total:" <+> ppr (allOps ops)

-- The SPL formula for which we generate code and count operations.
formula :: MonadSpiral m => [Flag] -> Int -> m (SPL (Exp (Complex Double)))
formula fs n =
  case fs of
    [Dif]                -> return $ dif n
    [Dit]                -> return $ dit n
    [SplitRadix]         -> runSearch () splitRadixSearch (DFT n)
    [ConjPairSplitRadix] -> runSearch () conjSplitRadixSearch (DFT n)
    [ImpSplitRadix]      -> runSearch () impSplitRadixSearch (DFT n)
    _                    -> fail "Must specify exactly on of --dif, --dit, --split-radix, conj-split-radix, or --improved-split-radix"
  where
    splitRadixSearch :: (Typeable a, Typed a, MonadSpiral m)
                     => SPL (Exp a)
                     -> S s m (SPL (Exp a))
    splitRadixSearch (F n w) = splitRadixBreakdown n w
    splitRadixSearch _       = mzero

    conjSplitRadixSearch :: (Typeable a, Typed a, MonadSpiral m)
                         => SPL (Exp a)
                         -> S s m (SPL (Exp a))
    conjSplitRadixSearch (F n w) = conjPairSplitRadixBreakdown n w
    conjSplitRadixSearch _       = mzero

    impSplitRadixSearch :: (Typeable a, Typed a, MonadSpiral m)
                        => SPL (Exp a)
                        -> S s m (SPL (Exp a))
    impSplitRadixSearch (F n w) = improvedSplitRadixBreakdown n w
    impSplitRadixSearch _       = mzero

data Flag = Dif
          | Dit
          | SplitRadix
          | ConjPairSplitRadix
          | ImpSplitRadix
  deriving (Eq, Ord, Show)

options :: [OptDescr Flag]
options =
    [ Option [] ["dif"] (NoArg Dif)                             "Use DIF"
    , Option [] ["dit"] (NoArg Dit)                             "Use DIT"
    , Option [] ["split-radix"] (NoArg SplitRadix)              "Use split radix"
    , Option [] ["conj-split-radix"] (NoArg ConjPairSplitRadix) "Use conjugate pair split radix"
    , Option [] ["imp-split-radix"] (NoArg ImpSplitRadix)       "Use improved split radix"
    ]


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE OverloadedStrings #-}

module Model where

import           Control.Monad.Bayes.Class
import           Control.Monad.Bayes.Sampler
import qualified Data.Text.IO as TIO
import           Control.Monad
import           Control.Monad.Bayes.Inference.PMMH
import           Control.Monad.State
import Control.Monad.Bayes.Weighted
import           Numeric.Log
import           Text.Megaparsec

import           DataParser
import           Statistics
import           Utils

data Params = Params {
    rho :: Double, -- ^ Rate of detection
    beta :: Double, -- ^ Mean contact rate between susceptible and infected people
    gamma :: Double -- ^ Mean recovery rate
} deriving Show

data FixedParams = FixedParams {
    numPop :: Int,
    timeSlices :: Int
}

data LatentState = LatentState {
    sus :: Int, -- ^ Number of people susceptible to infection
    inf :: Int, -- ^ Number of people currently infected
    recov :: Int -- ^ Number of people recovered from infection
}

{-
observation model: Poisson rho * I
-}

-- | Model for how we observe the number of infected people
observationModel :: MonadSample m => Params -> LatentState -> m InfectionCount
observationModel (Params rho _ _) (LatentState _ inf _) =
    poisson (rho * fromIntegral inf)

{-
dN_SI <-
-}
-- | Transition the model a single time slice
transitionModelSingleStep
    :: MonadSample m => FixedParams -> Params -> LatentState -> m LatentState
transitionModelSingleStep (FixedParams numPop timeSlices) (Params rho beta gamma) (LatentState sus inf recov)
    = do
        let dt = 1 / fromIntegral timeSlices
        dN_SI <- binomial
            sus
            (1 - exp ((-beta * dt * (fromIntegral inf)) / (fromIntegral numPop))
            )
        dN_IR <- binomial
            inf
            (1 - exp (-gamma * dt)
            )
        let sus'   = sus - dN_SI
        let inf'   = inf + dN_SI - dN_IR
        let recov' = recov + dN_IR
        return (LatentState sus' inf' recov')

-- | Transition the model for a full step
transitionModel :: MonadSample m => FixedParams -> Params -> LatentState -> m LatentState
transitionModel fixedParams params =
    transitionModelSingleStep fixedParams params

-- | Simulate a single step, returning the new latent state and appending the observed infection count to the state.
simulateStep
    :: (MonadSample m, MonadState [InfectionCount] m)
    => FixedParams
    -> Params
    -> LatentState
    -> m LatentState
simulateStep fixedParams params latent = do
    latent'        <- transitionModel fixedParams params latent
    infectionCount <- observationModel params latent'
    modify (++ [infectionCount])
    return latent'

-- | Simulate nsteps steps of an epidemic with the specified initial state and parameters
simulateEpidemic
    :: MonadSample m => LatentState -> FixedParams -> Params -> Int -> m Epidemic
simulateEpidemic initialState fixedParams params nsteps =
    Epidemic
        <$> execStateT
                (repeatFunction nsteps (simulateStep fixedParams params) initialState)
                []

-- | Execute a single simulation of an epidemic
generateSingleEpidemic :: LatentState -> FixedParams -> Params -> Int -> IO Epidemic
generateSingleEpidemic initialState fixedParams params nsteps =
    sampleIO $ simulateEpidemic initialState fixedParams params nsteps

generateEpidemics :: LatentState -> FixedParams -> Params -> Int -> Int -> IO [Epidemic]
generateEpidemics initialState fixedParams params nsteps nepidemics = replicateM nepidemics (generateSingleEpidemic initialState fixedParams params nsteps)

fixedParams :: FixedParams
fixedParams = FixedParams 763 1

params :: Params
params = Params 0.9 2.0 0.6

initialState :: LatentState
initialState = LatentState 762 1 0

scoreEpidemicToDatum
    :: MonadInfer m
    => FixedParams
    -> Epidemic
    -> LatentState
    -> Params
    ->  m Params
scoreEpidemicToDatum fixedParams dat  x params = do
    let obs lambda y = score (poissonPdf lambda y)
        go [] x = return x
        go (y:ys) x = do
            x' <- transitionModel fixedParams params x
            obs ((fromIntegral $ inf x') * (rho params)) y
            return x'
    (go $! (unwrapEpidemic dat)) $! x
    return params

unwrapEpidemic :: Epidemic -> [Int]
unwrapEpidemic (Epidemic xs) = xs

{-
ddprior <- function(params) {
    dgamma(params[[1]], 0.3, 10, log = TRUE) +
    dgamma(params[[3]], 1, 8, log=TRUE) +
    dbeta(params[[2]], 2,7, log=TRUE)
}
-}
paramsPrior :: MonadSample m => m Params
paramsPrior = do
    pBeta <- Control.Monad.Bayes.Class.gamma 2 1
    pRho <- Control.Monad.Bayes.Class.beta 2.0 7.0
    pGamma <- Control.Monad.Bayes.Class.gamma 1.0  (1 / 8.0)
    return (Params pRho pBeta pGamma )

testInferenceEpidemic :: Int -> Int -> IO [[(Params, Numeric.Log.Log Double)]]
testInferenceEpidemic nsteps nparticles = do
    ys <- parseFromFile epidemicParser "data/datafile"
    case ys of (Left _) -> error "naughty"
               (Right dat) -> sampleIO $ do
                       pmmhRes <- prior $ pmmh nsteps 14 nparticles paramsPrior (scoreEpidemicToDatum fixedParams dat initialState)
                       return pmmhRes

parseFromFile p file = runParser p file <$> TIO.readFile file

extractParams :: (Params -> Double) -> [[(Params, Numeric.Log.Log Double)]] -> [Double]
extractParams project samples = (project . fst . head) <$> samples

--scoreEpidemicToData :: (MonadSample m, MonadState [InfectionCount] m => Epidemic ->  Params -> LatentState -> m Epidemic
--scoreEpidemicToData data params initialState = do
--    let obs lambda y = score (poissonPdf lambda y)
    -- simulate a new x from old x (x is a latent state)
    -- calculate new lambda
    -- score observation at time using lambda
    -- store x in the monad state
    -- repeat

generateSamples :: [Double]
generateSamples = sampleSTfixed $ replicateM 1000 (normal 3 1)


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE RankNTypes #-}
-----------------------------------------------------------------------------
--
-- Module      :  AI.Network.RMM.LSTM
-- Copyright   :  (c) JP Moresmau
-- License     :  BSD3
--
-- Maintainer  :  JP Moresmau <jp@moresmau.fr>
-- Stability   :  experimental
-- Portability :
--
-- | Long Short Term Memory networks
--
-----------------------------------------------------------------------------

module AI.Network.RNN.LSTM where

import Control.DeepSeq

import Data.List

import Numeric.LinearAlgebra.HMatrix

import AI.Network.RNN.Types
import AI.Network.RNN.Util
import AI.Network.RNN.Expr
import Debug.Trace
import qualified Data.IntMap as I

import Numeric.AD

-- | The LSTM data type, with the different weights and states
data LSTMNetwork = LSTMNetwork
    { lstmSize :: !Int
    , lstmWeightsW :: !(Matrix Double)
    , lstmWeightsU :: !(Matrix Double)
    , lstmBias :: !(Vector Double)
    , lstmState :: !(Vector Double)
    , lstmOutput :: !(Vector Double)
    } deriving (Show,Read,Eq)

--instance Show LSTMNetwork where
--     show s = let
--        a= show $ (rnnsize s,toVector s)
--        in trace (show $ length a) a
--
--instance Read LSTMNetwork where
--    readsPrec d s =
--        let [(ok,r)] = readsPrec d s
--            dat = uncurry fromVector ok
--        in [(dat,r)]

-- | Force evaluation instance
instance NFData LSTMNetwork where
    rnf LSTMNetwork{..} = rnf (lstmSize,lstmWeightsW,lstmWeightsU,lstmBias,lstmState,lstmOutput)

-- | Implement the network evaluation and conversions functions
instance RNNEval LSTMNetwork where
    type Size LSTMNetwork = Int
    evalStep n@LSTMNetwork{..} is =
        let z = (lstmWeightsW #> is) + (lstmWeightsU #> lstmOutput) + lstmBias
            [i,f,c1,o] = takesV (replicate 4 lstmSize) z
            c2 = cmap tanh c1
            ns = (c2 * cmap sigmoid i) + (cmap sigmoid f * lstmState)
            no = cmap sigmoid o * cmap tanh ns
        in force (n{lstmState=ns,lstmOutput=no},no)
    fromVector sz vs =
        let
            msize = sz * sz
            [v1,v2,v3,v4,v5] = takesV [msize * 4,msize * 4, sz * 4,sz,sz] vs
            m1 = reshape sz v1
            m2 = reshape sz v2
        in LSTMNetwork sz m1 m2 v3 v4 v5
    toVector LSTMNetwork{..} = force $ vjoin
        [ flatten lstmWeightsW, flatten lstmWeightsU, lstmBias, lstmState, lstmOutput]
    rnnsize = lstmSize
    fullSize = lstmFullSize .lstmSize

-- | Full size of a network
lstmFullSize :: FullSize LSTMNetwork
lstmFullSize sz = (sz * sz) * 8 + sz* 4 + sz + sz

data LSTMIO = LSTMIO
    { lioInput :: Matrix Double
    , lioLstms :: [LSTMNetwork]
    , lioOutput :: Matrix Double
    , lioSer   :: Maybe (Vector Double)
    } deriving (Eq)

instance NFData LSTMIO where
  rnf LSTMIO{..} = rnf (lioInput,lioLstms,lioOutput)

instance RNNEval LSTMIO  where
    type Size LSTMIO = (Int,Int,Int,Int)
    evalStep io@LSTMIO{..} is =
        let inps = cmap sigmoid $ lioInput #> (vjoin [is,scalar 1])
            (lstm2,outs) = evalStep lioLstms inps
        in (io{lioLstms=lstm2,lioSer=Nothing},lioOutput #> outs)
    fromVector (is,lnnSize,lnnNumber,os) vs =
        let [v1,v2,v3] = takesV [(is+1)* lnnSize ,lnnNumber * (lstmFullSize lnnSize),os * lnnSize ] vs
            m1 = reshape (is+1) v1
            m2 = reshape lnnSize v3
        in LSTMIO m1 (fromVector (replicate lnnNumber lnnSize) v2) m2 $ Just vs
    toVector LSTMIO{..} = case lioSer of
        Just vs -> vs
        Nothing -> vjoin [flatten lioInput,toVector lioLstms,flatten lioOutput]
    rnnsize LSTMIO{..} = (cols lioInput - 1, rnnsize $ head lioLstms, length lioLstms,rows lioOutput)
    fullSize =lstmioFullSize . rnnsize
--        let (r1,c1)=size lioInput
--            (r2,c2)=size lioOutput
--        in (r1*c1) + (fullSize lioLstms) + (r2*c2)

instance Show LSTMIO where
    show l = show (rnnsize l,toVector l)

instance Read LSTMIO where
    readsPrec d s =
        let [(ok,r)] = readsPrec d s
            dat = uncurry fromVector ok
        in [(dat,r)]

lstmioFullSize :: FullSize LSTMIO
lstmioFullSize (is,lnnSize,lnnNumber,os) = (is+1) * lnnSize + lnnNumber * (lstmFullSize lnnSize) + os * lnnSize

data LSTMList = LSTMList
    { llSize :: (Int,Int,Int,Int)
    , llData :: [Double]
    } deriving (Show, Read, Eq)

instance NFData LSTMList where
  rnf LSTMList{..} = rnf (llSize,llData)

instance RNNEval LSTMList where
    type Size LSTMList = (Int,Int,Int,Int)
    rnnsize = llSize
    fullSize =lstmioFullSize . rnnsize
    toDList = llData
    fromDList sz = LSTMList sz
    toVector = fromList . toDList
    fromVector sz = fromDList sz . toList
    evalStep l@LSTMList{..} is =
        let lstmio :: LSTMIO = fromDList llSize llData
            (io2,out) = evalStep lstmio is
        in (l{llData=toDList io2},out)

-- | Implementation of the LSTM evaluation step without explicit matrices and vectors
-- just using lists, so we can use AD on it
lstmList :: (Num b,Floating b) => Int -> [b] -> [b] -> ([b],[b])
lstmList sz lstm is = let
    msize = sz * sz
    [mW,mU,vB,vS,vO] = takes [msize * 4,msize * 4, sz * 4,sz,sz] lstm
    z = zipWith3 (\a b c->a+b+c) (listMProd mW is) (listMProd mU vO) vB
    [i,f,c1,o] = takes (replicate 4 sz) z
    c2 = map tanh c1
    ns = zipWith (+) (zipWith (*) c2 (map sigmoid i)) (zipWith (*) (map sigmoid f) vS)
    no = zipWith (*) (map sigmoid o) (map tanh ns)
    in (mW++mU++vB++ns++no,no)

-- | Cost calculation using list representation for AD
cost' :: (Num b,Floating b,Fractional b) => Int -> [[b]] -> [[b]] -> [b] -> b
cost' sz is os lstm = let
    (_,res) = mapAccumL (lstmList sz) lstm is
--    in - (calcMeanList $ three (last res) (last os))
--    in - (calcMeanList $ concat $ zipWith (three) res os)
    in sum $ zipWith err os res
    where
      err :: (Num b,Floating b) => [b] -> [b] -> b
      err a b  = sum (zipWith (\c d -> (c- d)**2 ) a b)
--      one i o = zipWith (*) o (map log i)
--      two i o = zipWith (*) (map (\x->1 - x) o) (map (\x->1 - log x) i)
--      three i o= zipWith (+) (one i o) (two i o)

-- | Gradient descent learning
-- The third parameter is a call back function to monitor progress and stop the learning process if needed
learnGradientDescent :: (Monad m) => LSTMNetwork -> TrainData a -> (LSTMNetwork -> TrainData a -> Int -> m Bool) -> m LSTMNetwork
learnGradientDescent lstm td progressF =  go (toList $ toVector lstm) 0
    -- go2 gds 0
    where
      go2 (ls:lss) gen =  do
        let rnn::LSTMNetwork = fromVector (tdRecSize td) (fromList ls)
        cont <- progressF rnn td gen
        if cont
            then go2 lss (gen+1)
            else return rnn
      go ls gen = do
        let rnn::LSTMNetwork = fromVector (tdRecSize td) (fromList ls)
        cont <- progressF rnn td gen
        if cont
            then do
                let
                    gs= gf ls -- gradients using AD
                    ls2 = zipWith (\o g->o-g*0.1) ls gs
                go ls2 (gen+1)
            else return rnn
      lis = map toList (tdInputs td)
      los = map toList (tdOutputs td)
      gf = grad (cost' (tdRecSize td) (map (map auto) lis) (map (map auto) los))
      gds = gradientDescent (cost' (tdRecSize td) (map (map auto) lis) (map (map auto) los)) (toList $ toVector lstm)

-- | Gradient descent learning using symbolic differentiation
--   The goal was to calculate the derivative once and then just close and eval the result using the current data
--   However it is much slower than normal automatic differentiation
learnGradientDescentSym :: (Monad m) => LSTMNetwork -> TrainData a -> (LSTMNetwork -> TrainData a -> Int -> m Bool) -> m LSTMNetwork
learnGradientDescentSym lstm td progressF =  go (toList $ toVector lstm) 0
    where
      go ls gen = do
        let rnn::LSTMNetwork = fromVector (tdRecSize td) (fromList ls)
        cont <- progressF rnn td gen
        if cont
            then do
                let
                    i = I.fromList $ zip [0..] ls
                    cexpr = map (\g-> close g i) gf
                    gs = map eval cexpr
                    ls2 = zipWith (\o g->o-g*0.1) ls gs
                go ls2 (gen+1)
            else return rnn
      lis = map toList (tdInputs td)
      los = map toList (tdOutputs td)
      ls0 = toList $ toVector lstm
      gf = map fullSimplify $ grad (cost'
            (tdRecSize td)
            (map (map (\x -> autoEval (Lit $ Lit x) I.empty)) lis)
            (map (map (\x -> autoEval (Lit $ Lit x) I.empty)) los))
            (zipWith (\_ i->Var i) ls0 [0..])

-- | RMSProp learning, as far as I can make out
-- The third parameter is a call back function to monitor progress and stop the learning process if needed
learnRMSProp :: (Monad m) => LSTMNetwork -> TrainData a -> (LSTMNetwork -> TrainData a -> Int -> m Bool) -> m LSTMNetwork
learnRMSProp lstm td progressF = go ls0 (replicate myl 0) (replicate myl 0) (replicate myl 0) 0
    where
      go ls rgs rgs2 ugs gen = do
        let rnn::LSTMNetwork = fromVector (tdRecSize td) (fromList ls)
        cont <- progressF rnn td gen
        if cont
            then do
                let
                    gs= gf ls -- gradients using AD
                    rgup = zipWith (\rg g-> 0.95 * rg + 0.05 * g) rgs ls
                    rg2up = zipWith (\rg2 g-> 0.95 * rg2 + 0.05 * (g ** 2)) rgs2 ls
                    ugup = zipWith4 (\ud zg rg rg2 -> 0.9 * ud - 1e-4 * zg / sqrt(rg2 - rg ** 2 + 1e-4)) ugs gs rgup rg2up
                    ls2 = zipWith (+) ls ugup
                go (force ls2) (force rgup) (force rg2up) (force ugup) (gen+1)
            else return rnn
      lis = map toList (tdInputs td)
      los = map toList (tdOutputs td)
      gf = grad (cost' (tdRecSize td) (map (map auto) lis) (map (map auto) los))
      ls0 = toList $ toVector lstm
      myl= length ls0


{-# LANGUAGE CPP                        #-}
{-# LANGUAGE PolyKinds                  #-}
{-# LANGUAGE ConstraintKinds            #-}
{-# LANGUAGE DefaultSignatures          #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE RebindableSyntax           #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE TypeFamilies               #-}

module Data.Semimodule.Algebra (
  -- * Algebras
    Algebra(..)
  -- * Unital algebras
  , Unital(..)
  -- * Coalgebras
  , Coalgebra(..)
  -- * Unital coalgebras
  , Counital(..)
  -- * Bialgebras
  , Bialgebra
) where

import Control.Arrow
import Control.Applicative
import Control.Category (Category, (<<<), (>>>))
import Data.Bool
import Data.Complex
import Data.Finite (Finite, finites)
import Data.Fixed
import Data.Functor.Apply
import Data.Functor.Compose
import Data.Functor.Contravariant
import Data.Functor.Product
import Data.Functor.Rep
import Data.Functor.Rep
import Data.Int
import Data.Semiring
import Data.Semimodule
import Data.Sequence hiding (reverse,index)
import Data.Tuple (swap)
import Data.Word
import Foreign.C.Types (CFloat(..),CDouble(..))
import GHC.Real hiding (Fractional(..))
import GHC.TypeNats (KnownNat(..))
import Numeric.Natural
import Prelude (Ord, reverse)
import Prelude (fromInteger)
import Prelude hiding (Num(..), Fractional(..), negate, sum, product)
import qualified Control.Category as C
import qualified Control.Monad as M
import qualified Data.IntSet as IntSet
import qualified Data.Map as Map
import qualified Data.Sequence as Seq
import qualified Data.Set as Set

-------------------------------------------------------------------------------
-- Algebras
-------------------------------------------------------------------------------

-- | An < https://en.wikipedia.org/wiki/Algebra_over_a_field#Generalization:_algebra_over_a_ring algebra > over a semiring.
--
-- Note that the algebra < https://en.wikipedia.org/wiki/Non-associative_algebra needn't be associative >.
--
class Semiring a => Algebra a b where

    -- |
    --
    -- @
    -- 'joined' = 'Data.Semimodule.Free.over' 'Data.Semimodule.Free.diagonal' '.' 'uncurry'
    -- @
    --
    joined :: (b -> b -> a) -> b -> a

-- | A < https://en.wikipedia.org/wiki/Algebra_over_a_field#Unital_algebra unital algebra > over a semiring.
--
class Algebra a b => Unital a b where

    -- | Obtain a vector from the unit of a unital algebra.
    --
    -- @
    -- 'unital' = 'Data.Semimodule.Transform.over' 'initial' '.' 'const'
    -- @
    --
    unital :: a -> b -> a

-------------------------------------------------------------------------------
-- Coalgebras
-------------------------------------------------------------------------------

-- | A coalgebra over a semiring.
--
class Semiring a => Coalgebra a c where

    -- |
    --
    -- @
    -- 'cojoined' = 'curry' '.' 'Data.Semimodule.Free.over' 'Data.Semimodule.Free.codiagonal'
    -- @
    --
    cojoined :: (c -> a) -> c -> c -> a

-- | A counital coalgebra over a semiring.
--
class Coalgebra a c => Counital a c where

    -- |
    --
    -- @
    -- 'Data.Semimodule.Free.Cov' 'counital' = 'Data.Semimodule.Free.coover' 'Data.Semimodule.Free.coinitial' $ pure '()'
    -- @
    --
    counital :: (c -> a) -> a

-------------------------------------------------------------------------------
-- Bialgebras
-------------------------------------------------------------------------------

-- | A < https://en.wikipedia.org/wiki/Bialgebra bialgebra > over a semiring.
--
class (Unital a b, Counital a b) => Bialgebra a b

-------------------------------------------------------------------------------
-- Algebra instances
-------------------------------------------------------------------------------

instance Semiring a => Algebra a () where
  joined f = f ()

instance Semiring a => Unital a () where
  unital r () = r

--TODO: consider separating out n=3 cross product case
instance Semiring a => Algebra a (Finite n) where
  joined = M.join

instance Semiring a => Unital a (Finite n) where
  unital = const

instance (Algebra a b1, Algebra a b2) => Algebra a (b1, b2) where
  joined f (a,b) = joined (\a1 a2 -> joined (\b1 b2 -> f (a1,b1) (a2,b2)) b) a

instance (Unital a b1, Unital a b2) => Unital a (b1, b2) where
  unital r (a,b) = unital r a * unital r b

instance (Algebra a b1, Algebra a b2, Algebra a b3) => Algebra a (b1, b2, b3) where
  joined f (a,b,c) = joined (\a1 a2 -> joined (\b1 b2 -> joined (\c1 c2 -> f (a1,b1,c1) (a2,b2,c2)) c) b) a

instance (Unital a b1, Unital a b2, Unital a b3) => Unital a (b1, b2, b3) where
  unital r (a,b,c) = unital r a * unital r b * unital r c

-- | Tensor algebra on /b/.
--
-- >>> joined (<>) [1..3 :: Int]
-- [1,2,3,1,2,3,1,2,3,1,2,3]
--
-- >>> joined (\f g -> fold (f ++ g)) [1..3] :: Int
-- 24
--
instance Semiring a => Algebra a [b] where
  joined f = go [] where
    go ls rrs@(r:rs) = f (reverse ls) rrs + go (r:ls) rs
    go ls [] = f (reverse ls) []

instance Semiring a => Unital a [b] where
  unital a [] = a
  unital _ _ = zero

instance Semiring a => Algebra a (Seq b) where
  joined f = go Seq.empty where
    go ls s = case viewl s of
       EmptyL -> f ls s
       r :< rs -> f ls s + go (ls |> r) rs

instance Semiring a => Unital a (Seq b) where
  unital a b | Seq.null b = a
             | otherwise = zero

instance (Semiring a, Ord b) => Algebra a (Set.Set b) where
  joined f = go Set.empty where
    go ls s = case Set.minView s of
       Nothing -> f ls s
       Just (r, rs) -> f ls s + go (Set.insert r ls) rs

instance (Semiring a, Ord b) => Unital a (Set.Set b) where
  unital a b | Set.null b = a
           | otherwise = zero

instance Semiring a => Algebra a IntSet.IntSet where
  joined f = go IntSet.empty where
    go ls s = case IntSet.minView s of
       Nothing -> f ls s
       Just (r, rs) -> f ls s + go (IntSet.insert r ls) rs

instance Semiring a => Unital a IntSet.IntSet where
  unital a b | IntSet.null b = a
             | otherwise = zero

---------------------------------------------------------------------
-- Coalgebra instances
---------------------------------------------------------------------

--instance (Representable f, Algebra a (Rep (Co f))) => Coalgebra a (Co f a) where
--cojoined k (Co f) (Co g) = k (index f * index g)
--  cojoined k f g = k (f * g)

instance Semiring a => Coalgebra a () where
  cojoined = const

instance Semiring a => Counital a () where
  counital f = f ()

instance Semiring a => Coalgebra a (Finite n) where
  cojoined f i j = bool zero (f i) $ i == j

instance (KnownNat n, Semiring a) => Counital a (Finite n) where
  counital f = sum . fmap f $ finites

instance Algebra a b => Coalgebra a (b -> a) where
  cojoined k f g = k (f * g)

instance Unital a b => Counital a (b -> a) where
  counital f = f one

instance (Coalgebra a c1, Coalgebra a c2) => Coalgebra a (c1, c2) where
  cojoined f (a1,b1) (a2,b2) = cojoined (\a -> cojoined (\b -> f (a,b)) b1 b2) a1 a2

instance (Counital a c1, Counital a c2) => Counital a (c1, c2) where
  counital k = counital $ \a -> counital $ \b -> k (a,b)

instance (Coalgebra a c1, Coalgebra a c2, Coalgebra a c3) => Coalgebra a (c1, c2, c3) where
  cojoined f (a1,b1,c1) (a2,b2,c2) = cojoined (\a -> cojoined (\b -> cojoined (\c -> f (a,b,c)) c1 c2) b1 b2) a1 a2

instance (Counital a c1, Counital a c2, Counital a c3) => Counital a (c1, c2, c3) where
  counital k = counital $ \a -> counital $ \b -> counital $ \c -> k (a,b,c)

-- | The tensor coalgebra on /c/.
--
instance Semiring a => Coalgebra a [c] where
  cojoined f as bs = f (mappend as bs)

instance Semiring a => Counital a [c] where
  counital f = f []

instance Semiring a => Coalgebra a (Seq c) where
  cojoined f as bs = f (mappend as bs)

instance Semiring a => Counital a (Seq c) where
  counital f = f Seq.empty

-- | The free commutative band coalgebra
instance (Semiring a, Ord c) => Coalgebra a (Set.Set c) where
  cojoined f as bs = f (Set.union as bs)

instance (Semiring a, Ord c) => Counital a (Set.Set c) where
  counital f = f Set.empty

-- | The free commutative band coalgebra over Int
instance Semiring a => Coalgebra a IntSet.IntSet where
  cojoined f as bs = f (IntSet.union as bs)

instance Semiring a => Counital a IntSet.IntSet where
  counital f = f IntSet.empty

-- | The free commutative coalgebra over a set and a given semigroup
instance (Semiring r, Ord a, Semigroup b) => Coalgebra r (Map.Map a b) where
  cojoined f as bs = f (Map.unionWith (<>) as bs)

instance (Semiring r, Ord a, Semigroup b) => Counital r (Map.Map a b) where
  counital f = f Map.empty

{-

-- | The free commutative coalgebra over a set and a given semigroup
instance (Semiring r, Ord a, Commutative b) => Coalgebra r (Map a b) where
  cojoined f as bs = f (Map.unionWith (+) as bs)
  counital k = k (Map.empty)

-- | The free commutative coalgebra over a set and Int
instance (Semiring r, Commutative b) => Coalgebra r (IntMap b) where
  cojoined f as bs = f (IntMap.unionWith (+) as bs)
  counital k = k (IntMap.empty)
-}

---------------------------------------------------------------------
-- Bialgebra instances
---------------------------------------------------------------------

instance Semiring a => Bialgebra a () where

instance (KnownNat n, Semiring a) => Bialgebra a (Finite n) where

instance (Bialgebra a b1, Bialgebra a b2) => Bialgebra a (b1, b2) where

instance (Bialgebra a b1, Bialgebra a b2, Bialgebra a b3) => Bialgebra a (b1, b2, b3) where

instance Semiring a => Bialgebra a [b]

instance Semiring a => Bialgebra a (Seq b)


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
module LSI.FrequencyResponse where

import Data.Complex
import Data.Vector as V hiding (map, zip, foldr)
import Linear.V as V

import LSI.RationalFunction

computeFR :: RealFloat f
          => RationalFunction d f
          -> V d f
          -> Complex f
computeFR rf = (computeFR' rf) . V.toList .V.toVector
  where computeFR' f =
          case f of
            Monomial c es ->
              let fs = map toInteger . V.toList . V.toVector $ es in
                \vs -> {-# SCC monom #-}
                  let factors = map iexp $ zip zs fs
                      iexp (z, e) | e < 0 = 1/(z^(-e))
                      iexp (z, e) = z^e
                      zs = map (\w -> mkPolar 1.0 w) vs in
                    (c :+ 0.0) * (foldr (*) 1 factors)

            Add rf1 rf2 ->
              let fr1 = computeFR' rf1
                  fr2 = computeFR' rf2 in
                \vs -> (fr1 vs) + (fr2 vs)

            Mul rf1 rf2 ->
              let fr1 = computeFR' rf1
                  fr2 = computeFR' rf2 in
                \vs -> (fr1 vs) * (fr2 vs)

            Div rf1 rf2 ->
              let fr1 = computeFR' rf1
                  fr2 = computeFR' rf2 in
                \vs -> (fr1 vs) / (fr2 vs)


{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TupleSections         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE FlexibleContexts      #-}

import           Control.Applicative
import           Control.Monad
import           Control.Monad.Random

import           Codec.Compression.GZip ( decompress )
import           Data.Serialize ( Get )
import qualified Data.Serialize as Serialize
import qualified Data.ByteString.Lazy as B

import           Data.List ( foldl' )
import           Data.List.Split ( chunksOf )
import           Data.Maybe ( fromMaybe )
#if ! MIN_VERSION_base(4,13,0)
import           Data.Semigroup ( (<>) )
#endif
import           Data.Word ( Word32 , Word8 )
import qualified Data.Vector.Storable as V

import           Numeric.LinearAlgebra ( maxIndex )
import qualified Numeric.LinearAlgebra.Static as SA

import           Options.Applicative
import           System.FilePath ( (</>) )

import           Grenade
import           Grenade.Utils.OneHot

#ifdef COMPLEX
-- It's logistic regression!
--
-- This network is used to show how we can embed a Network as a layer in the larger MNIST
-- type.
--
type FL i o =
  Network
    '[ FullyConnected i o, Logit ]
    '[ 'D1 i, 'D1 o, 'D1 o ]

-- The definition of our convolutional neural network.
-- In the type signature, we have a type level list of shapes which are passed between the layers.
-- One can see that the images we are inputing are two dimensional with 28 * 28 pixels.

-- It's important to keep the type signatures, as there's many layers which can "squeeze" into the gaps
-- between the shapes, so inference can't do it all for us.

-- With the mnist data from Kaggle normalised to doubles between 0 and 1, learning rate of 0.01 and 15 iterations,
-- this network should get down to about a 1.3% error rate.
--
-- /NOTE:/ This model is actually too complex for MNIST, and one should use the type given belowe instead.
--         This one is just here to demonstrate Inception layers in use.
--
type MNIST =
  Network
    '[ Reshape,
       Concat ('D3 28 28 1) Trivial ('D3 28 28 14) (InceptionMini 28 28 1 5 9),
       Pooling 2 2 2 2, Relu,
       Concat ('D3 14 14 3) (Convolution 15 3 1 1 1 1) ('D3 14 14 15) (InceptionMini 14 14 15 5 10), Crop 1 1 1 1, Pooling 3 3 3 3, Relu,
       Reshape, FL 288 80, FL 80 10 ]
    '[ 'D2 28 28, 'D3 28 28 1,
       'D3 28 28 15, 'D3 14 14 15, 'D3 14 14 15, 'D3 14 14 18,
       'D3 12 12 18, 'D3 4 4 18, 'D3 4 4 18,
       'D1 288, 'D1 80, 'D1 10 ]

#else

-- The simpler network can just be dropped in without changing any of the other code
--

type MNIST
  = Network
    '[ Convolution 1 10 5 5 1 1, Pooling 2 2 2 2, Relu
     , Convolution 10 16 5 5 1 1, Pooling 2 2 2 2, Reshape, Relu
     , FullyConnected 256 80, Logit, FullyConnected 80 10, Logit]
    '[ 'D2 28 28, 'D3 24 24 10, 'D3 12 12 10, 'D3 12 12 10
     , 'D3 8 8 16, 'D3 4 4 16, 'D1 256, 'D1 256
     , 'D1 80, 'D1 80, 'D1 10, 'D1 10]

-- ... and this is an even simpler network from  https://crypto.stanford.edu/~blynn/haskell/brain.html
{-
type MNIST
  = Network
      '[Reshape, FullyConnected 784 30, Relu, FullyConnected 30 10, Logit]
      '[ 'D2 28 28, 'D1 784, 'D1 30, 'D1 30, 'D1 10, 'D1 10]
-}

#endif

randomMnist :: MonadRandom m => m MNIST
randomMnist = randomNetwork

convTest :: Int -> FilePath -> Maybe Int -> LearningParameters -> IO ()
convTest iterations dataDir nSamples rate = do
  net0         <- randomMnist
  trainData    <- readMNIST (dataDir </> "train-images-idx3-ubyte.gz")
                            (dataDir </> "train-labels-idx1-ubyte.gz")
  validateData <- readMNIST (dataDir </> "t10k-images-idx3-ubyte.gz")
                            (dataDir </> "t10k-labels-idx1-ubyte.gz")

  foldM_ (runIteration (maybe trainData (`take` trainData) nSamples) validateData) net0 [1..iterations]

    where
  trainEach rate' !network (i, o) = train rate' network i o

  runIteration trainRows validateRows net i = do
    let trained' = foldl' (trainEach ( rate { learningRate = learningRate rate * 0.9 ^ i} )) net trainRows
    print trained'

    putStrLn "Checking..."
    let res      = fmap (\(rowP,rowL) -> (rowL,) $ runNet trained' rowP) validateRows
    let res'     = fmap (\(S1D label, S1D prediction) -> (maxIndex (SA.extract label), maxIndex (SA.extract prediction))) res
    let matched   = length $ filter ((==) <$> fst <*> snd) res'
    let total     = length res'
    let matchedpc = fromIntegral matched / fromIntegral total * 100.0 :: Float
    putStrLn $ "Iteration " ++ show i ++ ": matched " ++ show matched ++ " of " ++ show total ++ " (" ++ show matchedpc ++ "%)"
    return trained'

data MnistOpts = MnistOpts FilePath (Maybe Int) Int LearningParameters

mnist' :: Parser MnistOpts
mnist' = MnistOpts <$> argument str (metavar "DATADIR")
                       -- option to reduce the number of training samples used from 60,000
                       -- to avoid running out of memory
                   <*> option (Just <$> auto) (long "limit_samples_to" <> short 'l' <> value Nothing)
                   <*> option auto (long "iterations" <> short 'i' <> value 15)
                   <*> (LearningParameters
                       <$> option auto (long "train_rate" <> short 'r' <> value 0.01)
                       <*> option auto (long "momentum" <> value 0.9)
                       <*> option auto (long "l2" <> value 0.0005)
                       )

main :: IO ()
main = do
    MnistOpts dataDir nSamples iter rate <- execParser (info (mnist' <**> helper) idm)
    putStr "Training convolutional neural network with "
    putStr $ maybe "all" show nSamples
    putStrLn " samples..."

    convTest iter dataDir nSamples rate

-- Adapted from https://github.com/tensorflow/haskell/blob/master/tensorflow-mnist/src/TensorFlow/Examples/MNIST/Parse.hs
-- Could also have used Data.IDX, although that uses a different Vector variant from that need for fromStorable
readMNIST :: FilePath -> FilePath -> IO [(S ( 'D2 28 28), S ( 'D1 10))]
readMNIST iFP lFP = do
  labels  <- readMNISTLabels lFP
  samples <- readMNISTSamples iFP
  return $ zip
    (fmap (fromMaybe (error "bad samples") . fromStorable) samples)
    (fromMaybe (error "bad labels") . oneHot . fromIntegral <$> labels)

-- | Check's the file's endianess, throwing an error if it's not as expected.
checkEndian :: Get ()
checkEndian = do
  magic <- Serialize.getWord32be
  when (magic `notElem` ([2049, 2051] :: [Word32]))
    $ error "Expected big endian, but image file is little endian."

-- | Reads an MNIST file and returns a list of samples.
readMNISTSamples :: FilePath -> IO [V.Vector Double]
readMNISTSamples path = do
  raw <- decompress <$> B.readFile path
  either fail ( return . fmap (V.map normalize) ) $ Serialize.runGetLazy getMNIST raw
 where
  getMNIST :: Get [V.Vector Word8]
  getMNIST = do
    checkEndian
    -- Parse header data.
    cnt    <- fromIntegral <$> Serialize.getWord32be
    rows   <- fromIntegral <$> Serialize.getWord32be
    cols   <- fromIntegral <$> Serialize.getWord32be
    -- Read all of the data, then split into samples.
    pixels <- Serialize.getLazyByteString $ fromIntegral $ cnt * rows * cols
    return $ V.fromList <$> chunksOf (rows * cols) (B.unpack pixels)

  normalize :: Word8 -> Double
  normalize = (/ 255) . fromIntegral
  -- There are other normalization functions in the literature, such as
  -- normalize = (/ 0.3081) . (`subtract` 0.1307) . (/ 255) . fromIntegral
  -- but we need values in the range [0..1] for the showShape' pretty printer

-- | Reads a list of MNIST labels from a file and returns them.
readMNISTLabels :: FilePath -> IO [Word8]
readMNISTLabels path = do
  raw <- decompress <$> B.readFile path
  either fail return $ Serialize.runGetLazy getLabels raw
 where
  getLabels :: Get [Word8]
  getLabels = do
    checkEndian
    -- Parse header data.
    cnt <- fromIntegral <$> Serialize.getWord32be
    -- Read all of the labels.
    B.unpack <$> Serialize.getLazyByteString cnt


module Numeric.FFT.Execute ( execute ) where

import           Control.Monad               (when)
import qualified Control.Monad               as CM
import           Control.Monad.ST
import           Data.Complex
import qualified Data.IntMap.Strict          as IM
import           Data.STRef
import qualified Data.Vector                 as V
import           Data.Vector.Unboxed
import qualified Data.Vector.Unboxed.Mutable as MV
import           Prelude                     hiding (concatMap, foldr, length,
                                              map, mapM_, null, reverse, sum,
                                              zip, zipWith)

import           Numeric.FFT.Special
import           Numeric.FFT.Types
import           Numeric.FFT.Utils

-- | Main FFT plan execution driver.
execute :: Plan -> Direction -> VCD -> VCD
execute (Plan dlinfo perm base) dir h
  | n == 1 = h
  | V.null dlinfo = runST $ do
                          mhin <- case perm of
                            Nothing -> thaw h
                            Just p  -> unsafeThaw $ backpermute h p
                          mhout <- MV.replicate n 0
                          applyBase base sign mhin mhout
                          when (dir == Inverse) $ do
                            let s = 1.0 / fromIntegral n :+ 0
                            CM.forM_ [0..n-1] $ \i -> do
                              x <- MV.unsafeRead mhout i
                              MV.unsafeWrite mhout i $ s * x
                          unsafeFreeze mhout
  | otherwise=  fullfft
  where
    n = length h              -- Input vector length.
    bsize = baseSize base     -- Size of base transform.

    -- Root of unity sign.
    sign = case dir of
      Forward -> 1
      Inverse -> -1

    -- Apply Danielson-Lanczos steps and base transform to digit
    -- reversal ordered input vector.
    fullfft = runST $ do
      mhin <- case perm of
            Nothing -> thaw h
            Just p  -> unsafeThaw $ backpermute h p
      mhtmp <- MV.replicate n 0
      multBase mhin mhtmp
      mhr <- newSTRef (mhtmp, mhin)
      V.forM_ dlinfo $ \dlstep -> do
        (mh0, mh1) <- readSTRef mhr
        dl sign dlstep mh0 mh1
        writeSTRef mhr (mh1, mh0)
      mhs <- readSTRef mhr
      let vout = fst mhs
      when (dir == Inverse) $ do
        let s = 1.0 / fromIntegral n :+ 0
        CM.forM_ [0..n-1] $ \i -> do
          x <- MV.unsafeRead vout i
          MV.unsafeWrite vout i $ s * x
      unsafeFreeze vout

    -- Multiple base transform application for "bottom" of algorithm.
    multBase :: MVCD s -> MVCD s -> ST s ()
    multBase xmin xmout =
      V.zipWithM_ (applyBase base sign)
                  (slicemvecs bsize xmin) (slicemvecs bsize xmout)

-- | Monadic FFT plan execution driver -- used by Rader's algorithm
-- for convolutions.
executeM :: Plan -> Direction -> MVCD s -> MVCD s -> ST s ()
executeM (Plan dlinfo perm base) dir hin hout =
  if n == 1
  then MV.copy hout hin
  else do
    htmp <- MV.replicate n 0

    -- Input permutation.
    case perm of
      Nothing -> MV.copy htmp hin
      Just p  -> backpermuteM n p hin htmp

    -- Apply Danielson-Lanczos steps and base transform to digit
    -- reversal ordered input vector.
    multBase htmp hout
    mhr <- newSTRef (hout, htmp)
    V.forM_ dlinfo $ \dlstep -> do
      (mh0, mh1) <- readSTRef mhr
      dl sign dlstep mh0 mh1
      writeSTRef mhr (mh1, mh0)
    when (odd $ V.length dlinfo) $ MV.copy hout htmp

    -- Output scaling for inverse transform.
    when (dir == Inverse) $ do
      let s = 1.0 / fromIntegral n :+ 0
      forM_ (enumFromN 0 n) $ \i -> do
        x <- MV.unsafeRead hout i
        MV.unsafeWrite hout i $ s * x
  where
    n = MV.length hin         -- Input vector length.
    bsize = baseSize base     -- Size of base transform.

    -- Root of unity sign.
    sign = case dir of
      Forward -> 1
      Inverse -> -1

    -- Multiple base transform application for "bottom" of algorithm.
    multBase :: MVCD s -> MVCD s -> ST s ()
    multBase xmin xmout =
      V.zipWithM_ (applyBase base sign)
                  (slicemvecs bsize xmin) (slicemvecs bsize xmout)

-- | Single Danielson-Lanczos step: process all duplicates and
-- concatenate into a single vector.
dl :: Int -> (Int, Int, VVVCD, VVVCD) -> MVCD s -> MVCD s -> ST s ()
dl sign (wfac, split, dmatp, dmatm) mhin mhout =
  V.zipWithM_ doone (slicemvecs wfac mhin) (slicemvecs wfac mhout)
  where
    -- Twiddled diagonal entries in row r, column c (both
    -- zero-indexed), where each row and column if a wfac x wfac
    -- matrix.
    dmat = if sign == 1 then dmatp else dmatm
    d r c = (dmat V.! r) V.! c

    -- Size of each diagonal sub-matrix.
    ns = wfac `div` split

    -- Index vectors.
    nsidxs = enumFromN 0 ns
    splitidxs = enumFromN 1 (split-1)

    -- Process one duplicate by processing all rows and writing the
    -- results into a single output vector.
    doone :: MVCD s -> MVCD s -> ST s ()
    doone vin vout = do
      let vs = (slicemvecs ns vin, slicemvecs ns vout)
      mapM_ (single vs) $ enumFromN 0 split
      where
        -- Multiply a single block by its appropriate diagonal
        -- elements and accumulate the result.
        mult :: VMVCD s -> MVCD s -> Int -> Bool -> Int -> ST s ()
        mult vins vo r first c = do
          let vi = vins V.! c
              dvals = d r c
          forM_ nsidxs $ \i -> do
            xi <- MV.unsafeRead vi i
            xo <- if first then return 0 else MV.unsafeRead vo i
            MV.unsafeWrite vo i (xo + xi * dvals ! i)
        -- Multiply all blocks by the corresponding diagonal
        -- elements in a single row.
        single :: (VMVCD s, VMVCD s) -> Int -> ST s ()
        single (vis, vos) r = do
          mult vis (vos V.! r) r True 0
          mapM_ (mult vis (vos V.! r) r False) splitidxs
        -- single (vis, vos) r =
        --   let m = mult vis (vos V.! r) r
        --   in do
        --     m True 0
        --     mapM_ (m False) splitidxs

-- | Apply a base transform to a single vector.
applyBase :: BaseTransform -> Int -> MVCD s -> MVCD s -> ST s ()

-- Simple DFT algorithm.
applyBase (DFTBase sz wsfwd wsinv) sign mhin mhout = do
  h <- freeze mhin
  forM_ (enumFromN 0 sz) $ \i -> MV.unsafeWrite mhout i (doone h i)
  where ws = if sign == 1 then wsfwd else wsinv
        doone h i = sum $ zipWith (*) h $
                    generate sz (\k -> ws ! (i * k `mod` sz))

-- Special hard-coded cases.
applyBase (SpecialBase sz) sign mhin mhout =
  case IM.lookup sz specialBases of
    Just f  -> f sign mhin mhout
    Nothing -> error "invalid problem size for SpecialBase"

-- Rader prime-length FFT.
applyBase (RaderBase sz outperm bfwd binv csz cplan) sign mhin mhout = do
  -- Padding size.
  let pad = csz - (sz - 1)

  -- Permuted input vector padded to next greater power of two size
  -- for fast convolution.
  apad <- MV.replicate csz 0
  forM_ (enumFromN 0 csz) $ \i -> do
    val <- if i == 0 then MV.unsafeRead mhin 1
           else if i > pad
                then MV.unsafeRead mhin $ i - pad + 1
                else return 0
    MV.unsafeWrite apad i val

  -- FFT-based convolution calculation.
  convtmp <- MV.replicate csz 0
  executeM cplan Forward apad convtmp
  let bmult = if sign == 1 then bfwd else binv
  forM_ (enumFromN 0 csz) $ \i -> do
    x <- MV.unsafeRead convtmp i
    MV.unsafeWrite convtmp i $ x * (bmult ! i)
  executeM cplan Inverse convtmp apad
  conv <- unsafeFreeze apad

  -- Input vector sum.
  sumhref <- newSTRef 0
  forM_ (enumFromN 0 sz) $ \i -> do
    val <- MV.unsafeRead mhin i
    modifySTRef sumhref (+ val)
  sumh <- readSTRef sumhref

  -- Write output based on output generator index ordering.
  h0 <- MV.unsafeRead mhin 0
  forM_ (enumFromN 0 sz) $ \i -> do
    let (idx, val) = case i of
          0 -> (0, sumh)
          _ -> (outperm ! (i - 1), h0 + conv ! (i - 1))
    MV.unsafeWrite mhout idx val


{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes, GADTs, TypeFamilies, TypeOperators, ConstraintKinds #-}
{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, FlexibleContexts #-}
{-# LANGUAGE DataKinds, PolyKinds, ConstraintKinds #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE BangPatterns #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Language.Grappa.Distribution where

import Control.Monad
import Control.Monad.Trans
import GHC.Exts
import Data.Typeable
import qualified Data.Vector.Storable as V
--import Data.Number.LogFloat hiding (sum, product, log)

import Numeric.LinearAlgebra hiding (R, Uniform, (<>), (<#))
import qualified Numeric.LinearAlgebra as M

import qualified Numeric.Log as Log
import qualified Numeric.SpecFunctions as Gamma

import qualified Numeric.AD.Mode.Forward as ADF
import qualified Numeric.AD.Internal.Forward as ADF
import qualified Numeric.AD.Mode.Reverse as ADR
import qualified Numeric.AD.Internal.Reverse as ADR
import qualified Numeric.AD.Internal.Identity as ADR
import qualified Numeric.AD.Jacobian as ADR
import qualified Data.Reflection as ADR (Reifies)

-- | The R type alias
type R = Double

----------------------------------------------------------------------
-- * Supporting the Gamma and Beta Functions
----------------------------------------------------------------------

-- | Class for computing @ln (gamma x)@
class HasGamma a where
  logGamma :: a -> a
  digamma :: a -> a

instance HasGamma Double where
  logGamma = Gamma.logGamma
  digamma = Gamma.digamma

-- | The 'trigamma' function = the derivative of 'digamma' (which itself is the
-- derivative of 'logGamma'). This implementation was cribbed from
-- <https://stackoverflow.com/questions/2978979/how-to-improve-performance-of-this-numerical-computation-in-haskell/2983578 this>
-- Stack Overflow post
trigamma :: Double -> Double
trigamma x = go 0 (x + 5) $ computeP $ x + 6
  where
    go :: Int -> Double -> Double -> Double
    go !i !z !p
        | i >= 6    = p
        | otherwise = go (i+1) (z-1) (1 / (z*z) + p)
    invSq z = 1 / (z * z)
    computeP z = (((((5/66*p-1/30)*p+1/42)*p-1/30)*p+1/6)*p+1)/z+0.5*p
      where p = invSq z

----------------------------------------------------------------------
-- * Helper Definitions for Automatic Differentiation
----------------------------------------------------------------------

-- | Convert between two 'AD.Forward' types, using a conversion function which
-- should numerically be the identity, i.e., it should the representation of @r@
-- in type @a@ to the closest representable number to @r@ in @b@.
{-
unsafeConvertAD :: (a -> b) -> AD.Forward a -> AD.Forward b
unsafeConvertAD f (AD.Forward x y) =
  AD.Forward (f x) (f y)
unsafeConvertAD f (AD.Lift x) =
  AD.Lift (f x)
unsafeConvertAD _ AD.Zero = AD.Zero
-}

-- | Apply a function to an 'ADF.Forward' object that has a more efficient or
-- accurate version for the underlying type. For instance, log1p = log . (1+)
-- has a more accurate version in C code, but we still need to apply the above
-- mathematical definition to compute the derivative.
optimizedADOp :: Num a => (ADF.Forward a -> ADF.Forward a) -> (a -> a) ->
                 ADF.Forward a -> ADF.Forward a
optimizedADOp f f_opt x =
  ADF.bundle (f_opt $ ADF.primal x) (ADF.tangent $ f x)

-- Needed to use @'ADF.Forward' 'Double'@ Log
instance Log.Precise (ADF.Forward Double) where
  log1p = optimizedADOp Log.log1p (log . (1 +))
  expm1 = optimizedADOp Log.expm1 ((\x -> x - 1) . exp)

-- Needed to use @'ADR.Reverse' 'Double'@ Log
instance ADR.Reifies rs ADR.Tape => Log.Precise (ADR.Reverse rs Double) where
  log1p x = ADR.unary Log.log1p (ADR.Id (1 / (ADR.primal x + 1))) x
  expm1 x = ADR.unary Log.expm1 (ADR.Id ((ADR.primal x) * exp (ADR.primal x))) x

instance ADR.Reifies s ADR.Tape => HasGamma (ADR.Reverse s Double) where
  logGamma x =
    ADR.unary Gamma.logGamma (ADR.Id $ Gamma.digamma $ ADR.primal x) x
  digamma x =
    ADR.unary Gamma.digamma (ADR.Id $ trigamma $ ADR.primal x) x

----------------------------------------------------------------------
-- * Probabilities
----------------------------------------------------------------------

-- | We use the shorthand @Prob@ for the type of probabilities
newtype Prob = Prob { fromProb :: Log.Log Double }
             deriving (Num,Eq,Ord,Real,Fractional,Typeable)

instance Show Prob where
  show (Prob x) = show x

-- | Specialized method for testing if a 'Prob' is zero
probIsZero :: Prob -> Bool
probIsZero (Prob x) = x == 0

-- | Specialized method for testing if a 'Prob' is a NaN
probIsNaN :: Prob -> Bool
probIsNaN (Prob x) = isNaN $ Log.ln x

-- | Convert a (non-negative) real to a 'Prob'
rToProb :: R -> Prob
rToProb = Prob . Log.Exp . log

-- | Convert a real that is already in log-space to a 'Prob'
logRToProb :: R -> Prob
logRToProb = Prob . Log.Exp

-- | Convert a 'Prob' to a real (not in log space)
probToR :: Prob -> R
probToR = exp . Log.ln . fromProb

-- | Convert a 'Prob' to a real in log space
probToLogR :: Prob -> R
probToLogR = Log.ln . fromProb

-- | Lifting the 'Log.sum' method to 'Prob'
sumProb :: [Prob] -> Prob
sumProb ps = Prob $ Log.sum $ map fromProb ps

----------------------------------------------------------------------
-- * Matrices of reals and probabilities
----------------------------------------------------------------------

-- | A bundled up type for unboxed real-valued vectors
newtype RVector = RVector { unRVector :: Vector Double } deriving (Num,Eq,Show)

-- | Return the length of an 'RVector'
lengthV :: RVector -> Int
lengthV (RVector v) = V.length v

-- | Return the @i@th element of a 'RVector'
atV :: RVector -> Int -> Double
atV (RVector v) i = v V.! i

-- | Generate a probability vector of the given size
generateV :: Int -> (Int -> Double) -> RVector
generateV n f = RVector $ V.generate n f

-- | Map over an 'RVector'
mapV :: (R -> R) -> RVector -> RVector
mapV f (RVector v) = RVector (V.map f v)

-- | Fold over an 'RVector'
foldrV :: (R -> b -> b) -> b -> RVector -> b
foldrV f b (RVector v) = V.foldr' f b v

-- | Take the sum of an 'RVector'
sumV :: RVector -> Double
sumV (RVector v) = V.foldl' (+) 0 v

-- | Generate a vector from a list of its elements
fromListV :: [Double] -> RVector
fromListV = RVector . V.fromList

-- | Extract the elements of a vector as a list
toListV :: RVector -> [Double]
toListV = V.toList . unRVector

-- | The type of real-valued matrices
newtype RMatrix = RMatrix { unRMatrix :: Matrix Double } deriving (Num,Eq,Show)

-- | Return the number of rows of an 'RMatrix'
rowsM :: RMatrix -> Int
rowsM (RMatrix m) = rows m

-- | Return the number of columns of a probability matrix
colsM :: RMatrix -> Int
colsM (RMatrix m) = M.cols m

-- | Return the @(i,j)@th element of a 'RMatrix'
atM :: RMatrix -> Int -> Int -> Double
atM (RMatrix m) i j = atIndex m (i,j)

-- | Build a matrix from a list of its rows
fromRowsM :: [RVector] -> RMatrix
fromRowsM vs = RMatrix $ fromRows $ map unRVector vs

-- | Build a matrix of probabilities from a list of its columns
fromColsM :: [RVector] -> RMatrix
fromColsM vs = RMatrix $ fromColumns $ map unRVector vs

-- | Get the rows of a matrix as a list of vectors
toRowsM :: RMatrix -> [RVector]
toRowsM (RMatrix m) = map RVector $ toRows m

-- | Get the columns of a matrix as a list of vectors
toColsM :: RMatrix -> [RVector]
toColsM (RMatrix m) = map RVector $ toColumns m

-- | Matrix multiplication
mulM :: RMatrix -> RMatrix -> RMatrix
mulM (RMatrix m1) (RMatrix m2) = RMatrix (m1 M.<> m2)

-- | Matrix-vector multiplication
mulMV :: RMatrix -> RVector -> RVector
mulMV (RMatrix m) (RVector v) = RVector (m #> v)

-- | Vector-matrix multiplication
mulVM :: RVector -> RMatrix -> RVector
mulVM (RVector v) (RMatrix m) = RVector (v M.<# m)

-- | The type of (efficient storable) vectors of values in log space
newtype ProbVector = ProbVector { unProbVector :: Vector Double }
                   deriving (Eq, Show)

-- | Return the length of a 'ProbVector'
lengthPV :: ProbVector -> Int
lengthPV (ProbVector v) = V.length v

-- | Return the @i@th element of a 'ProbVector'
atPV :: ProbVector -> Int -> Prob
atPV (ProbVector v) i = Prob $ Log.Exp $ v V.! i

-- | Generate a probability vector of the given size
generatePV :: Int -> (Int -> Prob) -> ProbVector
generatePV n f = ProbVector $ V.generate n (Log.ln . fromProb . f)

-- | Generate a probability vector from a list of its elements
fromListPV :: [Prob] -> ProbVector
fromListPV = ProbVector . V.fromList . map probToLogR

-- | Extract the elements of a probability vector as a list
toListPV :: ProbVector -> [Prob]
toListPV = map logRToProb . V.toList . unProbVector

-- | Map over a 'ProbVector'
mapPV :: (Prob -> Prob) -> ProbVector -> ProbVector
mapPV f (ProbVector v) = ProbVector (V.map (probToLogR . f . logRToProb) v)

-- | Fold over a 'ProbVector'
foldrPV :: (Prob -> b -> b) -> b -> ProbVector -> b
foldrPV f b (ProbVector v) = V.foldr' (f . logRToProb) b v

-- | Apply a binary operator pointwise to two probability vectors
binOpPV :: (Prob -> Prob -> Prob) -> ProbVector -> ProbVector -> ProbVector
binOpPV f v1 v2 =
  generatePV (min (lengthPV v1) (lengthPV v2)) (\i -> f (atPV v1 i) (atPV v2 i))

instance Num ProbVector where
  (+) = binOpPV (+)
  (-) = binOpPV (-)
  (*) = binOpPV (*)
  abs = mapPV abs
  signum = mapPV signum
  fromInteger i = generatePV 1 (const $ fromInteger i)

-- | Take the sum of a 'ProbVector'. The algorithm for this is adapted from
-- 'Log.sum', though that function requires a 'Foldable' instance.
sumPV :: ProbVector -> Prob
sumPV (ProbVector v) | V.length v == 0 = Prob $ Log.Exp $ log 0
sumPV (ProbVector v) =
  let max_v = V.foldl1' max v in
  Prob $ Log.Exp $
  if isInfinite max_v then max_v else
    max_v + Log.log1p (V.foldl' (\r x -> r + Log.expm1 (x - max_v)) 0 v
                       + fromIntegral (V.length v - 1))

-- | The type of matrices of values in log space
newtype ProbMatrix = ProbMatrix (Matrix Double) deriving (Eq, Show)

-- | Return the number of rows of a probability matrix
rowsPM :: ProbMatrix -> Int
rowsPM (ProbMatrix m) = rows m

-- | Return the number of columns of a probability matrix
colsPM :: ProbMatrix -> Int
colsPM (ProbMatrix m) = M.cols m

-- | Return the @(i,j)@th element of a 'ProbMatrix'
atPM :: ProbMatrix -> Int -> Int -> Prob
atPM pm@(ProbMatrix m) i j =
  if i >= rowsPM pm then
    error ("atPM index out of bounds: row " ++ show i ++ " >= " ++ show (rowsPM pm))
    else
    if j >= colsPM pm then
      error ("atPM index out of bounds: column " ++ show j
             ++ " >= " ++ show (colsPM pm))
    else
      Prob $ Log.Exp $ atIndex m (i,j)

-- | Build a matrix of probabilities from a list of its rows
fromRowsPM :: [ProbVector] -> ProbMatrix
fromRowsPM vs = ProbMatrix $ fromRows $ map unProbVector vs

-- | Build a matrix of probabilities from a list of its columns
fromColsPM :: [ProbVector] -> ProbMatrix
fromColsPM vs = ProbMatrix $ fromColumns $ map unProbVector vs

-- | Get the rows of a matrix as a list of vectors
toRowsPM :: ProbMatrix -> [ProbVector]
toRowsPM (ProbMatrix m) = map ProbVector $ toRows m

-- | Get the columns of a matrix as a list of vectors
toColsPM :: ProbMatrix -> [ProbVector]
toColsPM (ProbMatrix m) = map ProbVector $ toColumns m

-- | Matrix-vector multiplication for probability matrices
mulPMV :: ProbMatrix -> ProbVector -> ProbVector
mulPMV m v =
  generatePV (rowsPM m) $ \i ->
  sumPV $ generatePV (min (colsPM m) (lengthPV v)) $ \j ->
  (atPM m i j) * (atPV v j)

-- | Vector-matrix multiplication
mulPVM :: ProbVector -> ProbMatrix -> ProbVector
mulPVM v m =
  generatePV (colsPM m) $ \j ->
  sumPV $ generatePV (min (rowsPM m) (lengthPV v)) $ \i ->
  (atPM m i j) * (atPV v i)

-- | Matrix-matrix multiplication for probability matrices
mulPM :: ProbMatrix -> ProbMatrix -> ProbMatrix
mulPM m1 (ProbMatrix m2) =
  ProbMatrix $ fromColumns $
  map (unProbVector . mulPMV m1 . ProbVector) $ toColumns m2

----------------------------------------------------------------------
-- * Distributions
----------------------------------------------------------------------

-- | This defines the support type of a distribution
type family Support (d :: *) :: *

-- | The base class for distributions, stating that we can evaluate the density
-- of any given distribution at a particular value.  For debugging purposes, we
-- also require distributions to be 'Show'able.
class PDFDist d where
  distDensity :: d -> Support d -> Prob

-- | This type class states that monad @m@ supports distribution type @d@ by
-- allowing distributions of type @d@ to be randomly sampled in @m@
class Monad m => SampleableIn m d where
  distSample :: d -> m (Support d)

instance (SampleableIn m d,
          Monad (t m), MonadTrans t) => SampleableIn (t m) d where
  distSample = lift . distSample

-- | This type family states that all distribution types in @ds@ satisfy @c@
type family DistsIn (ds :: [*]) (c :: * -> Constraint) :: Constraint where
  DistsIn '[] c = ()
  DistsIn (d ': ds) c = (c d, DistsIn ds c)

-- | This constraint says that the distribution type @d@ is continuous, i.e.,
-- isomorphic to the reals.
--
-- The class functions take @d@, not @Proxy d@, since e.g. we need
-- different implementations for @uniform 0 1 :: Uniform@ and @uniform
-- 1 2 :: Uniform@.
class Continuous d where
  -- | Note that 'toReal' need not be defined on all of @Support d@,
  -- and implementation should raise an error on out of bounds
  -- values. E.g., for @uniform 0 1@, we have @Support Uniform ~
  -- R@, but in fact 'toReal' is only defined on @[0,1]@. We
  -- don't use @Maybe R@ as the return type, since passing an out
  -- of bounds value to 'toReal' is probably a bug.
  toReal :: d -> Support d -> R
  fromReal :: d -> Double -> Support d

{- FIXME: this does not seem to be the derivative in the below, and is in fact
just fromReal in log-space. Figure this out or remove it!

  -- | The derivative of 'fromReal'.
  --
  -- In the multivariable case the derivative becomes the determinant
  -- of the Jacobian (matrix of partial derivatives).
  --
  -- We can factor this into a separate class if we end up with some
  -- continuous distributions for which we don't have Jacobian
  -- determinants.
  fromReal' :: d -> Double -> LogFloat
-}

-- | The PDF of a continuous dist transformed to be a dist on the
-- whole real line.
--
-- All of our single variable transforms are increasing, so I don't
-- think we have to worry about taking absolute values; indeed, the
-- @LogFloat@ is not even defined for negative values. (compare
-- formulas for single dimensional change of variables which uses
-- derivative and multiple dimension change of variables where the
-- *absolute value* of the Jacobian is used:
-- https://en.wikipedia.org/wiki/Integration_by_substitution).
{-
fromRealDistDensity :: (Continuous d, PDFDist d) => d -> Double -> Prob
fromRealDistDensity d x = distDensity d (fromReal d x) * Prob (fromReal' d x)
-}

----------------------------------------------------------------
-- * Helper functions for defining 'Continuous' instances
----------------------------------------------------------------

-- TODO(conathan): revisit these bijections once we have a better
-- understanding of any additional properties we want from them. We
-- plan to use these bijections to work in the reals and not have to
-- worry about bounds.
--
-- Potential trouble: consider the open interval @(lb,ub)@ bijection
-- below in @toRealLbUb@ and @fromRealLbUb@. The logit function @\p ->
-- log (p / (1 - p))@ is small for most of the reals, so if we do
-- anything where we choose random reals over a large range, nearly
-- all of them will cluster in the ends of underlying interval
-- (i.e. near @lb@ and @ub@). If we want more uniform choices in the
-- reals to map to more uniform choices in our underlying interval,
-- then we'll need a bijection that spreads the interval much wider.
--
-- Some bijections to consider (these need to be scaled to work with
-- @(lb,ub)@):
--
-- - @\x -> log (x / (1 - x))@:
--   https://www.wolframalpha.com/input/?i=plot+log+(x+/+(1+-+x))
--
-- - @\x -> (1 / (1 - x)) - (1 / x)@:
--   https://www.wolframalpha.com/input/?i=plot+(1/(1-x)+-+1/x)
--
-- - @\x -> tan x@:
--   https://www.wolframalpha.com/input/?i=plot+tan(x)
--
-- All of these can be made more uniform by scaling them by a
-- constant, e.g. @\x -> 100000 * tan x@.

{- FIXME: this is only used for fromReal'; see comments for it above
-- | The derivative of the identify function, returning @LogFloat@.
fromRealId' :: Double -> LogFloat
fromRealId' = logFloat
-}

-- | Biject the open interval @(lb,ub)@ onto the reals.
toRealLbUb :: Floating a => a -> a -> a -> a
toRealLbUb lb ub x = r
  where
    -- Biject @(lb,ub)@ onto @(0,1)@
    p = (x - lb) / (ub - lb)
    -- Biject @(0,1)@ onto the reals
    --r = log (p / (1 - p))
    r = (2*p - 1) / (p - p*p)

-- | Biject the reals on the open interval @(lb,ub)@
fromRealLbUb :: Floating a => a -> a -> a -> a
fromRealLbUb lb ub r = x
  where
    -- Inverse of @\p -> log (p / 1 - p)@, except that we have a special case
    -- for infinite values of exp(-r), to prevent us getting NaN values for AD
    -- exp_neg_r = exp (-r)
    -- p = if isInfinite exp_neg_r then 0 else 1 / (1 + exp_neg_r)
    p = (r - 2 + sqrt (r*r + 4)) / (2*r)
    -- Inverse of @\x -> (x - lb) / (ub - lb)@.
    x = p * (ub - lb) + lb

{-
-- | Derivative of 'fromRealLbUb'.
fromRealLbUb' :: R -> R -> R -> LogFloat
fromRealLbUb' lb ub r =
  -- The below is an expansion of
  -- @log ((ub - lb) * (exp (-r) / (1 + exp (-r))**2))@.
  --
  -- Here we use @log1p@ -- an implementation of @\p -> log (1 + p)@
  -- that's optimized for small @p@ -- from
  -- @Data.Numeric.LogFloat@. This implementation of @log1p@ will use
  -- an optimized C implementation if compiled with the @useffi@ flag,
  -- which is supposed to be enabled by default.
  --
  -- The @log-domain@ package that provides an native Haskell
  -- implementation of @log1p@, I think.
  logToLogFloat $ log (ub - lb) - r - 2 * log1p (exp (- r))
-}

-- | Biject the open interval @(0,\infty)@ onto the reals, using the piecewise
-- function that maps @x@ in @(0,1)@ to @log x@ and otherwise maps @x@ to @x-1@.
toRealZeroInftyPW :: Double -> Double
toRealZeroInftyPW x = if x < 1 then log x else x - 1

-- | Biject the reals on the open interval @(0,\infty)@, using the piecewise
-- function that maps negative @x@ to @exp x@ and otherwise maps @x@ to @x+1@.
fromRealZeroInftyPW :: Double -> Double
fromRealZeroInftyPW x = if x < 0 then exp x else x + 1

-- | Biject the open interval @(lb,\infty)@ onto the reals.
toRealLbInfty :: Double -> Double -> R
toRealLbInfty lb x = log (x - lb)

-- | Biject the reals onto the open interval @(lb,\infty)@.
fromRealLbInfty :: Double -> Double -> R
fromRealLbInfty lb r = exp r + lb

-- | Biject the reals onto the open interval @(lb,\infty)@ in log-space
fromRealLbInfty' :: Double -> Double -> Prob
fromRealLbInfty' lb r = logRToProb lb + logRToProb r

-- | Biject the open interval @(-\infty,ub)@ onto the reals.
toRealInftyUb :: Double -> Double -> R
toRealInftyUb ub x = - log (ub - x)

-- | Biject the reals onto the open interval @(-\infty,ub)@.
fromRealInftyUb :: Double -> Double -> R
fromRealInftyUb ub r = - (exp (- r) - ub)

-- | Biject the reals onto the open interval @(lb,\infty)@ in log-space
fromRealInftyUb' :: Double -> Double -> Prob
fromRealInftyUb' _ub _r =
  error "FIXME: fromRealInftyUb' is not right!"
  -- logRToProb ub + logRToProb (- r)

----------------------------------------------------------------------
-- * Example distribution types
----------------------------------------------------------------------

-- | Evaluate the density of a normal distribution at value @r@. This is
-- "unchecked" in that it assumes @sigma@ is non-negative.
normalDensityUnchecked :: Floating r => r -> r -> r -> Log.Log r
normalDensityUnchecked mu sigma x =
  Log.Exp $ - 0.5 * sq ((x - mu) / sigma) - log sigma - halfLogTwoPi
  where
    sq n = n * n
    halfLogTwoPi = 0.5 * log (2 * pi)

-- | Evaluate the density of a normal distribution at value @r@
normalDensityGen :: (Ord r, Floating r, Show r) => r -> r -> r -> Log.Log r
normalDensityGen mu sigma x =
  if sigma <= 0 then
    error $ "normalDensity: sigma is <= 0: sigma = " ++ show sigma
  else
    normalDensityUnchecked mu sigma x

-- | Evaluate the density of a normal distribution at value @r@
normalDensity :: Double -> Double -> Double -> Prob
normalDensity mu sigma x = Prob $ normalDensityGen mu sigma x

data StdNormal = StdNormal deriving Show
type instance Support StdNormal = R
instance PDFDist StdNormal where
  distDensity StdNormal x = logRToProb $ - 0.5 * sq x - halfLogTwoPi
    where
    sq n = n * n
    halfLogTwoPi = 0.5 * log (2 * pi)

instance Continuous StdNormal where
  toReal _ = id
  fromReal _ = id
  --fromReal' _ = fromRealId'

-- | The normal distribution
data Normal = Normal R R deriving Show
type instance Support Normal = R

instance PDFDist Normal where
  distDensity (Normal mu sigma) r = normalDensity mu sigma r

instance Continuous Normal where
  toReal _ = id
  fromReal _ = id
  --fromReal' _ = fromRealId'

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

-- | Multivariate normal distribution.
--
-- Our multivariate normal of @k@ dimensions is parameterized by the
-- mean vectur @mu = [mu1,...,muk]@ and the *Cholesky factor* @c@ of
-- the covariance matrix @Sigma@. I.e., for multivariate normal with
-- mean @mu@ and covariance matrix @Sigma@, we use
--
-- > MVNormal mu c
--
-- for
--
-- > transpose c * c = Sigma
--
-- The requirements for @c@ to be well formed are that
--
-- - @c@ is upper triangular.
--
-- - all diagonal entries of @c@ are positive.
--
-- There are several reasons to prefer the Cholesky factor @c@ to the
-- covariance matrix @Sigma@:
--
-- - computing the PDF is more efficient using @c@.
--
-- - the sampling algorithm uses @c@, so we need to compute it anyway.
--
-- - all upper triangular matrices with positive diagonal entries are
--   valid Cholesky factors, and a matrix is positive definite iff it
--   has a Cholesky factorization, and the Cholesky factorization is
--   unique. So, it's easy to uniformly sample the space of
--   multivariate normals by uniformly sampling the space upper
--   triangular matrices with positive diagonal entries. We care about
--   sampling multivariate normals e.g. when using multivariate
--   normals as a variational family in Black Box Variation Inference.
--
-- There is a more general definition of multivariate normal which
-- only requires non-negative entries on the diagonal of @c@, where
-- zero values correspond to zero variance random variables,
-- i.e. constants, but we don't support those here (the PDFs would be
-- Dirac deltas).
--
-- A covariance matrix @Sigma@ can be converted into a Cholesky factor
-- @c@ using the @Numeric.LinearAlgebra.chol@ function:
--
-- > c = chol Sigma
--
-- The covariance matrix @Sigma@ must be symmetric and positive
-- definite; you can use @Numeric.LinearAlgebra.trustsym@ to construct
-- it.
data MVNormal = MVNormal (Vector Double) (Matrix Double) deriving Show

type instance Support MVNormal = Vector Double

instance PDFDist MVNormal where
  distDensity (MVNormal mu c) x =
    mvnDensityMuC mu c x

instance (Monad m, SampleableIn m Normal) => SampleableIn m MVNormal where
  -- https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Drawing_values_from_the_distribution
  distSample (MVNormal mu c) = do
    z <- V.fromList <$> (replicateM (size mu) $ distSample (Normal 0 1))
    -- Our Cholesky factorization is @Sigma = transpose c * c@ and
    -- Wikipedia has @Sigma = c * transpose c@, so we need to
    -- transpose here.
    return $ mu + (tr c #> z)

-- | Multivariate normal density in terms of Cholesky factor.
mvnDensityMuC :: Vector Double -> Matrix Double -> Vector Double -> Prob
mvnDensityMuC mu c x =
  if not isValidCholeskyFactor
  then error "distDensity: Cholesky factor is malformed!"
  -- https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Likelihood_function.
  else logRToProb $ -0.5 * (logAbsDet + prod + k * log2pi)
  where
    log2pi = log (2 * pi)
    k = fromIntegral $ size x
    -- We have
    --
    -- > Sigma = transpose c * c
    --
    -- and so
    --
    -- > det Sigma = det c * det c .
    --
    -- We know @det c > 0@ since the diag is positive.
    logAbsDet = 2 * (log $ V.product (takeDiag c))
    -- Here @cholSolve@ solves a linear system for the matrix @Sigma@
    -- for which @c@ is the Cholesky factor. I.e. if
    --
    -- > u == cholSolve m v
    --
    -- then
    --
    -- > transpose c * c * v == u .
    prod = (x - mu) <.> flatten (cholSolve c (asColumn $ x - mu))

    isValidCholeskyFactor = all (> 0) (V.toList $ takeDiag c)

-- | Multivariate normal density in terms of covariance matrix @Sigma@.
--
-- Can use this to test the other implementation above; see
-- 'testMvDensity' below.
mvnDensityMuSigma :: Vector Double -> Matrix Double -> Vector Double -> Prob
mvnDensityMuSigma mu sigma x =
  if signDet < 0
  then error "distDensity: matrix is not positive definite!"
  -- https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Likelihood_function.
  else logRToProb $ -0.5 * (logAbsDet + prod + k * log2pi)
  where
    log2pi = log (2 * pi)
    k = fromIntegral $ size x
    (sigmaInv, (logAbsDet, signDet)) = invlndet sigma
    prod = (x - mu) <.> (sigmaInv #> (x - mu))

-- | Check that the two mv normal density functions agree.
--
-- They agree if this function returns a small number.
testMvnDensity :: Vector Double -> Matrix Double -> Vector Double -> Double
testMvnDensity mu c x =
   abs $ (probToR $ mvnDensityMuSigma mu (tr c <> c) x) -
         (probToR $ mvnDensityMuC mu c x)

-- | Should return a smalllllll number.
testMvnDensityExample :: Double
testMvnDensityExample = sum [ testMvnDensity mu c x | x <- xs ]
  where
    mu = V.fromList [1,2]
    c = chol (trustSym $ (2><2) [1,0.5,0.5,1])
    xs = [ V.fromList [x0,x1] | x0 <- [-5,-4.5..5] , x1 <- [-5,-4.5..5] ]

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

-- | The uniform distribution
data Uniform = Uniform Double R deriving Show
type instance Support Uniform = R

uniformDensityGen :: (Ord r, RealFloat r, Log.Precise r) =>
                     r -> r -> r -> Log.Log r
uniformDensityGen lb ub r
  | lb >= ub =
    error "Malformed uniform distribution! Lower bound >= upper bound!"
  | r >= lb && r < ub =
    -- The expression 1 / (ub - lb) in log-space
      Log.Exp $ negate $ log $ ub - lb
  | otherwise = 0

uniformDensity :: Double -> Double -> Double -> Prob
uniformDensity lb ub r
  | lb >= ub =
    error "Malformed uniform distribution! Lower bound >= upper bound!"
  | r >= lb && r < ub = 1 / rToProb (ub - lb)
  | otherwise = 0

instance PDFDist Uniform where
  distDensity (Uniform lb ub) = uniformDensity lb ub

instance Continuous Uniform where
  toReal (Uniform lb ub) = toRealLbUb lb ub
  fromReal (Uniform lb ub) = fromRealLbUb lb ub
  --fromReal' (Uniform lb ub) = fromRealLbUb' lb ub

data StudentT = StudentT Double deriving Show
type instance Support StudentT = R

instance Continuous StudentT where
  toReal _ = id
  fromReal _ = id
  --fromReal' _ = fromRealId'

data Cauchy = Cauchy deriving Show
type instance Support Cauchy = R
instance PDFDist Cauchy where
  distDensity Cauchy x = logRToProb $ - log (1 + sq x) - log pi
    where sq n = n * n

instance Continuous Cauchy where
  toReal _ = id
  fromReal _ = id
  --fromReal' _ = fromRealId'

data For d = forall t. For [t] (t -> d)
type instance Support (For d) = [Support d]
instance PDFDist d => PDFDist (For d) where
  distDensity (For xs f) ys =
    product [distDensity (f x) y | (x,y) <- zip xs ys]

data StdUniform = StdUniform deriving Show
type instance Support StdUniform = R

instance PDFDist StdUniform where
  distDensity _ = uniformDensity 0 1

instance Continuous StdUniform where
  toReal _ = toRealLbUb 0 1
  fromReal _ = fromRealLbUb 0 1
  --fromReal' _ = fromRealLbUb' 0 1

-- | The categorical distribution, that picks a natural number between 1 and @n@
-- from a list of probabilisties of each number. Note that the probabilities
-- need not be normalized, e.g., the list @[1,1,1]@ is treated the same as
-- @[1/3,1/3,1/3]@, i.e., that each of 0, 1, and 2 has equal probability.
data Categorical = Categorical [Prob] deriving Show
type instance Support Categorical = Int

instance PDFDist Categorical where
  distDensity (Categorical ws) n =
    let total_w = sum ws in
    if n >= length ws then 0 else ws!!n / total_w

data Exponential = Exponential R deriving Show
type instance Support Exponential = R

exponentialDensityUnchecked :: (RealFloat a, Log.Precise a) =>
                               a -> a -> Log.Log a
exponentialDensityUnchecked rate x = Log.Exp $ log rate - rate * x

exponentialDensity :: (Ord a, RealFloat a, Log.Precise a) => a -> a -> Log.Log a
exponentialDensity rate x =
  if rate > 0 then exponentialDensityUnchecked rate x else 0

instance PDFDist Exponential where
  distDensity (Exponential rate) = Prob . exponentialDensity rate

data Beta = Beta R R deriving Show
type instance Support Beta = R

betaDensityUnchecked :: (HasGamma a, Floating a) => a -> a -> a -> Log.Log a
betaDensityUnchecked alpha beta x =
  Log.Exp $
  ((alpha - 1) * log x) + ((beta - 1) * log (1 - x)) -
  (logGamma alpha + logGamma alpha - logGamma (alpha + beta))

betaDensity :: (Ord a, HasGamma a, RealFloat a, Log.Precise a) =>
               a -> a -> a -> Log.Log a
betaDensity alpha beta x =
  if x > 0 && x < 1 then betaDensityUnchecked alpha beta x else 0

-- | Compute the PDF of the Beta distribution where the argument is already in
-- log space
betaDensityLog :: (Ord a, HasGamma a, RealFloat a, Log.Precise a) =>
                  a -> a -> Log.Log a -> Log.Log a
betaDensityLog alpha beta x =
  if x > 0 && x < 1 then
    Log.Exp $
    ((alpha - 1) * Log.ln x) + ((beta - 1) * Log.ln (1 - x)) -
    (logGamma alpha + logGamma alpha - logGamma (alpha + beta))
  else 0

instance PDFDist Beta where
  distDensity (Beta alpha beta) = Prob . betaDensity alpha beta

data Gamma = Gamma R R deriving Show
type instance Support Gamma = R

-- | Calculate the PDF of the gamma distribution, given shape parameter @k@ and
-- scale parameter @theta@, but without the @1 / (gamma k * theta ** k)@
-- normalization constant
relativeGammaDensity :: Floating a => a -> a -> a -> Log.Log a
relativeGammaDensity k theta x =
  Log.Exp $ (k - 1) * log x - (x / theta)

-- | Calculate the PDF of the gamma distribution, assuming @x > 0@, given shape
-- parameter @k@ and scale parameter @theta@
gammaDensityUnchecked :: (HasGamma a, Floating a) => a -> a -> a -> Log.Log a
gammaDensityUnchecked k theta x =
  Log.Exp (Log.ln (relativeGammaDensity k theta x) - logGamma k - k * log theta)

-- | Calculate the PDF of the gamma distribution, given shape parameter @k@ and
-- scale parameter @theta@
gammaDensity :: (Ord a, HasGamma a, RealFloat a, Log.Precise a) =>
                a -> a -> a -> Log.Log a
gammaDensity k theta x =
  if x <= 0 then 0 else gammaDensityUnchecked k theta x

instance PDFDist Gamma where
  distDensity (Gamma k theta) = Prob . gammaDensity k theta

data Dirichlet = Dirichlet [R] deriving Show
type instance Support Dirichlet = [R]

-- | The log of the multivariate beta function
logMVBeta :: (HasGamma a, Floating a) => [a] -> Log.Log a
logMVBeta alphas =
  -- Gamma (a_1) * ... * Gamma (a_n) / Gamma (a_1 + ... + a_n) in log space
  Log.Exp $ sum (map logGamma alphas) - logGamma (sum alphas)

-- | Calculate the density of the Dirichlet distribution at any type that
-- supports the gamma function
dirichletDensity :: (HasGamma a, Floating a) => [a] -> [a] -> Log.Log a
dirichletDensity alphas xs =
  -- x_1 ** (alpha_1 - 1) * ... * x_n ** (alpha_n - 1) / Beta (alphas)
  Log.Exp $
  sum (zipWith (\alpha x -> (alpha - 1) * log x) alphas xs) -
  Log.ln (logMVBeta alphas)

-- | Calculate the density of the Dirichlet distribution over log space
dirichletDensityLog :: (HasGamma a, Floating a) => [a] -> [Log.Log a] ->
                       Log.Log a
dirichletDensityLog alphas xs =
  -- x_1 ** (alpha_1 - 1) * ... * x_n ** (alpha_n - 1) / Beta (alphas)
  Log.Exp $
  sum (zipWith (\alpha x -> (alpha - 1) * Log.ln x) alphas xs) -
  Log.ln (logMVBeta alphas)

-- | The log of the multivariate beta function
logMVBetaV :: RVector -> Log.Log Double
logMVBetaV (RVector alphas) =
  -- Gamma (a_1) * ... * Gamma (a_n) / Gamma (a_1 + ... + a_n) in log space
  Log.Exp $ V.foldl' (\r alpha -> r + logGamma alpha) 0 alphas

-- | Calculate the density of the Dirichlet distribution over an 'RVector'
dirichletDensityV :: RVector -> RVector -> Prob
dirichletDensityV (RVector alphas) (RVector xs) =
  -- x_1 ** (alpha_1 - 1) * ... * x_n ** (alpha_n - 1) / Beta (alphas)
  Prob $ Log.Exp $
  V.ifoldl' (\r i alpha -> r + (alpha - 1) * log (xs!i)) 0 alphas -
  Log.ln (logMVBetaV $ RVector alphas)

-- | Calculate the density of the Dirichlet distribution over a 'ProbVector'
dirichletDensityPV :: RVector -> ProbVector -> Prob
dirichletDensityPV (RVector alphas) (ProbVector xs) =
  -- x_1 ** (alpha_1 - 1) * ... * x_n ** (alpha_n - 1) / Beta (alphas)
  Prob $ Log.Exp $
  V.ifoldl' (\r i alpha -> r + (alpha - 1) * xs!i) 0 alphas -
  Log.ln (logMVBetaV $ RVector alphas)

instance PDFDist Dirichlet where
  distDensity (Dirichlet alphas) xs = Prob $ dirichletDensity alphas xs

data Bernoulli = Bernoulli Double deriving Show

data Binomial = Binomial Int Double deriving Show

data Wishart = Wishart (Matrix Double) Double deriving Show
type instance Support Wishart = Matrix Double
instance PDFDist Wishart where
  distDensity (Wishart _scaleChol _df) _x = undefined

-- | A dirac delta distribution, that has a 100 percent chance of returning a
-- given constant value
data Dirac a = Dirac a deriving Show
type instance Support (Dirac a) = a

instance Eq a => PDFDist (Dirac a) where
  -- FIXME: this density is not really correct for R; it should really be
  -- infinite for continuous functions and 1 for discrete ones
  distDensity (Dirac a) x = if a == x then 1 else 0

-- | A technically invalid distribution, that always returns the value @x@ but
-- always has a score of @1@
data DontCare a = DontCare a deriving Show
type instance Support (DontCare a) = a

instance PDFDist (DontCare a) where
  distDensity (DontCare _) _ = 1

----------------------------------------------------------------------
-- * Continuity
----------------------------------------------------------------------

class Continuous (d f) => ReprContinuous d f
instance Continuous (d f) => ReprContinuous d f

-- | This type class states that monad @m@ supports distribution type @d@ by
-- allowing distributions of type @d@ to be randomly sampled in @m@
class Monad m => ReprSampleableIn m d f where
  distSampleRepr :: d f -> m (f (ReprSupport d))

class DiffContinuous d f where
  toRealF   :: d f -> f (ReprSupport d) -> f R
  fromRealF :: d f -> f R -> f (ReprSupport d)

class DiffPDF d f where
  distDensityF :: d f -> f (ReprSupport d) -> f Prob

{-
instance DiffContinuous ReprNormal AD.Forward where
  toRealF _ d = d
  fromRealF _ d = d

instance DiffPDF ReprNormal AD.Forward where
  distDensityF (ReprNormal mu sigma) x =
    if sigma <= 0
      then error $ "distDensityAD Normal: sigma is <= 0: sigma = " ++ show sigma
           -- FIXME: instead of using exp below, figure out how to do an
           -- adToProb that takes a real in log space and preserves
      else adToProb $ exp $ -0.5 * sq((x - mu) / sigma) - log sigma - halfLogTwoPi
      where sq n = n * n
            halfLogTwoPi = 0.5 * log (2 * pi)

instance DiffContinuous ReprUniform AD.Forward where
  toRealF (ReprUniform lb ub) =
    toRealLbUbNum lb ub
  fromRealF (ReprUniform lb ub) =
    fromRealLbUbNum lb ub

instance DiffPDF ReprUniform AD.Forward where
  distDensityF (ReprUniform lb ub) _
    | lb >= ub = error "Malformed uniform distribution! Lower bound >= upper bound!"
    | otherwise =
      (1 / adToProb (lb - ub))
-}

----------------------------------------------------------------------
-- Entropy
----------------------------------------------------------------------

-- | FIXME: document this!
class Entropy d where
  entropy :: d -> Double

instance Entropy Normal where
  entropy (Normal _mu sigma) = k + log sigma
    where
    k = 0.5 * (1 + log (2 * pi))

----------------------------------------------------------------------
-- * The ReprDistribution stuff
----------------------------------------------------------------------

-- | This defines the support type of a distribution
type family ReprSupport (d :: (* -> *) -> *) :: *

data ReprNormal f = ReprNormal (f R) (f R)
type instance ReprSupport ReprNormal = R

data ReprUniform f = ReprUniform (f R) (f R)
type instance ReprSupport ReprUniform = R

data ReprCauchy f = ReprCauchy
type instance ReprSupport ReprCauchy = R

data ReprCategorical f = ReprCategorical [f Prob]
type instance ReprSupport ReprCategorical = Int

-- FIXME: this uses different vector and matrix types than MVNormal
data ReprMVNormal f = ReprMVNormal (f RMatrix) (f RMatrix)
type instance ReprSupport ReprMVNormal = RMatrix