text:
{-# OPTIONS_GHC -fno-warn-dodgy-exports #-}
{- | an incomplete extension of dimensional-tf to work with hmatrix and ad
Haddocks are organized to follow hmatrix

note: for subscripting, use the 'HNat' while the units still use 'N.NumType'

TODO:

*  Friendly syntax for introduction (more matD)

*  Friendly syntax for elimination (matD as pattern)

*  A pretty-printer that multiplies out the dimensions at each place?
   The current show instance makes you mentally multiply out row and column units
   (which helps you to see the big picture, but may be more work in other cases)

*  check that all types that could could be inferred are

*  default columns/rows to dimensionless?

*  missing operations (check export list comments)
-}
module DimMat.Internal  where

import Foreign.Storable (Storable)
import GHC.Exts (Constraint)

-- import Data.Type.Equality (type (==))
import qualified Prelude as P
import Prelude (Double)
import Numeric.Units.Dimensional hiding (Mul, Div)
import Numeric.Units.Dimensional.Prelude hiding (Mul, Div)
import qualified Numeric.Units.Dimensional as D

import qualified Numeric.LinearAlgebra as H
import qualified Numeric.LinearAlgebra.LAPACK as H

import Text.PrettyPrint.ANSI.Leijen
import Data.List (transpose)

import Data.HList.CommonMain hiding (MapFst)

-- * Replacements for Dimensional classes

-- | a version of Numeric.Units.Dimensional.'D.Mul' which
-- requires the arguments to include the 'D.Dim' type constructor
class (D.Mul a b c) => Mul a b c

instance (D.Mul a b c,
       a ~ Dim l m t i th n j,
       b ~ Dim l' m' t' i' th' n' j',
       c ~ Dim l'' m'' t'' i'' th'' n'' j'') =>
       Mul a b c

-- | a version of Numeric.Units.Dimensional.'D.Div' which
-- requires the arguments to include the 'D.Dim' type constructor
class (D.Div a b c) => Div a b c

instance (D.Div a b c,
       a ~ Dim l m t i th n j,
       b ~ Dim l' m' t' i' th' n' j',
       c ~ Dim l'' m'' t'' i'' th'' n'' j'') =>
       Div a b c

-- * AD
{- $ad
TODO: gradients, hessians, etc.

Types for derivative towers can see hlist's @HList\/Data\/HList\/broken\/Lazy.hs@,
but laziness doesn't really make much sense if the @take@ that is eventually used
to get a finite list for printing etc.

Complications include the fact that AD.grad needs a traversable,
but hmatrix stuff is not traversable (due needing Storable). In ipopt-hs
I got around this problem by copying data. Perhaps that is the solution?

> BROKEN
> grad :: (Num a, AreRecips i iinv, H.Element a, Storable a,
>           MapMultEq o iinv r) =>
>         (forall s. (AD.Mode s, H.Container H.Vector (AD.AD s a),
>                     Storable (AD.AD s a), H.Field (AD.AD s a))
>                 => DimMat '[i] (AD.AD s a)
>                 -> Quantity o (AD.AD s a))
>      -> DimMat '[i] a
>      -> DimMat '[r] a
> grad f (DimVec x) = DimMat (H.fromLists [AD.grad (unQty . f . DimVec . H.fromList) (H.toList x)])
>     where unQty (Dimensional a) = a
-}

#ifdef WITH_AD
{- |
>>> let ke velocity = velocity*velocity*(1*~kilo gram)
>>> diff ke (3 *~ (metre/second))
6.0 m^-1 kg^-1 s
-}

diff :: (Num a) =>
        (forall s. AD.Mode s => Dimensional v x (AD.AD s a)
                             -> Dimensional v y (AD.AD s a))
        -> Dimensional v x a -> Dimensional v (Div x y) a
diff f z = Dimensional $ AD.diff (unD . f . Dimensional) (unD z)
    where unD (Dimensional a) = a
#endif

-- * GADT for linear algebra with units

{- | Generalization of 'Dimensional' to matrices and vectors. Units
in each coordinate are known at compile-time. This wraps up HMatrix.

[@ 'DMat' @] the units at coordinate ij are @(r1 ': r)_i (DOne ': c)_j@

[@ 'DVec' @] the units at coordinate i are @(r1 ': r)_i@

[@ DScal @] is the same as Dimensional
-}
data D (sh :: ( *, [[ * ]])) e where
  DMat :: (H.Container H.Matrix e, H.Field e)
    => H.Matrix e -> D '(r1,[r, c]) e
  DVec :: (H.Container H.Vector e, H.Field e)
    => H.Vector e -> D '(r1, '[r]) e
  DScal :: (H.Field e) => e -> D '(r1,'[]) e

instance (Show a, PPUnits sh) => Pretty (D sh a) where
    pretty (DVec v) = case ppUnits (Proxy :: Proxy sh) of
        [rs] -> vcat
             [ dullgreen (string label) <+> string (show e)
                | (e,label) <- H.toList v `zip` pad rs ]
    pretty (DMat m) = case ppUnits (Proxy :: Proxy sh) of
        [rs,cs] ->
            vcat $
            map (hsep . onHead dullgreen) $
            transpose $ map (onHead dullgreen . map string . pad) $
            zipWith (\a b -> a:b)
                ((show (H.rows m) ++ "><"++ show (H.cols m)) : cs) $
            transpose $
            zipWith (\r e -> r : map show e) rs (H.toLists m)
        where
            onHead f (x:xs) = f x : xs
            onHead _ [] = []

instance Pretty (D sh a) => Show (D sh a) where
    showsPrec p x = showsPrec p (pretty x)

-- ** pretty instance
pad :: [String] -> [String]
pad [] = []
pad xs = let
    w = maximum (map length xs)
    in map (\x -> take w $ x ++ replicate w ' ') xs

class PPUnits (sh :: k) where
    ppUnits :: Proxy sh -> [[String]]

instance forall (r1 :: *) (r::[*]) (c :: [*]) l m t i th n j.
      (PPUnits [r1 ': r, c], Show (Quantity r1 Int), PPUnits' c, PPUnits' r,
       r1 ~ Dim l m t i th n j) => PPUnits '(r1, [r,c]) where
    ppUnits _ = ppUnits (Proxy :: Proxy [r1 ': r, DOne ': c])

instance (PPUnits' x, PPUnits xs) => PPUnits (x ': xs) where
    ppUnits _ = ppUnits' (Proxy :: Proxy x) : ppUnits (Proxy :: Proxy xs)
instance PPUnits '[] where
    ppUnits _ = []

class PPUnits' (sh :: [ * ]) where
    ppUnits' :: Proxy sh -> [String]
instance (PPUnits' xs) => PPUnits' (DOne ': xs) where
    ppUnits' _ = "1" : ppUnits' (Proxy :: Proxy xs)
instance (ShowDimSpec x, PPUnits' xs) => PPUnits' (x ': xs) where
    ppUnits' _ = showDimSpec (Proxy :: Proxy x) : ppUnits' (Proxy :: Proxy xs)
instance PPUnits' '[] where
    ppUnits' _ = []

class ShowDimSpec a where
    showDimSpec :: Proxy a -> String

instance (Show (Quantity a Int), Dim l m t i th n j ~ a) => ShowDimSpec a where
    showDimSpec _ = case drop 2 $ show (1 *~ (Dimensional 1 :: Unit a Int)) of
          "" -> "1"
          x -> x

-- * Constraints
{- $justification
A major theme in this library is that type inference goes in whichever direction
it can: in ordinary haskell it is very common for argument types to be determined
by the result types. For example see any code that uses 'Num' or 'Read'.

When we use type families, things look more convenient:

@
data Nat = S Nat | Z

type family Add (a :: Nat) (b :: Nat) :: Nat
type instance Add Z b = b
type instance Add (S a) b = Add a (S b)
@

But ghc is unable to deduce things like @a ~ Z@ given evidence such as @Add Z a ~ Z@.
One way around this is to use @ConstraintKinds@:

@
type AddT (a :: Nat) (b :: Nat) (c :: Nat)
          = (Add a b ~ c, Sub c a ~ b, Sub c b ~ a)
@

Which leads to functions like  @f :: AddT a b ab => ... @. This is bad for a couple reasons:

* the right-hand-side of the type can only mention type variables on the
  left-hand-side.
* the left hand side can only bind type variables. Working around this
  leads to many auxiliary type families such as Fst, Snd and Head, or
  leads to a @type family AddT@

So below many constraints expressed as classes, since they have less
of those limitations.

-}

-- | @a*b = c@ when any are lists
class MultEq (a :: k1) (b :: k2) (c :: k3)

-- instance (Zip3 MultEq aas bbs ccs) => MultEq aas bbs ccs
instance Zip3 Mul as bs cs => MultEq as bs cs
instance Zip1 Mul as b  c  => MultEq as b  c
instance Zip1 Mul bs a  c  => MultEq a  bs c
instance (Zip1 Mul cs aInv b,
          Mul a aInv DOne) => MultEq a  b  cs

instance (SameLength as bs,
        Zip2 Mul as bs c)  => MultEq as bs c
instance (SameLength as cs,
          Zip2 Div as cs bInv,
          Mul b bInv DOne) => MultEq as b  cs
instance (Zip2 Div bs cs aInv,
          SameLength bs cs,
          Mul a aInv DOne) => MultEq a  bs cs
instance Mul a b c => MultEq a b c

-- ** Zip

class (SameLength a b, SameLength b c) =>
    Zip3
      (op :: k -> k -> k -> Constraint)
      (a :: [k])
      (b :: [k])
      (c :: [k])

instance (SameLength aas bbs,
          SameLength ccs bbs,
          op a b c,
          (a ': as) ~ aas,
          (b ': bs) ~ bbs,
          (c ': cs) ~ ccs,
          Zip3 op as bs cs) => Zip3 op aas bbs ccs
instance Zip3 op '[] '[] '[]

class (SameLength a b) =>
    Zip2
      (op :: k -> k -> k -> Constraint)
      (a :: [k])
      (b :: [k])
      (c ::  k)

instance (SameLength aas bbs,
      op a b c,
      (a ': as) ~ aas,
      (b ': bs) ~ bbs,
      Zip2 op as bs c) => Zip2 op aas bbs c

instance Zip2 op '[] '[] c

class Zip1
      (op :: k -> k -> k -> Constraint)
      (a :: [k])
      (b ::  k)
      (c ::  k)

instance ((a ': as) ~ aas,
    op a b c,
    Zip1 op as b c) => Zip1 op aas b c

instance Zip1 op '[] b c

{- | given @ijs :: [[Quantity a]]@ (except the : and [] constructors are
actually (,) and (), ie. a HList that doesn't use the HList constructors),
calculate a @DimMat rowUnits colUnits@, where the outer product of rowUnits
and colUnits gives the units at each index in the ijs.  The first element
of colUnits is DOne.
-}
class (SameLength a ab) => Outer a b ab
instance Outer '[] b '[]
instance (SameLength aas ccs,
          (a ': as) ~ aas,
          (c ': cs) ~ ccs,
          MultEq a b c,
          Outer as b cs)
  => Outer aas b ccs

-- * DimMatFromTuple
class DimMatFromTuple ijs r1 r c e

type family TupleToHListU (a :: *) :: [*]
type instance TupleToHListU (a, b) = () ': TupleToHListU b
type instance TupleToHListU () = '[]

type family TuplesToHListU (a :: *) :: [[*]]
type instance TuplesToHListU (a, b) = TupleToHListU a ': TuplesToHListU b
type instance TuplesToHListU () = '[]

instance (Outer (r1 ': r) (DOne ': c) ijs',
      DMFromTuple1 e ijs ijs',
      SameLength (TuplesToHListU ijs) ijs') => DimMatFromTuple ijs r1 r c e

-- | helper for 'DimMatFromTuple'
type family DMFromTuple1 e b (b' :: [[*]]) :: Constraint
type family DMFromTuple2 e b (b' :: [*]) :: Constraint
type family DMFromTuple3 e b b' :: Constraint

type instance DMFromTuple3 e (Quantity b e') b' = (e ~ e', b ~ b')
type instance DMFromTuple1 e (x, xs) (x' ': xs') = (TupleToHListU x `SameLength` x',
                                                    DMFromTuple2 e x x', DMFromTuple1 e xs xs')
type instance DMFromTuple1 e () xs = (xs ~ '[])
type instance DMFromTuple2 e (x, xs) (x' ': xs') = (DMFromTuple3 e x x', DMFromTuple2 e xs xs')
type instance DMFromTuple2 e () xs = (xs ~ '[])

-- | just for types produced by the matD quasiquote
toDM :: DimMatFromTuple ijs r1 r c e => ijs -> Proxy (D '(r1, [r, c]) e)
toDM _ = Proxy

-- | @InnerCxt t a b = t ~ 'H.dot' a b@
type family InnerCxt (t :: k) (a :: [k]) (b :: [k]) :: Constraint
type instance InnerCxt t (a ': as) (b ': bs) = (MultEq a b t, InnerCxt t as bs)
type instance InnerCxt t '[] '[] = ()

class (SameLength a b, InnerCxt c a b) => Inner (a :: [*]) (b :: [*]) (c :: *)

instance (SameLength aas bbs, InnerCxt c aas bbs) => Inner aas bbs c

-- | @ProdEq a b@ is @product a ~ b@
class ProdEq a b
instance (ProdEq as b', Mul a b' b) => ProdEq (a ': as) b
instance (dOne ~ DOne) => ProdEq '[] dOne

-- | @RecipEq a aInv@ is @a*aInv ~ DOne@ (or a list of DOne)
class RecipEq (a :: k) (aInv :: k)
instance (MultEq as aInvs DOne) => RecipEq as aInvs

-- | @AtEq a n b m c@ calculates @(At a n \`Mult\` At b m) ~ c@,
-- but also can infer part of the `a` if the `b` and `c` are known
type family AtEq2 (a :: [k]) (n :: HNat) (b :: [k]) (m :: HNat) (c :: k) :: Constraint
type instance AtEq2  (a ': as) HZero (b ': bs) HZero c = (MultEq a b c)
type instance AtEq2  (a ': as) (HSucc n) bs m c = AtEq2 as n bs m c
type instance AtEq2  as HZero (b ': bs) (HSucc m) c = AtEq2 as HZero bs m c

type family AtEq (a :: [k]) (n :: HNat) (b :: k) :: Constraint
type instance AtEq (a ': as) HZero b = (a ~ b)
type instance AtEq (a ': as) (HSucc n) b = AtEq as n b

-- | Data.Packed.Vector.'H.@>'
(@>) :: (HNat2Integral i, AtEq r i ri, MultEq r1 ri u)
    => D '(r1,'[r]) a
    -> Proxy i
    -> Quantity u a
DVec v @> i = Dimensional (v H.@> hNat2Integral i)

-- | Data.Packed.Matrix.'H.@@>'
(@@>) :: (HNat2Integral i, HNat2Integral j, AtEq2 (r1 ': r) i (DOne ': c) j ty)
    => D '(r1, [r,c]) a
    -> (Proxy i, Proxy j)
    -> Quantity ty a
DMat m @@> (i,j) = Dimensional (m H.@@> (hNat2Integral i,hNat2Integral j))

pnorm :: (AllEq r1 r, AllEq DOne c)
         => H.NormType -> D '(r1, [r, c]) a -> Quantity r1 (H.RealOf a)
pnorm normType (DMat a) = Dimensional (H.pnorm normType a)

{- | @AllEq a xs@ is like @all (all (a==)) xs@, @all (a ==) xs@, @a == xs@:
whichever amount of [ ] is peeled off before making the comparison (with ~)
-}
class AllEq (a :: k1) (xs :: k2)

instance (a ~ x, AllEq a xs) => AllEq a (x ': xs)
instance AllEq a '[]
instance AllEq '[] xs
instance (AllEq a xs, AllEq as xs) => AllEq (a ': as) xs

{- | @c = a `dot` b@ is one of:

> c_ij = sum_j a_ij b_jk
> c_k  = sum_j a_j  b_jk
> c_i  = sum_j a_ij b_j
> c    = sum_j a_j  b_j

-}
class Dot a b c {- | a b -> c -} where
    dot :: H.Element e => D a e -> D b e -> D c e

instance
    ( MultEq (a ': ra) b (c ': rc),
      MultEq ca rb b,
      shA ~ '(a,[ra, ca]),
      shB ~ '(b, rb ': cb),
      shC ~ '(c, rc ': cb))
    => Dot shA shB shC where
    dot (DMat a) (DMat b) = DMat (H.multiply a b)
    dot (DMat a) (DVec b) = DVec (H.mXv a b)
    {-
    dot (DVec a) (DMat b) = DVec (H.vXm a b)
    dot (DVec a) (DVec b) = DScal (H.dot a b)
    -}

class DotS s a b

instance (rest' ~ rest, MultEq s (a ': as) (b ': bs))
    => DotS s '(a, as ': rest) '(b, bs ': rest')

class Trans a b where
    trans :: D a e -> D b e

instance (a ~ '(r1, [x,y]),
          b ~ '(r1, [y,x]))
    => Trans a b where
    trans (DMat a) = DMat (H.trans a)

{- | type for a pseudoinverse (and inverse):

The single instance comes from looking at inverses from a 2x2 matrix (let's call A):

> a b
> c d

and the inverse * determinant of the original

>  d  -b
> -c   a

In the product A * A^-1 the diagonal is dimensionless ('DOne').

That happens if the row and column type-level unit lists are reciprocals of
eachother ('AreRecips'), so the constraint on the instance of PInv encodes
this exactly (plus some constraints requiring that sh and sh' are at least
1x1)
-}
class PInv a b where
  pinv :: D a e -> D b e

type family LengthSndTwo (a :: k) :: Constraint
type instance LengthSndTwo '(a, as) = SameLength as '[(), ()]

type AreRecips a b = MultEq a b DOne
instance (MultEq ra cb a,
          MultEq ca rb b,
          AreRecips a b,
          bSh ~ '(b, [rb, cb]),
          aSh ~ '(a, [ra, ca]))
    => PInv aSh bSh where
  pinv (DMat a) = DMat (H.pinv a)

inv :: (PInv a b,
        b ~ '(t1, [t2, t3]))
  => D a e -> D b e
inv (DMat a) = DMat (H.inv a)

pinvTol :: (PInv a b,
           b ~ '(t1, [t2, t3]) )
  => Double -> D a e -> D b e
pinvTol tol (DMat a) = DMat (H.pinvTol tol a)

class Det a b where
    det :: D a e -> Quantity b e

instance (SameLength r c,
       ProdEq (r1 ': r) (pr :: *),
       ProdEq c (pc :: *),
       MultEq pr pc b,
       a ~ '(r1, [r, c])) =>
    Det a b where
  det (DMat a) = Dimensional (H.det a)

{- | Numeric.LinearAlgebra.Algorithms.'H.expm'

@y t = expm (scale t a) \`multiply\` y0@ solves the DE @y' = Ay@ where y0 is the
value of y at time 0

-}
expm :: (AreRecips r c)
    => D '(r1, [r, c]) a
    -> D '(r1, [r, c])  a
expm (DMat a) = DMat (H.expm a)

{- | Numeric.Container.'H.scale'

-}
class Scale a b c where
  scale :: Quantity a e -> D b e -> D c e

fromQty :: H.Field e => Quantity a e -> D '(a, '[]) e
fromQty (Dimensional a) = DScal a

toQty :: D '(a, '[]) e -> Quantity a e
toQty (DScal a) = Dimensional a

instance (MultEq a (r1 ': r) (r1' ': r'),
      b ~ '(r1, r ': rs),
      c ~ '(r1', r' ': rs)) =>
  Scale a b c where
  scale (Dimensional t) (DMat a) = DMat (H.scale t a)
  scale (Dimensional t) (DVec a) = DVec (H.scale t a)

{- | Numeric.Container.'H.scaleRecip'
-}
class ScaleRecip a b c where
  scaleRecip :: D '(a, '[]) e -> D b e -> D c e

class ScaleRecip1 (bool :: Bool) a b c where
  scaleRecip1 :: Proxy bool -> D '(a, '[]) e -> D b e -> D c e

instance
  (ScaleRecipCxt r1 r1' r r' rs rs' a b c
   , rs' ~ '[ t1 ]
   ) => ScaleRecip1 True a b c where
  scaleRecip1 _ (DScal t) (DMat a) = DMat (H.scaleRecip t a)

instance
  (ScaleRecipCxt r1 r1' r r' rs rs' a b c,
   rs' ~ '[]) =>
  ScaleRecip1 False a b c where
  scaleRecip1 _ (DScal t) (DVec a) = DVec (H.scaleRecip t a)

instance (HEq (HLength bs) (HSucc (HSucc HZero)) bool1,
          HEq (HLength cs) (HSucc (HSucc HZero)) bool1,
    ScaleRecip1 bool1 a '(b, bs) '(c, cs)) => ScaleRecip a '(b, bs) '(c, cs) where
  scaleRecip = scaleRecip1 (Proxy :: Proxy bool1)

type ScaleRecipCxt (r1 :: *) (r1' :: *) r r' rs rs' (a :: *) b c =
  (MultEq a (r1' ': r') (r1 ': r) ,
   MultEq rs rs' DOne,
   b ~ '(r1, r ': rs),
   c ~ '(r1', r' ': rs'))

-- | a shortcut for @scaleRecip (DScal 1)@
recipMat :: forall b c e. (H.Field e, ScaleRecip DOne b c) => D b e -> D c e
recipMat m = scaleRecip (DScal 1 :: D '(DOne, '[]) e) m

liftH2 ::
  ( forall h f. (H.Container f e, h ~ f e) => h -> h -> h) ->
    D a e -> D a e -> D a e
liftH2 f (DMat a) (DMat b) = DMat (f a b)
liftH2 f (DVec a) (DVec b) = DVec (f a b)

add a b = liftH2 H.add a b
sub a b = liftH2 H.sub a b

mulMat :: ( MultEq as bs cs, MultEq a b c, cs ~ '[t1 , t2] )
  => D '(a,as) e -> D '(b,bs) e -> D '(c,cs) e
mulMat (DMat a) (DMat b) = DMat (H.mul a b)

mulVec :: ( MultEq as bs cs, MultEq a b c, cs ~ '[t1] )
  => D '(a,as) e -> D '(b,bs) e -> D '(c,cs) e
mulVec (DVec a) (DVec b) = DVec (H.mul a b)

divideMat :: ( MultEq as cs bs, MultEq a c b, cs ~ '[t1 , t2] )
  => D '(a,as) e -> D '(b,bs) e -> D '(c,cs) e
divideMat (DMat a) (DMat b) = DMat (H.divide a b)

divideVec :: ( MultEq as cs bs, MultEq a c b, cs ~ '[t1] )
  => D '(a,as) e -> D '(b,bs) e -> D '(c,cs) e
divideVec (DVec a) (DVec b) = DVec (H.divide a b)

arctan2 :: (bs ~ MapMapConst DOne as) => D '(a,as) e -> D '(a,as) e -> D '(b,bs) e
arctan2 (DMat a) (DMat b) = DMat (H.arctan2 a b)
arctan2 (DVec a) (DVec b) = DVec (H.arctan2 a b)

equal :: D a e -> D a e -> Bool
equal (DMat a) (DMat b) = H.equal a b
equal (DVec a) (DVec b) = H.equal a b

{- | @cmap f m@ gives a matrix @m'@

@f@ is applied to

-}
class CMap f a b e e' where
    cmap :: f -> D a e -> D b e'

-- | Maybe there's a way to implement in terms of the real cmap (possibly
-- unsafeCoerce?)
instance
    (ToHLists sh e xs,
     FromHLists xs' sh' e',
     SameLength xs xs',
     HMapCxt HList (HMap f) xs xs') =>
    CMap f sh sh' e e' where
  cmap f m = fromHLists (HMap f `hMap` (toHLists m :: HList xs) :: HList xs')

type family AppendEq' (a :: [k]) (b :: [k]) (ab :: [k]) :: Constraint
type instance AppendEq' (a ': as) b (a' ': abs) = (a ~ a', AppendEq' as b abs)
type instance AppendEq' '[] b abs = (b ~ abs)

-- | a bit overkill?
--  @a ++ b = ab@
type AppendEq a b ab =
   (ab ~ HAppendR a b,
    AppendEq' a b ab,
    SameLength (DropPrefix a ab) b,
    SameLength (DropPrefix b ab) a)

type instance HAppendR (x ': xs) ys = x ': HAppendR xs ys
type instance HAppendR '[] ys = ys

type family DropPrefix (a :: [k]) (ab :: [k]) :: [k]
type instance DropPrefix (a ': as) (a' ': abs) = DropPrefix as abs
type instance DropPrefix '[] bs = bs

{- | the slightly involved type here exists because
ci1 and ci2 both start with DOne, but ci2's contribution
to ci3 does not need a DOne at the start. Another way to
read the constraints here is:

> map (*rem) (a11 : ri) = b11 : bi
> ci3 = ci1 ++ map (*rem) ci2

The same idea happens with vconcat.
-}
hconcat ::
    ( MultEq (rem :: *) a b,
      MultEq rem ra rb,
      MultEq rem (DOne ': cb) cb',
      AppendEq ca cb' cc ) =>
    D '(a, [ra,ca]) e -> D '(b, [rb, cb]) e -> D '(a, [ra, cc]) e
hconcat (DMat a) (DMat b) = DMat (H.fromBlocks [[a, b]])

vconcat :: (AppendEq ra (b ': rb) rc) =>
    D '(a, '[ra,ca]) e -> D '(b, '[rb,ca]) e -> D '(a, '[rc,ca]) e
vconcat (DMat a) (DMat b) = DMat (H.fromBlocks [[a],[b]])

rank, rows, cols :: D t a -> Int
rank (DMat a) = H.rank a
rows (DMat a) = H.rows a
cols (DMat a) = H.cols a

-- | H.'H.rows' except type-level
rowsNT :: D '(a, r ': c) e -> Proxy (HLength (a ': ri))
rowsNT _ = Proxy

-- | H.'H.cols' except type-level
colsNT :: D '(a, r ': c ': cs) e -> Proxy (HLength (DOne ': c))
colsNT _ = Proxy

-- | (m `hasRows` n) constrains the matrix/vector @m@ to have @n@ rows
hasRows :: (SameLength (HReplicateR n ()) r, -- forwards
            HLength r ~ n -- backwards
    ) => D '(a, ra ': ca) e -> Proxy (n :: HNat) -> D '(a, ra ': ca) e
hasRows x _ = x

-- | (m `hasRows` n) constrains the matrix/vector @m@ to have @n@ rows
hasCols :: (SameLength (HReplicateR n ()) ci, -- forwards
            HLength ci ~ n -- backwards
    ) => D '(a, ra ': ca ': rest) e -> Proxy (n :: HNat) -> D '(a, ra ': ca ': rest) e
hasCols x _ = x

-- | H.'H.scalar'
class (MapConst '[] as ~ as) => Scalar as where
    scalar :: D '(a, '[]) e -> D '(a, as) e

instance Scalar '[ '[] ] where
  scalar (DScal a) = DVec (H.scalar a)

instance Scalar '[ '[], '[] ] where
  scalar (DScal a) = DMat (H.scalar a)

{- | Numeric.Container.'H.konst', but the size is determined by the type.

>>> let n = hSucc (hSucc hZero) -- 2
>>> konst ((1::Double) *~ second) `hasRows` n `hasCols` n
2><2 1   1
s    1.0 1.0
s    1.0 1.0

-}
konst :: forall e a ra ca.
    (H.Field e,
     HNat2Integral (HLength (a ': ra)),
     HNat2Integral (HLength (DOne ': ca)),
     AllEq DOne ca,
     AllEq a ra)
    => D '(a, '[]) e -> D '(a, '[ra, ca]) e
konst (DScal a) = DMat (H.konst a
    (hNat2Integral (Proxy :: Proxy (HLength (a ': ra))),
     hNat2Integral (Proxy :: Proxy (HLength (DOne ': ca)))))

-- | identity matrix. The size is determined by the type.
ident :: forall ones e.
    (H.Field e, HNat2Integral (HLength (DOne ': ones))) =>
    D '(DOne, [ones, ones]) e
ident = DMat (H.ident (hNat2Integral (Proxy :: Proxy (HLength (DOne ': ones)))))

-- | zero matrix. The size and dimension is determined by the type.
zeroes :: forall c a r e. (H.Field e,
                        HNat2Integral (HLength (a ': r)),
                        HNat2Integral (HLength (DOne ': c)))
    => D '(a, '[r, c]) e
zeroes = DMat (H.konst 0
        (hNat2Integral (Proxy :: Proxy (HLength (a ': r))),
         hNat2Integral (Proxy :: Proxy (HLength (DOne ': c)))))

type family CanAddConst (a :: k) (m :: [[k]]) :: Constraint
type instance CanAddConst a [as, ones] = (AllEq a as, AllEq '[] ones)
type instance CanAddConst a '[as] = (AllEq a as)

addConstant :: (H.Field e, CanAddConst a sh)
    => D '(a, '[]) e
    -> D '(a, sh) e
    -> D '(a, sh) e
addConstant (DScal a) (DMat b) = DMat (H.addConstant a b)
addConstant (DScal a) (DVec b) = DVec (H.addConstant a b)

conj :: D sh a -> D sh a
conj (DMat a) = DMat (H.conj a)
conj (DVec a) = DVec (H.conj a)

-- | conjugate transpose
ctrans x = conj . trans $ x

diag :: (MapConst DOne r ~ c, SameLength r c)
  => D '(a, '[r]) t -> D '(a, '[r,c]) t
diag (DVec a) = DMat (H.diag a)

-- | 'H.blockDiag'. The blocks should be provided as:
--
-- @blockDiag $ 'hBuild' m1 m2 m3@
--
-- only available if hmatrix >= 0.15
diagBlock :: (HMapOut UnDimMat (b ': bs) (H.Matrix e),
              Num e, H.Field e, AppendShOf b bs (D '(a, sh) e),
              sh ~ '[r,c])
  => HList (b ': bs)
  -> D '(a, sh) e
diagBlock pairs = DMat (H.diagBlock (hMapOut UnDimMat pairs))

data UnDimMat = UnDimMat
instance (D sh a ~ x, H.Matrix a ~ y) => ApplyAB UnDimMat x y where
        applyAB _ (DMat x) = x

class DiagBlock (bs :: [*]) t

-- | @AppendShOf a [b,c,d] aas@ makes aas have the type of a matrix that
-- has a,b,c,d along the diagonal
class AppendShOf a (as :: [*]) aas
instance
 (e ~ f, f ~ g,
  AppendShOf (D ab e) ds z,
  AppendDims a b ab,

  -- constraints to force D in the type
  x ~ D a e,
  y ~ D b f,
  z ~ D c g) =>
  AppendShOf x (y ': ds) z
instance (x ~ z) => AppendShOf x '[] z

class AppendDims (a :: (*, [[*]])) (b :: (*, [[*]])) (c :: (*, [[*]]))
    | a b -> c, c a -> b, c b -> a
instance (c ~ a, AppendEq ra (b ': rb) rc, AppendEq ca cb cc) =>
  AppendDims '(a, [ra,ca]) '(b, [rb,cb]) '(c, [rc,cc])
-- how to handle vectors?
--type instance AppendDims '(a, '[ra]) '(b, '[rb]) = '(a, '[HAppendR ra (b ': rb)])

class ToHList sh e result where
    toHList :: D sh e -> HList result

-- | given a vector like @x = DimMat '[units] e@ this does something like
-- @[ (x \@> i) | i <- [1 .. n] ]@, if we had comprehensions for HLists
instance
    (HListFromList e e1,
     SameLength result e1,
     HMapCxt HList AddDimensional e1 result,
     -- HMapAddDimensional result e1,
     ToHListRow (r ': rs) e result) =>
   ToHList '(r, '[rs]) e result where
  toHList (DVec v) = case hListFromList (H.toList v) :: HList e1 of
      e1 -> hMap AddDimensional e1

class HMapAddDimensional (a :: [*]) (b :: [*]) | a -> b
instance HMapAddDimensional '[] '[]
instance (HMapAddDimensional as bs,
    dta ~ Quantity t b) => HMapAddDimensional (dta ': as) (b ': bs)

class HListFromList e e' | e' -> e where
        hListFromList :: [e] -> HList e'

instance HListFromList e '[e] where
        hListFromList (e : _) = HCons e HNil
        hListFromList _ = error "HListFromList: not enough input"
instance (e ~ e', HListFromList e es) => HListFromList e (e' ': es) where
        hListFromList (e : es) = e `HCons` hListFromList es

type family ToHListRow (a :: [*]) e (b :: [*]) :: Constraint
type instance ToHListRow (a ': as) e (b ': bs) = (Quantity a e ~ b, ToHListRow as e bs)

data AddDimensional = AddDimensional
instance (Quantity t x ~ y) => ApplyAB AddDimensional x y where
        applyAB _ x = Dimensional x

class FromHList list sh e where
  fromHList :: HList list -> D sh e

instance
    (H.Field e,
     HMapOut RmDimensional list e,
     ToHListRow (r ': rs) e list) =>
  FromHList list '(r, '[rs]) e where
  fromHList xs = DVec (H.fromList (hMapOut RmDimensional xs))

data RmDimensional = RmDimensional
instance (x ~ Quantity d y) => ApplyAB RmDimensional x y where
        applyAB _ (Dimensional a) = a

class FromHLists lists sh e where
  fromHLists :: HList lists -> D sh e

-- | [[Dim e unit]] -> DimMat units e
instance
  (ToHListRows' (r1 ': r) c e lists,
   HMapOut (HMapOutWith RmDimensional) lists [e],
   H.Field e) =>
  FromHLists lists '(r1, [r,c]) e where
    fromHLists xs = DMat (H.fromLists (hMapOut (HMapOutWith RmDimensional) xs))

newtype HMapOutWith f = HMapOutWith f
instance (HMapOut f l e, es ~ [e], HList l ~ hl) => ApplyAB (HMapOutWith f) hl es where
    applyAB (HMapOutWith f) = hMapOut f

class ToHListRows' (r :: [*]) (c :: [*]) (e :: *) (rows :: [*])
instance ToHListRows' '[] c e '[]

instance (ToHListRows' r c e rows,
          MultEq r c c',
          HMapCxt HList (AddQty e) c' row')
  => ToHListRows' (r1 ': r) c e (hListRow ': rows)

data AddQty u
instance (qty ~ Quantity u e) => ApplyAB (AddQty u) e qty where
    applyAB _ = Dimensional

class ToHLists sh e xs where
    toHLists :: D sh e -> HList xs

-- | DimMat units e -> [[Dim e unit]]
instance
    (HListFromList e e1,
     HListFromList (HList e1) e2,
     HMapCxt HList (HMap AddDimensional) e2 xs,
     ToHListRows' ri ci e xs,
     SameLength e2 xs,
     (r1 ': r) ~ ri, (DOne ': c) ~ ci )
  => ToHLists '(r1, [r,c]) e xs where
  toHLists (DMat m) = case hListFromList (map hListFromList (H.toLists m) :: [HList e1]) :: HList e2 of
    e2 -> hMap (HMap AddDimensional) e2

{- still bad

class PairsToList a t where
        pairsToList :: a -> [H.Matrix t]
instance PairsToList () t where
        pairsToList _ = []
instance (PairsToList b t, t' ~ t) => PairsToList (DimMat sh t',b) t where
        pairsToList (DimMat a,b) = a : pairsToList b

class EigV (sh :: [[ [DimSpec *] ]])
           (eigenValue  :: [[DimSpec *]])
           (eigenVector :: [[[DimSpec *]]])

instance
  ( SameLengths [r,c,r',c',rinv,cinv,eigval,erinv],
    -- ZipWithMul r c eigval,
    MapConst '[] r ~ eigval,
    PInv [r',c'] [rinv,cinv],
    -- AreRecips r' cinv,
    -- AreRecips c' rinv,
    cinv ~ c,
    c ~ ('[] ': _1),
    c' ~ ('[] ': _2),
    ZipWithMul eigval rinv erinv,
    MultiplyCxt [r',c'] erinv r,
    sh ~ [r,c],
    sh' ~ [r',c'])
    =>  EigV sh eigval sh'
-- | when no eigenvectors are needed
type family EigE (sh :: [[ [DimSpec *] ]]) (eigenValue  :: [ [DimSpec *] ]) :: Constraint
type instance EigE [r,c] eigval = ( SameLengths [r,c,eigval], ZipWithMul r c eigval)

{- $eigs

The Hmatrix eig factors A into P and D where A = P D inv(P) and D is diagonal.

The units for eigenvalues can be figured out:

>               _____
>      -1       |  c
> P D P  = A =  |r
>               |

>       _______
>       |   d
> P   = |c
>       |

>       _______
>       |   -1
>       |  c
>  -1   |
> P   = | -1
>       |d

So we can see that the dimension labeled `d-1` in P inverse is actually the
same `c` in `A`. The actual units of `d` don't seem to matter because the
`inv(d)` un-does any units that the `d` adds. So `d` can be all DOne. But
another choice, such as 1/c would be more appropriate, since then you can
expm your eigenvectors (not that that seems to be something people do)?

To get the row-units of A to match up, sometimes `D` will have units.
The equation ends up as D/c = r

Please ignore the type signatures on 'eig' 'eigC' etc. instead look at the type of
'wrapEig' 'wrapEigOnly' together with the hmatrix documentation (linked).

Perhaps the convenience definitions `eig m = wrapEig H.eig m` should be in
another module.
-}

{-
-- | 'wrapEig' H.'H.eig'
eig m = wrapEig H.eig m
-- | 'wrapEig' H.'H.eigC'
eigC m = wrapEig H.eigC m
-- | 'wrapEig' H.'H.eigH'
eigH m = wrapEig H.eigH m
-- | 'wrapEig' H.'H.eigH''
eigH' m = wrapEig H.eigH' m
-- | 'wrapEig' H.'H.eigR'
eigR m = wrapEig H.eigR m
-- | 'wrapEig' H.'H.eigS'
eigS m = wrapEig H.eigS m
-- | 'wrapEig' H.'H.eigS''
eigS' m = wrapEig H.eigS' m
-- | 'wrapEig' H.'H.eigSH'
eigSH m = wrapEig H.eigSH m
-- | 'wrapEig' H.'H.eigSH''
eigSH' m = wrapEig H.eigSH' m

-- | 'wrapEigOnly' H.'H.eigOnlyC'
eigOnlyC m = wrapEigOnly H.eigOnlyC m
-- | 'wrapEigOnly' H.'H.eigOnlyH'
eigOnlyH m = wrapEigOnly H.eigOnlyH m
-- | 'wrapEigOnly' H.'H.eigOnlyR'
eigOnlyR m = wrapEigOnly H.eigOnlyR m
-- | 'wrapEigOnly' H.'H.eigOnlyS'
eigOnlyS m = wrapEigOnly H.eigOnlyS m
-- | 'wrapEigOnly' H.'H.eigenvalues'
eigenvalues m = wrapEigOnly H.eigenvalues m
-- | 'wrapEigOnly' H.'H.eigenvaluesSH'
eigenvaluesSH m = wrapEigOnly H.eigenvaluesSH m
-- | 'wrapEigOnly' H.'H.eigenvaluesSH''
eigenvaluesSH' m = wrapEigOnly H.eigenvaluesSH' m
-}

wrapEig :: (c' ~ ('[] ': _1),
            EigV [r,c] eigVal [r',c'],
    H.Field y, H.Field z)
    => (H.Matrix x -> (H.Vector y, H.Matrix z)) ->
    DimMat [r,c] x ->
    (DimMat '[eigVal] y, DimMat [r',c'] z)
wrapEig hmatrixFun (DimMat a) = case hmatrixFun a of
    (e,v) -> (DimVec e, DimMat v)

wrapEigOnly :: (EigE [r,c] eigVal, H.Field y)
    => (H.Matrix x -> H.Vector y) ->
    DimMat [r,c] x -> DimMat '[eigVal] y
wrapEigOnly hmatrixFun (DimMat a) = case hmatrixFun a of
    (e) -> DimVec e

-}

-- | @\\a xs -> map (map (const a)) xs@
type family MapMapConst (a::k) (xs :: [[l]]) :: [[k]]
type instance MapMapConst a (x ': xs) = MapConst a x ': MapMapConst a xs
type instance MapMapConst a '[] = '[]

-- | @\\a xs -> map (const a) xs@
type family MapConst (a :: k) (xs :: [l]) :: [k]
type instance MapConst a (x ': xs) = a ': MapConst a xs
type instance MapConst a '[] = '[]

-- | convert from (a,(b,(c,(d,())))) to '[a,b,c,d]
type family FromPairs (a :: *) :: [*]
type instance FromPairs (a,b) = a ': FromPairs b
type instance FromPairs () = '[]

{-
-- | @map fst@
type family MapFst (a :: *) :: [*]
type instance MapFst ((a,_t) , as) = a ': MapFst as
type instance MapFst () = '[]

-- | @\\a xs -> map (/a) xs@
type family MapDiv (a :: k) (xs :: [k]) :: [k]
type instance MapDiv a (x ': xs) = (x @- a) ': MapDiv a xs
type instance MapDiv a '[] = '[]

type family UnDQuantity (a :: [*]) :: [ [*] ]
type instance UnDQuantity (x ': xs) = UnDQuantity1 x ': UnDQuantity xs
type instance UnDQuantity '[] = '[]

type family UnDQuantity1 (a :: *) :: [*]
type instance UnDQuantity1 (Unit t x) = x

type family DimMatFromTuple ijs :: * -> *
type instance DimMatFromTuple ijs =
        DimMat [UnDQuantity (MapFst ijs),
               '[] ': MapDiv (UnDQuantity1 (Fst (Fst ijs)))
               (UnDQuantity (FromPairs (Snd (Fst ijs))))]
type family Append (a :: [k]) (b :: [k]) :: [k]
type instance Append (a ': as) b = a ': Append as b
type instance Append '[] b = b
-}

max_stars_repo_name: aavogt/DimMat

max_issues_repo_name: aavogt/DimMat

max_forks_repo_name: aavogt/DimMat


text:
{-# LANGUAGE DeriveGeneric #-}

-- import BDSCOD.Conditioning
import BDSCOD.Aggregation
import qualified BDSCOD.InhomogeneousBDSLlhd as InhomBDSLlhd
import BDSCOD.Llhd
import BDSCOD.Types
import BDSCOD.Utility
import Control.Monad (replicateM)
import Data.List (sort)
import Data.Maybe (fromJust, isJust)
import Data.Tuple (swap)
import qualified Data.Vector.Unboxed as Unboxed
import qualified Epidemic as EpiSim
import qualified Epidemic.BirthDeathSampling as EpiBDS
import qualified Epidemic.BDSCOD as EpiBDSCOD
import Epidemic.Types.Events
import Epidemic.Types.Parameter
import Epidemic.Types.Population
import qualified Epidemic.Utility as EpiUtil
import GHC.Generics
import Generic.Random (genericArbitraryU)
import Numeric.LinearAlgebra.HMatrix
import qualified System.Random.MWC as MWC
import Test.Hspec
import Test.Hspec.QuickCheck
import Test.QuickCheck
import Test.QuickCheck.Gen
import Test.Hspec.Core.QuickCheck (modifyMaxDiscardRatio)
import Data.Either.Combinators (fromRight', isLeft)

-- | Check if @y@ is withing @delta@ of @x@
withinDeltaOf :: (Ord a, Num a)
              => a -- ^ delta
              -> a -- ^ y
              -> a -- ^ x
              -> Bool
withinDeltaOf delta y x = abs (y - x) < delta

-- | Apply the @withinDeltaOf@ function to two lists.
allWithinDeltaOf :: (Ord a, Num a) => a -> [a] -> [a] -> Bool
allWithinDeltaOf _ [] [] = True
allWithinDeltaOf delta [y] [x] = withinDeltaOf delta y x
allWithinDeltaOf delta (y:ys) (x:xs) = withinDeltaOf delta y x && allWithinDeltaOf delta ys xs
allWithinDeltaOf _ _ _ = False

unsafeNBFromMAndV x = fromRight' $ nbFromMAndV x

-- | Approximate the derivative of @f@ at @x@ with a step of size @h@.
finiteDifference :: Fractional a
                 => a         -- ^ h
                 -> (a -> a)  -- ^ f
                 -> a         -- ^ x
                 -> a
finiteDifference h f x = (f (x+h) - f (x-h)) / (2*h)

testTestingHelpers =
  describe "Testing the helper functions for the testing suite" $ do
    context "withinDeltaOf" $ do
      it "sanity" $ do
        withinDeltaOf 0.2 1 2 `shouldBe` False
        withinDeltaOf 0.2 1 1.1 `shouldBe` True
      it "identity" $ property $
        \x -> withinDeltaOf 0.2 x (x :: Double)
      it "range" $ property $
        \x -> withinDeltaOf 1 1 (1 + (x :: Double)) || not (x < 1 && x > (-1))
    context "finiteDifference" $ do
      it "sine" $ property $
        \x -> let fFD = finiteDifference 0.01 sin
                  f' = cos
                in withinDeltaOf 1e-3 (fFD (x :: Double)) (f' x)
      it "polynomial" $ property $
        \x -> let fFD = finiteDifference 0.01 (\z -> z + 0.5 * z ** 2)
                  f' = \z -> 1 + z
                in withinDeltaOf 1e-3 (fFD (x :: Double)) (f' x)

testLogSumExp =
  describe "Test the log-sum-exp function" $
    it "equivalence to unsafe implementation" $ property $
      \x -> (not $ null x) ==> withinDeltaOf 1e-6 (logSumExp x) (log (sum [exp x' | x' <- x :: [Double]]))

testNbPGF = do
  describe "Test nbPGF: 1" $ do
    it "known value of PGF is correct 1" $
      nbPGF Zero 0.0 `shouldBe` 1

    it "known value of PGF is correct 2" $
      nbPGF Zero 0.5 `shouldBe` 1

    it "known value of PGF is correct 3" $
      nbPGF Zero 1.0 `shouldBe` 1

    it "known value of PGF is correct 4" $
      nbPGF (NegBinomSizeProb 1 0.5) 0.0 `shouldBe` 0.5

    it "known value of PGF is correct 5" $
      nbPGF (NegBinomSizeProb 1 0.5) 0.5 `shouldSatisfy` (withinDeltaOf 1e-6 (2.0 / 3.0))

    it "known value of PGF is correct 5" $
      nbPGF (NegBinomSizeProb 1 0.5) 1.0 `shouldBe` 1.0

    it "PGF partial derivative seems correct 1" $
      nbPGF' (NegBinomSizeProb 1 0.5) 1.0 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-3 (\x -> nbPGF (NegBinomSizeProb 1 0.5) x) 1.0))

    it "PGF partial derivative seems correct 2" $
      nbPGF' (NegBinomSizeProb 1 0.5) 0.5 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-3 (\x -> nbPGF (NegBinomSizeProb 1 0.5) x) 0.5))

    it "PGF partial derivative seems correct 3" $
      nbPGF' (NegBinomSizeProb 1 0.5) 0.0 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-3 (\x -> nbPGF (NegBinomSizeProb 1 0.5) x) 0.0))

    it "PGF second partial derivative seems correct 1" $
      nbPGF'' (NegBinomSizeProb 1 0.5) 1.0 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-3 (\x -> nbPGF' (NegBinomSizeProb 1 0.5) x) 1.0))

    it "PGF second partial derivative seems correct 2" $
      nbPGF'' (NegBinomSizeProb 1 0.5) 0.5 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-5 (\x -> nbPGF' (NegBinomSizeProb 1 0.5) x) 0.5))

    it "PGF second partial derivative seems correct 3" $
      nbPGF'' (NegBinomSizeProb 1 0.5) 0.0 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-5 (\x -> nbPGF' (NegBinomSizeProb 1 0.5) x) 0.0))

  describe "Test nbPGF: 2" $ do
    it "test pochhammer and logPochhammer" $ do
      pochhammer 2 5 > 2 `shouldBe` True
      let pochhammersWorking (a,b) = withinDeltaOf 1e-5 (log $ pochhammer a b) (logPochhammer a b)
      all pochhammersWorking [(a,b) | a <- [1..10], b <- [1..10], a <= b] `shouldBe` True
      all pochhammersWorking [(a+0.1,b) | a <- [1..10], b <- [1..10], a <= b] `shouldBe` True

    it "test nbPGFdash and logNbPGFdash" $ do
      let nbPGFdashWorking (j,r,p,z) = withinDeltaOf 1e-5 (log $ nbPGFdash j (NegBinomSizeProb r p) z) (logNbPGFdash j (NegBinomSizeProb r p) z)
      all nbPGFdashWorking [(j,r,p,z) | j <- [2..50], r <- [2..50], p <- [0.1,0.3,0.5,0.7,0.9], z <- [0.1,0.3,0.5,0.7,0.9]] `shouldBe` True

    it "test nbPGF and logNbPGF" $ do
      let nbPGFWorking (r,p,z) = withinDeltaOf 1e-5 (log $ nbPGF (NegBinomSizeProb r p) z) (logNbPGF (NegBinomSizeProb r p) z)
      all nbPGFWorking [(r,p,z) | r <- [2..50], p <- [0.1,0.3,0.5,0.7,0.9], z <- [0.1,0.3,0.5,0.7,0.9]] `shouldBe` True

    it "test nbPGF' and logNbPGF'" $ do
      let nbPGFWorking (r,p,z) = withinDeltaOf 1e-5 (log $ nbPGF' (NegBinomSizeProb r p) z) (logNbPGF' (NegBinomSizeProb r p) z)
      all nbPGFWorking [(r,p,z) | r <- [2..50], p <- [0.1,0.3,0.5,0.7,0.9], z <- [0.1,0.3,0.5,0.7,0.9]] `shouldBe` True

    it "test nbPGF'' and logNbPGF''" $ do
      let nbPGFWorking (r,p,z) = withinDeltaOf 1e-5 (log $ nbPGF'' (NegBinomSizeProb r p) z) (logNbPGF'' (NegBinomSizeProb r p) z)
      all nbPGFWorking [(r,p,z) | r <- [2..50], p <- [0.1,0.3,0.5,0.7,0.9], z <- [0.1,0.3,0.5,0.7,0.9]] `shouldBe` True

testLogPdeGF1 = do
    describe "test pdeGF and logPdeGF" $ do
      modifyMaxDiscardRatio (const 1000) $
        it "test pdeGF and logPdeGF" $ property $
        \(z
         , delay
         , lam
         , mu) -> (z > 0) &&
                  (delay < 30) &&
                  (delay > 0) &&
                  (lam < 30) &&
                  (lam > 0) &&
                  (mu < 30) &&
                  (mu > 0) ==>
                  withinDeltaOf 1e-3 (log $ pdeGF (params lam mu) (TimeDelta delay) pdeSol z) (logPdeGF (params lam mu) (TimeDelta delay) pdeSol z)
                  where params lam mu = (Parameters (lam, mu, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed []))
                        pdeSol = (PDESol Zero 1)

testLogPdeGF2 = do
    describe "test pdeGF and logPdeGF again" $ do
      modifyMaxDiscardRatio (const 10000) $
        it "test pdeGF and logPdeGF again" $ property $
        \(z
         , lam
         , delay
         , nbMean
         , nbVar) -> (z <= 1) && (z > 0) &&
                     (lam < 200) &&
                     (lam > 0) &&
                     (delay < 200) &&
                     (delay > 0) &&
                     (nbMean > 0) &&
                     (nbVar > nbMean) ==>
                     withinDeltaOf 1e-3 (log $ pdeGF (params lam) (scaledDelay delay) (pdeSol (nbMean,nbVar)) z) (logPdeGF (params lam) (scaledDelay delay) (pdeSol (nbMean,nbVar)) z)
                     where params lam = (Parameters (lam / 10, 0.3, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed []))
                           pdeSol nbStats = (PDESol (unsafeNBFromMAndV nbStats) 1)
                           scaledDelay d = TimeDelta $ d / 10

testLogPdeGFDash1 = do
    describe "test pdeGF' and logPdeGF'" $ do
      modifyMaxDiscardRatio (const 1000) $
        it "test pdeGF' and logPdeGF'" $ property $
        \(z100
         , delay
         , lam
         , mu) -> (z100 <= 100) &&
                  (z100 > 0) &&
                  (delay < 30) &&
                  (delay > 0) &&
                  (lam < 30) &&
                  (lam > 0) &&
                  (mu < 30) &&
                  (mu > 0) ==>
                  withinDeltaOf 1e-3 (log $ pdeGF' (params lam mu) (TimeDelta delay) pdeSol (z100 / 100)) (logPdeGF' (params lam mu) (TimeDelta delay) pdeSol (z100 / 100))
                  where params lam mu = (Parameters (lam, mu, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed []))
                        pdeSol = (PDESol Zero 1)

testLogPdeGFDash2 = do
    describe "test pdeGF' and logPdeGF' again" $ do
      modifyMaxDiscardRatio (const 10000) $
        it "test pdeGF' and logPdeGF' again" $ property $
        \(z
         , lam
         , delay
         , nbMean
         , nbVar) -> (z <= 1) && (z > 0) &&
                     (lam < 200) &&
                     (lam > 0) &&
                     (delay < 200) &&
                     (delay > 0) &&
                     (nbMean > 0) &&
                     (nbVar > nbMean) ==>
                     withinDeltaOf 1e-3 (log $ pdeGF' (params lam) (scaledDelay delay) (pdeSol (nbMean,nbVar)) z) (logPdeGF' (params lam) (scaledDelay delay) (pdeSol (nbMean,nbVar)) z)
                     where params lam = (Parameters (lam / 10, 0.3, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed []))
                           pdeSol nbStats = (PDESol (unsafeNBFromMAndV nbStats) 1)
                           scaledDelay d = TimeDelta $ d / 50

testLogPdeGFDashDash1 = do
    describe "test pdeGF'' and logPdeGF''" $ do
      modifyMaxDiscardRatio (const 1000) $
        it "test pdeGF'' and logPdeGF''" $ property $
        \(z100
         , delay
         , lam
         , mu) -> (z100 <= 100) &&
                  (z100 > 0) &&
                  (delay < 30) &&
                  (delay > 0) &&
                  (lam < 30) &&
                  (lam > 0) &&
                  (mu < 30) &&
                  (mu > 0) ==>
                  withinDeltaOf 1e-3 (log $ pdeGF'' (params lam mu) (TimeDelta delay) pdeSol (z100 / 100)) (logPdeGF'' (params lam mu) (TimeDelta delay) pdeSol (z100 / 100))
                  where params lam mu = (Parameters (lam, mu, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed []))
                        pdeSol = (PDESol Zero 1)

testLogPdeGFDashDash2 = do
    describe "test pdeGF'' and logPdeGF'' again" $ do
      modifyMaxDiscardRatio (const 10000) $
        it "test pdeGF'' and logPdeGF'' again" $ property $
        \(z
         , lam
         , delay
         , nbMean
         , nbVar) -> (z <= 1) && (z > 0) &&
                     (lam < 200) &&
                     (lam > 0) &&
                     (delay < 200) &&
                     (delay > 0) &&
                     (nbMean > 0) &&
                     (nbVar > nbMean) ==>
                     withinDeltaOf 1e-3 (log $ pdeGF'' (params lam) (scaledDelay delay) (pdeSol (nbMean,nbVar)) z) (logPdeGF'' (params lam) (scaledDelay delay) (pdeSol (nbMean,nbVar)) z)
                     where params lam = (Parameters (lam / 10, 0.3, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed []))
                           pdeSol nbStats = (PDESol (unsafeNBFromMAndV nbStats) 1)
                           scaledDelay d = TimeDelta $ d / 50

testLogPdeStatistics = do
    describe "test pdeStatistics and logPdeStatistics" $ do
      modifyMaxDiscardRatio (const 10000) $
        it "test pdeStatistics and logPdeStatistics" $ property $
        \( lam
         , delay
         , nbMean
         , nbVar) -> (lam < 200) &&
                     (lam > 6) &&
                     (delay < 200) &&
                     (delay > 0) &&
                     (nbMean > 1) &&
                     (nbVar > nbMean) ==>
                     withinDeltaOf 1e-3 (log . fst' $ fooUnlogged lam (TimeDelta delay) nbMean nbVar) (fst' $ fooLogged lam (TimeDelta delay) nbMean nbVar) &&
                     withinDeltaOf 1e-3 (log . snd' $ fooUnlogged lam (TimeDelta delay) nbMean nbVar) (snd' $ fooLogged lam (TimeDelta delay) nbMean nbVar) &&
                     withinDeltaOf 1e-3 (log . thd' $ fooUnlogged lam (TimeDelta delay) nbMean nbVar) (thd' $ fooLogged lam (TimeDelta delay) nbMean nbVar)
                     where params lam = (Parameters (lam / 10, 0.3, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed []))
                           pdeSol nbStats = (PDESol (unsafeNBFromMAndV nbStats) 1)
                           scaledDelay d = TimeDelta $ d / 50
                           fooUnlogged lam (TimeDelta delay) nbMean nbVar = pdeStatistics (params lam) (scaledDelay delay) (pdeSol (nbMean,nbVar))
                           fooLogged lam (TimeDelta delay) nbMean nbVar  = (logPdeStatistics (params lam) (scaledDelay delay) (pdeSol (nbMean,nbVar)))
                           fst' (a,_,_) = a
                           snd' (_,a,_) = a
                           thd' (_,_,a) = a

testp0 = do
  describe "Test p0" $ do
    it "Initial condition 1" $ do
      p0 (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 0.0001) 0.2 `shouldSatisfy` (withinDeltaOf 1e-3 0.2)
      p0 (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 0.0001) 1.0 `shouldSatisfy` (withinDeltaOf 1e-3 1.0)

    it "Evolution 1" $
      let a = p0 (Parameters (2.3, 0.5, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) 0.2
          b = p0 (Parameters (2.3, 0.5, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.1) 0.2
       in a > b `shouldBe` True

    it "Long term condition 1" $ do
      p0 (Parameters (0.001, 5.9, 0.001, Timed [(AbsoluteTime 1000, 0.001)], 0.01, Timed [])) (TimeDelta 100.0) 0.2 `shouldSatisfy` (withinDeltaOf 1e-2 1.0)
      p0 (Parameters (0.001, 5.9, 0.001, Timed [(AbsoluteTime 1000, 0.001)], 0.01, Timed [])) (TimeDelta 100.0) 0.9 `shouldSatisfy` (withinDeltaOf 1e-2 1.0)

    it "Long term condition 2" $ do
      p0 (Parameters (10.1, 0.1, 9.0, Timed [(AbsoluteTime 1000, 1.0)], 0.01, Timed [])) (TimeDelta 100.0) 0.2 `shouldSatisfy` (withinDeltaOf 1e-2 0.0)
      p0 (Parameters (10.1, 0.1, 9.0, Timed [(AbsoluteTime 1000, 1.0)], 0.01, Timed [])) (TimeDelta 100.0) 0.9 `shouldSatisfy` (withinDeltaOf 1e-2 0.0)

    it "First partial derivative seems correct 1" $ do
      p0' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) 0.7 `shouldSatisfy` (withinDeltaOf 1e-2 (finiteDifference 1e-5 (\z -> p0 (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) z) 0.7))
      p0' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) 0.9 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-5 (\z -> p0 (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) z) 0.9))
      p0' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) 0.7 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-5 (\z -> p0 (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) z) 0.7))
      p0' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) 0.9 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-5 (\z -> p0 (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) z) 0.9))

    it "Second partial derivative seems correct 1" $ do
      p0'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) 0.7 `shouldSatisfy` (withinDeltaOf 1e-2 (finiteDifference 1e-5 (\z -> p0' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) z) 0.7))
      p0'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) 0.9 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-5 (\z -> p0' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) z) 0.9))
      p0'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) 0.7 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-5 (\z -> p0' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) z) 0.7))
      p0'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) 0.9 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-5 (\z -> p0' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) z) 0.9))

testRr = do
  describe "Test rr" $ do
    it "Initial condition 1" $ do
      rr (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 0.0001) 0.2 `shouldSatisfy` (withinDeltaOf 1e-3 (0.8/0.8))
      rr (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 0.0001) 0.9 `shouldSatisfy` (withinDeltaOf 1e-3 (0.1/0.1))

    it "Evolution 1" $
      let a = rr (Parameters (0.001, 0.001, 0.001, Timed [(AbsoluteTime 1000,0.9)], 0.001, Timed [])) (TimeDelta 0.01) 0.3
          b = rr (Parameters (0.001, 0.001, 0.001, Timed [(AbsoluteTime 1000,0.9)], 0.001, Timed [])) (TimeDelta 0.01) 0.2
       in a > b `shouldBe` True

    it "Long term condition 1" $ do
      rr (Parameters (3, 0.9, 0.01, Timed [(AbsoluteTime 1000,0.1)], 0.1, Timed [])) (TimeDelta 100.0) 0.2 `shouldSatisfy` (withinDeltaOf 1e-2 0.0)
      rr (Parameters (3, 0.9, 0.01, Timed [(AbsoluteTime 1000,0.1)], 0.1, Timed [])) (TimeDelta 100.0) 0.9 `shouldSatisfy` (withinDeltaOf 1e-2 0.0)

    it "Long term condition 2" $
      let a = rr (Parameters (3, 0.1, 0.01, Timed [(AbsoluteTime 1000,0.1)], 0.1, Timed [])) (TimeDelta 1.1) (0.9/(1-0.9))
          b = rr (Parameters (3, 0.1, 0.01, Timed [(AbsoluteTime 1000,0.1)], 0.1, Timed [])) (TimeDelta 1.1) (0.8/(1-0.8))
          c = rr (Parameters (3, 0.1, 0.01, Timed [(AbsoluteTime 1000,0.1)], 0.1, Timed [])) (TimeDelta 1.1) (0.3/(1-0.3))
       in do
        a < b `shouldBe` True
        b < c `shouldBe` True

    it "First partial derivative seems correct 1" $ do
      rr' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) 0.7 `shouldSatisfy` (withinDeltaOf 1e-2 (finiteDifference 1e-5 (\z -> rr (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) z) 0.7))
      rr' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) 0.9 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-5 (\z -> rr (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) z) 0.9))
      rr' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) 0.7 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-5 (\z -> rr (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) z) 0.7))
      rr' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) 0.9 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-5 (\z -> rr (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) z) 0.9))

    it "Second partial derivative seems correct 1" $ do
      rr'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) 0.7 `shouldSatisfy` (withinDeltaOf 1e-2 (finiteDifference 1e-5 (\z -> rr' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) z) 0.7))
      rr'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) 0.9 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-5 (\z -> rr' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) z) 0.9))
      rr'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) 0.7 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-5 (\z -> rr' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) z) 0.7))
      rr'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) 0.9 `shouldSatisfy` (withinDeltaOf 1e-5 (finiteDifference 1e-5 (\z -> rr' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) z) 0.9))

testPdeGF = do
  describe "Test pdeGF" $ do
    it "First partial derivative seems correct 1" $ do
      pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol Zero 1) 0.7 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 ((\z -> pdeGF (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol Zero 1) z)) 0.7))
      pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol Zero 1) 0.9 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 (\z -> pdeGF (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol Zero 1) z) 0.9))
      pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol Zero 1) 0.7 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 (\z -> pdeGF (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol Zero 1) z) 0.7))
      pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol Zero 1) 0.9 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 (\z -> pdeGF (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol Zero 1) z) 0.9))
      pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) 0.7 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 (\z -> pdeGF (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) z) 0.7))
      pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) 0.9 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 (\z -> pdeGF (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) z) 0.9))
      pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) 0.7 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 (\z -> pdeGF (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) z) 0.7))
      pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) 0.9 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 (\z -> pdeGF (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) z) 0.9))

    it "Second partial derivative seems correct 1" $ do
      pdeGF'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol Zero 1) 0.7 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 (\z -> pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol Zero 1) z) 0.7))
      pdeGF'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol Zero 1) 0.9 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 (\z -> pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol Zero 1) z) 0.9))
      pdeGF'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol Zero 1) 0.7 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 (\z -> pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol Zero 1) z) 0.7))
      pdeGF'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol Zero 1) 0.9 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 (\z -> pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol Zero 1) z) 0.9))
      pdeGF'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) 0.7 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 (\z -> pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) z) 0.7))
      pdeGF'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) 0.9 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-6 (\z -> pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 1.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) z) 0.9))
      pdeGF'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) 0.7 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 (\z -> pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) z) 0.7))
      pdeGF'' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) 0.9 `shouldSatisfy` (withinDeltaOf 1e-4 (finiteDifference 1e-5 (\z -> pdeGF' (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) (TimeDelta 2.0) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 1) z) 0.9))

testPdeStatistics = do
  describe "Test pdeStatistics" $ do
    it "Properties 1" $
      let (c,m,v) = pdeStatistics (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1) (PDESol Zero 1)
          (c',m',v') = pdeStatistics (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2) (PDESol Zero 1)
          (c'',m'',v'') = pdeStatistics (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 20) (PDESol Zero 1)
       in do
        c > c' `shouldBe` True
        c' > c'' `shouldBe` True
        c'' `shouldSatisfy` (withinDeltaOf 1e-6 0.0)
        m < m' `shouldBe` True
        m' < m'' `shouldBe` True
        v < v' `shouldBe` True
        v' < v'' `shouldBe` True
        m < v `shouldBe` True
        m' < v' `shouldBe` True
        m'' < v'' `shouldBe` True

    it "Properties 2" $
      let (c,m,v) = pdeStatistics (Parameters (2.3, 1.2, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 1) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 2)
          (c',m',v') = pdeStatistics (Parameters (2.3, 1.2, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 2) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 2)
          (c'',m'',v'') = pdeStatistics (Parameters (2.3, 1.2, 0.3, Timed [(AbsoluteTime 1000,0.5)], 0.6, Timed [])) (TimeDelta 20) (PDESol (unsafeNBFromMAndV (3.0,9.0)) 2)
       in do
        c > c' `shouldBe` True
        c' > c'' `shouldBe` True
        c'' `shouldSatisfy` (withinDeltaOf 1e-6 0.0)
        m < m' `shouldBe` True
        m' < m'' `shouldBe` True
        v < v' `shouldBe` True
        v' < v'' `shouldBe` True
        m < v `shouldBe` True
        m' < v' `shouldBe` True
        m'' < v'' `shouldBe` True

    it "Properties 3" $
      let (c,m,v) = pdeStatistics (Parameters (2.0,1.0,0.5, Timed [(AbsoluteTime 1000,0.5)],0.4,Timed [])) (TimeDelta 2.0) (PDESol (NegBinomSizeProb 3.9 0.5) 1.0)
       in do
        c > 0 `shouldBe` True
        m > 0 `shouldBe` True
        v > m `shouldBe` True

testLlhd = do
  describe "Test llhd" $ do
    it "Manceau example" $
      let obs = [ (TimeDelta 1.0, OBirth)
                , (TimeDelta 1.0, OOccurrence)
                , (TimeDelta 1.0, OBirth)
                , (TimeDelta 1.0, OBirth)
                , (TimeDelta 1.0, ObsUnscheduledSequenced)
                , (TimeDelta 1.0, OOccurrence)
                , (TimeDelta 1.0, OCatastrophe 3) ]
          params lam = Parameters (lam,1.0,0.3, Timed [(AbsoluteTime 7.0,0.5)],0.6,Timed [])
          (llhdVal1,_) = unsafeLlhdAndNB obs (params 1.1) initLlhdState
          (llhdVal2,_) = unsafeLlhdAndNB obs (params 1.2) initLlhdState
          (llhdVal3,_) = unsafeLlhdAndNB obs (params 1.3) initLlhdState
          (llhdVal9,_) = unsafeLlhdAndNB obs (params 1.9) initLlhdState
       in do
        llhdVal1 `shouldSatisfy` (withinDeltaOf 1e-1 (-40.5))
        llhdVal2 `shouldSatisfy` (withinDeltaOf 2e-1 (-41.0))
        llhdVal3 `shouldSatisfy` (withinDeltaOf 2e-1 (-41.5))
        llhdVal9 `shouldSatisfy` (withinDeltaOf 2e-1 (-46.0))

testInhomBDSLlhd = do
  describe "Test inhomogeneous BDS LLHD" $ do
    it "Check for constant parameters it looks right" $
      let obs = [(TimeDelta 1.0,OBirth),(TimeDelta 1.0,OBirth),(TimeDelta 1.0,ObsUnscheduledSequenced),(TimeDelta 1.0,ObsUnscheduledSequenced),(TimeDelta 1.0,ObsUnscheduledSequenced)]
          tlams = fromJust $ asTimed [(AbsoluteTime 0,1.2)]
          tlams' = fromJust $ asTimed [(AbsoluteTime 0,1.2),(AbsoluteTime 10,5.0)]
          tlams'' = fromJust $ asTimed [(AbsoluteTime 0,1.3),(AbsoluteTime 10,5.0)]
          tlams''' = fromJust $ asTimed [(AbsoluteTime 0,1.3),(AbsoluteTime 0.5,1.4)]
          tlams'''' = fromJust $ asTimed [(AbsoluteTime 0,1.3),(AbsoluteTime 1.5,1.3),(AbsoluteTime 2.5,1.3)]
          lam = fromJust $ cadlagValue tlams (AbsoluteTime 0.1)
          lam'' = fromJust $ cadlagValue tlams'' (AbsoluteTime 0.1)
          (llhdValXXX1,_) = InhomBDSLlhd.inhomLlhdAndNB obs (InhomBDSLlhd.InhomParams (tlams,1.0,0.3)) InhomBDSLlhd.initLlhdState
          (llhdValXXX2,_) = InhomBDSLlhd.inhomLlhdAndNB obs (InhomBDSLlhd.InhomParams (tlams',1.0,0.3)) InhomBDSLlhd.initLlhdState
          (llhdValXXX3,_) = InhomBDSLlhd.inhomLlhdAndNB obs (InhomBDSLlhd.InhomParams (tlams'',1.0,0.3)) InhomBDSLlhd.initLlhdState
          (llhdValXXX4,_) = InhomBDSLlhd.inhomLlhdAndNB obs (InhomBDSLlhd.InhomParams (tlams''',1.0,0.3)) InhomBDSLlhd.initLlhdState
          (llhdValXXX5,_) = InhomBDSLlhd.inhomLlhdAndNB obs (InhomBDSLlhd.InhomParams (tlams'''',1.0,0.3)) InhomBDSLlhd.initLlhdState
          (llhdValYYY1,_) = unsafeLlhdAndNB obs (Parameters (lam,1.0,0.3,Timed [],0.0,Timed [])) initLlhdState
          (llhdValYYY2,_) = unsafeLlhdAndNB obs (Parameters (lam'',1.0,0.3,Timed [],0.0,Timed [])) initLlhdState
          (llhdValYYY3,_) = unsafeLlhdAndNB obs (Parameters (lam'' + 0.1,1.0,0.3,Timed [],0.0,Timed [])) initLlhdState
       in do
        llhdValXXX1 `shouldSatisfy` (withinDeltaOf 1e-3 (llhdValYYY1))
        llhdValXXX2 `shouldSatisfy` (withinDeltaOf 1e-3 (llhdValYYY1))
        llhdValXXX3 `shouldSatisfy` (withinDeltaOf 1e-3 (llhdValYYY2))
        llhdValXXX3 `shouldSatisfy` (\l -> not $ withinDeltaOf 1e-3 llhdValYYY1 l)
        (if llhdValYYY3 < llhdValYYY2 then llhdValXXX3 > llhdValXXX4 else llhdValXXX3 < llhdValXXX4) `shouldBe` True
        llhdValXXX5 `shouldSatisfy` (withinDeltaOf 1e-1 (llhdValYYY2)) -- exposes limitation of approximation!!!
    describe "Check values are finite when sensible" $
      let infParams = (InhomBDSLlhd.InhomParams (fromJust $ asTimed [(AbsoluteTime 0.0,1.0),(AbsoluteTime 1.0,1.0)],0.4,0.4)) :: InhomBDSLlhd.InhomParams
          (InhomBDSLlhd.InhomParams (tlams,_,_)) = infParams
          obs = [(TimeDelta 0.3,OBirth),(TimeDelta 0.5,OBirth),(TimeDelta 0.19,ObsUnscheduledSequenced)]
          llhdVal = fst $ InhomBDSLlhd.inhomLlhdAndNB obs infParams InhomBDSLlhd.initLlhdState
          obs' = [(TimeDelta 0.3,OBirth),(TimeDelta 0.5,OBirth),(TimeDelta 0.20,ObsUnscheduledSequenced)]
          llhdVal' = fst $ InhomBDSLlhd.inhomLlhdAndNB obs' infParams InhomBDSLlhd.initLlhdState
          obs'' = [(TimeDelta 0.3,OBirth),(TimeDelta 0.5,OBirth),(TimeDelta 0.21,ObsUnscheduledSequenced)]
          llhdVal'' = fst $ InhomBDSLlhd.inhomLlhdAndNB obs'' infParams InhomBDSLlhd.initLlhdState
       in do
        it "Check cadlagValue" $ do
          cadlagValue tlams (AbsoluteTime 1.1) == Just 1.0 `shouldBe` True
          cadlagValue tlams (AbsoluteTime 1.0) == Just 1.0 `shouldBe` True
          cadlagValue tlams (AbsoluteTime 0.9) == Just 1.0 `shouldBe` True
        it "Check nextTime" $ do
          isJust (nextTime tlams (AbsoluteTime 1.1)) `shouldBe` True
          isJust (nextTime tlams (AbsoluteTime 1.0)) `shouldBe` True
          isJust (nextTime tlams (AbsoluteTime 0.9)) `shouldBe` True
        it "Check llhdValue" $ do
          isInfinite llhdVal `shouldBe` False
          isInfinite llhdVal' `shouldBe` False
          isInfinite llhdVal'' `shouldBe` False

testConversion = do
  describe "Test conversion between event types" $ do
    it "Demonstration data set 1" $
      let p1 = Person (Identifier 1)
          p2 = Person (Identifier 2)
          p4 = Person (Identifier 4)
          p5 = Person (Identifier 5)
          p6 = Person (Identifier 6)
          simObsEvents =
            [ Infection (AbsoluteTime 1) p1 p2
            , Sampling (AbsoluteTime 3) p1
            , Infection (AbsoluteTime 4) p2 p4
            , Sampling (AbsoluteTime 6) p4
            , Occurrence (AbsoluteTime 8) p2
            , Occurrence (AbsoluteTime 11) p6
            , Sampling (AbsoluteTime 12) p5
            ]
          llhdObsEvents =
            [ (TimeDelta 1.0, OBirth)
            , (TimeDelta 2.0, ObsUnscheduledSequenced)
            , (TimeDelta 1.0, OBirth)
            , (TimeDelta 2.0, ObsUnscheduledSequenced)
            , (TimeDelta 2.0, OOccurrence)
            , (TimeDelta 3.0, OOccurrence)
            , (TimeDelta 1.0, ObsUnscheduledSequenced)
            ]
       in eventsAsObservations simObsEvents `shouldSatisfy` (== llhdObsEvents)

testImpossibleParameters = do
  describe "Test correct handling of impossible parameters" $ do
    it "Test negative birth rate is impossible" $
      let obs = [(TimeDelta 1.0,OBirth),(TimeDelta 1.0,OOccurrence),(TimeDelta 1.0,OBirth),(TimeDelta 1.0,OBirth),(TimeDelta 1.0,ObsUnscheduledSequenced),(TimeDelta 1.0,OOccurrence)]
          llhd1 = fst $ unsafeLlhdAndNB obs (Parameters (0.0000000001,1.0,0.3,Timed [],0.6,Timed [])) initLlhdState
          llhd2 = llhdAndNB obs (Parameters (0.0000000000,1.0,0.3,Timed [],0.6,Timed [])) initLlhdState
          llhd3 = llhdAndNB obs (Parameters (-0.0000000001,1.0,0.3,Timed [],0.6,Timed [])) initLlhdState
       in do
        isNaN llhd1 `shouldBe` False
        isInfinite llhd1 `shouldBe` False
        llhd1 < 0 `shouldBe` True
        isLeft llhd2 `shouldBe` True
        isLeft llhd3 `shouldBe` True
    it "Test negative sampling rate is impossible" $
      let obs1 = [(TimeDelta 1.0,OBirth),(TimeDelta 1.0,OOccurrence),(TimeDelta 1.0,OBirth),(TimeDelta 1.0,OBirth),(TimeDelta 1.0,ObsUnscheduledSequenced),(TimeDelta 1.0,OOccurrence)]
          obs2 = [(TimeDelta 1.0,OBirth),(TimeDelta 1.0,OOccurrence),(TimeDelta 1.0,OBirth),(TimeDelta 1.0,OBirth),(TimeDelta 1.0,OOccurrence)]
          llhd11 = fst $ unsafeLlhdAndNB obs1 (Parameters (1.0,1.0,0.1,Timed [],0.6,Timed [])) initLlhdState
          llhd12 = llhdAndNB obs1 (Parameters (1.0,1.0,0.0,Timed [],0.6,Timed [])) initLlhdState
          llhd13 = llhdAndNB obs1 (Parameters (1.0,1.0,-0.1,Timed [],0.6,Timed [])) initLlhdState
          llhd21 = fst $ unsafeLlhdAndNB obs2 (Parameters (1.0,1.0,0.1,Timed [],0.6,Timed [])) initLlhdState
          llhd22 = fst $ unsafeLlhdAndNB obs2 (Parameters (1.0,1.0,0.0,Timed [],0.6,Timed [])) initLlhdState
          llhd23 = llhdAndNB obs2 (Parameters (1.0,1.0,-0.1,Timed [],0.6,Timed [])) initLlhdState
       in do
        isNaN llhd11 `shouldBe` False
        isLeft llhd12 `shouldBe` True
        isLeft llhd13 `shouldBe` True
        isNaN llhd21 `shouldBe` False
        isNaN llhd22 `shouldBe` False
        isLeft llhd23 `shouldBe` True
        isInfinite llhd11 `shouldBe` False
        isInfinite llhd21 `shouldBe` False
        isInfinite llhd22 `shouldBe` False

tmpIsSampling :: EpidemicEvent -> Bool
tmpIsSampling e = case e of
  Sampling{} -> True
  _ -> False

testParameterUpdate :: SpecWith ()
testParameterUpdate =
  let params1 = (Parameters (2.4, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) :: Parameters
      params2 = (Parameters (2.3, 1, 0.3, Timed [(AbsoluteTime 1000, 0.5)], 0.6, Timed [])) :: Parameters
      params3 = (Parameters (2.4, 1, 0.3, Timed [(AbsoluteTime 1000, 0.6)], 0.6, Timed [])) :: Parameters
    in do describe "Testing parameter update function" $ do
            it "test lambda update" $ do
              (params1 /= params2) `shouldBe` True
              (params1 == putLambda params2 2.4) `shouldBe` True
            it "test rhos update" $ do
              (params1 /= params3) `shouldBe` True
              (params1 == putRhos params3 (Timed [(AbsoluteTime 1000, 0.5)])) `shouldBe` True

-- | This test case looks at how to set the seed when using @mwc-random@. We
-- care about this because we want to be able to generate reproducible
-- simulations without being restricted to the fixed seed that is the default
-- value.
testMWCSeeding :: SpecWith ()
testMWCSeeding = do
  describe "Testing MWC seeding" $ do
    it "test create works as expected" $ do
      gen <- MWC.create
      x1 <- (MWC.uniform gen :: IO Double)
      x2 <- (MWC.uniform gen :: IO Double)
      gen' <- MWC.create
      x1' <- (MWC.uniform gen' :: IO Double)
      (x1 /= x2) `shouldBe` True
      (x1 == x1') `shouldBe` True
    it "test initialise works as expected" $ do
      xGen <- MWC.create
      x1 <- (MWC.uniform xGen :: IO Double)
      yGen <- MWC.initialize (Unboxed.fromList [1,2,3])
      y1 <- (MWC.uniform yGen :: IO Double)
      (x1 /= y1) `shouldBe` True
      y2 <- (MWC.uniform yGen :: IO Double)
      (y1 /= y2) `shouldBe` True
      zGen <- MWC.initialize (Unboxed.fromList [1,2,3])
      z1 <- (MWC.uniform zGen :: IO Double)
      (y1 == z1) `shouldBe` True
      wGen <- MWC.initialize (Unboxed.fromList [3,2,1])
      w1 <- (MWC.uniform wGen :: IO Double)
      (z1 /= w1) `shouldBe` True

-- | Generate a random @NumLineages@
qcRandomNumLineages :: Gen NumLineages
qcRandomNumLineages = do
  kDouble <- choose (1 + 1e-6, 1e2 :: Double)
  return . fromIntegral $ round kDouble

-- | Generate a random @Rate@
qcRandomRate :: Gen Rate
qcRandomRate = choose (1e-3, 1e1)

-- | Generate a random @Probability@
qcRandomProbability :: Gen Probability
qcRandomProbability = choose (0, 1)

-- | Generate a random @NegativeBinomial@ using @NegBinomSizeProb@ /not/ any
-- other constructor.
qcRandomNegBinomSizeProb :: Gen NegativeBinomial
qcRandomNegBinomSizeProb = do
  r <- qcRandomRate
  p <- qcRandomProbability
  return $ NegBinomSizeProb r p

-- | Generate a random @TimeDelta@
qcRandomTimeDelta :: Gen TimeDelta
qcRandomTimeDelta = do
  td <- qcRandomRate
  return $ TimeDelta td

-- | Generate a random __small__ @TimeDelta@
qcRandomSmallTimeDelta :: Gen TimeDelta
qcRandomSmallTimeDelta = do
  td <- choose (1e-3,2)
  return $ TimeDelta td

-- | Generate a list of absolute times that occur after the origin within the
-- duration.
qcRandomOrderedAbsTimes :: AbsoluteTime -> TimeDelta -> Gen [AbsoluteTime]
qcRandomOrderedAbsTimes (AbsoluteTime a) (TimeDelta d) = do
  times <- listOf1 $ choose (a, a + d)
  return [AbsoluteTime t | t <- sort times]

-- | Generate a random @Timed x@ from the start time and the duration and the
-- constant value to store at each time.
qcRandomTimedX :: Num x => AbsoluteTime -> TimeDelta -> x -> Gen (Timed x)
qcRandomTimedX originTime duration x = do
  absTimes <- qcRandomOrderedAbsTimes originTime duration
  case asTimed (zip absTimes (repeat x)) of
    Just timedVals -> return timedVals
    Nothing -> qcRandomTimedX originTime duration x

-- | Generate a random @Parameters@
qcRandomParameters :: AbsoluteTime -> TimeDelta -> Gen Parameters
qcRandomParameters originTime duration = do
  -- randLambda <- qcRandomRate
  randMu <- qcRandomRate
  randPsi <- qcRandomRate
  randRho <- qcRandomProbability
  randOmega <- qcRandomRate
  tmp <- choose (0.5,1.0)
  let randLambda = (randMu + randPsi + randOmega) / tmp
  randNu <- qcRandomProbability
  randTimedRho <- qcRandomTimedX originTime duration randRho
  randTimedNu <- qcRandomTimedX originTime duration randNu
  return $
    Parameters
      (randLambda, randMu, randPsi, randTimedRho, randOmega, randTimedNu)

-- | Generate a random @ObservedEvent@
qcRandomObservedEvent :: Gen ObservedEvent
qcRandomObservedEvent = do
  isUnscheduled <- chooseAny
  if isUnscheduled
    then elements [OBirth, ObsUnscheduledSequenced, OOccurrence]
    else do
      isSequenced <- chooseAny
      numLineages <- suchThat chooseAny (> 0)
      if isSequenced
        then return (OCatastrophe numLineages)
        else return (ODisaster numLineages)

-- | Generate a random list of @observation@ values
qcRandomObservations :: Gen [Observation]
qcRandomObservations = do
  durationDouble <- suchThat chooseAny (> 0) :: Gen Double
  eventAbsTimesDoubles <- listOf1 (choose (0, durationDouble))
  let duration = TimeDelta durationDouble
      eventAbsTimes = [AbsoluteTime t | t <- eventAbsTimesDoubles]
  let eats = sort ((AbsoluteTime 0) : eventAbsTimes)
      timeDeltas = [timeDelta a b | (a, b) <- zip (init eats) (tail eats)]
      numEvents = length eventAbsTimes
  eventTypes <- vectorOf numEvents qcRandomObservedEvent
  return $ zip timeDeltas eventTypes

-- | A copy of the @withinDeltaOf@ function specialised to
-- @AggregatedObservations@ to test for approximate equality of aggregated
-- observations.
withinDeltaOfAggObs :: Double
                    -> AggregatedObservations
                    -> AggregatedObservations
                    -> Bool
withinDeltaOfAggObs delta (AggregatedObservations aggTimes obs) (AggregatedObservations aggTimes' obs') =
  withinDeltaOfAggTimes delta aggTimes aggTimes' &&  allWithinDeltaOfObs delta obs obs'

withinDeltaOfAbsoluteTimes :: Double -> AbsoluteTime -> AbsoluteTime -> Bool
withinDeltaOfAbsoluteTimes delta (AbsoluteTime a) (AbsoluteTime b) = withinDeltaOf delta a b

allWithinDeltaOfAbsoluteTimes :: Double -> [AbsoluteTime] -> [AbsoluteTime] -> Bool
allWithinDeltaOfAbsoluteTimes delta absTimes1 absTimes2 =
  let f (a,b) = withinDeltaOfAbsoluteTimes delta a b
      absTimePairs = zip absTimes1 absTimes2
  in all f absTimePairs

withinDeltaOfTimeDeltas :: Double -> TimeDelta -> TimeDelta -> Bool
withinDeltaOfTimeDeltas delta (TimeDelta a) (TimeDelta b) = withinDeltaOf delta a b

allWithinDeltaOfTimeDeltas :: Double -> [TimeDelta] -> [TimeDelta] -> Bool
allWithinDeltaOfTimeDeltas delta timeDels1 timeDels2 =
  let f (a,b) = withinDeltaOfTimeDeltas delta a b
      timeDelPairs = zip timeDels1 timeDels2
  in all f timeDelPairs

withinDeltaOfAggTimes :: Double -> AggregationTimes -> AggregationTimes -> Bool
withinDeltaOfAggTimes delta (AggTimes ts) (AggTimes ts') =
  let times = map fst ts
      obsEvents = map snd ts
      times' = map fst ts'
      obsEvents' = map snd ts'
      timesWithinDelta = allWithinDeltaOfAbsoluteTimes delta times times'
      observedEventsEqual = all (uncurry (withinDeltaOfObsEvent delta)) (zip obsEvents obsEvents')
  in timesWithinDelta && observedEventsEqual

withinDeltaOfObsEvent :: Double -> ObservedEvent -> ObservedEvent -> Bool
withinDeltaOfObsEvent delta (ODisaster nl) (ODisaster nl') = withinDeltaOf delta nl nl'
withinDeltaOfObsEvent delta (OCatastrophe nl) (OCatastrophe nl') = withinDeltaOf delta nl nl'
withinDeltaOfObsEvent _ OBirth OBirth = True
withinDeltaOfObsEvent _ ObsUnscheduledSequenced ObsUnscheduledSequenced = True
withinDeltaOfObsEvent _ OOccurrence OOccurrence = True
withinDeltaOfObsEvent _ _ _ = False

withinDeltaOfObs :: Double -> Observation -> Observation -> Bool
withinDeltaOfObs delta (t,oe) (t',oe') = withinDeltaOfTimeDeltas delta t t' && withinDeltaOfObsEvent delta oe oe'

allWithinDeltaOfObs :: Double -> [Observation] -> [Observation] -> Bool
allWithinDeltaOfObs _ [] [] = True
allWithinDeltaOfObs delta [y] [x] = withinDeltaOfObs delta y x
allWithinDeltaOfObs delta (y:ys) (x:xs) = withinDeltaOfObs delta y x && allWithinDeltaOfObs delta ys xs
allWithinDeltaOfObs _ _ _ = False

testAggregation :: SpecWith ()
testAggregation =
  describe "Testing Aggregation" $ do
    let smallDelta = 1e-4
        tinyDelta = TimeDelta 1e-6
        duration obs = AbsoluteTime $ sum [t | (TimeDelta t, _) <- obs]
        multiplyAbsTime a (AbsoluteTime x) = AbsoluteTime (a * x)
        propertyRemoveSeq obs =
          let dur = duration obs
              ats =
                fromJust $
                maybeAggregationTimes [timeAfterDelta dur tinyDelta] []
              (AggregatedObservations _ obs') =
                aggregateUnscheduledObservations ats obs
           in not $ any isUnscheduledSequenced obs'
        propertyRemoveUnseq obs =
          let dur = duration obs
              ats =
                fromJust $
                maybeAggregationTimes [] [timeAfterDelta dur tinyDelta]
              (AggregatedObservations _ obs') =
                aggregateUnscheduledObservations ats obs
           in not $ any isOccurrence obs'
        propertyRemoveUnsched1 obs =
          let dur = duration obs
              ats =
                fromJust $
                maybeAggregationTimes
                  [timeAfterDelta dur tinyDelta]
                  [ timeAfterDelta
                      (timeAfterDelta dur tinyDelta)
                      (TimeDelta 1.0)
                  ]
              (AggregatedObservations _ obs') =
                aggregateUnscheduledObservations ats obs
           in not (any isOccurrence obs') &&
              not (any isUnscheduledSequenced obs')
        propertyRemoveUnsched2 obs =
          let dur = duration obs
              ats =
                fromJust $
                maybeAggregationTimes
                  [multiplyAbsTime 0.4 dur]
                  [multiplyAbsTime 0.5 dur]
              (AggregatedObservations _ obs') =
                aggregateUnscheduledObservations ats obs
           in not (any isOccurrence obs') &&
              not (any isUnscheduledSequenced obs')
        propertyBirthsRemain obs =
          let dur = duration obs
              numBs = length $ filter isBirth obs
              ats =
                fromJust $
                maybeAggregationTimes
                  [timeAfterDelta dur tinyDelta]
                  [ timeAfterDelta
                      (timeAfterDelta dur tinyDelta)
                      (TimeDelta 1.0)
                  ]
              (AggregatedObservations _ obs') =
                aggregateUnscheduledObservations ats obs
              numBs' = length $ filter isBirth obs'
           in numBs == numBs'
        propertyLineagesConst os =
          let dur = duration os
              numSeq = sum $ map numSequenced os
              numUnseq = sum $ map numUnsequenced os
              ats =
                fromJust $
                maybeAggregationTimes
                  [timeAfterDelta dur tinyDelta]
                  [ timeAfterDelta
                      (timeAfterDelta dur tinyDelta)
                      (TimeDelta 1.0)
                  ]
              (AggregatedObservations _ os') =
                aggregateUnscheduledObservations ats os
              numSeq' = sum $ map numSequenced os'
              numUnseq' = sum $ map numUnsequenced os'
           in withinDeltaOf smallDelta numSeq numSeq' &&
              withinDeltaOf smallDelta numUnseq numUnseq'
    it "sequenced aggregation removes all such unscheduled observations" $
      forAll qcRandomObservations propertyRemoveSeq
    it "unsequenced aggregation removes all such unscheduled observations" $
      forAll qcRandomObservations propertyRemoveUnseq
    it "aggregating both removes all relevent observations 1" $
      forAll qcRandomObservations propertyRemoveUnsched1
    it "aggregating both removes all relevent observations 2" $
      forAll qcRandomObservations propertyRemoveUnsched2
    it "aggregating leaves birth observations unchanged" $
      forAll qcRandomObservations propertyBirthsRemain
    it "aggregating leaves the number of observed lineages unchanged" $
      forAll qcRandomObservations propertyLineagesConst

testIntervalLlhd :: SpecWith ()
testIntervalLlhd =
  describe "Testing the intervalLlhd function" $ do
    let propertyNBNotNaN (params, delay, k, nb) =
          let (_, nb') = fromRight' $ intervalLlhd params delay k nb
              (m, _) = mAndVFromNb nb'
           in not $ isNaN m
        absTimeZero = AbsoluteTime 0
        qcIntervalLlhdArgs ::
             Gen (Parameters, TimeDelta, NumLineages, NegativeBinomial)
        qcIntervalLlhdArgs = do
          totalDuration <- qcRandomTimeDelta
          params <- qcRandomParameters absTimeZero totalDuration
          delay <- qcRandomSmallTimeDelta
          k <- qcRandomNumLineages
          nb <- qcRandomNegBinomSizeProb
          return (params, delay, k, nb)
        propertyPDEStatsNonNaN (params, delay, k, nb) =
          let (_,logm,_) = logPdeStatistics params delay (PDESol nb k)
          in not $ isNaN logm
        propertyLogCNotNaN (params, delay, k, nb) =
          let pdeSol = PDESol nb k
              logmGF = logPdeGF params delay pdeSol
              logC = logmGF 1
          in not $ isNaN logC
        propertyLogmTermNotNaN (params, delay, k, nb) =
          let pdeSol = PDESol nb k
              logmGF' = logPdeGF' params delay pdeSol
              logmTerm = logmGF' 1
          in not $ isNaN logmTerm
        propertyOdeHelperNotNaN (params, delay, _, _) =
          let (x1,x2,disc,expFact) = odeHelpers params delay
              notNaN = not . isNaN
          in notNaN x1 && notNaN x2 && 0 < disc && 1 > expFact && x2 > x1
    it "resulting odeHelpers are not NaN" $ forAll qcIntervalLlhdArgs propertyOdeHelperNotNaN
    it "resulting logmTerm from logPdeGF' is not NAN" $ forAll qcIntervalLlhdArgs propertyLogmTermNotNaN
    it "resulting logC from logPdeGF is not NAN" $ forAll qcIntervalLlhdArgs propertyLogCNotNaN
    it "resulting log(mean) from logPdeStatistics is not NAN" $ forAll qcIntervalLlhdArgs propertyPDEStatsNonNaN
    it "resulting NB is not NAN" $ forAll qcIntervalLlhdArgs propertyNBNotNaN

testLogP0Dash :: SpecWith ()
testLogP0Dash = do
  describe "Testing the logP0' function" $ do
    let absTimeZero = AbsoluteTime 0
        qcP0Args :: Gen (Parameters, TimeDelta, Probability)
        qcP0Args = do
          totalDuration <- qcRandomTimeDelta
          params_ <- qcRandomParameters absTimeZero totalDuration
          delay_ <- qcRandomSmallTimeDelta
          z_ <- qcRandomProbability
          return (params_, delay_, z_)
        propertyLogValNotNaN (params_, delay_, z_) =
          let logValue = logP0' params_ delay_ z_
          in not $ isNaN logValue
        propertyApproximateEquality (params, delay, z) =
          let linearValue = p0' params delay z
              logValue = logP0' params delay z
          in withinDeltaOf 1e-10 (exp logValue) linearValue -- fails for smaller delta :)
    it "log version is not nan" $ forAll qcP0Args propertyLogValNotNaN
    it "approximate equality to p0'" $ forAll qcP0Args propertyApproximateEquality
  describe "Testing the logP0'' function" $ do
    let absTimeZero = AbsoluteTime 0
        qcP0Args :: Gen (Parameters, TimeDelta, Probability)
        qcP0Args = do
          totalDuration <- qcRandomTimeDelta
          params_ <- qcRandomParameters absTimeZero totalDuration
          delay_ <- qcRandomSmallTimeDelta
          z_ <- qcRandomProbability
          return (params_, delay_, z_)
        propertyLogValNotNaN (params_, delay_, z_) =
          let logValue = logP0'' params_ delay_ z_
          in not $ isNaN logValue
        propertyApproximateEquality (params, delay, z) =
          let linearValue = p0'' params delay z
              logValue = logP0'' params delay z
          in withinDeltaOf 1e-10 (exp logValue) linearValue -- fails for smaller delta :)
    it "log version is not nan" $ forAll qcP0Args propertyLogValNotNaN
    it "approximate equality to p0''" $ forAll qcP0Args propertyApproximateEquality

-- | Two of the tests have been commented out because they are broken by changes
-- to the LLHD interface rather than the underlying code that does the
-- computation.
main :: IO ()
main = hspec $ do
  -- ** slow tests **
  testNbPGF
  -- ** fast tests **
  testTestingHelpers
  testPdeStatistics
  testp0
  testRr
  testPdeGF
  -- testLlhd -- broken because we guard against R0 < 1 in current version.
  testConversion
  -- testImpossibleParameters -- broken because we use either monad instead of infinite values now.
  -- testInhomBDSLlhd -- not relevant to the first manuscript.
  testParameterUpdate
  testMWCSeeding
  testLogPdeGF1
  testLogPdeGF2
  testLogPdeGFDash1
  testLogPdeGFDash2
  testLogPdeGFDashDash1
  testLogPdeGFDashDash2
  testLogSumExp
  testLogPdeStatistics
  testAggregation
  testIntervalLlhd
  testLogP0Dash

ext: hs

max_stars_repo_name: aezarebski/timtam

max_issues_repo_name: aezarebski/timtam

max_forks_repo_name: aezarebski/timtam


text:
module Statistics
    ( runJSONFold
    , receivedCreatedF
    , module Statistics.Block
    , module Statistics.Chart
    , module Statistics.CSV
    , module Statistics.Focus
    , module Statistics.Graph
    , module Statistics.MemPool
    , module Statistics.Report
    , module Statistics.Throughput
    , module Statistics.Tx
    ) where

import           Control.Foldl (Fold (..))
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as M

import           JSONLog
import           Statistics.Block
import           Statistics.Chart
import           Statistics.CSV
import           Statistics.Focus
import           Statistics.Graph
import           Statistics.MemPool
import           Statistics.Report
-- import Statistics.Relay
import           Statistics.Throughput
import           Statistics.Tx
import           Types
import           Universum
import           Util.Pipes (fold')

runJSONFold :: FilePath -> Fold IndexedJLTimedEvent a -> IO a
runJSONFold logDir fd = runParseLogs logDir $ fold' fd

receivedCreatedF :: Fold IndexedJLTimedEvent (Map TxHash (Maybe Timestamp))
receivedCreatedF = f <$> txFirstReceivedF <*> inBlockChainF
  where
    f :: Map TxHash Timestamp -> Map TxHash Timestamp -> Map TxHash (Maybe Timestamp)
    f rm cm = M.mapWithKey g rm
      where
        g :: TxHash -> Timestamp -> Maybe Timestamp
        g tx ts = maybe Nothing (\ts' -> Just $ ts' - ts) $ M.lookup tx cm

max_stars_repo_name: Davdegreat/cardano-sl

max_issues_repo_name: Davdegreat/cardano-sl

max_forks_repo_name: Davdegreat/cardano-sl