{-# 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


{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE MagicHash            #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE UnboxedTuples        #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}

-- |
-- Module      :  Internal.Vector
-- Copyright   :  (c) Alberto Ruiz 2007-15
-- License     :  BSD3
-- Maintainer  :  Alberto Ruiz
-- Stability   :  provisional
--

module Internal.Vector(
    I,Z,R,C,
    fi,ti,
    Vector, fromList, unsafeToForeignPtr, unsafeFromForeignPtr, unsafeWith,
    createVector, avec, inlinePerformIO,
    toList, dim, (@>), at', (|>),
    vjoin, subVector, takesV, idxs,
    buildVector,
    asReal, asComplex,
    toByteString,fromByteString,
    zipVector, unzipVector, zipVectorWith, unzipVectorWith,
    foldVector, foldVectorG, foldVectorWithIndex, foldLoop,
    mapVector, mapVectorM, mapVectorM_,
    mapVectorWithIndex, mapVectorWithIndexM, mapVectorWithIndexM_
) where

import           Data.Complex
import           Data.Int                      (Int64)
import           Data.Vector.Storable          (Vector, fromList, unsafeFromForeignPtr,
                                                unsafeToForeignPtr, unsafeWith)
import qualified Data.Vector.Storable          as Vector
import           Foreign.C.Types               (CInt)
import           Foreign.ForeignPtr
import           Foreign.Marshal.Array
import           Foreign.Ptr
import           Foreign.Storable
import           GHC.Base                      (IO (IO), realWorld#, when)
import           GHC.ForeignPtr                (mallocPlainForeignPtrBytes)
import           System.IO.Unsafe              (unsafePerformIO)

import           Control.Monad                 (replicateM)
import           Data.Binary
import           Data.Binary.Put
import qualified Data.ByteString.Internal      as BS
import           Data.Vector.Storable.Internal (updPtr)

type I = CInt
type Z = Int64
type R = Float
type C = Complex Float

-- | specialized fromIntegral
fi :: Int -> CInt
fi = fromIntegral

-- | specialized fromIntegral
ti :: CInt -> Int
ti = fromIntegral

-- | Number of elements
dim :: (Storable t) => Vector t -> Int
dim = Vector.length
{-# INLINE dim #-}

-- C-Haskell vector adapter
{-# INLINE avec #-}
avec :: Storable a => Vector a -> (f -> IO r) -> ((CInt -> Ptr a -> f) -> IO r)
avec v f g = unsafeWith v $ \ptr -> f (g (fromIntegral (Vector.length v)) ptr)

-- allocates memory for a new vector
createVector :: Storable a => Int -> IO (Vector a)
createVector n = do
    when (n < 0) $ error ("trying to createVector of negative dim: "++show n)
    fp <- doMalloc undefined
    return $ unsafeFromForeignPtr fp 0 n
  where
    --
    -- Use the much cheaper Haskell heap allocated storage
    -- for foreign pointer space we control
    --
    doMalloc :: Storable b => b -> IO (ForeignPtr b)
    doMalloc dummy = do
        mallocPlainForeignPtrBytes (n * sizeOf dummy)

{- | creates a Vector from a list:

@> fromList [2,3,5,7]
4 |> [2.0,3.0,5.0,7.0]@

-}

safeRead :: Storable a => Vector a -> (Ptr a -> IO c) -> c
safeRead v = inlinePerformIO . unsafeWith v
{-# INLINE safeRead #-}

inlinePerformIO :: IO a -> a
inlinePerformIO (IO m) = case m realWorld# of (# _, r #) -> r
{-# INLINE inlinePerformIO #-}

{- extracts the Vector elements to a list

>>> toList (linspace 5 (1,10))
[1.0,3.25,5.5,7.75,10.0]

-}
toList :: Storable a => Vector a -> [a]
toList v = safeRead v $ peekArray (dim v)

{- | Create a vector from a list of elements and explicit dimension. The input
     list is truncated if it is too long, so it may safely
     be used, for instance, with infinite lists.

>>> 5 |> [1..]
[1.0,2.0,3.0,4.0,5.0]
it :: (Enum a, Num a, Foreign.Storable.Storable a) => Vector a

-}
(|>) :: (Storable a) => Int -> [a] -> Vector a
infixl 9 |>
n |> l
    | length l' == n = fromList l'
    | otherwise      = error "list too short for |>"
  where
    l' = take n l

-- | Create a vector of indexes, useful for matrix extraction using '(??)'
idxs :: [Int] -> Vector I
idxs js = fromList (map fromIntegral js) :: Vector I

{- | takes a number of consecutive elements from a Vector

>>> subVector 2 3 (fromList [1..10])
[3.0,4.0,5.0]
it :: (Enum t, Num t, Foreign.Storable.Storable t) => Vector t

-}
subVector :: Storable t => Int       -- ^ index of the starting element
                        -> Int       -- ^ number of elements to extract
                        -> Vector t  -- ^ source
                        -> Vector t  -- ^ result
subVector = Vector.slice
{-# INLINE subVector #-}

{- | Reads a vector position:

>>> fromList [0..9] @> 7
7.0

-}
(@>) :: Storable t => Vector t -> Int -> t
infixl 9 @>
v @> n
    | n >= 0 && n < dim v = at' v n
    | otherwise = error "vector index out of range"
{-# INLINE (@>) #-}

-- | access to Vector elements without range checking
at' :: Storable a => Vector a -> Int -> a
at' v n = safeRead v $ flip peekElemOff n
{-# INLINE at' #-}

{- | concatenate a list of vectors

>>> vjoin [fromList [1..5::Float], konst 1 3]
[1.0,2.0,3.0,4.0,5.0,1.0,1.0,1.0]
it :: Vector Float

-}
vjoin :: Storable t => [Vector t] -> Vector t
vjoin [] = fromList []
vjoin [v] = v
vjoin as = unsafePerformIO $ do
    let tot = sum (map dim as)
    r <- createVector tot
    unsafeWith r $ \ptr ->
        joiner as tot ptr
    return r
  where joiner [] _ _ = return ()
        joiner (v:cs) _ p = do
            let n = dim v
            unsafeWith v $ \pb -> copyArray p pb n
            joiner cs 0 (advancePtr p n)

{- | Extract consecutive subvectors of the given sizes.

>>> takesV [3,4] (linspace 10 (1,10::Float))
[[1.0,2.0,3.0],[4.0,5.0,6.0,7.0]]
it :: [Vector Float]

-}
takesV :: Storable t => [Int] -> Vector t -> [Vector t]
takesV ms w | sum ms > dim w = error $ "takesV " ++ show ms ++ " on dim = " ++ (show $ dim w)
            | otherwise = go ms w
    where go [] _ = []
          go (n:ns) v = subVector 0 n v
                      : go ns (subVector n (dim v - n) v)

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

-- | transforms a complex vector into a real vector with alternating real and imaginary parts
asReal :: (RealFloat a, Storable a) => Vector (Complex a) -> Vector a
asReal v = unsafeFromForeignPtr (castForeignPtr fp) (2*i) (2*n)
    where (fp,i,n) = unsafeToForeignPtr v

-- | transforms a real vector into a complex vector with alternating real and imaginary parts
asComplex :: (RealFloat a, Storable a) => Vector a -> Vector (Complex a)
asComplex v = unsafeFromForeignPtr (castForeignPtr fp) (i `div` 2) (n `div` 2)
    where (fp,i,n) = unsafeToForeignPtr v

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

-- | map on Vectors
mapVector :: (Storable a, Storable b) => (a-> b) -> Vector a -> Vector b
mapVector f v = unsafePerformIO $ do
    w <- createVector (dim v)
    unsafeWith v $ \p ->
        unsafeWith w $ \q -> do
            let go (-1) = return ()
                go !k = do x <- peekElemOff p k
                           pokeElemOff      q k (f x)
                           go (k-1)
            go (dim v -1)
    return w
{-# INLINE mapVector #-}

-- | zipWith for Vectors
zipVectorWith :: (Storable a, Storable b, Storable c) => (a-> b -> c) -> Vector a -> Vector b -> Vector c
zipVectorWith f u v = unsafePerformIO $ do
    let n = min (dim u) (dim v)
    w <- createVector n
    unsafeWith u $ \pu ->
        unsafeWith v $ \pv ->
            unsafeWith w $ \pw -> do
                let go (-1) = return ()
                    go !k = do x <- peekElemOff pu k
                               y <- peekElemOff pv k
                               pokeElemOff      pw k (f x y)
                               go (k-1)
                go (n -1)
    return w
{-# INLINE zipVectorWith #-}

-- | unzipWith for Vectors
unzipVectorWith :: (Storable (a,b), Storable c, Storable d)
                   => ((a,b) -> (c,d)) -> Vector (a,b) -> (Vector c,Vector d)
unzipVectorWith f u = unsafePerformIO $ do
      let n = dim u
      v <- createVector n
      w <- createVector n
      unsafeWith u $ \pu ->
          unsafeWith v $ \pv ->
              unsafeWith w $ \pw -> do
                  let go (-1) = return ()
                      go !k   = do z <- peekElemOff pu k
                                   let (x,y) = f z
                                   pokeElemOff      pv k x
                                   pokeElemOff      pw k y
                                   go (k-1)
                  go (n-1)
      return (v,w)
{-# INLINE unzipVectorWith #-}

foldVector :: Storable a => (a -> b -> b) -> b -> Vector a -> b
foldVector f x v = unsafePerformIO $
    unsafeWith v $ \p -> do
        let go (-1) s = return s
            go !k !s = do y <- peekElemOff p k
                          go (k-1::Int) (f y s)
        go (dim v -1) x
{-# INLINE foldVector #-}

-- the zero-indexed index is passed to the folding function
foldVectorWithIndex :: Storable a => (Int -> a -> b -> b) -> b -> Vector a -> b
foldVectorWithIndex f x v = unsafePerformIO $
    unsafeWith v $ \p -> do
        let go (-1) s = return s
            go !k !s = do y <- peekElemOff p k
                          go (k-1::Int) (f k y s)
        go (dim v -1) x
{-# INLINE foldVectorWithIndex #-}

foldLoop :: (Int -> t -> t) -> t -> Int -> t
foldLoop f s0 d = go (d - 1) s0
     where
       go 0 s   = f (0::Int) s
       go !j !s = go (j - 1) (f j s)

foldVectorG :: Storable t1 => (Int -> (Int -> t1) -> t -> t) -> t -> Vector t1 -> t
foldVectorG f s0 v = foldLoop g s0 (dim v)
    where g !k !s = f k (safeRead v . flip peekElemOff) s
          {-# INLINE g #-} -- Thanks to Ryan Ingram (http://permalink.gmane.org/gmane.comp.lang.haskell.cafe/46479)
{-# INLINE foldVectorG #-}

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

-- | monadic map over Vectors
--    the monad @m@ must be strict
mapVectorM :: (Storable a, Storable b, Monad m) => (a -> m b) -> Vector a -> m (Vector b)
mapVectorM f v = do
    w <- return $! unsafePerformIO $! createVector (dim v)
    mapVectorM' w 0 (dim v -1)
    return w
    where mapVectorM' w' !k !t
              | k == t               = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f x
                                       return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
              | otherwise            = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f x
                                       _ <- return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
                                       mapVectorM' w' (k+1) t
{-# INLINE mapVectorM #-}

-- | monadic map over Vectors
mapVectorM_ :: (Storable a, Monad m) => (a -> m ()) -> Vector a -> m ()
mapVectorM_ f v = do
    mapVectorM' 0 (dim v -1)
    where mapVectorM' !k !t
              | k == t            = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    f x
              | otherwise         = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    _ <- f x
                                    mapVectorM' (k+1) t
{-# INLINE mapVectorM_ #-}

-- | monadic map over Vectors with the zero-indexed index passed to the mapping function
--    the monad @m@ must be strict
mapVectorWithIndexM :: (Storable a, Storable b, Monad m) => (Int -> a -> m b) -> Vector a -> m (Vector b)
mapVectorWithIndexM f v = do
    w <- return $! unsafePerformIO $! createVector (dim v)
    mapVectorM' w 0 (dim v -1)
    return w
    where mapVectorM' w' !k !t
              | k == t               = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f k x
                                       return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
              | otherwise            = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f k x
                                       _ <- return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
                                       mapVectorM' w' (k+1) t
{-# INLINE mapVectorWithIndexM #-}

-- | monadic map over Vectors with the zero-indexed index passed to the mapping function
mapVectorWithIndexM_ :: (Storable a, Monad m) => (Int -> a -> m ()) -> Vector a -> m ()
mapVectorWithIndexM_ f v = do
    mapVectorM' 0 (dim v -1)
    where mapVectorM' !k !t
              | k == t            = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    f k x
              | otherwise         = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    _ <- f k x
                                    mapVectorM' (k+1) t
{-# INLINE mapVectorWithIndexM_ #-}

mapVectorWithIndex :: (Storable a, Storable b) => (Int -> a -> b) -> Vector a -> Vector b
--mapVectorWithIndex g = head . mapVectorWithIndexM (\a b -> [g a b])
mapVectorWithIndex f v = unsafePerformIO $ do
    w <- createVector (dim v)
    unsafeWith v $ \p ->
        unsafeWith w $ \q -> do
            let go (-1) = return ()
                go !k = do x <- peekElemOff p k
                           pokeElemOff      q k (f k x)
                           go (k-1)
            go (dim v -1)
    return w
{-# INLINE mapVectorWithIndex #-}

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

-- a 64K cache, with a Float taking 13 bytes in Bytestring,
-- implies a chunk size of 5041
chunk :: Int
chunk = 5000

chunks :: Int -> [Int]
chunks d = let c = d `div` chunk
               m = d `mod` chunk
           in if m /= 0 then reverse (m:(replicate c chunk)) else (replicate c chunk)

putVector :: (Storable t, Binary t) => Vector t -> Data.Binary.Put.PutM ()
putVector v = mapM_ put $! toList v

getVector :: (Storable a, Binary a) => Int -> Get (Vector a)
getVector d = do
              xs <- replicateM d get
              return $! fromList xs

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

toByteString :: Storable t => Vector t -> BS.ByteString
toByteString v = BS.PS (castForeignPtr fp) (sz*o) (sz * dim v)
  where
    (fp,o,_n) = unsafeToForeignPtr v
    sz = sizeOf (v@>0)

fromByteString :: Storable t => BS.ByteString -> Vector t
fromByteString (BS.PS fp o n) = r
  where
    r = unsafeFromForeignPtr (castForeignPtr (updPtr (`plusPtr` o) fp)) 0 n'
    n' = n `div` sz
    sz = sizeOf (r@>0)

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

instance (Binary a, Storable a) => Binary (Vector a) where

    put v = do
            let d = dim v
            put d
            mapM_ putVector $! takesV (chunks d) v

    -- put = put . v2bs

    get = do
          d <- get
          vs <- mapM getVector $ chunks d
          return $! vjoin vs

    -- get = fmap bs2v get

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

{- | creates a Vector of the specified length using the supplied function to
     to map the index to the value at that index.

@> buildVector 4 fromIntegral
4 |> [0.0,1.0,2.0,3.0]@

-}
buildVector :: Storable a => Int -> (Int -> a) -> Vector a
buildVector len f =
    fromList $ map f [0 .. (len - 1)]

-- | zip for Vectors
zipVector :: (Storable a, Storable b, Storable (a,b)) => Vector a -> Vector b -> Vector (a,b)
zipVector = zipVectorWith (,)

-- | unzip for Vectors
unzipVector :: (Storable a, Storable b, Storable (a,b)) => Vector (a,b) -> (Vector a,Vector b)
unzipVector = unzipVectorWith id

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


--
-- ConvLayer : convolution layer
--

module CNN.ConvLayer (
  initFilterC
, convolve
, deconvolve
, reverseConvFilter
, updateConvFilter
) where

import Control.Monad hiding (msum)
import Data.List hiding (transpose)
import Data.Maybe
import Debug.Trace
import Numeric.LinearAlgebra
import System.Random.Mersenne as MT

import CNN.Algebra
import CNN.Image
import CNN.LayerType

{- |
initFilterC

  IN : #kernel
       #channel
       input size (x * y)
       kernel size (n * n)
       pool size (m * m)

  OUT: filter of convolution layer

-}

initFilterC :: Int -> Int -> Int -> Int -> Int -> Int -> IO [FilterC]
initFilterC k c x y n m = do
  f <- forM [1..k] $ \i -> do
    w <- initKernel c n
    return (w, 0.0)
  return f
  where
    ri = fromIntegral (n * n * c)   -- resolution (input)
    ro = fromIntegral (n * n * k) / (fromIntegral (m * m))
    a = sqrt (6.0 / (ri + ro))
    initKernel :: Int -> Int -> IO Kernel
    initKernel c n = do
      let
        sz = n * n
      w <- forM [1..c] $ \i -> do
        w' <- forM [1..sz] $ \j -> do
          r <- MT.randomIO :: IO Double
          return ((r * 2.0 - 1.0) * a)
        return w'
      return $ fromLists w

--

{- |
convolve

  IN: kernel size
      filter
      image

test: filter k=3, c=2, s=2
>>> let filter = [(fromLists [[1.0,2.0,3.0,4.0],[5.0,6.0,7.0,8.0]],0.25),(fromLists [[3.0,4.0,5.0,6.0],[7.0,8.0,1.0,2.0]],0.5),(fromLists [[5.0,6.0,7.0,8.0],[1.0,2.0,3.0,4.0]],0.75)] :: [FilterC]
>>> let img1 = [fromLists [[1.0,2.0,3.0],[4.0,5.0,6.0],[7.0,8.0,9.0]], fromLists [[4.0,5.0,6.0],[7.0,8.0,9.0],[1.0,2.0,3.0]]] :: Image
>>> let img2 = [fromLists [[1.0,2.0,3.0,4.0],[4.0,5.0,6.0,7.0],[7.0,8.0,9.0,10.0],[10.0,11.0,12.0,13.0]],fromLists [[4.0,5.0,6.0,7.0],[7.0,8.0,9.0,10.0],[1.0,2.0,3.0,4.0],[4.0,5.0,6.0,7.0]]] :: Image
>>> convolve 2 filter img1
[(2><2)
 [ 200.25, 236.25
 , 173.25, 209.25 ],(2><2)
 [ 152.5, 188.5
 , 233.5, 269.5 ],(2><2)
 [ 152.75, 188.75
 , 197.75, 233.75 ]]
>>> convolve 2 filter img2
[(3><3)
 [ 200.25, 236.25, 272.25
 , 173.25, 209.25, 245.25
 , 182.25, 218.25, 254.25 ],(3><3)
 [ 152.5, 188.5, 224.5
 , 233.5, 269.5, 305.5
 , 206.5, 242.5, 278.5 ],(3><3)
 [ 152.75, 188.75, 224.75
 , 197.75, 233.75, 269.75
 , 278.75, 314.75, 350.75 ]]

-}

--
-- USE corr2 ??
--

convolve :: Int -> [FilterC] -> Image -> Image
convolve s fs im = map (convolveImage x iss) fs
  where
    pl = head im
    x  = cols pl - (s - 1)
    y  = rows pl - (s - 1)
    ps = [(i, j) | i <- [0..(x-1)], j <- [0..(y-1)]]
    iss = map subImage im
    subImage :: Plain -> Matrix R
    subImage i = fromColumns $ map (\p -> flatten $ subMatrix p (s, s) i) ps

convolveImage :: Int -> [Matrix R] -> FilterC -> Plain
convolveImage x iss (k, b) = reshape x $ cmap (+ b) $ vsum vs
  where
    vs = zipWith (<#) (toRows k) iss

-- back prop

deconvolve :: Int -> [FilterC] -> Image -> Delta -> (Delta, Maybe Layer)
deconvolve s fs im d = (delta, Just (ConvLayer s dw))
  where
    delta = convolve s fs $ addBorder s d
    sim = slideImage s im
    dw = map (hadamard sim) d

{- |
addBorder

  IN : kernel size
       delta (Image)
  OUT: delta (Image)

>>> let s = [fromLists [[1.0, 2.0],[3.0,4.0]], fromLists [[5.0,6.0],[7.0,8.0]]]
>>> addBorder 2 s
[(4><4)
 [ 0.0, 0.0, 0.0, 0.0
 , 0.0, 1.0, 2.0, 0.0
 , 0.0, 3.0, 4.0, 0.0
 , 0.0, 0.0, 0.0, 0.0 ],(4><4)
 [ 0.0, 0.0, 0.0, 0.0
 , 0.0, 5.0, 6.0, 0.0
 , 0.0, 7.0, 8.0, 0.0
 , 0.0, 0.0, 0.0, 0.0 ]]

-}

addBorder :: Int -> Delta -> Delta
addBorder s d = map ab d
  where
    s' = s - 1
    ab :: Plain -> Plain
    ab p = yb ||| ((xb === (p === xb)) ||| yb)
      where
        x = cols p
        y = rows p
        xb = (s'><x) $ repeat 0.0
        yb = ((y + 2 * s')><s') $ repeat 0.0

{- |
slideImage

>>> let p = [fromLists [[1.0,2.0,3.0,4.0],[5.0,6.0,7.0,8.0],[9.0,10.0,11.0,12.0],[13.0,14.0,15.0,16.0]]] :: Image
>>> slideImage 2 p
[[(3><3)
 [ 1.0,  2.0,  3.0
 , 5.0,  6.0,  7.0
 , 9.0, 10.0, 11.0 ],(3><3)
 [  2.0,  3.0,  4.0
 ,  6.0,  7.0,  8.0
 , 10.0, 11.0, 12.0 ],(3><3)
 [  5.0,  6.0,  7.0
 ,  9.0, 10.0, 11.0
 , 13.0, 14.0, 15.0 ],(3><3)
 [  6.0,  7.0,  8.0
 , 10.0, 11.0, 12.0
 , 14.0, 15.0, 16.0 ]]]

-}

slideImage :: Int -> Image -> [[Plain]]
slideImage s im = map (slidePlain s) im

{- |
slidePlain

>>> let p = fromLists [[1.0,2.0,3.0,4.0],[5.0,6.0,7.0,8.0],[9.0,10.0,11.0,12.0],[13.0,14.0,15.0,16.0]] :: Plain
>>> slidePlain 2 p
[(3><3)
 [ 1.0,  2.0,  3.0
 , 5.0,  6.0,  7.0
 , 9.0, 10.0, 11.0 ],(3><3)
 [  2.0,  3.0,  4.0
 ,  6.0,  7.0,  8.0
 , 10.0, 11.0, 12.0 ],(3><3)
 [  5.0,  6.0,  7.0
 ,  9.0, 10.0, 11.0
 , 13.0, 14.0, 15.0 ],(3><3)
 [  6.0,  7.0,  8.0
 , 10.0, 11.0, 12.0
 , 14.0, 15.0, 16.0 ]]

>>> slidePlain 3 p
[(2><2)
 [ 1.0, 2.0
 , 5.0, 6.0 ],(2><2)
 [ 2.0, 3.0
 , 6.0, 7.0 ],(2><2)
 [ 3.0, 4.0
 , 7.0, 8.0 ],(2><2)
 [ 5.0,  6.0
 , 9.0, 10.0 ],(2><2)
 [  6.0,  7.0
 , 10.0, 11.0 ],(2><2)
 [  7.0,  8.0
 , 11.0, 12.0 ],(2><2)
 [  9.0, 10.0
 , 13.0, 14.0 ],(2><2)
 [ 10.0, 11.0
 , 14.0, 15.0 ],(2><2)
 [ 11.0, 12.0
 , 15.0, 16.0 ]]

-}

slidePlain :: Int -> Plain -> [Plain]
slidePlain s m
  | s < 1     = []
  | s == 1    = [m]
  | otherwise = map (\p -> subMatrix p (x, y) m) ps
  where
    x = cols m - (s - 1)
    y = rows m - (s - 1)
    ps = [(i, j) | i <- [0..s-1], j <- [0..s-1]]

{- |
hadamard

  IN : slided images
       delta
  OUT: (Kernel, Bias)

>>> let p = [fromLists [[1.0,2.0,3.0,4.0],[5.0,6.0,7.0,8.0],[9.0,10.0,11.0,12.0],[13.0,14.0,15.0,16.0]], fromLists [[16.0,15.0,14.0,13.0],[12.0,11.0,10.0,9.0],[8.0,7.0,6.0,5.0],[4.0,3.0,2.0,1.0]]] :: Image
>>> let ds = fromLists [[0.1,0.2,0.3],[0.4,0.5,0.6],[0.7,0.8,0.9]]
>>> hadamard (slideImage 2 p) ds
((2><4)
 [ 34.800000000000004, 39.3, 52.800000000000004, 57.3
 ,               41.7, 37.2, 23.700000000000006, 19.2 ],4.5)

-}

hadamard :: [[Plain]] -> Plain -> FilterC
hadamard sim ds = (w, b)
  where
    b = sumElements ds
    d' = flatten ds
    w = fromLists $ map (map ((<.>) d' . flatten)) sim

{- |
reverseConvFilter

  IN : kernel size
       original filter
  OUT: reversed filter

>>> let fs = [(fromLists [[0.1,0.2,0.3,0.4],[0.5,0.6,0.7,0.8]],1.0),(fromLists [[0.1,0.3,0.5,0.7],[0.2,0.4,0.6,0.8]],2.0),(fromLists [[0.1,0.2,0.5,0.6],[0.3,0.4,0.7,0.8]],3.0)] :: [FilterC]
>>> let ConvLayer s fs' = reverseConvFilter 2 fs
>>> s
2
>>> fs'
[((3><4)
 [ 0.4, 0.3, 0.2, 0.1
 , 0.7, 0.5, 0.3, 0.1
 , 0.6, 0.5, 0.2, 0.1 ],0.0),((3><4)
 [ 0.8, 0.7, 0.6, 0.5
 , 0.8, 0.6, 0.4, 0.2
 , 0.8, 0.7, 0.4, 0.3 ],0.0)]

-}

reverseConvFilter :: Int -> [FilterC] -> Layer
reverseConvFilter s fs = ConvLayer s (zip rv (repeat 0.0))
  where
    (k, _) = unzip fs
    k'     = map (toLists . fliprl) k
    --r      = transpose $ map (map reverse) k'
    r      = transpose k'
    rv     = map fromLists r

-- update

updateConvFilter :: Int -> [FilterC] -> Double -> [Maybe Layer] -> Layer
updateConvFilter s fs lr dl
  | dl' == [] = ConvLayer s fs
  | otherwise = ConvLayer s $ zip ks' (toList bs')
  where
    dl' = catMaybes dl
    sc = lr / fromIntegral (length dl')
    (ks , bs ) = unzip fs
    (kss, bss) = unzip $ map unzip $ strip dl'
    dbs = vsum $ map fromList bss
    dks = map (mscale sc . msum) $ transpose kss
    bs' = (fromList bs) - (vscale sc dbs)
    ks' = zipWith (-) ks dks

strip :: [Layer] -> [[FilterC]]
strip [] = []
strip (ConvLayer s fs:ds) = fs:strip ds
strip (_:ds) = strip ds  -- if not ConvLayer


-- | Data structure of compiler. https://en.wikibooks.org/wiki/Write_Yourself_a_Scheme_in_48_Hours/

module LispVal
       ( LispVal(..)
       , Env
       , unwordsList
       ) where

import           Data.Complex
import           Data.IORef
import           Data.Ratio                 (denominator, numerator)
import           ErrorCheckingAndExceptions
import           System.IO                  (Handle)

data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal
             | Number Integer
             | Float Double
             | Ratio Rational
             | Complex (Complex Double)
             | Character Char
             | String String
             | Bool Bool
             | PrimitiveFunc ([LispVal] -> ThrowsError LispVal)
             | Func { params :: [String], vararg :: Maybe String, body :: [LispVal], closure :: Env}
             | IOFunc ([LispVal] -> IOThrowsError LispVal)
             | Port Handle

showVal :: LispVal -> String
showVal (Atom x) = x
showVal (Bool True) = "#t"
showVal (Bool False) = "#f"
showVal (Ratio x) = show (numerator x) ++ ('/' : show (denominator x))
showVal (Complex x) = show (realPart x) ++ ('+' : show (imagPart x) ++ "i")
showVal (Float x) = show x
showVal (Number x) = show x
showVal (Character x) = "#\\" ++ show x
showVal (String x) = "\"" ++ x ++ "\""
showVal (List x) = "(" ++ unwordsList x ++ ")"
showVal (DottedList first rest)= "(" ++ unwordsList first ++ " . " ++ showVal rest ++ ")"
showVal (PrimitiveFunc _) = "<primitive>"
showVal (Func {params = args, vararg = varargs, body = body, closure = env}) =
  "(lambda (" ++ unwords (map show args) ++
  (case varargs of
     Nothing -> ""
     Just arg -> " . " ++ arg) ++ ") ...)"
showVal (Port _) = "<IO port>"
showVal (IOFunc _) = "<IO primitive>"

unwordsList :: [LispVal] -> String
unwordsList = unwords . map showVal

instance Show LispVal where
  show = showVal

type Env = IORef [(String, IORef LispVal)]


{-# LANGUAGE ScopedTypeVariables, TypeFamilies, TupleSections, FlexibleContexts #-}

module Fuml.Unsupervised where

import Fuml.Core
import qualified Fuml.Base.PCA as PCA
import Fuml.Base.KNN (euclideanDistance)
import Numeric.LinearAlgebra
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector as V
import Data.List (foldl1', sortBy, nub)
import Data.Ord (comparing)
import Control.Monad.Identity

type Unsupervisor m p a = Supervisor m () p a

with :: b -> [a] -> [(a,b)]
with y = map (,y)

data PCA = PCA { pcaComponents :: Int, pcaStat :: PCA.Stat }

pca :: Int -> Unsupervisor Identity PCA.Stat (Vector Double)
pca ncomp  = Supervisor $ \_ theData ->
  let stat@(m,_,v) = PCA.statVs $ map fst theData
  in return $ Predict stat $ \x -> takeRows ncomp v #> (x - m)

findNearestCentroidIx :: [Vector Double] -> Vector Double -> Int
findNearestCentroidIx ctrs x =
  let vdixs = zip [0..] $ map (\c -> euclideanDistance c x) ctrs
  in fst $ head $ sortBy (comparing snd) vdixs

centroid :: [Vector Double] -> Vector Double
centroid vs =
  let vadd = VS.zipWith (+)
      n = realToFrac $ length vs
  in VS.map (/n) $ foldl1' vadd vs

{-kmeans :: Int -> Unsupervisor Identity [Vector Double] Int
kmeans nclus
  =  Supervisor $ \_ theData ->
        let clus = KM.kmeans (VU.convert) KM.euclidSq nclus (map fst theData)
            ctrs = map (centroid . KM.elements) $ V.toList clus
            pr v = findNearestCentroidIx ctrs v
        in return $ Predict ctrs pr

cluster :: Eq b => [a] -> (a-> Vector Double) -> Unsupervisor Identity p b -> [[a]]
cluster xs f unsup =
  let p = runIdentity $ runSupervisor unsup Nothing $ with () $ map (f) xs
      withClus x = (x,predict p $ f x)
      withCluss = map withClus xs
      cluss = nub $ map snd withCluss
      getElems clus = map fst $ filter ((==clus) . snd) withCluss
  in map getElems cluss -}


{-|
Module : Mnist
|-}
module Mnist
  ( Image(..)
  , readTrainingData
  , readTestData
  , getData
  ) where

import qualified Codec.Compression.GZip as GZip
import qualified Data.ByteString.Lazy as BL
import Data.Binary.Get
import Data.Word
import qualified Data.List.Split as S
import Numeric.LinearAlgebra

-- | data representing an image
data Image = Image {
      iRows :: Int -- ^ Number of rows
    , iColumns :: Int -- ^ Number of columns
    , iPixels :: [Word8] -- ^ Pixels
    } deriving (Eq, Show)

-- | Create column vector of Image pixels
toColVector :: Image
         -> Matrix Double
toColVector image = (rc><1) p :: Matrix Double
  where rc = iRows image * iColumns image
        p = normalisedData image

-- | Normalise Image pixels
normalisedData :: Image -> [Double]
normalisedData image = normalisePixel <$> iPixels image

-- | Normalise pixel
normalisePixel :: Word8 -> Double
normalisePixel p = fromIntegral p / 255.0

{-
[offset] [type]          [value]          [description]
0000     32 bit integer  0x00000801(2049) magic number (MSB first)
0004     32 bit integer  60000            number of items
0008     unsigned byte   ??               label
0009     unsigned byte   ??               label
 ........
xxxx     unsigned byte   ??               label
The labels values are 0 to 9.
-}

-- | Extract information from labels file
deserialiseLabels :: Get (Word32, Word32, [Word8])
deserialiseLabels = do
  magicNumber <- getWord32be
  count <- getWord32be
  labelData <- getRemainingLazyByteString
  let labels = BL.unpack labelData
  return (magicNumber, count, labels)

-- | Read labels from file
readLabels :: FilePath -> IO [Int]
readLabels filename = do
  content <- GZip.decompress <$> BL.readFile filename
  let (_, _, labels) = runGet deserialiseLabels content
  return (fromIntegral <$> labels)

{-
[offset] [type]          [value]          [description]
0000     32 bit integer  0x00000803(2051) magic number
0004     32 bit integer  60000            number of images
0008     32 bit integer  28               number of rows
0012     32 bit integer  28               number of columns
0016     unsigned byte   ??               pixel
0017     unsigned byte   ??               pixel
 ........
xxxx     unsigned byte   ??               pixel
Pixels are organized row-wise. Pixel values are 0 to 255. 0 means background (white), 255 means foreground (black).
-}

-- | Extract information from image file
deserialiseHeader :: Get (Word32, Word32, Word32, Word32, [[Word8]])
deserialiseHeader = do
  magicNumber <- getWord32be
  imageCount <- getWord32be
  r <- getWord32be
  c <- getWord32be
  packedData <- getRemainingLazyByteString
  let len = fromIntegral (r * c)
  let unpackedData = S.chunksOf len (BL.unpack packedData)
  return (magicNumber, imageCount, r, c, unpackedData)

-- | Read Images form file
readImages :: FilePath -> IO [Image]
readImages filename = do
  content <- GZip.decompress <$> BL.readFile filename
  let (_, _, r, c, unpackedData) = runGet deserialiseHeader content
  return (Image (fromIntegral r) (fromIntegral c) <$> unpackedData)

-- | Create column vector of label
vectorizedLabels :: Int -> Matrix Double
vectorizedLabels x = (10><1) p :: Matrix Double
  where p = [i | j <- [0 .. 9], let i = if j == x then 1.0 else 0.0]

-- | Get training data
readTrainingData :: IO [(Matrix Double, Matrix Double)]
readTrainingData = do
  trainingLabels <- readLabels "train-labels-idx1-ubyte.gz"
  trainingImages <- readImages "train-images-idx3-ubyte.gz"
  return (zip (map toColVector trainingImages) (vectorizedLabels <$> trainingLabels))

-- | Get test data
readTestData :: IO [(Matrix Double, Int)]
readTestData = do
  testLabels <- readLabels "t10k-labels-idx1-ubyte.gz"
  testImages <- readImages "t10k-images-idx3-ubyte.gz"
  return (zip (map toColVector testImages) testLabels)

-- | Get training and test data
getData :: IO ([(Matrix Double, Matrix Double)], [(Matrix Double, Int)])
getData = do
  trainingData <- readTrainingData
  testData <- readTestData
  return (trainingData, testData)


{-# LANGUAGE FlexibleContexts #-}

module Data.AdaptiveCoordinateDescent.Internal
  ( normalize
  , perturb
  , makePerturbed
  , fromVector
  , generate
  , plus
  , minus
  , sphere
  , rosen )
  where

import Control.Monad
import Control.Monad.Primitive
import Control.Monad.State.Strict
import Control.Monad.ST
import Data.Foldable
import Data.Traversable
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Unboxed.Mutable as VM
import Numeric.LinearAlgebra hiding ( toList )
import Pipes hiding ( for )
import System.Random.MWC

normalize :: Matrix Double -> Matrix Double
normalize mat = runST $ do
  stds <- VM.replicate columns (0, 0)
  for_ [0..columns-1] $ \column -> do
    for_ [0..rows-1] $ \row -> do
      (avg, std) <- VM.read stds column
      VM.write stds column (avg + mat `atIndex` (row, column),
                            std + (mat `atIndex` (row, column))**2)
    (avg, std) <- VM.read stds column
    VM.write stds column (avg / fromIntegral rows
                         ,sqrt $ (std / fromIntegral rows) - (avg / fromIntegral rows)**2)

  fstds <- V.unsafeFreeze stds

  return $ (rows><columns) $ generate (rows*columns) $ \offset ->
             let row = offset `div` columns
                 column = offset `mod` columns

                 (avg, _var) = fstds V.! column -- variance ignored for now

              in ((mat `atIndex` (row, column)) - avg)
 where
  (rows, columns) = size mat

perturb :: (PrimMonad m, Traversable f)
        => f Double
        -> Gen (PrimState m)
        -> m (f Double)
perturb thing rng = for thing $ \val -> do
  perturbance <- uniformR (-0.1, 0.1) rng
  return $ val + perturbance

makePerturbed :: (PrimMonad m, Traversable f)
              => f Double
              -> Gen (PrimState m)
              -> Int
              -> Producer (f Double) m ()
makePerturbed params rng num_items =
  replicateM_ num_items $ do
    new_values <- lift $ perturb params rng
    yield new_values

fromVector :: (Container Vector a, Traversable f) => Vector a -> f void -> f a
fromVector vec structure =
  flip evalState 0 $ for structure $ \_ -> do
    idx <- get
    put (idx+1)
    return $ vec `atIndex` idx
{-# INLINE fromVector #-}

generate' :: Int -> (Int -> a) -> [a]
generate' 0 _ = []
generate' n fun = fun (n-1):generate' (n-1) fun

generate :: Int -> (Int -> a) -> [a]
generate x fun = reverse $ generate' x fun

plus :: (Num a, Traversable f) => f a -> f a -> f a
plus tr1 tr2 = flip evalState (toList tr2) $ for tr1 $ \item -> do
  (x:rest) <- get
  put rest
  return $ item + x

minus :: (Num a, Traversable f) => f a -> f a -> f a
minus tr1 tr2 = flip evalState (toList tr2) $ for tr1 $ \item -> do
  (x:rest) <- get
  put rest
  return $ item - x

sphere :: [Double] -> Double
sphere [x, y] = sqrt $ x*x + y*y
sphere _ = error "sphere takes two arguments."

rosen :: [Double] -> Double
rosen [x, y] = (1 - x)**2 + 100*(y - x**2)**2
rosen _ = error "rosen takes two arguments."


-- {-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE DataKinds, GADTs, TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FunctionalDependencies #-}

module UnitBLAS.Level2(unitTestLevel2BLAS) where

import Test.HUnit
--import Numerical.Array.Shape as S
import Prelude as P
import Test.Tasty
import Test.Tasty.HUnit

import qualified Data.Vector.Storable as SV
import qualified Data.Vector.Storable.Mutable as SMV

import Data.Complex

import  Numerical.HBLAS.MatrixTypes as Matrix
import  Numerical.HBLAS.BLAS as BLAS

--unitTestShape = testGroup "Shape Unit tests"
--    [ testCase "foldl on shape" $ ( S.foldl (+) 0 (1:* 2:* 3 :* Nil )  @?=  ( P.foldl   (+) 0  [1,2,3])  )
--    , testCase "foldr on shape" $ ( S.foldr (+) 0 (1:* 2:* 3 :* Nil )  @?=  ( P.foldr  (+) 0  [1,2,3])  )
--    , testCase "scanr1 on shape" (S.scanr1 (+) 0 (1:* 1 :* 1:* Nil )   @?=  (3:* 2:* 1 :* Nil ) )
--    , testCase "scanl1 on shape" (S.scanl1 (+) 0 (1:* 1 :* 1:* Nil )   @?=  (1:* 2:* 3:* Nil ) )
--    ]

matmatTest1SGEMV:: IO ()
matmatTest1SGEMV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2) (\_ -> (1.0::Float))
    right <- Matrix.generateMutableDenseVector 2 (\_ -> (1.0 :: Float))
    res  <- Matrix.generateMutableDenseVector  2 (\_ -> (0.0 :: Float))
    BLAS.sgemv Matrix.NoTranspose  1.0 1.0 left right res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList @?= [2,2]

matmatTest1DGEMV:: IO ()
matmatTest1DGEMV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2) (\_ -> (1.0))
    right <- Matrix.generateMutableDenseVector 2 (\_ -> (1.0 ))
    res  <- Matrix.generateMutableDenseVector 2  (\_ -> (0.0 ))
    BLAS.dgemv Matrix.NoTranspose 1.0 1.0 left right res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList @?= [2.0,2.0]

matmatTest1CGEMV:: IO ()
matmatTest1CGEMV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2) (\_ -> (1.0))
    right <- Matrix.generateMutableDenseVector 2 (\_ -> (1.0 ))
    res  <- Matrix.generateMutableDenseVector  2 (\_ -> (0.0 ))
    BLAS.cgemv Matrix.NoTranspose  1.0 1.0 left right res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList @?= [2.0,2.0]

matmatTest1ZGEMV:: IO ()
matmatTest1ZGEMV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2) (\_ -> (1.0))
    right <- Matrix.generateMutableDenseVector 2 (\_ -> (1.0 ))
    res  <- Matrix.generateMutableDenseVector 2 (\_ -> (0.0 ))
    BLAS.zgemv Matrix.NoTranspose  1.0 1.0 left right res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList @?= [2.0,2.0]

----
----

matmatTest1SGER :: IO ()
matmatTest1SGER = do
  res <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (2,2) (\_ -> 1.0)
  x <- Matrix.generateMutableDenseVector 2 (\_ -> 2.0)
  y <- Matrix.generateMutableDenseVector 2 (\_ -> 3.0)
  BLAS.sger 2.0 x y res
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat res
  resList @?= [13.0,13.0,13.0,13.0]

matmatTest1DGER :: IO ()
matmatTest1DGER = do
  res <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (2,2) (\_ -> 1.0)
  x <- Matrix.generateMutableDenseVector 2 (\_ -> 2.0)
  y <- Matrix.generateMutableDenseVector 2 (\_ -> 3.0)
  BLAS.sger 2.0 x y res
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat res
  resList @?= [13.0,13.0,13.0,13.0]

----
----

matmatTest1STRSV:: IO ()
matmatTest1STRSV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2)
            (\(i,j) -> if i >= j then (1.0::Float) else 0 )

    res  <- Matrix.generateMutableDenseVector  2 (\i -> if i == 0 then 3 else 1)
    BLAS.strsv MatUpper NoTranspose MatUnit left res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList @?= [2,1]

matmatTest1DTRSV:: IO ()
matmatTest1DTRSV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2)
                (\(i,j) -> if i >= j then (1.0::Double) else 0 )

    res  <- Matrix.generateMutableDenseVector  2 (\i -> if i == 0 then 3 else 1)
    BLAS.dtrsv MatUpper NoTranspose MatUnit left res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList @?= [2,1]

matmatTest1CTRSV:: IO ()
matmatTest1CTRSV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2)
                (\(i,j) -> if i >= j then (1.0::(Complex Float)) else 0 )

    res  <- Matrix.generateMutableDenseVector  2 (\i -> if i == 0 then 3 else 1)
    BLAS.ctrsv MatUpper NoTranspose MatUnit left res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList @?= [2,1]

matmatTest1ZTRSV:: IO ()
matmatTest1ZTRSV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2)
                (\(i,j) -> if i >= j then (1.0::(Complex Double )) else 0 )
    res  <- Matrix.generateMutableDenseVector  2 (\i -> if i == 0 then 3 else 1)
    BLAS.ztrsv MatUpper NoTranspose MatUnit left res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList @?= [2,1]

unitTestLevel2BLAS = testGroup "BLAS Level 2 tests " [
----- gemv tests
    testCase "sgemv on 2x2 all 1s"    matmatTest1SGEMV
    ,testCase "dgemv  on 2x2 all 1s " matmatTest1DGEMV
    ,testCase "cgemv  on 2x2 all 1s" matmatTest1CGEMV
    ,testCase "zgemv on 2x2 all 1s" matmatTest1ZGEMV
----- trsv tests
    ,testCase "strsv on 2x2 upper 1s" matmatTest1STRSV
    ,testCase "dtrsv on 2x2 upper 1s" matmatTest1DTRSV
    ,testCase "ctrsv on 2x2 upper 1s" matmatTest1CTRSV
    ,testCase "ztrsv on 2x2 upper 1s" matmatTest1ZTRSV
---- ger tests
    ,testCase "sger on 2x2 all 1s" matmatTest1SGER
    ,testCase "dger on 2x2 all 1s" matmatTest1DGER
    ]