{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE TypeApplications #-}
module Numeric.ADDual.Examples where

import Control.DeepSeq
import Control.Monad (replicateM)
import Data.Maybe (catMaybes)
import qualified Data.Vector as V
import GHC.Generics (Generic)
import Hedgehog (Gen, Size)
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import qualified Hedgehog.Internal.Gen as HI.Gen
import qualified Hedgehog.Internal.Seed as HI.Seed
import qualified Hedgehog.Internal.Tree as HI.Tree


type Matrix s = V.Vector s

data FNeural a = FNeural [(Matrix a, V.Vector a)] (V.Vector a)
  deriving (Show, Eq, Functor, Foldable, Traversable, Generic)

instance NFData a => NFData (FNeural a)

fneural :: (Floating a, Ord a) => FNeural a -> a
fneural (FNeural layers input) =
  let dotp v1 v2 = V.sum (V.zipWith (*) v1 v2)

      mat @. vec =
        let n = V.length vec
            m = V.length mat `div` n
        in V.fromListN m $ map (\i -> dotp (V.slice (n*i) n mat) vec) [0 .. m-1]
      (+.) = V.zipWith (+)

      relu x = if x >= 0.0 then x else 0.0
      safeSoftmax vec = let m = V.maximum vec
                            factor = V.sum (V.map (\z -> exp (z - m)) vec)
                        in V.map (\z -> exp (z - m) / factor) vec
      forward [] x = safeSoftmax x
      forward ((weights, bias) : lys) x =
        let x' = V.map relu ((weights @. x) +. bias)
        in forward lys x'
  in V.sum $ forward layers input

makeNeuralInput :: Int -> FNeural Double
makeNeuralInput scale = sampleGenPure 100 (genNeuralInput scale)

genNeuralInput :: Int -> Gen (FNeural Double)
genNeuralInput scale = do
  let genScalar = Gen.double (Range.linearFracFrom 0 (-1) 1)
      genMatrix nin nout = V.fromListN (nin*nout) <$> replicateM (nin*nout) genScalar
      genVector nout = V.fromListN nout <$> replicateM nout genScalar
  nIn <- Gen.integral (Range.linear 1 scale)
  n1 <- Gen.integral (Range.linear 1 scale)
  n2 <- Gen.integral (Range.linear 1 scale)
  m1 <- genMatrix nIn n1; v1 <- genVector n1
  m2 <- genMatrix n1 n2; v2 <- genVector n2
  inp <- genVector nIn
  pure $ FNeural [(m1, v1), (m2, v2)] inp


sampleGenPure :: Size -> Gen a -> a
sampleGenPure size gen =
  HI.Tree.treeValue $ head $ catMaybes
    [HI.Gen.evalGen size (HI.Seed.from n) gen | n <- [42..]]