{-# LANGUAGE DataKinds #-}
{-# LANGUAGE BangPatterns #-}

module Main where

import Control.Category
import Criterion.Main
import Data.MyPrelude
import Data.Utils
import Data.Void
import Numeric.Neural
import Prelude           hiding (id, (.))

main :: IO ()
main = defaultMain 
    [ bgroup "white" 
        [ bench "10/200"     $ whnf (w   10)   200
        , bench "10/2000"    $ whnf (w   10)  2000
        , bench "10/20000"   $ whnf (w   10) 20000
        , bench "100/200"    $ whnf (w  100)   200
        , bench "100/2000"   $ whnf (w  100)  2000
        , bench "100/20000"  $ whnf (w  100) 20000
        , bench "1000/200"   $ whnf (w 1000)   200
        , bench "1000/2000"  $ whnf (w 1000)  2000
        , bench "1000/20000" $ whnf (w 1000) 20000
        ]
    , env setupEnv $ \ ~(m, xss) -> bgroup "linear"
        [ l m xss  1 5 
        , l m xss  5 5
        , l m xss 10 5
        ]
    ]

w :: Int -> Int -> Double
w sampleCount testCount = flip evalRand (mkStdGen 123456) $ do
    stats   <- mkStats'
    samples <- replicateM sampleCount $ mkSample stats
    let m = whiten model' samples
    xss <- replicateM testCount $ mkSample stats
    return $ sum [model m xs | xs <- xss]

  where

    mkStats' :: MonadRandom m => m (Vector Width (Double, Double))
    mkStats' = sequenceA (pure $ (,) <$> getRandomR (-100, 100) <*> getRandomR (0.1, 20))

    mkSample :: MonadRandom m => Vector Width (Double, Double) -> m (Vector Width Double)
    mkSample = mapM $ uncurry boxMuller'

    model' :: Model (Vector Width) Identity Void (Vector Width Double) Double
    model' = Model (cArr $ Diff $ Identity . sum) absurd id runIdentity

type Width = 10

l :: M -> [Vector Width' Double] -> Int -> Int -> Benchmark
l m xss batchSize steps = bench (printf "%d/%d" batchSize steps) $ whnf l' steps where

    l' :: Int -> Double
    l' steps' = 
        let m' = loop steps' m 
            xs = pure 0
        in  modelError m' [(xs, xs)]

    loop :: Int -> M -> M
    loop 0  m' = m'
    loop !n m' =
        let m'' = m' `deepseq` snd $ descent m' 0.01 [(xs, xs) | xs <- take batchSize xss]
        in  loop (pred n) m''

setupEnv :: IO (M, [Vector Width' Double])
setupEnv = return $ flip evalRand (mkStdGen 987654) $ do
    let e xs = Diff $ Identity . sqDiff (fromDouble <$> xs)
    m   <- modelR $ mkStdModel linearLayer e id id
    xss <- replicateM 100 $ let r = getRandomR (-5, 5) in sequence $ pure r
    return (m, xss)

type M = StdModel (Vector Width') (Vector Width') (Vector Width' Double) (Vector Width' Double)

type Width' = 100