{-# LANGUAGE DataKinds, RankNTypes, TypeFamilies #-}
module TestMnistFCNN
  ( testTrees, shortTestForCITrees, mnistTestCase2T, mnistTestCase2D
  ) where

import Prelude

import           Control.DeepSeq
import           Control.Monad (foldM, when)
import           Data.Array.Internal (valueOf)
import           Data.Coerce (coerce)
import           Data.List.Index (imap)
import           Data.Proxy (Proxy (Proxy))
import           Data.Time.Clock.POSIX (POSIXTime, getPOSIXTime)
import qualified Data.Vector.Generic as V
import           GHC.TypeLits (KnownNat)
import qualified Numeric.LinearAlgebra as HM
import           System.IO (hFlush, hPutStrLn, stderr)
import           System.Random
import           Test.Tasty
import           Test.Tasty.HUnit hiding (assert)
import           Test.Tasty.QuickCheck hiding (label, shuffle)
import           Text.Printf

import HordeAd
import HordeAd.Core.OutdatedOptimizer
import HordeAd.Tool.MnistTools

import TestCommon

testTrees :: [TestTree]
testTrees = [ dumbMnistTests
            , bigMnistTests
            , vectorMnistTests
            , matrixMnistTests
            , fusedMnistTests
            ]

shortTestForCITrees :: [TestTree]
shortTestForCITrees = [ dumbMnistTests
                      , shortCIMnistTests
                      ]

sgdShow :: HasDelta r
        => r
        -> (a -> DualNumberVariables 'DModeGradient r -> DualMonadGradient r (DualNumber 'DModeGradient r))
        -> [a]  -- ^ training data
        -> Domain0 r  -- ^ initial parameters
        -> r
sgdShow gamma f trainData params0Init =
  let result =
        fst $ sgd gamma f trainData (params0Init, V.empty, V.empty, V.empty)
  in snd $ dReverse 1 (f $ head trainData) result

sgdTestCase :: String
            -> IO [a]
            -> (Int
                -> Int
                -> a
                -> DualNumberVariables 'DModeGradient Double
                -> DualMonadGradient Double
                                     (DualNumber 'DModeGradient Double))
            -> Double
            -> Double
            -> TestTree
sgdTestCase prefix trainDataIO trainWithLoss gamma expected =
  let widthHidden = 250
      widthHidden2 = 50
      nParams0 = fcnnMnistLen0 widthHidden widthHidden2
      vec = HM.randomVector 33 HM.Uniform nParams0 - HM.scalar 0.5
      name = prefix ++ " "
             ++ unwords [show widthHidden, show nParams0, show gamma]
  in testCase name $ do
       trainData <- trainDataIO
       sgdShow gamma (trainWithLoss widthHidden widthHidden2)
                     trainData vec
         @?= expected

mnistTestCase2
  :: String
  -> Int
  -> Int
  -> (Int
      -> Int
      -> MnistData Double
      -> DualNumberVariables 'DModeGradient Double
      -> DualMonadGradient Double (DualNumber 'DModeGradient Double))
  -> Int
  -> Int
  -> Double
  -> Double
  -> TestTree
mnistTestCase2 prefix epochs maxBatches trainWithLoss widthHidden widthHidden2
               gamma expected =
  let nParams0 = fcnnMnistLen0 widthHidden widthHidden2
      params0Init = HM.randomVector 44 HM.Uniform nParams0 - HM.scalar 0.5
      name = prefix ++ ": "
             ++ unwords [ show epochs, show maxBatches
                        , show widthHidden, show widthHidden2
                        , show nParams0, show gamma ]
  in testCase name $ do
       hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
              prefix epochs maxBatches
       trainData <- loadMnistData trainGlyphsPath trainLabelsPath
       testData <- loadMnistData testGlyphsPath testLabelsPath
       -- Mimic how backprop tests and display it, even though tests
       -- should not print, in principle.
       let runBatch :: Domain0 Double -> (Int, [MnistData Double])
                    -> IO (Domain0 Double)
           runBatch !params0 (k, chunk) = do
             let f = trainWithLoss widthHidden widthHidden2
                 (!res, _, _, _) =
                   fst $ sgd gamma f chunk (params0, V.empty, V.empty, V.empty)
                 !trainScore = fcnnMnistTest0 (Proxy @Double)
                                         widthHidden widthHidden2 chunk res
                 !testScore  = fcnnMnistTest0 (Proxy @Double)
                                         widthHidden widthHidden2 testData res
                 !lenChunk = length chunk
             hPutStrLn stderr $ printf "\n%s: (Batch %d with %d points)" prefix k lenChunk
             hPutStrLn stderr $ printf "%s: Training error:   %.2f%%" prefix ((1 - trainScore) * 100)
             hPutStrLn stderr $ printf "%s: Validation error: %.2f%%" prefix ((1 - testScore ) * 100)
             return res
       let runEpoch :: Int -> Domain0 Double -> IO (Domain0 Double)
           runEpoch n params0 | n > epochs = return params0
           runEpoch n params0 = do
             hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData
                 chunks = take maxBatches
                          $ zip [1 ..] $ chunksOf 5000 trainDataShuffled
             !res <- foldM runBatch params0 chunks
             runEpoch (succ n) res
       res <- runEpoch 1 params0Init
       let testErrorFinal = 1 - fcnnMnistTest0 (Proxy @Double) widthHidden widthHidden2 testData res
       testErrorFinal @?= expected

mnistTestCase2V
  :: String
  -> Int
  -> Int
  -> (Int
      -> Int
      -> MnistData Double
      -> DualNumberVariables 'DModeGradient Double
      -> DualMonadGradient Double (DualNumber 'DModeGradient Double))
  -> Int
  -> Int
  -> Double
  -> Double
  -> TestTree
mnistTestCase2V prefix epochs maxBatches trainWithLoss widthHidden widthHidden2
                gamma expected =
  let (nParams0, nParams1, _, _) = fcnnMnistLen1 widthHidden widthHidden2
      params0Init = HM.randomVector 44 HM.Uniform nParams0 - HM.scalar 0.5
      params1Init = V.fromList $
        imap (\i nPV -> HM.randomVector (44 + nPV + i) HM.Uniform nPV
                        - HM.scalar 0.5)
             nParams1
      name = prefix ++ ": "
             ++ unwords [ show epochs, show maxBatches
                        , show widthHidden, show widthHidden2
                        , show nParams0, show (length nParams1)
                        , show (sum nParams1 + nParams0), show gamma ]
  in testCase name $ do
       hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
              prefix epochs maxBatches
       trainData <- loadMnistData trainGlyphsPath trainLabelsPath
       testData <- loadMnistData testGlyphsPath testLabelsPath
       -- Mimic how backprop tests and display it, even though tests
       -- should not print, in principle.
       let runBatch :: (Domain0 Double, Domain1 Double)
                    -> (Int, [MnistData Double])
                    -> IO (Domain0 Double, Domain1 Double)
           runBatch (!params0, !params1) (k, chunk) = do
             let f = trainWithLoss widthHidden widthHidden2
                 (resS, resV, _, _) =
                   fst $ sgd gamma f chunk (params0, params1, V.empty, V.empty)
                 res = (resS, resV)
                 !trainScore = fcnnMnistTest1
                                         widthHidden widthHidden2 chunk res
                 !testScore = fcnnMnistTest1
                                        widthHidden widthHidden2 testData res
                 !lenChunk = length chunk
             hPutStrLn stderr $ printf "\n%s: (Batch %d with %d points)" prefix k lenChunk
             hPutStrLn stderr $ printf "%s: Training error:   %.2f%%" prefix ((1 - trainScore) * 100)
             hPutStrLn stderr $ printf "%s: Validation error: %.2f%%" prefix ((1 - testScore ) * 100)
             return res
       let runEpoch :: Int -> (Domain0 Double, Domain1 Double)
                    -> IO (Domain0 Double, Domain1 Double)
           runEpoch n params2 | n > epochs = return params2
           runEpoch n params2 = do
             hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData
                 chunks = take maxBatches
                          $ zip [1 ..] $ chunksOf 5000 trainDataShuffled
             !res <- foldM runBatch params2 chunks
             runEpoch (succ n) res
       res <- runEpoch 1 (params0Init, params1Init)
       let testErrorFinal =
             1 - fcnnMnistTest1 widthHidden widthHidden2 testData res
       testErrorFinal @?= expected

fcnnMnistLossTanh :: DualMonad 'DModeGradient Double m
                => Int
                -> Int
                -> MnistData Double
                -> DualNumberVariables 'DModeGradient Double
                -> m (DualNumber 'DModeGradient Double)
fcnnMnistLossTanh widthHidden widthHidden2 (xs, targ) vec = do
  res <- fcnnMnist0 tanhAct softMaxAct widthHidden widthHidden2 xs vec
  lossCrossEntropy targ res

fcnnMnistLossRelu :: DualMonad 'DModeGradient Double m
                => Int
                -> Int
                -> MnistData Double
                -> DualNumberVariables 'DModeGradient Double
                -> m (DualNumber 'DModeGradient Double)
fcnnMnistLossRelu widthHidden widthHidden2 (xs, targ) vec = do
  res <- fcnnMnist0 reluAct softMaxAct widthHidden widthHidden2 xs vec
  lossCrossEntropy targ res

mnistTestCase2L
  :: String
  -> Int
  -> Int
  -> (MnistData Double
      -> DualNumberVariables 'DModeGradient Double
      -> DualMonadGradient Double (DualNumber 'DModeGradient Double))
  -> Int
  -> Int
  -> Double
  -> Double
  -> TestTree
mnistTestCase2L prefix epochs maxBatches trainWithLoss widthHidden widthHidden2
                gamma expected =
  let ((nParams0, nParams1, nParams2, _), totalParams, range, parameters0) =
        initializerFixed 44 0.5 (fcnnMnistLen2 widthHidden widthHidden2)
      name = prefix ++ ": "
             ++ unwords [ show epochs, show maxBatches
                        , show widthHidden, show widthHidden2
                        , show nParams0, show nParams1, show nParams2
                        , show totalParams, show gamma, show range]
  in testCase name $ do
       hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
              prefix epochs maxBatches
       trainData <- loadMnistData trainGlyphsPath trainLabelsPath
       testData <- loadMnistData testGlyphsPath testLabelsPath
       -- Mimic how backprop tests and display it, even though tests
       -- should not print, in principle.
       let runBatch :: Domains Double
                    -> (Int, [MnistData Double])
                    -> IO (Domains Double)
           runBatch (!params0, !params1, !params2, !paramsX) (k, chunk) = do
             let f = trainWithLoss
                 res = fst $ sgd gamma f chunk
                                 (params0, params1, params2, paramsX)
                 !trainScore = fcnnMnistTest2 @Double chunk res
                 !testScore = fcnnMnistTest2 @Double testData res
                 !lenChunk = length chunk
             hPutStrLn stderr $ printf "\n%s: (Batch %d with %d points)" prefix k lenChunk
             hPutStrLn stderr $ printf "%s: Training error:   %.2f%%" prefix ((1 - trainScore) * 100)
             hPutStrLn stderr $ printf "%s: Validation error: %.2f%%" prefix ((1 - testScore ) * 100)
             return res
       let runEpoch :: Int
                    -> Domains Double
                    -> IO (Domains Double)
           runEpoch n params2 | n > epochs = return params2
           runEpoch n params2 = do
             hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData
                 chunks = take maxBatches
                          $ zip [1 ..] $ chunksOf 5000 trainDataShuffled
             !res <- foldM runBatch params2 chunks
             runEpoch (succ n) res
       res <- runEpoch 1 parameters0
       let testErrorFinal = 1 - fcnnMnistTest2 testData res
       testErrorFinal @?= expected

mnistTestCase2T
  :: Bool
  -> String
  -> Int
  -> Int
  -> (MnistData Double
      -> DualNumberVariables 'DModeGradient Double
      -> DualMonadGradient Double (DualNumber 'DModeGradient Double))
  -> Int
  -> Int
  -> Double
  -> Double
  -> TestTree
mnistTestCase2T reallyWriteFile
                prefix epochs maxBatches trainWithLoss widthHidden widthHidden2
                gamma expected =
  let ((nParams0, nParams1, nParams2, _), totalParams, range, !parameters0) =
        initializerFixed 44 0.5 (fcnnMnistLen2 widthHidden widthHidden2)
      name = prefix ++ " "
             ++ unwords [ show epochs, show maxBatches
                        , show widthHidden, show widthHidden2
                        , show nParams0, show nParams1, show nParams2
                        , show totalParams, show gamma, show range]
  in testCase name $ do
       hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
              prefix epochs maxBatches
       trainData0 <- loadMnistData trainGlyphsPath trainLabelsPath
       testData <- loadMnistData testGlyphsPath testLabelsPath
       let !trainData = force $ shuffle (mkStdGen 6) trainData0
       -- Mimic how backprop tests and display it, even though tests
       -- should not print, in principle.
       let runBatch :: (Domains Double, [(POSIXTime, Double)])
                    -> (Int, [MnistData Double])
                    -> IO (Domains Double, [(POSIXTime, Double)])
           runBatch ((!params0, !params1, !params2, !paramsX), !times)
                    (k, chunk) = do
             when (k `mod` 100 == 0) $ do
               hPutStrLn stderr $ printf "%s: %d " prefix k
               hFlush stderr
             let f = trainWithLoss
                 (!params0New, !value) =
                   sgd gamma f chunk (params0, params1, params2, paramsX)
             time <- getPOSIXTime
             return (params0New, (time, value) : times)
       let runEpoch :: Int
                    -> (Domains Double, [(POSIXTime, Double)])
                    -> IO (Domains Double, [(POSIXTime, Double)])
           runEpoch n params2times | n > epochs = return params2times
           runEpoch n (!params2, !times2) = do
             hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
             let !trainDataShuffled =
                   if n > 1
                   then shuffle (mkStdGen $ n + 5) trainData
                   else trainData
                 chunks = take maxBatches
                          $ zip [1 ..] $ chunksOf 1 trainDataShuffled
             res <- foldM runBatch (params2, times2) chunks
             runEpoch (succ n) res
       timeBefore <- getPOSIXTime
       (res, times) <- runEpoch 1 (parameters0, [])
       let ppTime (t, l) = init (show (t - timeBefore)) ++ " " ++ show l
       when reallyWriteFile $
         writeFile "walltimeLoss.txt" $ unlines $ map ppTime times
       let testErrorFinal = 1 - fcnnMnistTest2 testData res
       testErrorFinal @?= expected

mnistTestCase2D
  :: Bool
  -> Int
  -> Bool
  -> String
  -> Int
  -> Int
  -> (MnistData Double
      -> DualNumberVariables 'DModeGradient Double
      -> DualMonadGradient Double (DualNumber 'DModeGradient Double))
  -> Int
  -> Int
  -> Double
  -> Double
  -> TestTree
mnistTestCase2D reallyWriteFile miniBatchSize decay
                prefix epochs maxBatches trainWithLoss widthHidden widthHidden2
                gamma0 expected =
  let np = fcnnMnistLen2 widthHidden widthHidden2
      ((nParams0, nParams1, nParams2, _), totalParams, range, !parameters0) =
        initializerFixed 44 0.5 np
      name = prefix ++ " "
             ++ unwords [ show epochs, show maxBatches
                        , show widthHidden, show widthHidden2
                        , show nParams0, show nParams1, show nParams2
                        , show totalParams, show gamma0, show range]
  in testCase name $ do
       hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
              prefix epochs maxBatches
       trainData0 <- loadMnistData trainGlyphsPath trainLabelsPath
       testData <- loadMnistData testGlyphsPath testLabelsPath
       let !trainData = force $ shuffle (mkStdGen 6) trainData0
       -- Mimic how backprop tests and display it, even though tests
       -- should not print, in principle.
       let runBatch :: (Domains Double, [(POSIXTime, Double)])
                    -> (Int, [MnistData Double])
                    -> IO (Domains Double, [(POSIXTime, Double)])
           runBatch ((!params0, !params1, !params2, !paramsX), !times)
                    (k, chunk) = do
             when (k `mod` 100 == 0) $ do
               hPutStrLn stderr $ printf "%s: %d " prefix k
               hFlush stderr
             let f = trainWithLoss
                 gamma = if decay
                         then gamma0 * exp (- fromIntegral k * 1e-4)
                         else gamma0
                 (!params0New, !value) =
                   sgdBatchForward (33 + k * 7) miniBatchSize gamma f chunk
                                   (params0, params1, params2, paramsX) np
             time <- getPOSIXTime
             return (params0New, (time, value) : times)
       let runEpoch :: Int
                    -> (Domains Double, [(POSIXTime, Double)])
                    -> IO (Domains Double, [(POSIXTime, Double)])
           runEpoch n params2times | n > epochs = return params2times
           runEpoch n (!params2, !times2) = do
             hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
             let !trainDataShuffled =
                   if n > 1
                   then shuffle (mkStdGen $ n + 5) trainData
                   else trainData
                 chunks = take maxBatches
                          $ zip [1 ..]
                          $ chunksOf miniBatchSize trainDataShuffled
             res <- foldM runBatch (params2, times2) chunks
             runEpoch (succ n) res
       timeBefore <- getPOSIXTime
       (res, times) <- runEpoch 1 (parameters0, [])
       let ppTime (t, l) = init (show (t - timeBefore)) ++ " " ++ show l
       when reallyWriteFile $
         writeFile "walltimeLoss.txt" $ unlines $ map ppTime times
       let testErrorFinal = 1 - fcnnMnistTest2 testData res
       testErrorFinal @?= expected

mnistTestCase2F
  :: Bool
  -> Int
  -> Bool
  -> String
  -> Int
  -> Int
  -> (MnistData Double
      -> DualNumberVariables 'DModeDerivative Double
      -> DualMonadForward Double (DualNumber 'DModeDerivative Double))
  -> Int
  -> Int
  -> Double
  -> Double
  -> TestTree
mnistTestCase2F reallyWriteFile miniBatchSize decay
                prefix epochs maxBatches trainWithLoss widthHidden widthHidden2
                gamma0 expected =
  let np = fcnnMnistLen2 widthHidden widthHidden2
      ((nParams0, nParams1, nParams2, _), totalParams, range, !parameters0) =
        initializerFixed 44 0.5 np
      name = prefix ++ " "
             ++ unwords [ show epochs, show maxBatches
                        , show widthHidden, show widthHidden2
                        , show nParams0, show nParams1, show nParams2
                        , show totalParams, show gamma0, show range]
  in testCase name $ do
       hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
              prefix epochs maxBatches
       trainData0 <- loadMnistData trainGlyphsPath trainLabelsPath
       testData <- loadMnistData testGlyphsPath testLabelsPath
       let !trainData = coerce $ force $ shuffle (mkStdGen 6) trainData0
       -- Mimic how backprop tests and display it, even though tests
       -- should not print, in principle.
       let runBatch :: (Domains Double, [(POSIXTime, Double)])
                    -> (Int, [MnistData Double])
                    -> IO (Domains Double, [(POSIXTime, Double)])
           runBatch ((!params0, !params1, !params2, !paramsX), !times)
                    (k, chunk) = do
             when (k `mod` 100 == 0) $ do
               hPutStrLn stderr $ printf "%s: %d " prefix k
               hFlush stderr
             let f = trainWithLoss
                 gamma = if decay
                         then gamma0 * exp (- fromIntegral k * 1e-4)
                         else gamma0
                 (!params0New, !value) =
                   sgdBatchFastForward (33 + k * 7) miniBatchSize gamma f chunk
                                       (params0, params1, params2, paramsX) np
             time <- getPOSIXTime
             return (params0New, (time, value) : times)
       let runEpoch :: Int
                    -> (Domains Double, [(POSIXTime, Double)])
                    -> IO (Domains Double, [(POSIXTime, Double)])
           runEpoch n params2times | n > epochs = return params2times
           runEpoch n (!params2, !times2) = do
             hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
             let !trainDataShuffled =
                   if n > 1
                   then shuffle (mkStdGen $ n + 5) trainData
                   else trainData
                 chunks = take maxBatches
                          $ zip [1 ..]
                          $ chunksOf miniBatchSize trainDataShuffled
             res <- foldM runBatch (params2, times2) chunks
             runEpoch (succ n) res
       timeBefore <- getPOSIXTime
       (res, times) <- runEpoch 1 (parameters0, [])
       let ppTime (t, l) = init (show (t - timeBefore)) ++ " " ++ show l
       when reallyWriteFile $
         writeFile "walltimeLoss.txt" $ unlines $ map ppTime times
       let testErrorFinal = 1 - fcnnMnistTest2 testData res
       testErrorFinal @?= expected

mnistTestCase2S
  :: forall widthHidden widthHidden2.
     (KnownNat widthHidden, KnownNat widthHidden2)
  => Proxy widthHidden -> Proxy widthHidden2
  -> String
  -> Int
  -> Int
  -> (forall d r m. DualMonad d r m
      => Proxy widthHidden -> Proxy widthHidden2
      -> MnistData r -> DualNumberVariables d r -> m (DualNumber d r))
  -> Double
  -> Double
  -> TestTree
mnistTestCase2S proxy proxy2
                prefix epochs maxBatches trainWithLoss gamma expected =
  let ((_, _, _, nParamsX), totalParams, range, parametersInit) =
        initializerFixed 44 0.5 (fcnnMnistLenS @widthHidden @widthHidden2)
      name = prefix ++ ": "
             ++ unwords [ show epochs, show maxBatches
                        , show (valueOf @widthHidden :: Int)
                        , show (valueOf @widthHidden2 :: Int)
                        , show nParamsX, show totalParams
                        , show gamma, show range ]
  in testCase name $ do
    hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
           prefix epochs maxBatches
    trainData <- loadMnistData trainGlyphsPath trainLabelsPath
    testData <- loadMnistData testGlyphsPath testLabelsPath
    let runBatch :: Domains Double
                 -> (Int, [MnistData Double])
                 -> IO (Domains Double)
        runBatch (!params0, !params1, !params2, !paramsX) (k, chunk) = do
          let f = trainWithLoss proxy proxy2
              res = fst $ sgd gamma f chunk
                              (params0, params1, params2, paramsX)
              !trainScore = fcnnMnistTestS @widthHidden @widthHidden2
                                          chunk res
              !testScore = fcnnMnistTestS @widthHidden @widthHidden2
                                         testData res
              !lenChunk = length chunk
          hPutStrLn stderr $ printf "\n%s: (Batch %d with %d points)" prefix k lenChunk
          hPutStrLn stderr $ printf "%s: Training error:   %.2f%%" prefix ((1 - trainScore) * 100)
          hPutStrLn stderr $ printf "%s: Validation error: %.2f%%" prefix ((1 - testScore ) * 100)
          return res
    let runEpoch :: Int
                 -> Domains Double
                 -> IO (Domains Double)
        runEpoch n params2 | n > epochs = return params2
        runEpoch n params2 = do
          hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
          let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData
              chunks = take maxBatches
                       $ zip [1 ..] $ chunksOf 5000 trainDataShuffled
          !res <- foldM runBatch params2 chunks
          runEpoch (succ n) res
    res <- runEpoch 1 parametersInit
    let testErrorFinal = 1 - fcnnMnistTestS @widthHidden @widthHidden2
                                            testData res
    testErrorFinal @?= expected

dumbMnistTests :: TestTree
dumbMnistTests = testGroup "Dumb MNIST tests"
  [ testCase "1pretty-print in grey 3 2" $ do
      let (nParams0, lParams1, lParams2, _) = fcnnMnistLen2 4 3
          vParams1 = V.fromList lParams1
          vParams2 = V.fromList lParams2
          params0 = V.replicate nParams0 (1 :: Float)
          params1 = V.map (`V.replicate` 2) vParams1
          params2 = V.map (HM.konst 3) vParams2
          blackGlyph = V.replicate sizeMnistGlyph 4
          blackLabel = V.replicate sizeMnistLabel 5
          trainData = (blackGlyph, blackLabel)
          output = prettyPrintDf False (fcnnMnistLoss2 trainData)
                                 (params0, params1, params2, V.empty)
      -- printf "%s" output
      length output @?= 13348
  , testCase "2pretty-print in grey 3 2 fused" $ do
      let (nParams0, lParams1, lParams2, _) = fcnnMnistLen2 4 3
          vParams1 = V.fromList lParams1
          vParams2 = V.fromList lParams2
          params0 = V.replicate nParams0 (1 :: Float)
          params1 = V.map (`V.replicate` 2) vParams1
          params2 = V.map (HM.konst 3) vParams2
          blackGlyph = V.replicate sizeMnistGlyph 4
          blackLabel = V.replicate sizeMnistLabel 5
          trainData = (blackGlyph, blackLabel)
          output = prettyPrintDf True (fcnnMnistLossFused2 trainData)
                                 (params0, params1, params2, V.empty)
      --- printf "%s" output
      length output @?= 12431
  , testCase "3pretty-print on testset 3 2" $ do
      let (_, _, _, parameters0) = initializerFixed 44 0.5 (fcnnMnistLen2 4 3)
      testData <- loadMnistData testGlyphsPath testLabelsPath
      let trainDataItem = head testData
          output = prettyPrintDf True (fcnnMnistLoss2 trainDataItem) parameters0
      -- printf "%s" output
      length output @?= 16449
  , let blackGlyph = V.replicate sizeMnistGlyph 0
        blackLabel = V.replicate sizeMnistLabel 0
        trainData = replicate 10 (blackGlyph, blackLabel)
    in sgdTestCase "black"
         (return trainData) fcnnMnistLoss0 0.02 (-0.0)
  , let whiteGlyph = V.replicate sizeMnistGlyph 1
        whiteLabel = V.replicate sizeMnistLabel 1
        trainData = replicate 20 (whiteGlyph, whiteLabel)
    in sgdTestCase "white"
         (return trainData) fcnnMnistLoss0 0.02 23.02585095418536
  , let blackGlyph = V.replicate sizeMnistGlyph 0
        whiteLabel = V.replicate sizeMnistLabel 1
        trainData = replicate 50 (blackGlyph, whiteLabel)
    in sgdTestCase "black/white"
         (return trainData) fcnnMnistLoss0 0.02 23.025850929940457
  , let glyph = V.unfoldrExactN sizeMnistGlyph (uniformR (0, 1))
        label = V.unfoldrExactN sizeMnistLabel (uniformR (0, 1))
        trainData = map ((\g -> (glyph g, label g)) . mkStdGen) [1 .. 100]
    in sgdTestCase "random 100"
         (return trainData) fcnnMnistLoss0 0.02 11.089140063760212
  , sgdTestCase "first 100 trainset samples only"
      (take 100 <$> loadMnistData trainGlyphsPath trainLabelsPath)
      fcnnMnistLoss0 0.02 3.233123290489956
  , testCase "fcnnMnistTest0 on 0.1 params0 300 100 width 10k testset" $ do
      let nParams0 = fcnnMnistLen0 300 100
          params0 = V.replicate nParams0 0.1
      testData <- loadMnistData testGlyphsPath testLabelsPath
      (1 - fcnnMnistTest0 (Proxy @Double) 300 100 testData params0)
        @?= 0.902
  , testCase "fcnnMnistTest2VV on 0.1 params0 300 100 width 10k testset" $ do
      let (nParams0, nParams1, _, _) = fcnnMnistLen1 300 100
          params0 = V.replicate nParams0 0.1
          params1 = V.fromList $ map (`V.replicate` 0.1) nParams1
      testData <- loadMnistData testGlyphsPath testLabelsPath
      (1 - fcnnMnistTest1 300 100 testData (params0, params1))
        @?= 0.902
  , testCase "fcnnMnistTest2LL on 0.1 params0 300 100 width 10k testset" $ do
      let (nParams0, lParams1, lParams2, _) = fcnnMnistLen2 300 100
          vParams1 = V.fromList lParams1
          vParams2 = V.fromList lParams2
          params0 = V.replicate nParams0 0.1
          params1 = V.map (`V.replicate` 0.1) vParams1
          params2 = V.map (HM.konst 0.1) vParams2
      testData <- loadMnistData testGlyphsPath testLabelsPath
      (1 - fcnnMnistTest2 testData
                      (params0, params1, params2, V.empty))
        @?= 0.902
  , testProperty "Compare two forward derivatives and gradient for Mnist0" $
      \seed seedDs ->
      forAll (choose (1, 300)) $ \widthHidden ->
      forAll (choose (1, 100)) $ \widthHidden2 ->
      forAll (choose (0.01, 10)) $ \range ->
      forAll (choose (0.01, 10)) $ \rangeDs ->
        let createRandomVector n seedV = HM.randomVector seedV HM.Uniform n
            glyph = createRandomVector sizeMnistGlyph seed
            label = createRandomVector sizeMnistLabel seedDs
            mnistData :: MnistData Double
            mnistData = (glyph, label)
            nParams0 = fcnnMnistLen0 widthHidden widthHidden2
            paramShape = (nParams0, [], [], [])
            (_, _, _, parameters) = initializerFixed seed range paramShape
            (_, _, _, ds) = initializerFixed seedDs rangeDs paramShape
            (_, _, _, parametersPerturbation) =
              initializerFixed (seed + seedDs) 1e-7 paramShape
            f :: forall d r m. (DualMonad d r m, r ~ Double)
              => DualNumberVariables d r -> m (DualNumber d r)
            f = fcnnMnistLoss0 widthHidden widthHidden2 mnistData
        in
            qcPropDom f parameters ds parametersPerturbation 1
  , testProperty "Compare two forward derivatives and gradient for Mnist1" $
      \seed seedDs ->
      forAll (choose (1, 2000)) $ \widthHidden ->
      forAll (choose (1, 5000)) $ \widthHidden2 ->
      forAll (choose (0.01, 0.5)) $ \range ->  -- large nn, so NaNs fast
      forAll (choose (0.01, 10)) $ \rangeDs ->
        let createRandomVector n seedV = HM.randomVector seedV HM.Uniform n
            glyph = createRandomVector sizeMnistGlyph seed
            label = createRandomVector sizeMnistLabel seedDs
            mnistData :: MnistData Double
            mnistData = (glyph, label)
            paramShape = fcnnMnistLen1 widthHidden widthHidden2
            (_, _, _, parameters) = initializerFixed seed range paramShape
            (_, _, _, ds) = initializerFixed seedDs rangeDs paramShape
            (_, _, _, parametersPerturbation) =
              initializerFixed (seed + seedDs) 1e-7 paramShape
            f :: forall d r m. (DualMonad d r m, r ~ Double)
              => DualNumberVariables d r -> m (DualNumber d r)
            f = fcnnMnistLoss1 widthHidden widthHidden2 mnistData
        in
            qcPropDom f parameters ds parametersPerturbation 1
  , testProperty "Compare two forward derivatives and gradient for Mnist2" $
      \seed ->
      forAll (choose (0, sizeMnistLabel - 1)) $ \seedDs ->
      forAll (choose (1, 5000)) $ \widthHidden ->
      forAll (choose (1, 1000)) $ \widthHidden2 ->
      forAll (choose (0.01, 1)) $ \range ->
      forAll (choose (0.01, 10)) $ \rangeDs ->
        let createRandomVector n seedV = HM.randomVector seedV HM.Uniform n
            glyph = createRandomVector sizeMnistGlyph seed
            label = createRandomVector sizeMnistLabel seedDs
            labelOneHot = HM.konst 0 sizeMnistLabel V.// [(seedDs, 1)]
            mnistData, mnistDataOneHot :: MnistData Double
            mnistData = (glyph, label)
            mnistDataOneHot = (glyph, labelOneHot)
            paramShape = fcnnMnistLen2 widthHidden widthHidden2
            (_, _, _, parameters) = initializerFixed seed range paramShape
            (_, _, _, ds) = initializerFixed seedDs rangeDs paramShape
            (_, _, _, parametersPerturbation) =
              initializerFixed (seed + seedDs) 1e-7 paramShape
            f, fOneHot, fFused
              :: forall d r m. (DualMonad d r m, r ~ Double)
                 => DualNumberVariables d r -> m (DualNumber d r)
            f = fcnnMnistLoss2 mnistData
            fOneHot = fcnnMnistLoss2 mnistDataOneHot
            fFused = fcnnMnistLossFused2 mnistDataOneHot
        in
            qcPropDom f       parameters ds parametersPerturbation 1 .&&.
            qcPropDom fOneHot parameters ds parametersPerturbation 1 .&&.
            qcPropDom fFused  parameters ds parametersPerturbation 1 .&&.
            cmpTwoSimple fOneHot fFused  parameters ds
  ]

bigMnistTests :: TestTree
bigMnistTests = testGroup "MNIST tests with a 2-hidden-layer nn"
  [ mnistTestCase2 "1 epoch, 1 batch" 1 1 fcnnMnistLoss0 300 100 0.02
                   0.1269
  , mnistTestCase2 "tanh: 1 epoch, 1 batch" 1 1 fcnnMnistLossTanh 300 100 0.02
                   0.6406000000000001
  , mnistTestCase2 "relu: 1 epoch, 1 batch" 1 1 fcnnMnistLossRelu 300 100 0.02
                   0.7248
  , mnistTestCase2 "1 epoch, 1 batch, wider" 1 1 fcnnMnistLoss0 500 150 0.02
                   0.1269
  , mnistTestCase2 "2 epochs, but only 1 batch" 2 1 fcnnMnistLoss0 300 100 0.02
                   9.809999999999997e-2
  , mnistTestCase2 "artificial 1 2 3 4 5" 1 2 fcnnMnistLoss0 3 4 5
                   0.8972
  , mnistTestCase2 "artificial 5 4 3 2 1" 5 4 fcnnMnistLoss0 3 2 1
                   0.8991
  ]

vectorMnistTests :: TestTree
vectorMnistTests = testGroup "MNIST VV tests with a 2-hidden-layer nn"
  [ mnistTestCase2V "1 epoch, 1 batch" 1 1 fcnnMnistLoss1 300 100 0.02
                    0.12960000000000005
  , mnistTestCase2V "1 epoch, 1 batch, wider" 1 1 fcnnMnistLoss1 500 150 0.02
                    0.13959999999999995
  , mnistTestCase2V "2 epochs, but only 1 batch" 2 1 fcnnMnistLoss1 300 100 0.02
                    0.10019999999999996
  , mnistTestCase2V "1 epoch, all batches" 1 99 fcnnMnistLoss1 300 100 0.02
                    5.389999999999995e-2
  , mnistTestCase2V "artificial 1 2 3 4 5" 1 2 fcnnMnistLoss1 3 4 5
                    0.8972
  , mnistTestCase2V "artificial 5 4 3 2 1" 5 4 fcnnMnistLoss1 3 2 1
                    0.7756000000000001
  ]

matrixMnistTests :: TestTree
matrixMnistTests = testGroup "MNIST LL tests with a 2-hidden-layer nn"
  [ mnistTestCase2L "1 epoch, 1 batch" 1 1 fcnnMnistLoss2 300 100 0.02
                    0.12339999999999995
  , mnistTestCase2L "1 epoch, 1 batch, wider" 1 1 fcnnMnistLoss2 500 150 0.02
                    0.15039999999999998
  , mnistTestCase2L "2 epochs, but only 1 batch" 2 1 fcnnMnistLoss2 300 100 0.02
                    8.879999999999999e-2
  , mnistTestCase2L "1 epoch, all batches" 1 99 fcnnMnistLoss2 300 100 0.02
                    5.1100000000000034e-2
  , mnistTestCase2L "artificial 1 2 3 4 5" 1 2 fcnnMnistLoss2 3 4 5
                    0.8972
  , mnistTestCase2T False
                    "artificial TL 5 4 3 2 1" 5 4 fcnnMnistLoss2 3 2 1
                    0.8865
  , mnistTestCase2D False 1 False
                    "artificial DL 5 4 3 2 1" 5 4 fcnnMnistLoss2 3 2 1
                    0.8991
  , mnistTestCase2F False 1 False
                    "artificial FL 5 4 3 2 1" 5 4 fcnnMnistLoss2 3 2 1
                    0.8991
--  , mnistTestCase2T True False
--                    "2 epochs, all batches, TL, wider, to file"
--                    2 60000 fcnnMnistLoss2 500 150 0.02
--                    4.290000000000005e-2
--  , mnistTestCase2D True 1 False
--                    "2 epochs, all batches, DL, wider, to file"
--                    2 60000 fcnnMnistLoss2 500 150 0.02
--                    0.9079
--  , mnistTestCase2D True 64 False
--                    "2 epochs, all batches, DL, wider, to file"
--                    2 60000 fcnnMnistLoss2 500 150 0.02
--                    0.9261
--  , mnistTestCase2D True 64 True
--                    "2 epochs, all batches, DL, wider, to file"
--                    2 60000 fcnnMnistLoss2 500 150 0.02
--                    0.8993
--  , mnistTestCase2D True 64 True
--                    "2 epochs, all batches, DL, wider, to file"
--                    2 60000 fcnnMnistLoss2 500 150 2e-5
--                    0.9423
--  , mnistTestCase2D True 64 True
--                    "2 epochs, all batches, DL, wider, to file"
--                    2 60000 fcnnMnistLoss2 500 150 2e-4
--                    0.8714
--  , mnistTestCase2F True 64 True
--                    "2 epochs, all batches, FL, wider, to file"
--                    2 60000 fcnnMnistLoss2 500 150 2e-4
--                    0.8714
--  , mnistTestCase2D True 64 True
--                    "2 epochs, all batches, DL, wider, to file"
--                    2 60000 fcnnMnistLossFusedRelu2 1024 1024 2e-4
--                    0.902
--  , mnistTestCase2D False 64 True
--                    "2 epochs, all batches, 1024DL"
--                    2 60000 fcnnMnistLoss2 1024 1024 2e-4
--                    0.7465999999999999
--  , mnistTestCase2F False 64 True
--                    "2 epochs, all batches, 1024FL"
--                    2 60000 fcnnMnistLoss2 1024 1024 2e-4
--                    0.7465999999999999
  ]

fusedMnistTests :: TestTree
fusedMnistTests = testGroup "MNIST fused LL tests with a 2-hidden-layer nn"
  [ mnistTestCase2L "1 epoch, 1 batch" 1 1 fcnnMnistLossFused2 300 100 0.02
                    0.12339999999999995
  , mnistTestCase2L "1 epoch, 1 batch, wider" 1 1
                    fcnnMnistLossFused2 500 150 0.02
                    0.15039999999999998
  , mnistTestCase2L "2 epochs, but only 1 batch" 2 1
                    fcnnMnistLossFused2 300 100 0.02
                    8.879999999999999e-2
  , mnistTestCase2L "1 epoch, all batches" 1 99 fcnnMnistLossFused2 300 100 0.02
                    5.1100000000000034e-2
  , mnistTestCase2L "artificial 1 2 3 4 5" 1 2 fcnnMnistLossFused2 3 4 5
                    0.8972
  , mnistTestCase2L "artificial 5 4 3 2 1" 5 4 fcnnMnistLossFused2 3 2 1
                    0.7033
  , mnistTestCase2S (Proxy @300) (Proxy @100)
                    "S 1 epoch, 1 batch" 1 1 fcnnMnistLossFusedS 0.02
                    0.1311
  , mnistTestCase2S (Proxy @500) (Proxy @150)
                    "S 1 epoch, 1 batch, wider" 1 1 fcnnMnistLossFusedS 0.02
                    0.12470000000000003
  , mnistTestCase2S (Proxy @300) (Proxy @100)
                    "S 2 epochs, but only 1 batch" 2 1 fcnnMnistLossFusedS 0.02
                    9.630000000000005e-2
  , mnistTestCase2S (Proxy @300) (Proxy @100)
                    "S 1 epoch, all batches" 1 99 fcnnMnistLossFusedS 0.02
                    5.620000000000003e-2
  , mnistTestCase2S (Proxy @3) (Proxy @4)
                    "S artificial 1 2 3 4 5" 1 2 fcnnMnistLossFusedS 5
                    0.8972
  , mnistTestCase2S (Proxy @3) (Proxy @2)
                    "S artificial 5 4 3 2 1" 5 4 fcnnMnistLossFusedS 1
                    0.8246
  , mnistTestCase2S (Proxy @300) (Proxy @100)
                    "SR 1 epoch, 1 batch" 1 1 fcnnMnistLossFusedReluS 0.02
                    0.7068
  , mnistTestCase2S (Proxy @500) (Proxy @150)
                    "SR 1 epoch, 1 batch, wider" 1 1
                    fcnnMnistLossFusedReluS 0.02
                    0.8874
  , mnistTestCase2S (Proxy @300) (Proxy @100)
                    "SR 2 epochs, but 1 batch" 2 1 fcnnMnistLossFusedReluS 0.02
                    0.8352999999999999
  , mnistTestCase2S (Proxy @300) (Proxy @100)
                    "SR 1 epoch, all batches" 1 99 fcnnMnistLossFusedReluS 0.02
                    0.6415
  , mnistTestCase2S (Proxy @3) (Proxy @4)
                    "SR artificial 1 2 3 4 5" 1 2 fcnnMnistLossFusedReluS 5
                    0.8972
  , mnistTestCase2S (Proxy @3) (Proxy @2)
                    "SR artificial 5 4 3 2 1" 5 4 fcnnMnistLossFusedReluS 1
                    0.8991
  ]

shortCIMnistTests :: TestTree
shortCIMnistTests = testGroup "Short CI MNIST tests"
  [ mnistTestCase2 "2 artificial 1 2 3 4 5" 1 2 fcnnMnistLoss0 3 4 5
                   0.8972
  , mnistTestCase2 "2 artificial 5 4 3 2 1" 5 4 fcnnMnistLoss0 3 2 1
                   0.8991
  , mnistTestCase2V "VV 1 epoch, 1 batch" 1 1 fcnnMnistLoss1 300 100 0.02
                    0.12960000000000005
  , mnistTestCase2V "VV artificial 1 2 3 4 5" 1 2 fcnnMnistLoss1 3 4 5
                    0.8972
  , mnistTestCase2V "VV artificial 5 4 3 2 1" 5 4 fcnnMnistLoss1 3 2 1
                    0.7756000000000001
  , mnistTestCase2L "LL 1 epoch, 1 batch" 1 1 fcnnMnistLoss2 300 100 0.02
                    0.12339999999999995
  , mnistTestCase2L "LL artificial 1 2 3 4 5" 1 2 fcnnMnistLoss2 3 4 5
                    0.8972
  , mnistTestCase2L "LL artificial 5 4 3 2 1" 5 4 fcnnMnistLoss2 3 2 1
                    0.8085
  , mnistTestCase2L "fused LL 1/1 batch" 1 1 fcnnMnistLossFused2 300 100 0.02
                    0.12339999999999995
  , mnistTestCase2L "fused LL artificial 1 2 3 4 5" 1 2
                    fcnnMnistLossFused2 3 4 5
                    0.8972
  , mnistTestCase2T False
                    "fused TL artificial 5 4 3 2 1" 5 4
                    fcnnMnistLossFused2 3 2 1
                    0.8865
  , mnistTestCase2D False 1 False
                    "fused DL artificial 5 4 3 2 1" 5 4
                    fcnnMnistLossFused2 3 2 1
                    0.8991
  , mnistTestCase2S (Proxy @300) (Proxy @100)
                    "S 1 epoch, 1 batch" 1 1 fcnnMnistLossFusedS 0.02
                    0.1311
  , mnistTestCase2S (Proxy @3) (Proxy @4)
                    "S artificial 1 2 3 4 5" 1 2 fcnnMnistLossFusedS 5
                    0.8972
  , mnistTestCase2S (Proxy @3) (Proxy @2)
                    "S artificial 5 4 3 2 1" 5 4 fcnnMnistLossFusedS 1
                    0.8246
  , mnistTestCase2S (Proxy @3) (Proxy @4)
                    "SR artificial 1 2 3 4 5" 1 2 fcnnMnistLossFusedReluS 5
                    0.8972
  , mnistTestCase2S (Proxy @3) (Proxy @2)
                    "SR artificial 5 4 3 2 1" 5 4 fcnnMnistLossFusedReluS 1
                    0.8991
  ]


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

{-# LANGUAGE CPP                  #-}
{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE StandaloneDeriving   #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}
{-|
Module      : Grenade.Core.Shape
Description : Dependently typed shapes of data which are passed between layers of a network
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental

-}
module Grenade.Core.Shape (
    S (..)
  , Shape (..)
#if MIN_VERSION_singletons(2,6,0)
  , SShape (..)
#else
  , Sing (..)
#endif

  , randomOfShape
  , fromStorable
  , fromStorableMatrix
  , nk
  , visualise2D
  , splitChannels
  , combineChannels
  ) where

#if MIN_VERSION_singletons(2,6,0)
import           Data.Kind                    (Type)
#endif

import           Control.DeepSeq              (NFData (..))
import           Data.List.Split              (chunksOf)
import           Data.Maybe                   (fromJust)
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons
import           Data.Singletons.TypeLits
import           Data.Vector.Storable         (Vector)
import qualified Data.Vector.Storable         as V
import           GHC.TypeLits                 hiding (natVal)
import           Numeric.LinearAlgebra        (Matrix)
import qualified Numeric.LinearAlgebra        as NLA
import qualified Numeric.LinearAlgebra.Data   as NLAD
import           Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra.Static as H

import           System.Random.MWC

import           Grenade.Types

-- | The current shapes we accept.
--   at the moment this is just one, two, and three dimensional
--   Vectors/Matricies.
--
--   These are only used with DataKinds, as Kind `Shape`, with Types 'D1, 'D2, 'D3.
data Shape
  = D1 Nat
  -- ^ One dimensional vector
  | D2 Nat Nat
  -- ^ Two dimensional matrix. Row, Column.
  | D3 Nat Nat Nat
  -- ^ Three dimensional matrix. Row, Column, Channels.
  | D4 Nat Nat Nat Nat
  -- ^ Four dimensional matrix. Depth, Channels, Row, Column

-- | Concrete data structures for a Shape.
--
--   All shapes are held in contiguous memory.
--   3D is held in a matrix (usually row oriented) which has height depth * rows.
data S (n :: Shape) where
  S1D :: ( KnownNat len )
      => R len
      -> S ('D1 len)

  S2D :: ( KnownNat rows, KnownNat columns )
      => L rows columns
      -> S ('D2 rows columns)

  S3D :: ( KnownNat rows
         , KnownNat columns
         , KnownNat channels
         , KnownNat (rows * channels) )
      => L (rows * channels) columns
      -> S ('D3 rows columns channels)

  S4D :: ( KnownNat rows
         , KnownNat columns
         , KnownNat channels
         , KnownNat depth)
      => L (depth * channels * rows) columns
      -> S ('D4 depth channels rows columns)

deriving instance Show (S n)

-- Singleton instances.
--
-- These could probably be derived with template haskell, but this seems
-- clear and makes adding the KnownNat constraints simple.
-- We can also keep our code TH free, which is great.
#if MIN_VERSION_singletons(2,6,0)
-- In singletons 2.6 Sing switched from a data family to a type family.
type instance Sing = SShape

-- | Datatype to allow "pattern matching" on shapes at the type level through
--   the use the sing function from Singletons.
data SShape :: Shape -> Type where
  D1Sing :: Sing a -> SShape ('D1 a)
  D2Sing :: Sing a -> Sing b -> SShape ('D2 a b)
  D3Sing :: KnownNat (a * c) => Sing a -> Sing b -> Sing c -> SShape ('D3 a b c)
  D4Sing :: KnownNat (a * c * d) => Sing a -> Sing b -> Sing c -> Sing d -> SShape ('D4 a b c d)
#else
data instance Sing (n :: Shape) where
  D1Sing :: Sing a -> Sing ('D1 a)
  D2Sing :: Sing a -> Sing b -> Sing ('D2 a b)
  D3Sing :: KnownNat (a * c) => Sing a -> Sing b -> Sing c -> Sing ('D3 a b c)
  D4Sing :: KnownNat (a * c * d) => Sing a -> Sing b -> Sing c -> Sing d -> Sing ('D4 a b c d)
#endif

instance KnownNat a => SingI ('D1 a) where
  sing = D1Sing sing
instance (KnownNat a, KnownNat b) => SingI ('D2 a b) where
  sing = D2Sing sing sing
instance (KnownNat a, KnownNat b, KnownNat c, KnownNat (a * c)) => SingI ('D3 a b c) where
  sing = D3Sing sing sing sing
instance (KnownNat a, KnownNat b, KnownNat c, KnownNat d) => SingI ('D4 a b c d) where
  sing = D4Sing sing sing sing sing

instance SingI x => Num (S x) where
  (+) = n2 (+)
  (-) = n2 (-)
  (*) = n2 (*)
  abs = n1 abs
  signum = n1 signum
  fromInteger x = nk (fromInteger x)

instance SingI x => Fractional (S x) where
  (/) = n2 (/)
  recip = n1 recip
  fromRational x = nk (fromRational x)

instance SingI x => Floating (S x) where
  pi = nk pi
  exp = n1 exp
  log = n1 log
  sqrt = n1 sqrt
  (**) = n2 (**)
  logBase = n2 logBase
  sin = n1 sin
  cos = n1 cos
  tan = n1 tan
  asin = n1 asin
  acos = n1 acos
  atan = n1 atan
  sinh = n1 sinh
  cosh = n1 cosh
  tanh = n1 tanh
  asinh = n1 asinh
  acosh = n1 acosh
  atanh = n1 atanh

--
-- I haven't made shapes strict, as sometimes they're not needed
-- (the last input gradient back for instance)
--
instance NFData (S x) where
  rnf (S1D x) = rnf x
  rnf (S2D x) = rnf x
  rnf (S3D x) = rnf x
  rnf (S4D x) = rnf x

-- | Generate random data of the desired shape
randomOfShape :: forall x . (SingI x) => IO (S x)
randomOfShape = do
  seed :: Int <- withSystemRandom . asGenST $ \gen -> uniform gen
  return $ case (sing :: Sing x) of
    D1Sing SNat ->
        S1D (H.randomVector seed H.Uniform * 2 - 1)

    D2Sing SNat SNat ->
        S2D (H.uniformSample seed (-1) 1)

    D3Sing SNat SNat SNat ->
        S3D (H.uniformSample seed (-1) 1)

    D4Sing SNat SNat SNat SNat ->
        S4D (H.uniformSample seed (-1) 1)

-- | Generate a shape from a Storable Vector.
--
--   Returns Nothing if the vector is of the wrong size.
fromStorable :: forall x. SingI x => Vector RealNum -> Maybe (S x)
fromStorable xs = case sing :: Sing x of
    D1Sing SNat ->
      S1D <$> H.create xs

    D2Sing SNat SNat ->
      S2D <$> mkL xs

    D3Sing SNat SNat SNat ->
      S3D <$> mkL xs

    D4Sing SNat SNat SNat SNat ->
      S4D <$> mkL xs
  where
    mkL :: forall rows columns. (KnownNat rows, KnownNat columns)
        => Vector RealNum -> Maybe (L rows columns)
    mkL v =
      let rows    = fromIntegral $ natVal (Proxy :: Proxy rows)
          columns = fromIntegral $ natVal (Proxy :: Proxy columns)
      in  if rows * columns == V.length v
             then H.create $ NLA.reshape columns v
             else Nothing

-- | Generate a shape from a Storable Matrix.
--
--   Returns Nothing if the matrix is of the wrong size (or one dimensional)
fromStorableMatrix :: forall x. SingI x => Matrix RealNum -> Maybe (S x)
fromStorableMatrix xs = case sing :: Sing x of
    D1Sing SNat ->
      Nothing

    D2Sing SNat SNat ->
      S2D <$> mkL xs

    D3Sing SNat SNat SNat ->
      S3D <$> mkL xs

    D4Sing SNat SNat SNat SNat ->
      S4D <$> mkL xs
  where
    mkL :: forall rows columns. (KnownNat rows, KnownNat columns)
        => Matrix RealNum -> Maybe (L rows columns)
    mkL m =
      let rows    = fromIntegral $ natVal (Proxy :: Proxy rows)
          columns = fromIntegral $ natVal (Proxy :: Proxy columns)
          (h, w)  = NLA.size m
      in  if rows == h && columns == w
             then H.create m
             else Nothing

instance SingI x => Serialize (S x) where
  put i = (case i of
            (S1D x) -> putListOf put . NLA.toList . H.extract $ x
            (S2D x) -> putListOf put . NLA.toList . NLA.flatten . H.extract $ x
            (S3D x) -> putListOf put . NLA.toList . NLA.flatten . H.extract $ x
            (S4D x) -> putListOf put . NLA.toList . NLA.flatten . H.extract $ x
          ) :: PutM ()

  get = do
    Just i <- fromStorable . V.fromList <$> getListOf get
    return i

-- Helper function for creating the number instances
n1 :: ( forall a. Floating a => a -> a ) -> S x -> S x
n1 f (S1D x) = S1D (f x)
n1 f (S2D x) = S2D (f x)
n1 f (S3D x) = S3D (f x)
n1 f (S4D x) = S4D (f x)

-- Helper function for creating the number instances
n2 :: ( forall a. Floating a => a -> a -> a ) -> S x -> S x -> S x
n2 f (S1D x) (S1D y) = S1D (f x y)
n2 f (S2D x) (S2D y) = S2D (f x y)
n2 f (S3D x) (S3D y) = S3D (f x y)
n2 f (S4D x) (S4D y) = S4D (f x y)

-- | Helper function for creating the number instances
nk :: forall x. SingI x => RealNum -> S x
nk x = case (sing :: Sing x) of
  D1Sing SNat ->
    S1D (H.konst x)

  D2Sing SNat SNat ->
    S2D (H.konst x)

  D3Sing SNat SNat SNat ->
    S3D (H.konst x)

  D4Sing SNat SNat SNat SNat ->
    S4D (H.konst x)

-- | Prints out the contents of a matrix with entries between zero and the max value
visualise2D :: S ('D2 a b)  -- ^ input matrix
            -> RealNum      -- ^ maximum element value
            -> String       -- ^ pretty printed matrix
visualise2D (S2D mm) max =
  let m  = H.extract mm
      ms = NLAD.toLists m
      render n' | n' <= 0.2 * max  = ' '
                | n' <= 0.4 * max  = '.'
                | n' <= 0.6 * max  = '-'
                | n' <= 0.8 * max  = '='
                | otherwise =  '#'
      px = (fmap . fmap) render ms
  in unlines px

-- | TODO Theo
splitChannels :: forall rows columns channels.
                 (KnownNat rows, KnownNat columns)
                 => S ('D3 rows columns channels) -> [S ('D2 rows columns)]
splitChannels (S3D x)
  = let r   = fromIntegral $ natVal (Proxy :: Proxy rows)
        rs  = NLA.toRows $ H.extract x
        rs' = chunksOf r rs
        ms  = map (S2D . fromJust . H.create . NLA.fromRows) rs' :: [S ('D2 rows columns)]
    in ms

-- | TODO Theo
combineChannels :: forall rows columns channels.
                 (KnownNat rows, KnownNat columns, KnownNat channels)
                 => [S ('D2 rows columns)] -> S ('D3 rows columns channels)
combineChannels xs
  = let xs' = map (\(S2D x) -> NLA.toRows $ H.extract x) xs :: [[Vector RealNum]]
        xs'' = concat xs'
    in  S3D $ fromJust . H.create . NLA.fromRows $ xs''


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE BangPatterns #-}
-- | A music theory library for just intonation and other mathematically pure ideas.
module Boopadoop
  (module Boopadoop
  ,module Boopadoop.Diagram
  ,module Boopadoop.Rhythm
  ,module Boopadoop.Interval
  ,module Boopadoop.Discrete
  ) where

import qualified Data.WAVE as WAVE
import Control.Monad.ST
import Control.Monad
import Control.Applicative
import Boopadoop.Diagram
import Boopadoop.Rhythm
import Boopadoop.Interval
import Boopadoop.Discrete
import Data.List
import Data.Int
import Data.Complex
import qualified Data.Vector.Fixed as FV
import qualified Data.Primitive.ByteArray as BA
import qualified Data.Vector.Unboxed as Vector

-- | A 'Waveform' is a function (of time) that we can later sample.
newtype Waveform t a = Waveform
  {sample :: t -> a -- ^ 'sample' the 'Waveform' at a specified time
  }

instance Functor (Waveform t) where
  fmap f w = sampleFrom $ f . sample w

instance SummaryChar (Waveform t a) where
  sumUp _ = '~'

-- | A 'Double' valued wave with time also in terms of 'Double'.
-- This models a real-valued waveform which typically has values in @[-1,1]@ and
-- is typically supported on either the entire real line ('sinWave') or on a compact subset ('compactWave')
type DWave = Waveform Double Double

-- | Show a waveform by pretty printing some of the actual waveform in dot matrix form.
instance Show (Waveform Double Double) where
  show w = intercalate "\n" . transpose $ map sampleToString waveSamples
    where
      sampleToString k = if k <= quantLevel && k >= -quantLevel
        then replicate (quantLevel - k) '.' ++ "x" ++ replicate (quantLevel + k) '.'
        else let m = "k = " ++ show k in m ++ replicate (quantLevel * 2 + 1 - length m) ' '
      waveSamples = map (floor . (* realToFrac quantLevel) . sample w . (/sampleRate)) [0 .. 115]
      quantLevel = 15 :: Int
      sampleRate = 6400

instance Show (Waveform Double Discrete) where
  show w = intercalate "\n" . transpose $ map sampleToString waveSamples
    where
      sampleToString k = if k <= quantLevel && k >= -quantLevel
        then replicate (quantLevel - k) '.' ++ "x" ++ replicate (quantLevel + k) '.'
        else let m = "k = " ++ show k in m ++ replicate (quantLevel * 2 + 1 - length m) ' '
      waveSamples = map ((`quotRoundUp` (1 + (discFactor `quot` quantLevel))) . fromIntegral . unDiscrete . sample w . (/sampleRate)) [0 .. 115]
      quantLevel = 15 :: Int
      sampleRate = 6400

instance Show (Waveform Tick Discrete) where
  show w = intercalate "\n" . transpose $ map sampleToString waveSamples
    where
      sampleToString k = if k <= quantLevel && k >= -quantLevel
        then replicate (quantLevel - k) '.' ++ "x" ++ replicate (quantLevel + k) '.'
        else let m = "k = " ++ show k in m ++ replicate (quantLevel * 2 + 1 - length m) ' '
      waveSamples = map ((`quotRoundUp` (1 + (discFactor `quot` quantLevel))) . fromIntegral . unDiscrete . sample (skipTicks 1 w)) [0 .. 115]
      quantLevel = 15 :: Int

instance Show (Waveform Tick Double) where
  show w = intercalate "\n" . transpose $ map sampleToString waveSamples
    where
      sampleToString k = if k <= quantLevel && k >= -quantLevel
        then replicate (quantLevel - k) '.' ++ "x" ++ replicate (quantLevel + k) '.'
        else let m = "k = " ++ show k in m ++ replicate (quantLevel * 2 + 1 - length m) ' '
      waveSamples = map (floor . (* realToFrac quantLevel) . sample (skipTicks 1 w)) [0 .. 115]
      quantLevel = 15 :: Int

-- | A version of @'quot'@ that rounds away from zero instead of towards it.
quotRoundUp :: Int -> Int -> Int
quotRoundUp a b = if a `mod` b == 0 then a `quot` b else (signum a * signum b) + (a `quot` b)

-- | Build a 'Waveform' by sampling the given function.
sampleFrom :: (t -> a) -> Waveform t a
sampleFrom f = Waveform $ \t -> {-t `seq`-} f t

-- | Sample a 'Waveform' at specified time. @'sampleAt' = 'flip' 'sample'@
sampleAt :: t -> Waveform t a -> a
sampleAt = flip sample

-- | Pure sine wave of the given frequency
--sinWave :: Floating a => a -> Waveform a a
sinWave :: Double -> DWave
sinWave f = if sinHack -- speedy fast sin hack from https://www.gamedev.net/forums/topic/621589-extremely-fast-sin-approximation/ due to nightcracker
  then sampleFrom $ \t -> let
    x = pi * ((2 * f * t) - fromIntegral k)
    x2 = x * x
    k = spookyFastTruncate (2 * f * t)
    in (if odd k then negate else id) (x * (c + x2 * (b + a * x2)))
  else sampleFrom $ \t -> sin (freq * t)
  where
    sinHack = True
    !freq = 2 * pi * f
    a = 0.00735246819687011731341356165096815
    b = -0.16528911397014738207016302002888890
    c = 0.99969198629596757779830113868360584

-- | Pure sine wave of the given frequency and initial phase
sinWaveWithPhase :: Floating a => a -> a -> Waveform a a
sinWaveWithPhase f phi0 = sampleFrom $ \t -> let !freq = 2 * pi * f in sin (freq * t + phi0)

-- | Sine wave that is optimized to store only a small @'CompactWavetable'@.
-- The period is given in ticks because non-tick-multiple period sin waves will have aliasing under this optimization!
fastSin :: Tick -> Wavetable
fastSin per = exploitPeriodicity per $ tickTable (fromIntegral per) $ discretize $ sinWave 1

-- | @'compactWave' (l,h)@ is a wave which is @'True'@ on @[l,h)@ and @'False'@ elsewhere
compactWave :: (Ord t,Num t) => (t,t) -> Waveform t Bool
compactWave (low,high) = sampleFrom $ \t -> t >= low && t < high

-- | @'muting' 'True'@ is @'id'@ while @'muting' 'False'@ is @'const' 0@.
muting :: Num a => Bool -> a -> a
muting b s = if b then s else 0

-- | Modulate one wave with another according to the given function pointwise.
-- This means you can't implement 'phaseModulate' using only this combinator because phase modulation
-- requires information about the target wave at times other than the current time.
modulate :: (a -> b -> c) -> Waveform t a -> Waveform t b -> Waveform t c
modulate f a b = sampleFrom $ \t -> f (sample a t) (sample b t)

-- | Modulate the amplitude of one wave with another. This is simply pointwise multiplication:
-- @
--  'amplitudeModulate' = 'modulate' ('*')
-- @
amplitudeModulate :: Num a => Waveform t a -> Waveform t a -> Waveform t a
amplitudeModulate = modulate (*)

-- | Control rate of speed of one waveform with another
freqModulate :: Waveform t' p -> Waveform p a -> Waveform t' a
freqModulate phaseSelect w = sampleFrom $ \t -> sample w (sample phaseSelect t)

-- | First argument is valued in sample time steps, so @1/stdtr@ is a good starting point.
emuVCO :: Waveform Tick Double -> Waveform Double a -> [a]
emuVCO freqSelect w = go 0 0
  where
    go !i sampPoint = let skip = sample freqSelect i in sample w sampPoint : go (i + 1) (sampPoint + skip)

emuVCO' :: [Double] -> Waveform Double a -> [a]
emuVCO' fs w = go fs 0
  where
    go freqSelect !sampPoint = let (skip:ss) = freqSelect in sample w sampPoint : go ss (sampPoint + skip)

niceVCO :: Double -> [Double] -> Waveform Double a -> [a]
niceVCO amp fs = emuVCO' (fmap (\f -> (amp ** f)/stdtr) fs)

oscAboveZero :: Double -> Double -> Double
oscAboveZero amp = oscAbout (amp/2) amp

oscAbout :: Double -> Double -> Double -> Double
oscAbout z0 amp z = (amp * z) + z0

-- | Piecewise linear reconstruction (interpolation) of a sampled signal
linearResample :: Double -> Waveform Tick Double -> Waveform Double Double
linearResample sampleRate w = sampleFrom $ \t -> do
  let t' = t * sampleRate
  let (low,high) = (floor t',ceiling t')
  (fromIntegral high - t') * sample w low + (t' - fromIntegral low) * sample w high

recordStream :: [a] -> Waveform Tick a
recordStream xs = sampleFrom ((xs!!) . fromIntegral)

-- | The output wavetables here have @sample (toConvolutionTableStream ds !! t) k = ds !! (t - k)@ for @k < windowSize@
toConvolutionTableStream :: Tick -> [Discrete] -> [Wavetable]
toConvolutionTableStream windowSize xs = fmap toWavetable . scanl stepSt (startTable xs) $ xs
  where
    toWavetable :: (BA.ByteArray,Int) -> Wavetable
    toWavetable (v,b) = sampleFrom $ \k -> Discrete $ BA.indexByteArray v ((b - k) `mod` windowSize)
    startTable :: [Discrete] -> (BA.ByteArray,Int)
    startTable ds = (BA.byteArrayFromListN windowSize (fmap unDiscrete $ take windowSize ds),0)
    -- Table, breakpoint
    stepSt :: (BA.ByteArray, Int) -> Discrete -> (BA.ByteArray, Int)
    stepSt (!t,!b) x = runST $ do
      v <- BA.unsafeThawByteArray t
      BA.writeByteArray v (b `mod` windowSize) (unDiscrete x)
      v' <- BA.unsafeFreezeByteArray v
      pure (v',(b+1) `mod` windowSize)

streamWavetable :: Waveform Tick a -> [a]
streamWavetable w = map (sample w) [0..]

memorizeStream :: Int -> [Discrete] -> Wavetable
memorizeStream len = fromCompact . CompactWavetable . Vector.fromList . map unDiscrete . take len

-- | Very slow
integrateWave :: Num a => Waveform Tick a -> Waveform Tick a
integrateWave w = sampleFrom $ \t -> sum . map (sample w) $ [0 .. t]

-- | Modulate the phase of one wave with another. Used in synthesis.
-- @
--  'phaseModulate' beta ('setVolume' 0.2 $ 'sinWave' 'concertA') ('setVolume' 0.38 $ 'triWave' 'concertA')
-- @
-- (try beta=0.0005)
phaseModulate :: Num t
              => t -- ^ Tuning parameter. Modulation signal is @'amplitudeModulate'@d by @('const' beta)@
              -> Waveform t t -- ^ Modulation signal. Outputs the phase shift to apply
              -> Waveform t a -- ^ Target wave to be modulated
              -> Waveform t a
phaseModulate beta modulation target = sampleFrom $ \t -> sample target (t + beta * sample modulation t)

-- | Smoothly transition to playing a wave back at a different speed after some time
changeSpeed :: (Ord a,Fractional a) => a -> a -> a -> Waveform a a -> Waveform a a
changeSpeed startTime lerpTime newSpeed wave = sampleFrom $ \t -> sample wave $ if t < startTime
  then t
  else if t > startTime + lerpTime
    then startTime + newSpeed * t
    -- Lerp between sampling at 1 and sampling at newSpeed
    else startTime + (1 + ((t - startTime)/lerpTime) * (newSpeed - 1)) * t

-- | Play several waves on top of each other, normalizing so that e.g. playing three notes together doesn't triple the volume.
balanceChord :: Fractional a => [Waveform t a] -> Waveform t a
balanceChord notes = sampleFrom $ \t -> foldr (\x s -> s + factor * sample x t) 0 $ notes
  where
    factor = realToFrac . recip @Double . fromIntegral . length $ notes

-- | Play several waves on top of each other, without worrying about the volume. See 'balanceChord' for
-- a normalized version.
mergeWaves :: Num a => [Waveform t a] -> Waveform t a
mergeWaves notes = sampleFrom $ \t -> foldr (\x s -> s + sampleAt t x) 0 notes
  -- Average Frequency
  --,frequency = fmap (/(fromIntegral $ length notes)) . foldl (liftA2 (+)) (Just 0) . map frequency $ notes

-- | @'waveformToWAVE' outputLength@ gives a @'WAVE'@ file object by sampling the given @'DWave'@ at @44100Hz@.
-- May disbehave or clip based on behavior of @'doubleToSample'@ if the DWave takes values outside of @[-1,1]@.
waveformToWAVE :: Tick -> Int -> Wavetable -> WAVE.WAVE
waveformToWAVE outTicks sampleRate w = WAVE.WAVE
  {WAVE.waveHeader = WAVE.WAVEHeader
    {WAVE.waveNumChannels = 1
    ,WAVE.waveFrameRate = sampleRate
    ,WAVE.waveBitsPerSample = 32
    ,WAVE.waveFrames = Just $ fromIntegral outTicks
    }
  ,WAVE.waveSamples = [map (unDiscrete . sample w) [0 .. outTicks - 1]]
  }

-- | Take a specified number of samples at the given sample rate from the stream and put them in a @'WAVE.WAVE'@.
wavestreamToWAVE :: Tick -> Int -> [Discrete] -> WAVE.WAVE
wavestreamToWAVE outTicks sampleRate ws = WAVE.WAVE
  {WAVE.waveHeader = WAVE.WAVEHeader
    {WAVE.waveNumChannels = 1
    ,WAVE.waveFrameRate = sampleRate
    ,WAVE.waveBitsPerSample = 32
    ,WAVE.waveFrames = Just $ fromIntegral outTicks
    }
  ,WAVE.waveSamples = [take outTicks $ map unDiscrete ws]
  }

-- | Triangle wave of the given frequency
triWave :: Double -> Waveform Double Double
triWave f = sampleFrom $ \t -> let r = t * f - fromIntegral (floor (t * f) :: Int) in if r < 0.25
  then 4 * r
  else if r < 0.75
    then 2 - (4 * r)
    else -4 + (4 * r)

-- | Digital saw wave with given frequency
sawWave :: Double -> Waveform Double Double
sawWave f = sampleFrom $ \t -> 2 * (t * f - fromIntegral (floor (t * f) :: Int)) - 1

-- | Ramp from 0 to 1 over a time period
ramp :: Double -> Double -> Double
ramp = rampFrom 0

-- | Ramp from @x0@ to 1 from time zero to the max time
rampFrom :: Double -> Double -> (Double -> Double)
rampFrom x0 time t = if t < time
  then x0 + (1 - x0) * (max t 0)/time
  else 1

-- | Arbitrarily chosen standard tick rate, used in @'testWave'@
stdtr :: Num a => a
stdtr = 48000 --32000

-- | Output the first @len@ seconds of the given @'Wavetable'@ to a @.wav@ file at the given path for testing.
-- The volume is also attenuated by 50% to not blow out your eardrums.
-- Also pretty prints the wave.
testWave :: Double -> String -> Wavetable -> IO ()
testWave len fp w = print w >> pure w >>= WAVE.putWAVEFile (fp ++ ".wav") . waveformToWAVE (floor $ len*stdtr) stdtr . amplitudeModulate (sampleFrom $ const 0.5)

-- | Output the first @len@ seconds of the given wave stream to a @.wav@ file at the given path for testing.
testWaveStream :: Double -> String -> [Discrete] -> IO ()
testWaveStream len fp = WAVE.putWAVEFile (fp ++ ".wav") . wavestreamToWAVE (floor $ len*stdtr) stdtr

-- | Concert A4 frequency is 440Hz
concertA :: Num a => a
concertA = 440

-- | An envelope that suspends for @noteDuration@ before releasing
susEnvelope :: Envelope Tick Discrete -> Tick -> Wavetable
susEnvelope (Envelope del att hol dec sus rel) noteDuration = sampleFrom $ \t -> if t < del
  then 0
  else if t - del < att
    then doubleToDiscrete $ (t - del) `divTD` att
    else if t - del - att < hol
      then 1
      else if t - del - att - hol < dec
        then doubleToDiscrete $ 1 + ((t - del - att - hol) `divTD` dec) * (discreteToDouble sus - 1)
        else if t < noteDuration
          then sus
          else if t - noteDuration < rel
            then sus * doubleToDiscrete (1 - (t - noteDuration) `divTD` rel)
            else 0
  where
    divTD a b = fromIntegral a / fromIntegral b

-- | An envelope that suspends forever instead of releasing
suspendVelope :: Double -> Double -> Double -> Double -> Double -> DWave
suspendVelope del att hol dec sus = sampleFrom $ \t -> if t < del
  then 0
  else if t - del < att
    then (t - del) / att
    else if t - del - att < hol
      then 1
      else if t - del - att - hol < dec
        then 1 + (t - del - att - hol)/dec * (sus - 1)
        else sus

-- | Synthesize an envelope to a 'DWave'
envelope :: Envelope Double Double -> DWave
envelope (Envelope del att hol dec sus rel) = sampleFrom $ \t -> if t < del
  then 0
  else if t - del < att
    then (t - del) / att
    else if t - del - att < hol
      then 1
      else if t - del - att - hol < dec
        then 1 + (t - del - att - hol)/dec * (sus - 1)
        else if t - del - att - hol - dec < rel
          then sus * (1 - (t - del - att - hol - dec)/rel)
          else 0

-- | DAHDSR envelope with times @t@ and sustain level @b@
data Envelope t b = Envelope t t t t b t

-- | Discretize an envelope by converting times into ticks
discretizeEnvelope :: Double -> Envelope Double Double -> Envelope Tick Discrete
discretizeEnvelope tickRate (Envelope del att hol dec sus rel) = Envelope (dd del) (dd att) (dd hol) (dd dec) (doubleToDiscrete sus) (dd rel)
  where
    dd = floor . (*tickRate)

-- | Shift a wave in time to start at the specified time after its old start time
timeShift :: Num t => t -> Waveform t a -> Waveform t a
timeShift dt = sampleFrom . (. subtract dt) . sample

-- | Shift a wave in time such that the new zero is at the specified position
seekTo :: Num t => t -> Waveform t a -> Waveform t a
seekTo dt = sampleFrom . (. (+dt)) . sample

-- | Modify the amplitude of a wave by a constant multiple
setVolume :: Num a => a -> Waveform t a -> Waveform t a
setVolume = amplitudeModulate . sampleFrom . const

-- | The empty wave that is always zero when sampled
emptyWave :: Num a => Waveform t a
emptyWave = sampleFrom $ const 0

-- | Convolve with explicit discrete filter kernel weights.
discreteConvolve :: (Num a, Num t) => Waveform t [(t,a)] -> Waveform t a -> Waveform t a
discreteConvolve profile w = sampleFrom $ \t -> sum . map (\(dt,amp) -> amp * sample w (t + dt)) $ sample profile t

-- | This operation is not convolution, but something kind of like it. Use for creative purposes? Should be fast!
-- @wackyNotConvolution modf profile w = sampleFrom $ \t -> sample (modulate modf (sample profile t) w) t@
wackyNotConvolution :: (a -> b -> c) -> Waveform t (Waveform t a) -> Waveform t b -> Waveform t c
wackyNotConvolution modf profile w = sampleFrom $ \t -> sample (modulate modf (sample profile t) w) t

-- | Perform a discrete convolution. The output waveform is @f(t) = \int_{t-tickRadius}^{t+tickRadius} (kernel(t))(x) * w(t+x) dx@
-- but is discretized such that @x@ is always a multiple of @skipRate@.
tickConvolution :: (Show a,Fractional a)
                => Tick -- ^ @tickRadius@
                -> Tick -- ^ @skipRate@
                -> Waveform Tick (Waveform Tick a) -- ^ The kernel of the convolution at each @'Tick'@
                -> Waveform Tick a -- ^ w(t)
                -> Waveform Tick a
tickConvolution tickRadius skipRate profile w = sampleFrom $ \t -> let !kern = sample profile t in sum . map (\dt -> (*0.01) . (*sample w (t + dt)) . sample kern $ dt) $ sampleDeltas
  where
    sampleDeltas = map (*skipRate) [-stepsPerSide.. stepsPerSide]
    stepsPerSide = tickRadius `div` skipRate
    -- !stepModifier = realToFrac . recip . fromIntegral $ stepsPerSide :: Double

--accuCount :: IORef Integer
--accuCount = unsafePerformIO $ newIORef 0

-- | This computes a solid slice of a convolution much faster than sampling each point individually by using a convolution cache window.
fastTickConvolutionFixedKern :: Double -> Tick -> Tick -> Wavetable -> Wavetable -> [Discrete]
fastTickConvolutionFixedKern tickRate tickStart tickRadius kern w = kern `seq` go (vi,0) tickStart
  where
    vi = runST $ do
      v <- BA.newByteArray (windowSize * 4)
      forM_ [-tickRadius .. (tickRadius - 1)] $ \dt -> BA.writeByteArray v (1 + tickRadius + dt) (unDiscrete $ sample w (tickStart + dt))
      BA.unsafeFreezeByteArray v
    -- Table, breakpoint, offset
    stepSt :: (BA.ByteArray, Int) -> Tick -> (Discrete, BA.ByteArray)
    stepSt (!t,!b) o = runST $ do
      v <- BA.unsafeThawByteArray t
      BA.writeByteArray v (b `mod` windowSize) (unDiscrete $ sample w (o + tickRadius))
      let
        accu !tot (dt:dts) = do
          --let !() = unsafePerformIO $ readIORef accuCount >>= \x -> (if x `mod` 10000 == 0 then putStrLn ("accu: " ++ show x) else pure ()) >> writeIORef accuCount (x + 1)
          s <- BA.readByteArray v ((b + tickRadius + dt + 1) `mod` windowSize)
          let !toAdd = (0.01 * Discrete (BA.indexByteArray kernel dt)) * Discrete s
          accu (tot + toAdd) dts
        accu tot [] = pure tot
      k <- accu 0 sampleDeltas
      !v' <- BA.unsafeFreezeByteArray v
      pure (k,v')
    go (!t,!b) o = let (s,t') = stepSt (t,b) o in s `seq` (s : go (t',(b+1) `mod` windowSize) (o + 1))
    --accu :: BA.MutableByteArray s -> Int -> Tick -> Discrete -> Tick -> ST s Discrete
    windowSize = tickRadius * 2 + 1
    sampleDeltas = [-tickRadius .. tickRadius]
    !kernel = runST $ do
      v <- BA.newByteArray (windowSize * 4)
      forM_ [-tickRadius .. (tickRadius - 1)] $ \dt -> BA.writeByteArray v (1 + tickRadius + dt) (unDiscrete . (* doubleToDiscrete (recip tickRate)) $ sample kern dt)
      BA.unsafeFreezeByteArray v

-- | Same as @'tickConvolution'@ but for arbitarily valued waveforms. Works on @'DWave'@ for example.
sampledConvolution :: (RealFrac t, Fractional a)
                   => t -- ^ @convolutionSampleRate@, controls sampling for @x@
                   -> t -- ^ @convolutionRadius@, continuous analogue of @tickRadius@
                   -> Waveform t (Waveform t a) -- ^ Kernel of convolution for each time
                   -> Waveform t a -> Waveform t a
sampledConvolution convolutionSampleRate convolutionRadius profile w = sampleFrom $ \t -> sum . map (\dt -> (*(realToFrac . recip $ convolutionSampleRate * convolutionRadius)) . (* sample w (t + dt)) . sample (sample profile t) $ dt) $ sampleDeltas
  where
    sampleDeltas = map ((/convolutionSampleRate) . realToFrac) [-samplesPerSide .. samplesPerSide]
    samplesPerSide = floor (convolutionRadius * convolutionSampleRate) :: Int
    --sampleCount = 2 * samplesPerSide + 1

-- | Makes a filter which selects frequencies near @bandCenter@ with tuning parameter @bandSize@.
-- Try: @'optimizeFilter' 200 . 'tickTable' 'stdtr' $ 'bandpassFilter' 'concertA' 100@
bandpassFilter :: Fractional a
               => Double -- ^ @bandCenter@
               -> Double -- ^ @bandSize@
               -> Waveform Double a
bandpassFilter bandCenter bandSize = sampleFrom $ \t -> if t == 0 then 1 else realToFrac $ (sin (bandFreq * t)) / (bandFreq * t) * (cos (centerFreq * t))
  where
    !bandFreq = 2 * pi * bandSize
    !centerFreq = 2 * pi * bandCenter

-- | Ladder filter, adapted from VCV Rack
ladderFilter :: Double -> [Double] -> [Double] -> [Double] -> [(Double,Double,Double,Double)]
ladderFilter sampleRate reso cutoffs inputs = fmap (FV.convert :: FV.VecList 4 Double -> (Double,Double,Double,Double)) $ rkSolveStepSize (1/sampleRate) rkf (FV.mk4 1e-6 1e-6 1e-6 1e-6) (zip3 reso cutoffs inputs)
  where
    rkf (r,c,i) x = let cut = 2 * pi * c in FV.mk4 (cut * (i - r * (x FV.! 2) - (x FV.! 0))) (cut * (x FV.! 0 - x FV.! 1)) (cut * (x FV.! 1 - x FV.! 2)) (cut * (x FV.! 2 - x FV.! 3))

-- | Use Runge-Kutta method to solve the diffeq
rkSolve
  :: Applicative f => (i -> f Double -> f Double) -- ^ User suplied function mapping inputs to a coefficient matrix
  -> f Double -- ^ Initial value
  -> [i] -- ^ Stream of inputs
  -> [f Double] -- ^ Stream of outputs
rkSolve rkf = scanl stepSt
  where
    stepSt x i = let
      k1 = rkf i x
      k2 = rkf i (liftA2 (\k xi -> xi + k/2) k1 x)
      k3 = rkf i (liftA2 (\k xi -> xi + k/2) k2 x)
      k4 = rkf i (liftA2 (+) k3 x)
      in liftA2 (\xi ki -> xi + ki/6) x $ liftA2 (+) (liftA2 (+) (fmap (*2) k2) (fmap (*2) k3)) (liftA2 (+) k1 k4)

-- | Like @'rkSolve'@ but using non-unit step size by multiplying the coefficient matrix by a factor
rkSolveStepSize :: Applicative f => Double -> (i -> f Double -> f Double) -> f Double -> [i] -> [f Double]
rkSolveStepSize h rkf = rkSolve ((fmap (* h) .) . rkf)

-- | Drive scaling function adapted from VCV Rack
scaledDrive :: Double -> Double
scaledDrive drv = (1 + drv) ** 5

-- | Uses @'ladderFilter'@ to low pass filter the input with the given resonance and cutoff
lowPassFilter :: Double -> [Double] -> [Double] -> [Double] -> [Double]
lowPassFilter rate res cut = fmap (\(_,_,_,d) -> d) . ladderFilter rate res cut

-- | Synthesize from a continuous profile by sampling frequencies.
-- Doesn't work very well; prefer 'synthFromDisreteProfile'.
synthFromFreqProfile :: (Double,Double) -> Double -> DWave -> DWave
synthFromFreqProfile (f0,f1) fSampRate prof = sampleFrom $ \t -> max (-1) $ min 1 $ sum . fmap ((/fromIntegral (length sampPoints)) . sampleAt t . getComponent) $ sampPoints
  where
    getComponent f = fmap (*sample prof f) $ sinWaveWithPhase f (pi/2)
    sampPoints = [f0, f0 + 1/fSampRate .. f1]
    --jitterFactor k = (fromIntegral (floor (((k - f0) / (f1 - f0)) * 81083) `mod` 115)) / 115
    --jitter k = k + (2/fSampRate) * jitterFactor k

-- | Continuous saxopohone profile adapted from https://asa.scitation.org/doi/10.1121/1.396474
saxProfile :: DWave
saxProfile = sampleFrom $ \f -> let x = f/fb in sqrt (n * x/(1 + x**7))
  where
    fb = 618
    n = 1.088910

-- | Synthesize a wavetable from the specified relative strengths of various frequencies.
-- Essentially an inverse fourier transform
synthFromDiscreteProfile :: [(Double,Double)] -> Wavetable
synthFromDiscreteProfile fs = {-solidSlice 0 (truncate $ stdtr/(1 :: Double)) .-} tickTable stdtr . discretize . mergeWaves . fmap (\(f,amp) -> fmap (*(normFactor * amp)) $ sinWaveWithPhase f f) $ fs
  where
    normFactor = 0.9/(sum $ fmap snd fs)

-- | Produces a discrete profile of relative harmonic strengths from the given functions.
-- Try @'harmonicEquationToDiscreteProfile' (\x -> 0.1894 / (x ** 1.02)) (\x -> 0.0321 / (x ** 0.5669))@
harmonicEquationToDiscreteProfile :: (Double -> Double) -> (Double -> Double) -> Double -> [(Double,Double)]
harmonicEquationToDiscreteProfile oddF evenF f0 = go (15 :: Int)
  where
    go 0 = []
    go k = (fromIntegral k * f0,(if even k then evenF else oddF) $ fromIntegral k) : go (k-1)

-- | Roughly resemebles the spectrum of a saxophone
discSaxProfile :: [(Double,Double)]
discSaxProfile = f1 ++ f2
  where
  f1 =
    [(588.6,0.0911)
    ,(1179,0.05463)
    ,(1767,0.03471)
    ,(2357,0.01109)
    ,(2946,0.01625)
    ] -- y=63/x ish
  f2 =
    [(294.9,0.03032)
    ,(884,0.02445)
    ,(1474,0.01172)
    ,(2062,0.008224)
    ,(2651,0.00881)
    ,(3241,0.005567)
    ]

-- | Roughly resemebles the spectrum of a saxophone
genSaxProfile' :: Double -> [(Double,Double)]
genSaxProfile' f0 =
  [(1 * f0,0.0474475510012406)
  ,(2 * f0,0.0141427241744339)
  ,(3 * f0,0.00838229600425522)
  ,(4 * f0,0.00452452580892121)
  ,(5 * f0,0.0129662632033567)
  ,(6 * f0,0.0102954061023286)
  ,(7 * f0,0.00246665413386754)
  ,(8 * f0,0.00309631105551752)
  ,(9 * f0,0.00313836031805197)
  ,(10 * f0,0.00549688733850934)
  ,(11 * f0,0.00826471159302181)
  ,(12 * f0,0.00280485189554740)
  ]

-- | Roughly resemebles the spectrum of a saxophone
genSaxProfile :: Double -> [(Double,Double)]
genSaxProfile f0 = f1 ++ f2
  where
  f1 =
    [(2 * f0,0.0911)
    ,(4 * f0,0.05463)
    ,(4 * f0,0.03471)
    ,(6 * f0,0.01109)
    ,(8 * f0,0.01625)
    ] -- y=63/x ish
  f2 =
    [(f0,0.03032)
    ,(3 * f0,0.02445)
    ,(5 * f0,0.01172)
    ,(7 * f0,0.008224)
    ,(9 * f0,0.00881)
    ,(11 * f0,0.005567)
    ]

-- | Rough saxophone timbre from @'gen
fakedSaxTimbre :: Double -> Wavetable
fakedSaxTimbre = synthFromDiscreteProfile . genSaxProfile'

-- | A timbre that consists of several sin wave voices at fixed chord intervals from each other
chordSinTimbre :: ChordVoicing PitchFactorDiagram -> Double -> Wavetable
chordSinTimbre c r = discretize . tickTable stdtr . balanceChord . fmap (sinWave . flip intervalOf r) $ getVoiceList c

-- | Discretize the output of a @'Double'@ producing waveform
discretize :: Functor f => f Double -> f Discrete
discretize = fmap doubleToDiscrete

-- | Discretize the input to a @'Double'@ consuming waveform
tickTable :: Double -- ^ Sample rate. Each tick is @1/sampleRate@ seconds
          -> Waveform Double a -> Waveform Tick a
tickTable tickrate w = sampleFrom $ \t -> sample w (fromIntegral t/tickrate)

-- | Discretize the input to a @'Double'@ consuming waveform
tickTablePer :: Double -- ^ Sample period. Sampling rate is @1/period@
          -> Waveform Double a -> Waveform Tick a
tickTablePer per w = sampleFrom $ \t -> sample w (fromIntegral t * per)

-- | A domain- and codomain-discretized @'Waveform'@ suitable for writing to a WAVE file.
-- See @'waveformToWAVE'@.
type Wavetable = Waveform Tick Discrete

-- | A data structure for storing the results of a @'Wavetable'@ on some subset of its domain.
-- Used internally.
data CompactWavetable = CompactWavetable {getWavetable :: Vector.Vector Int32}

fromCompact :: CompactWavetable -> Wavetable
fromCompact cwt = sampleFrom $ \t -> case getWavetable cwt Vector.!? t of
  Just d -> Discrete d
  Nothing -> error "fromCompact Wavetable sampled outside size of compact table"

-- | Optimize a @'Wavetable'@ by storing its values in a particular range.
-- Uses @(tickEnd - tickStart + 1) * sizeOf (_ :: 'Discrete')@ bytes of memory to do this.
solidSlice :: Tick -> Tick -> Wavetable -> Wavetable
solidSlice tickStart tickEnd w = sampleFrom $ \t -> case getWavetable cwt Vector.!? (t-tickStart) of
  Just d -> Discrete d
  Nothing -> sample w t
  where
    cwt = CompactWavetable {getWavetable = Vector.generate (tickEnd - tickStart + 1) (unDiscrete . sample w . (+tickStart))}

-- | Optimize a filter by doing @'solidSlice'@ around @t=0@ since those values are sampled repeatedly in a filter
optimizeFilter :: Tick -> Wavetable -> Wavetable
optimizeFilter tickRadius = solidSlice (-tickRadius) tickRadius

-- | Take the Fourier Transform of a complex valued @'Tick'@ sampled waveform
fourierTransform :: Tick -> Double -> Waveform Tick (Complex Double) -> Waveform Double (Complex Double)
fourierTransform tickRadius fTickRate x = sampleFrom $ \f -> sum . map (\n -> sample x n / (fromIntegral tickRadius) * cis (2 * pi * f * (fromIntegral n / fTickRate))) $ [-tickRadius .. tickRadius]

-- | Take the Fourier Transform of a @'Wavetable'@
realDFT :: Tick -- ^ Radius of Fourier Transform window in @'Tick'@s. Try 200
        -> Double -- ^ Sampling rate to use for the Fourier transform. Try the sample sample rate as the @'Wavetable'@
        -> Wavetable -> Wavetable
realDFT tickRadius fTickRate w = discretize $ tickTable 1 $ fmap ((min 1) . magnitude) $ fourierTransform tickRadius fTickRate ((\x -> discreteToDouble x :+ 0) <$> solidSlice (-tickRadius) tickRadius w)

-- | Skip every @n@ ticks in the in the given @'Waveform'@.
-- @'sample' ('skipTicks' n w) k = 'sample' w (n*k)@
skipTicks :: Tick -- ^ @n@
          -> Waveform Tick a -> Waveform Tick a
skipTicks skipRate w = sampleFrom $ \t -> sample w (skipRate * t)

skipStream :: Tick -> [Discrete] -> [Discrete]
skipStream skipRate (x:xs) = x : skipStream skipRate (drop skipRate xs)
skipStream _ [] = []

-- | Optimize a @'Wavetable'@ that we know to be periodic by storing it's values on one period.
-- Takes @period * sizeOf (_ :: 'Discrete')@ bytes of memory to do this.
--
-- Warning: Causes clicking on period boundaries if the period isn't exactly the given value in ticks.
exploitPeriodicity :: Tick -- ^ Period in @'Tick'@s of the @'Wavetable'@.
                   -> Wavetable -> Wavetable
exploitPeriodicity period x = sampleFrom $ \t -> case getWavetable cwt Vector.!? (t `mod` period) of
  Just d -> Discrete d
  Nothing -> sample x t
  where
    cwt = CompactWavetable {getWavetable = Vector.generate period (unDiscrete . sample x)}

-- | Attempts to do a fast fourier transform, but the units of the domain of the output are highly suspect.
-- May be unreliable, use with caution.
usingFFT :: Tick -> Wavetable -> Wavetable
usingFFT tickRadius w = sampleFrom $ \t -> if t < (fromIntegral $ length l)
  then (!! t) . fmap (doubleToDiscrete . magnitude) $ l
  else 0
  where
    l = fft (map ((\x -> discreteToDouble x :+ 0) . sample w) [-tickRadius .. tickRadius])

-- | Cooley-Tukey fft
fft :: [Complex Double] -> [Complex Double]
fft [] = []
fft [x] = [x]
fft xs = zipWith (+) ys ts ++ zipWith (-) ys ts
    where n = length xs
          ys = fft evens
          zs = fft odds
          (evens, odds) = split xs
          split [] = ([], [])
          split [x] = ([x], [])
          split (x:y:xss) = (x:xt, y:yt) where (xt, yt) = split xss
          ts = zipWith (\z k -> exp' k n * z) zs [0..]
          exp' :: Int -> Int -> Complex Double
          exp' k a = cis $ -2 * pi * (fromIntegral k) / (fromIntegral a)


{-# LANGUAGE CPP #-}
{-# LANGUAGE DefaultSignatures, TypeOperators, FlexibleContexts, TypeSynonymInstances, FlexibleInstances #-}
module Pure.Data.Default where

import Control.Applicative

import Data.Complex
import Data.Int
import Data.Monoid
import Data.Ratio
import Data.Word

import GHC.Generics

class Default a where
  def :: a
  default def :: (Generic a, GDefault (Rep a)) => a
  def = to gdef

instance Default Bool where def = False -- like other languages with construction defaults

instance Default () where def = ()
instance Default Ordering where def = EQ
instance Default Any where def = Any False
instance Default All where def = All True
instance Default (Last a) where def = Last Nothing
instance (Num a) => Default (Sum a) where def = Sum 0
instance (Num a) => Default (Product a) where def = Product 1
instance Default (Endo a) where def = Endo id
instance Default a => Default (Const a b) where def = Const def
instance Default (Maybe a) where def = Nothing
instance Default [a] where def = []

-- Note that (def :: Try a) /= mempty
-- instance Default (Try a) where def = Trying

-- instance Default Txt where def = mempty

-- instance (Eq a, Hashable a) => Default (HashMap.HashMap a b) where
--   def = mempty

-- instance Default Value where
--   def =
-- #ifdef __GHCJS__
--     nullValue
-- #else
--     Null
-- #endif

-- instance Default Obj where def = mempty

-- instance Default Micros where def = 0
-- instance Default Millis where def = 0
instance Default Int where def = 0
instance Default Int8 where def = 0
instance Default Int16 where def = 0
instance Default Int32 where def = 0
instance Default Int64 where def = 0
instance Default Word where def = 0
instance Default Word8 where def = 0
instance Default Word16 where def = 0
instance Default Word32 where def = 0
instance Default Word64 where def = 0
instance Default Integer where def = 0
instance Default Float where def = 0
instance Default Double where def = 0
instance (Integral a) => Default (Ratio a) where def = 0
instance (Default a,RealFloat a) => Default (Complex a) where def = def :+ def
instance {-# OVERLAPPABLE #-} Default r => Default (x -> r) where def = const def
instance {-# OVERLAPPING #-} Default (a -> a) where def = id
instance Default a => Default (IO a) where def = return def
instance Default a => Default (Dual a) where def = Dual def

instance (Default a, Default b) => Default (a, b) where def = (def, def)
instance (Default a, Default b, Default c) => Default (a, b, c) where def = (def, def, def)
instance (Default a, Default b, Default c, Default d) => Default (a, b, c, d) where def = (def, def, def, def)
instance (Default a, Default b, Default c, Default d, Default e) => Default (a, b, c, d, e) where def = (def, def, def, def, def)
instance (Default a, Default b, Default c, Default d, Default e, Default f) => Default (a, b, c, d, e, f) where def = (def, def, def, def, def, def)
instance (Default a, Default b, Default c, Default d, Default e, Default f, Default g) => Default (a, b, c, d, e, f, g) where def = (def, def, def, def, def, def, def)

-- Inspired by Lukas Mai's data-default-class with a default instance for
-- sum types based on lexicographical order - similar to the Ord and Enum
-- instances
class GDefault f where
  gdef :: f a

instance GDefault V1 where
  gdef = undefined

instance GDefault U1 where
  gdef = U1

instance (Default a) => GDefault (K1 i a) where
  gdef = K1 def

instance (GDefault a, GDefault b) => GDefault (a :*: b) where
  gdef = gdef :*: gdef

instance (GDefault a, GDefault b) => GDefault (a :+: b) where
  gdef = L1 gdef

instance (GDefault a) => GDefault (M1 i c a) where
  gdef = M1 gdef


-- | See Hoffman, Gelman (2011) The No U-Turn Sampler: Adaptively Setting Path
--   Lengths in Hamiltonian Monte Carlo.

module Numeric.MCMC.NUTS where

import Control.Monad
import Control.Monad.Loops
import Control.Monad.Primitive
import System.Random.MWC
import System.Random.MWC.Distributions hiding (gamma)
import Statistics.Distribution.Normal

type Parameters = [Double]
type Density    = Parameters -> Double
type Gradient   = Parameters -> Parameters
type Particle   = (Parameters, Parameters)
type StateSpecs = ([Double], [Double], [Double], [Double], [Double], Int, Int)

newtype StateTree = StateTree {
    getStateTree :: StateSpecs
  }

instance Show StateTree where
  show (StateTree (tm, rm, tp, rp, t', n, s)) =
       "\n" ++ "tm: " ++ show tm
    ++ "\n" ++ "tp: " ++ show tp
    ++ "\n" ++ "t': " ++ show t'
    ++ "\n" ++ "n : " ++ show n
    ++ "\n" ++ "s : " ++ show s

data DualAveragingParameters = DualAveragingParameters {
    mAdapt :: Int
  , delta  :: Double
  , mu     :: Double
  , gamma  :: Double
  , tau0   :: Double
  , kappa  :: Double
  } deriving Show

-- | The NUTS sampler.
nuts
  :: PrimMonad m
  => Density
  -> Gradient
  -> Int
  -> Double
  -> Parameters
  -> Gen (PrimState m)
  -> m [Parameters]
nuts lTarget glTarget n e t g = go t 0 []
  where go position j acc
          | j >= n    = return acc
          | otherwise = do
              nextPosition <- nutsKernel lTarget glTarget e position g
              go nextPosition (succ j) (nextPosition : acc)

-- | The NUTS sampler with dual averaging.
nutsDualAveraging
  :: PrimMonad m
  => Density
  -> Gradient
  -> Int
  -> Int
  -> Parameters
  -> Gen (PrimState m)
  -> m [Parameters]
nutsDualAveraging lTarget glTarget n nAdapt t g = do
    e0 <- findReasonableEpsilon lTarget glTarget t g
    let daParams = basicDualAveragingParameters e0 nAdapt
    chain <- unfoldrM (kernel daParams) (1, e0, 1, 0, t)
    return $ drop nAdapt chain
  where
    kernel params (m, e, eAvg, h, t0) = do
      (eNext, eAvgNext, hNext, tNext) <-
        nutsKernelDualAvg lTarget glTarget e eAvg h m params t0 g
      return $ if   m > n + nAdapt
               then Nothing
               else Just (t0, (succ m, eNext, eAvgNext, hNext, tNext))

-- | Default DA parameters, given a base step size and burn in period.
basicDualAveragingParameters :: Double -> Int -> DualAveragingParameters
basicDualAveragingParameters step burnInPeriod = DualAveragingParameters {
    mu     = log (10 * step)
  , delta  = 0.5
  , mAdapt = burnInPeriod
  , gamma  = 0.05
  , tau0   = 10
  , kappa  = 0.75
  }

-- | A single iteration of dual-averaging NUTS.
nutsKernelDualAvg
  :: PrimMonad m
  => Density
  -> Gradient
  -> Double
  -> Double
  -> Double
  -> Int
  -> DualAveragingParameters
  -> [Double]
  -> Gen (PrimState m)
  -> m (Double, Double, Double, Parameters)
nutsKernelDualAvg lTarget glTarget e eAvg h m daParams t g = do
  r0 <- replicateM (length t) (normal 0 1 g)
  z0 <- exponential 1 g
  let logu = log (auxilliaryTarget lTarget t r0) - z0

  let go (tn, tp, rn, rp, tm, j, n, s, a, na) g
        | s == 1 = do
            vj <- symmetricCategorical [-1, 1] g
            z  <- uniform g

            (tnn, rnn, tpp, rpp, t1, n1, s1, a1, na1) <-
              if   vj == -1
              then do
                (tnn', rnn', _, _, t1', n1', s1', a1', na1') <-
                  buildTreeDualAvg lTarget glTarget g tn rn logu vj j e t r0
                return (tnn', rnn', tp, rp, t1', n1', s1', a1', na1')
              else do
                (_, _, tpp', rpp', t1', n1', s1', a1', na1') <-
                  buildTreeDualAvg lTarget glTarget g tp rp logu vj j e t r0
                return (tn, rn, tpp', rpp', t1', n1', s1', a1', na1')

            let accept = s1 == 1 && (min 1 (fi n1 / fi n :: Double)) > z

                n2 = n + n1
                s2 = s1 * stopCriterion tnn tpp rnn rpp
                j1 = succ j
                t2 | accept    = t1
                   | otherwise = tm

            go (tnn, tpp, rnn, rpp, t2, j1, n2, s2, a1, na1) g

        | otherwise = return (tm, a, na)

  (nextPosition, alpha, nalpha) <- go (t, t, r0, r0, t, 0, 1, 1, 0, 0) g

  let (hNext, eNext, eAvgNext) =
          if   m <= mAdapt daParams
          then (hm, exp logEm, exp logEbarM)
          else (h, eAvg, eAvg)
        where
          eta = 1 / (fromIntegral m + tau0 daParams)
          hm  = (1 - eta) * h
              + eta * (delta daParams - alpha / fromIntegral nalpha)

          zeta = fromIntegral m ** (- (kappa daParams))

          logEm    = mu daParams - sqrt (fromIntegral m) / gamma daParams * hm
          logEbarM = (1 - zeta) * log eAvg + zeta * logEm

  return (eNext, eAvgNext, hNext, nextPosition)

-- | A single iteration of NUTS.
nutsKernel
  :: PrimMonad m
  => Density
  -> Gradient
  -> Double
  -> Parameters
  -> Gen (PrimState m)
  -> m Parameters
nutsKernel lTarget glTarget e t g = do
  r0   <- replicateM (length t) (normal 0 1 g)
  z0   <- exponential 1 g
  let logu = log (auxilliaryTarget lTarget t r0) - z0

  let go (tn, tp, rn, rp, tm, j, n, s) g
        | s == 1 = do
            vj <- symmetricCategorical [-1, 1] g
            z  <- uniform g

            (tnn, rnn, tpp, rpp, t1, n1, s1) <-
              if   vj == -1
              then do
                (tnn', rnn', _, _, t1', n1', s1') <-
                  buildTree lTarget glTarget g tn rn logu vj j e
                return (tnn', rnn', tp, rp, t1', n1', s1')
              else do
                (_, _, tpp', rpp', t1', n1', s1') <-
                  buildTree lTarget glTarget g tp rp logu vj j e
                return (tn, rn, tpp', rpp', t1', n1', s1')

            let accept = s1 == 1 && (min 1 (fi n1 / fi n :: Double)) > z

                n2 = n + n1
                s2 = s1 * stopCriterion tnn tpp rnn rpp
                j1 = succ j
                t2 | accept    = t1
                   | otherwise = tm

            go (tnn, tpp, rnn, rpp, t2, j1, n2, s2) g

        | otherwise = return tm

  go (t, t, r0, r0, t, 0, 1, 1) g

-- | Build the 'tree' of candidate states.
buildTree
  :: PrimMonad m
  => Density
  -> Gradient
  -> Gen (PrimState m)
  -> Parameters
  -> Parameters
  -> Double
  -> Double
  -> Int
  -> Double
  -> m StateSpecs
buildTree lTarget glTarget g t r logu v 0 e = do
  let (t0, r0) = leapfrog glTarget (t, r) (v * e)
      joint    = log $ auxilliaryTarget lTarget t0 r0
      n        = indicate (logu < joint)
      s        = indicate (logu - 1000 < joint)
  return (t0, r0, t0, r0, t0, n, s)

buildTree lTarget glTarget g t r logu v j e = do
  z <- uniform g
  (tn, rn, tp, rp, t0, n0, s0) <-
    buildTree lTarget glTarget g t r logu v (pred j) e

  if   s0 == 1
  then do
    (tnn, rnn, tpp, rpp, t1, n1, s1) <-
      if   v == -1
      then do
        (tnn', rnn', _, _, t1', n1', s1') <-
          buildTree lTarget glTarget g tn rn logu v (pred j) e
        return (tnn', rnn', tp, rp, t1', n1', s1')
      else do
        (_, _, tpp', rpp', t1', n1', s1') <-
          buildTree lTarget glTarget g tp rp logu v (pred j) e
        return (tn, rn, tpp', rpp', t1', n1', s1')

    let accept = (fi n1 / max (fi (n0 + n1)) 1) > (z :: Double)
        n2     = n0 + n1
        s2     = s0 * s1 * stopCriterion tnn tpp rnn rpp
        t2     | accept    = t1
               | otherwise = t0

    return (tnn, rnn, tpp, rpp, t2, n2, s2)
  else return (tn, rn, tp, rp, t0, n0, s0)

-- | Determine whether or not to stop doubling the tree of candidate states.
stopCriterion :: (Integral a, Num b, Ord b) => [b] -> [b] -> [b] -> [b] -> a
stopCriterion tn tp rn rp =
      indicate (positionDifference `innerProduct` rn >= 0)
    * indicate (positionDifference `innerProduct` rp >= 0)
  where
    positionDifference = tp .- tn

-- | Build the tree of candidate states under dual averaging.
buildTreeDualAvg
  :: PrimMonad m
  => Density
  -> Gradient
  -> Gen (PrimState m)
  -> Parameters
  -> Parameters
  -> Double
  -> Double
  -> Int
  -> Double
  -> Parameters
  -> Parameters
  -> m ([Double], [Double], [Double], [Double], [Double], Int, Int, Double, Int)
buildTreeDualAvg lTarget glTarget g t r logu v 0 e t0 r0 = do
  let (t1, r1) = leapfrog glTarget (t, r) (v * e)
      joint    = log $ auxilliaryTarget lTarget t1 r1
      n        = indicate (logu < joint)
      s        = indicate (logu - 1000 <  joint)
      a        = min 1 (acceptanceRatio lTarget t0 t1 r0 r1)
  return (t1, r1, t1, r1, t1, n, s, a, 1)

buildTreeDualAvg lTarget glTarget g t r logu v j e t0 r0 = do
  z <- uniform g
  (tn, rn, tp, rp, t1, n1, s1, a1, na1) <-
    buildTreeDualAvg lTarget glTarget g t r logu v (pred j) e t0 r0

  if   s1 == 1
  then do
    (tnn, rnn, tpp, rpp, t2, n2, s2, a2, na2) <-
      if   v == -1
      then do
        (tnn', rnn', _, _, t1', n1', s1', a1', na1') <-
          buildTreeDualAvg lTarget glTarget g tn rn logu v (pred j) e t0 r0
        return (tnn', rnn', tp, rp, t1', n1', s1', a1', na1')
      else do
        (_, _, tpp', rpp', t1', n1', s1', a1', na1') <-
          buildTreeDualAvg lTarget glTarget g tp rp logu v (pred j) e t0 r0
        return (tn, rn, tpp', rpp', t1', n1', s1', a1', na1')

    let p      = fi n2 / max (fi (n1 + n2)) 1
        accept = p > (z :: Double)
        n3     = n1 + n2
        a3     = a1 + a2
        na3    = na1 + na2
        s3     = s1 * s2 * stopCriterion tnn tpp rnn rpp

        t3  | accept    = t2
            | otherwise = t1

    return (tnn, rnn, tpp, rpp, t3, n3, s3, a3, na3)
  else return (tn, rn, tp, rp, t1, n1, s1, a1, na1)

-- | Heuristic for initializing step size.
findReasonableEpsilon
  :: PrimMonad m
  => Density
  -> Gradient
  -> Parameters
  -> Gen (PrimState m)
  -> m Double
findReasonableEpsilon lTarget glTarget t0 g = do
  r0 <- replicateM (length t0) (normal 0 1 g)
  let (t1, r1) = leapfrog glTarget (t0, r0) 1.0
      a        = 2 * indicate (acceptanceRatio lTarget t0 t1 r0 r1 > 0.5) - 1

      go j e t r
        | j <= 0 = e -- no need to shrink this excessively
        | (acceptanceRatio lTarget t0 t r0 r) ^^ a > 2 ^^ (-a) =
            let (tn, rn) = leapfrog glTarget (t, r) e
            in  go (pred j) (2 ^^ a * e) tn rn
        | otherwise = e

  return $ go 10 1.0 t1 r1

-- | Simulate a single step of Hamiltonian dynamics.
leapfrog :: Gradient -> Particle -> Double -> Particle
leapfrog glTarget (t, r) e = (tf, rf)
  where
    rm = adjustMomentum glTarget e t r
    tf = adjustPosition e rm t
    rf = adjustMomentum glTarget e tf rm

-- | Adjust momentum.
adjustMomentum :: Fractional c => (t -> [c]) -> c -> t -> [c] -> [c]
adjustMomentum glTarget e t r = r .+ ((e / 2) .* glTarget t)

-- | Adjust position.
adjustPosition :: Num c => c -> [c] -> [c] -> [c]
adjustPosition e r t = t .+ (e .* r)

-- | The MH acceptance ratio for a given proposal.
acceptanceRatio :: Floating a => (t -> a) -> t -> t -> [a] -> [a] -> a
acceptanceRatio lTarget t0 t1 r0 r1 = auxilliaryTarget lTarget t1 r1
                                    / auxilliaryTarget lTarget t0 r0

-- | The negative potential.
auxilliaryTarget :: Floating a => (t -> a) -> t -> [a] -> a
auxilliaryTarget lTarget t r = exp (lTarget t - 0.5 * innerProduct r r)

-- | Simple inner product.
innerProduct :: Num a => [a] -> [a] -> a
innerProduct xs ys = sum $ zipWith (*) xs ys

-- | Vectorized multiplication.
(.*) :: Num b => b -> [b] -> [b]
z .* xs = map (* z) xs

-- | Vectorized subtraction.
(.-) :: Num a => [a] -> [a] -> [a]
xs .- ys = zipWith (-) xs ys

-- | Vectorized addition.
(.+) :: Num a => [a] -> [a] -> [a]
xs .+ ys = zipWith (+) xs ys

-- | Indicator function.
indicate :: Integral a => Bool -> a
indicate True  = 1
indicate False = 0

-- | A symmetric categorical (discrete uniform) distribution.
symmetricCategorical :: PrimMonad m => [a] -> Gen (PrimState m) -> m a
symmetricCategorical [] _ = error "symmetricCategorical: no candidates"
symmetricCategorical zs g = do
  j <- uniformR (0, length zs - 1) g
  return $ zs !! j

-- | Alias for fromIntegral.
fi :: (Integral a, Num b) => a -> b
fi = fromIntegral