{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE NoFieldSelectors #-}

-- | A port of Karpathy's microgpt.py (vendored in microgpt/microgpt.py).
--
-- * test suite (@cabal test@): checks the port against reference values
--   produced by microgpt/reference.py
-- * benchmark (@cabal bench@): same workloads as microgpt/bench.py
-- * executable (@cabal run microgpt@): the full microgpt.py script
--   (shuffle, 1000 training steps, 20 samples)
module Main (main, benchMain, microGPT) where

import Control.DeepSeq (NFData, deepseq, force)
import Control.Exception (evaluate)
import Control.Monad (foldM, foldM_, forM, forM_, replicateM, zipWithM, zipWithM_, (>=>))
import Control.Monad.State.Strict (State, evalState, get, put, runState, state)
import qualified Criterion.Main as Criterion
import qualified Data.Char as Char
import qualified Data.Foldable as Foldable
import qualified Data.IntMap.Strict as IntMap
import qualified Data.List as List
import qualified Data.Map.Strict as Map
import qualified Data.Maybe as Maybe
import qualified Data.Set as Set
import qualified Data.Traversable as Traversable
import GHC.Generics (Generic)
import qualified Numeric
import System.IO (hFlush, stdout)
import System.Random (StdGen)
import qualified System.Random as Random
import Test.Hspec

main :: IO ()
main = hspec spec

spec :: Spec
spec = do
  describe "Autograd Engine" $ do
    it "computes gradients for f(a, b) = a * b + a" $ do
      let graph = do
            a <- constant 2.0
            b <- constant 3.0
            ab <- mul a b
            add ab a

          (finalNode, tape) = runState graph initialTape
          grads = backward finalNode tape

      -- Forward pass result
      finalNode.node.val `shouldBe` 8.0

      let (NodeIdxMap gMap) = grads
      IntMap.lookup 0 gMap `shouldBe` Just 4.0 -- df/da
      IntMap.lookup 1 gMap `shouldBe` Just 2.0 -- df/db
  describe "MicroGPT Training Step" $ do
    it "computes loss and gradients for a single word" $ do
      let tok = mkTokenizer [['a' .. 'z']]
          cfg = microConfig tok.vocabSize
          w = deterministicWeights cfg

          (loss, grads) = trainingStep cfg w (encode tok "hello")

      loss `shouldSatisfy` (> 0.0)
      -- Verify gradient map is not empty
      let (NodeIdxMap gMap) = grads
      IntMap.size gMap `shouldSatisfy` (> 0)
    it "handles empty documents (BOS -> BOS) and degenerate single-token inputs" $ do
      let tok = mkTokenizer [['a' .. 'z']]
          cfg = microConfig tok.vocabSize
          w = deterministicWeights cfg
          (emptyDocLoss, _) = trainingStep cfg w (encode tok "")
          (singleTokLoss, _) = trainingStep cfg w [tok.bos]
      emptyDocLoss `shouldSatisfy` (> 0.0)
      singleTokLoss `shouldBe` 0.0
  describe "Equivalence with microgpt.py (values from microgpt/reference.py)" $ do
    let tok = mkTokenizer [['a' .. 'z']]
        cfg = microConfig tok.vocabSize
        w = deterministicWeights cfg
    it "has the same vocabulary and number of parameters" $ do
      tok.vocabSize `shouldBe` 27
      length w `shouldBe` 4192
    it "computes the same loss" $ do
      let (loss, _) = trainingStep cfg w (encode tok "emma")
      loss `shouldBeApprox` 3.506449207715624
    it "computes the same gradients" $ do
      let (_, grads) = trainingStep cfg w (encode tok "emma")
          paramGrads = Foldable.toList (parameterGradients grads w)
      forM_ referenceGradients $ \(i, expected) ->
        Maybe.listToMaybe (drop i paramGrads)
          `shouldSatisfy` maybe False (approxEq expected)
    it "follows the same training trajectory (Adam with linear decay)" $ do
      let docs = map (encode tok) ["emma", "olivia", "ava", "isabella", "sophia"]
          (losses, final) = train cfg microHyperparams (initTrainState w) docs
      length losses `shouldBe` 5
      zipWithM_
        shouldBeApprox
        losses
        [3.506449207715624, 3.2498071277112257, 2.7546668273691433, 3.496728826446422, 3.056047453583676]
      fmap (.value) (Maybe.listToMaybe (Foldable.toList final.params))
        `shouldSatisfy` maybe False (approxEq (-0.00267390585765643))
  describe "Inference" $ do
    it "samples documents made of known tokens, at most block_size long" $ do
      let tok = mkTokenizer [['a' .. 'z']]
          cfg = microConfig tok.vocabSize
          (sample, _) = generate cfg tok.bos 0.5 (deterministicWeights cfg) (Random.mkStdGen 42)
      length sample `shouldSatisfy` (<= cfg.blockSize)
      sample `shouldSatisfy` all (\t -> t >= 0 && t < tok.bos)

-- | (parameter index, gradient) for "emma", from microgpt/reference.py
referenceGradients :: [(Int, Double)]
referenceGradients =
  [ (64, 0.5013475135105153),
    (432, -0.10409231751964869),
    (700, 0.12256221981304923),
    (1200, -0.003786914738571059),
    (1500, -0.0012771713881593858),
    (1700, 0.015920736906842876),
    (2000, 0.018551333841752494),
    (2500, -0.0005752369491939385),
    (4191, 0.014184402670265658)
  ]

approxEq :: Double -> Double -> Bool
approxEq expected actual = abs (actual - expected) <= 1e-9 * max 1 (abs expected)

shouldBeApprox :: Double -> Double -> Expectation
shouldBeApprox actual expected = actual `shouldSatisfy` approxEq expected

-- | Same file as microgpt.py (makemore's names.txt)
datasetPath :: FilePath
datasetPath = "microgpt/input.txt"

-- | Loads documents like microgpt.py: one per non-blank line, stripped
loadDocs :: FilePath -> IO [String]
loadDocs path = do
  contents <- readFile path
  let strip = List.dropWhileEnd Char.isSpace . dropWhile Char.isSpace
  pure $ filter (not . null) (map strip (lines contents))

-- | Same workloads as microgpt/bench.py: the first X documents (unshuffled),
-- with deterministic weights
benchMain :: IO ()
benchMain = do
  docs <- loadDocs datasetPath
  let tok = mkTokenizer docs
      cfg = microConfig tok.vocabSize
      n = 100
  w <- evaluate (force (deterministicWeights cfg))
  batch <- evaluate (force (map (encode tok) (take n docs)))

  Criterion.defaultMain
    [ Criterion.bgroup
        "microGPT"
        [ Criterion.bench ("forward (" <> show n <> " docs)") $
            Criterion.nf (map (forwardLoss cfg w)) batch,
          Criterion.bench ("forward+backward (" <> show n <> " docs)") $
            Criterion.nf (map (trainingStep cfg w)) batch,
          Criterion.bench "train (10 steps)" $
            Criterion.nf
              ( \ds ->
                  let (losses, final) = train cfg microHyperparams (initTrainState w) ds
                   in (losses, fmap (.value) final.params)
              )
              batch
        ]
    ]

-- 1. Scalar AD engine

newtype NodeId
  = NodeId Int
  deriving newtype (Eq, Ord, Num, Enum, Show)

initialNodeId :: NodeId
initialNodeId = NodeId 0

nextNodeId :: NodeId -> NodeId
nextNodeId (NodeId x) = NodeId $ x + 1

newtype LocalPartialDerivative a
  = LocalPartialDerivative a
  deriving newtype (Eq, Ord, Num, Show)

newtype NodeIdxMap a
  = NodeIdxMap (IntMap.IntMap a)
  deriving newtype (NFData)

nimEmpty :: NodeIdxMap a
nimEmpty = NodeIdxMap IntMap.empty

nimInsert :: NodeId -> a -> NodeIdxMap a -> NodeIdxMap a
nimInsert (NodeId idx) x (NodeIdxMap m) = NodeIdxMap $ IntMap.insert idx x m

nimInsertWith :: (a -> a -> a) -> NodeId -> a -> NodeIdxMap a -> NodeIdxMap a
nimInsertWith f (NodeId k) v (NodeIdxMap m) = NodeIdxMap $ IntMap.insertWith f k v m

nimLookup :: NodeId -> NodeIdxMap a -> Maybe a
nimLookup (NodeId k) (NodeIdxMap m) = IntMap.lookup k m

data Node a = Node
  { val :: a,
    deps :: [(NodeId, LocalPartialDerivative a)]
  }

data NodeIdx a = NodeIdx
  { idx :: NodeId,
    node :: Node a
  }

data Tape a = Tape
  { nodes :: NodeIdxMap (Node a),
    nextId :: NodeId
  }

initialTape :: Tape a
initialTape = Tape {nodes = nimEmpty, nextId = initialNodeId}

type Gradients = NodeIdxMap Double

nadd :: (Num a) => NodeIdx a -> NodeIdx a -> Node a
nadd x y =
  Node
    { val = x.node.val + y.node.val,
      deps = [(x.idx, LocalPartialDerivative 1), (y.idx, LocalPartialDerivative 1)]
    }

nmul :: (Num a) => NodeIdx a -> NodeIdx a -> Node a
nmul x y =
  Node
    { val = x.node.val * y.node.val,
      deps = [(x.idx, nvLpd y), (y.idx, nvLpd x)]
    }
  where
    nvLpd z = LocalPartialDerivative z.node.val

npow :: (Floating a) => NodeIdx a -> a -> Node a
npow x p =
  Node
    { val = x.node.val ** p,
      deps = [(x.idx, LocalPartialDerivative $ p * (x.node.val ** (p - 1)))]
    }

nexp :: (Floating a) => NodeIdx a -> Node a
nexp x =
  Node
    { deps = [(x.idx, LocalPartialDerivative val)],
      ..
    }
  where
    val = exp x.node.val

nlog :: (Floating a) => NodeIdx a -> Node a
nlog x =
  Node
    { val = log x.node.val,
      deps = [(x.idx, LocalPartialDerivative $ 1 / x.node.val)]
    }

nrelu :: (Ord a, Num a) => NodeIdx a -> Node a
nrelu x =
  Node
    { val = max 0 x.node.val,
      deps = [(x.idx, LocalPartialDerivative $ if x.node.val > 0 then 1 else 0)]
    }

type GraphBuilder = State (Tape Double)

-- | Automatically handles ID assignment and tape insertion
pushNode :: Node Double -> GraphBuilder (NodeIdx Double)
pushNode node = do
  tape <- get
  let currentId = tape.nextId
      newIdx = NodeIdx {idx = currentId, node = node}
  put $
    tape
      { nodes = nimInsert currentId node tape.nodes,
        nextId = nextNodeId currentId
      }
  pure newIdx

constant :: Double -> GraphBuilder (NodeIdx Double)
constant x = pushNode Node {val = x, deps = []}

-- | Registers a learnable parameter. In the graph, it acts as a leaf node.
parameter :: Double -> GraphBuilder (NodeIdx Double)
parameter = constant

add :: NodeIdx Double -> NodeIdx Double -> GraphBuilder (NodeIdx Double)
add x y = pushNode (nadd x y)

mul :: NodeIdx Double -> NodeIdx Double -> GraphBuilder (NodeIdx Double)
mul x y = pushNode (nmul x y)

pow :: NodeIdx Double -> Double -> GraphBuilder (NodeIdx Double)
pow x p = pushNode (npow x p)

exp' :: NodeIdx Double -> GraphBuilder (NodeIdx Double)
exp' x = pushNode (nexp x)

log' :: NodeIdx Double -> GraphBuilder (NodeIdx Double)
log' x = pushNode (nlog x)

relu :: NodeIdx Double -> GraphBuilder (NodeIdx Double)
relu x = pushNode (nrelu x)

backward :: NodeIdx Double -> Tape Double -> Gradients
backward finalNode tape =
  foldl' step initialGrads [finalNode.idx, finalNode.idx - 1 .. 0]
  where
    -- \partial L / \partial L = 1.0
    initialGrads = nimInsertWith (+) finalNode.idx 1.0 nimEmpty

    step :: Gradients -> NodeId -> Gradients
    step grads currentId =
      let -- Get the accumulated gradient for this node (default 0.0 if unused)
          currentGrad = Maybe.fromMaybe 0.0 (nimLookup currentId grads)
       in case nimLookup currentId tape.nodes of
            Nothing -> grads
            Just node ->
              -- Distribute currentGrad to all dependencies using the chain rule
              foldl' (accum currentGrad) grads node.deps

    accum :: Double -> Gradients -> (NodeId, LocalPartialDerivative Double) -> Gradients
    accum g grads (depId, LocalPartialDerivative lpd) =
      nimInsertWith (+) depId (g * lpd) grads

-- 2. Vector & Matrix Primitives

type Vector = [NodeIdx Double]

-- | Stored as [nout][nin], like microgpt's state_dict
type Matrix = [Vector]

-- | Sums a list of nodes into a single node (like Python's sum, the empty sum is 0)
sumNodes :: [NodeIdx Double] -> GraphBuilder (NodeIdx Double)
sumNodes = \case
  [] -> constant 0
  x : xs -> foldM add x xs

dot :: Vector -> Vector -> GraphBuilder (NodeIdx Double)
dot xs ys = do
  prods <- zipWithM mul xs ys
  sumNodes prods

-- | microgpt's linear: one dot product per output row, no bias
linear :: Vector -> Matrix -> GraphBuilder Vector
linear x w = forM w (`dot` x)

-- | Numerically stable softmax (subtracts max as a non-differentiable constant, matching microgpt.py)
softmax :: Vector -> GraphBuilder Vector
softmax = \case
  [] -> pure []
  xs@(x0 : rest) -> do
    let maxVal = foldl' max x0.node.val (map (.node.val) rest)
    negMax <- constant (negate maxVal)
    exps <- mapM (\x -> add x negMax >>= exp') xs
    total <- sumNodes exps
    invTotal <- pow total (-1.0)
    mapM (`mul` invTotal) exps

-- | Root Mean Square normalization (no learnable parameters)
rmsnorm :: Vector -> GraphBuilder Vector
rmsnorm = \case
  [] -> pure []
  x -> do
    squares <- mapM (\xi -> mul xi xi) x
    sumSquares <- sumNodes squares
    ms <- mul sumSquares =<< constant (fromIntegral (length x) ** (-1))
    eps <- constant 1e-5
    scale <- add ms eps >>= (`pow` (-0.5))
    mapM (`mul` scale) x

-- | Element-wise vector addition
addVectors :: Vector -> Vector -> GraphBuilder Vector
addVectors = zipWithM add

-- 3. Model: GPT-2 with rmsnorm, no biases and ReLU (like microgpt)

data Config = Config
  { vocabSize :: Int,
    nLayer :: Int,
    nEmbd :: Int,
    blockSize :: Int,
    nHead :: Int
  }

microConfig :: Int -> Config
microConfig vocab =
  Config
    { vocabSize = vocab,
      nLayer = 1,
      nEmbd = 16,
      blockSize = 16,
      nHead = 4
    }

headDim :: Config -> Int
headDim cfg = cfg.nEmbd `div` cfg.nHead

-- | Parameters of a transformer layer, matrices are [nout][nin]
data Layer a = Layer
  { attnWq :: [[a]],
    attnWk :: [[a]],
    attnWv :: [[a]],
    attnWo :: [[a]],
    mlpFc1 :: [[a]],
    mlpFc2 :: [[a]]
  }
  deriving stock (Generic, Functor, Foldable, Traversable)
  deriving anyclass (NFData)

-- | The root microGPT structure.
-- The 'Foldable' order matches microgpt's flattened @params@ list.
data Weights a = Weights
  { wte :: [[a]],
    wpe :: [[a]],
    lmHead :: [[a]],
    layers :: [Layer a]
  }
  deriving stock (Generic, Functor, Foldable, Traversable)
  deriving anyclass (NFData)

-- | Shape of microgpt's state_dict
weightsShape :: Config -> Weights ()
weightsShape cfg =
  Weights
    { wte = matrix cfg.vocabSize cfg.nEmbd,
      wpe = matrix cfg.blockSize cfg.nEmbd,
      lmHead = matrix cfg.vocabSize cfg.nEmbd,
      layers = replicate cfg.nLayer layer
    }
  where
    matrix nout nin = replicate nout (replicate nin ())
    layer =
      Layer
        { attnWq = matrix cfg.nEmbd cfg.nEmbd,
          attnWk = matrix cfg.nEmbd cfg.nEmbd,
          attnWv = matrix cfg.nEmbd cfg.nEmbd,
          attnWo = matrix cfg.nEmbd cfg.nEmbd,
          mlpFc1 = matrix (4 * cfg.nEmbd) cfg.nEmbd,
          mlpFc2 = matrix cfg.nEmbd (4 * cfg.nEmbd)
        }

-- | Gaussian initialization (std 0.08), like microgpt
randomWeights :: Config -> StdGen -> (Weights Double, StdGen)
randomWeights cfg = runState (traverse (const (gauss 0 0.08)) (weightsShape cfg))

-- | Box-Muller Gaussian sampler
gauss :: Double -> Double -> State StdGen Double
gauss mu sigma = do
  u1 <- state Random.random
  u2 <- state Random.random
  pure $ mu + sigma * cos (2 * pi * u1) * sqrt (-(2 * log (max 1e-300 (1 - u2))))

-- | Deterministic initialization shared with microgpt/harness.py to compare both implementations
deterministicWeights :: Config -> Weights Double
deterministicWeights cfg =
  snd $ Traversable.mapAccumL (\i () -> (i + 1, 0.08 * sin (fromIntegral i * 0.37 + 0.11))) (0 :: Int) (weightsShape cfg)

-- | Embeds pure Double weights into the Autograd tape as parameter nodes.
-- On a fresh tape, the n-th parameter gets NodeId n (see 'withGradients').
embedWeights :: Weights Double -> GraphBuilder (Weights (NodeIdx Double))
embedWeights = traverse parameter

-- | KV-Cache of a layer: historical Key and Value vectors
data KVCache = KVCache
  { keys :: [Vector],
    values :: [Vector]
  }

emptyCaches :: Config -> [KVCache]
emptyCaches cfg = replicate cfg.nLayer KVCache {keys = [], values = []}

-- | Safe row lookup into an embedding matrix
lookupRow :: [[a]] -> Int -> [a]
lookupRow rows idx = Maybe.fromMaybe [] (Maybe.listToMaybe (drop idx rows))

-- | Multi-head causal self-attention for one token, updating the cache
attention :: Config -> Layer (NodeIdx Double) -> KVCache -> Vector -> GraphBuilder (Vector, KVCache)
attention cfg layer cache x = do
  q <- linear x layer.attnWq
  k <- linear x layer.attnWk
  v <- linear x layer.attnWv

  let keys = cache.keys <> [k]
      values = cache.values <> [v]
      hd = headDim cfg
      slice h = take hd . drop (h * hd)

  invScale <- constant (sqrt (fromIntegral hd) ** (-1))

  heads <- forM [0 .. cfg.nHead - 1] $ \h -> do
    let qH = slice h q
        kH = map (slice h) keys
        vH = map (slice h) values
    attnLogits <- forM kH (dot qH >=> (`mul` invScale))
    attnWeights <- softmax attnLogits
    forM (List.transpose vH) (dot attnWeights)

  out <- linear (concat heads) layer.attnWo
  pure (out, KVCache {keys = keys, values = values})

-- | A transformer layer: attention and MLP blocks, both pre-normalized and residual
transformerLayer :: Config -> Layer (NodeIdx Double) -> KVCache -> Vector -> GraphBuilder (Vector, KVCache)
transformerLayer cfg layer cache x = do
  xNorm <- rmsnorm x
  (attnOut, cache') <- attention cfg layer cache xNorm
  x' <- addVectors attnOut x

  xNorm' <- rmsnorm x'
  hidden <- linear xNorm' layer.mlpFc1 >>= mapM relu
  mlpOut <- linear hidden layer.mlpFc2
  x'' <- addVectors mlpOut x'

  pure (x'', cache')

-- | microgpt's gpt: logits of the next token, given a token, its position and the caches
forwardPassMonadic :: Config -> Weights (NodeIdx Double) -> Int -> Int -> [KVCache] -> GraphBuilder (Vector, [KVCache])
forwardPassMonadic cfg w tokenId posId caches = do
  x0 <- addVectors (lookupRow w.wte tokenId) (lookupRow w.wpe posId)
  x1 <- rmsnorm x0

  (x, caches') <-
    foldM
      ( \(h, acc) (layer, cache) -> do
          (h', cache') <- transformerLayer cfg layer cache h
          pure (h', acc <> [cache'])
      )
      (x1, [])
      (zip w.layers caches)

  logits <- linear x w.lmHead
  pure (logits, caches')

-- 4. Tokenizer

data Tokenizer = Tokenizer
  { uchars :: [Char],
    tokenIds :: Map.Map Char Int,
    bos :: Int,
    vocabSize :: Int
  }

mkTokenizer :: [String] -> Tokenizer
mkTokenizer docs =
  Tokenizer
    { uchars = uchars,
      tokenIds = Map.fromList (zip uchars [0 ..]),
      bos = length uchars,
      vocabSize = length uchars + 1
    }
  where
    uchars = Set.toAscList (Set.fromList (concat docs))

-- | Surrounds the document with BOS on both sides (unknown characters are dropped)
encode :: Tokenizer -> String -> [Int]
encode tok doc = [tok.bos] <> Maybe.mapMaybe (`Map.lookup` tok.tokenIds) doc <> [tok.bos]

decode :: Tokenizer -> [Int] -> String
decode tok = Maybe.mapMaybe (\t -> Maybe.listToMaybe (drop t tok.uchars))

-- 5. Training

-- | Mean cross-entropy loss over a document (at most block_size predictions)
sequenceLoss :: Config -> Weights (NodeIdx Double) -> [Int] -> GraphBuilder (NodeIdx Double)
sequenceLoss cfg w tokens =
  case take cfg.blockSize (zip tokens (drop 1 tokens)) of
    [] -> constant 0
    pairs -> do
      losses <- go 0 (emptyCaches cfg) pairs
      totalLoss <- sumNodes losses
      mul totalLoss =<< constant (1 / fromIntegral (length pairs))
  where
    go posId caches = \case
      [] -> pure []
      (tokenId, targetId) : rest -> do
        (logits, caches') <- forwardPassMonadic cfg w tokenId posId caches
        probs <- softmax logits
        loss <- case Maybe.listToMaybe (drop targetId probs) of
          Nothing -> constant 0
          Just targetProb -> log' targetProb >>= (\lp -> mul lp =<< constant (-1.0))
        (loss :) <$> go (posId + 1) caches' rest

-- | Loss of a document, forward pass only
forwardLoss :: Config -> Weights Double -> [Int] -> Double
forwardLoss cfg pureWeights tokens =
  (evalState (embedWeights pureWeights >>= \w -> sequenceLoss cfg w tokens) initialTape).node.val

-- | Runs a full forward and backward pass for a sequence of tokens
trainingStep :: Config -> Weights Double -> [Int] -> (Double, Gradients)
trainingStep cfg pureWeights tokens =
  let graph = do
        w <- embedWeights pureWeights
        sequenceLoss cfg w tokens

      (lossNode, tape) = runState graph initialTape
      grads = backward lossNode tape
   in (lossNode.node.val, grads)

-- | Pairs each parameter with its gradient (parameters are the first nodes of the tape)
withGradients :: Gradients -> Weights a -> Weights (Double, a)
withGradients grads =
  snd . Traversable.mapAccumL (\i x -> (i + 1, (Maybe.fromMaybe 0 (nimLookup (NodeId i) grads), x))) 0

parameterGradients :: Gradients -> Weights a -> Weights Double
parameterGradients grads = fmap fst . withGradients grads

data Hyperparams = Hyperparams
  { learningRate :: Double,
    beta1 :: Double,
    beta2 :: Double,
    epsilon :: Double,
    numSteps :: Int
  }

microHyperparams :: Hyperparams
microHyperparams =
  Hyperparams
    { learningRate = 0.01,
      beta1 = 0.85,
      beta2 = 0.99,
      epsilon = 1e-8,
      numSteps = 1000
    }

-- | A parameter with its Adam first and second moment buffers
data Param = Param
  { value :: Double,
    m :: Double,
    v :: Double
  }
  deriving stock (Generic)
  deriving anyclass (NFData)

data TrainState = TrainState
  { params :: Weights Param,
    step :: Int
  }
  deriving stock (Generic)
  deriving anyclass (NFData)

initTrainState :: Weights Double -> TrainState
initTrainState w =
  TrainState
    { params = fmap (\x -> Param {value = x, m = 0, v = 0}) w,
      step = 0
    }

-- | Updates a parameter and its Adam state, given the learning rate and the step (from 1)
updateParam :: Hyperparams -> Double -> Int -> (Double, Param) -> Param
updateParam hp lr t (grad, p) =
  let m = hp.beta1 * p.m + (1 - hp.beta1) * grad
      v = hp.beta2 * p.v + (1 - hp.beta2) * (grad ** 2)
      mHat = m / (1 - (hp.beta1 ** fromIntegral t))
      vHat = v / (1 - (hp.beta2 ** fromIntegral t))
   in Param {value = p.value - lr * mHat / (sqrt vHat + hp.epsilon), m = m, v = v}

-- | Forward, backward and Adam update on a document, returns the loss before the update
trainStep :: Config -> Hyperparams -> TrainState -> [Int] -> (Double, TrainState)
trainStep cfg hp st tokens =
  let (loss, grads) = trainingStep cfg (fmap (.value) st.params) tokens
      lr = hp.learningRate * (1 - fromIntegral st.step / fromIntegral hp.numSteps)
      params' = fmap (updateParam hp lr (st.step + 1)) (withGradients grads st.params)
   in (loss, TrainState {params = params', step = st.step + 1})

-- | One training step per document
train :: Config -> Hyperparams -> TrainState -> [[Int]] -> ([Double], TrainState)
train cfg hp = go
  where
    go st = \case
      [] -> ([], st)
      doc : rest ->
        let (loss, st') = trainStep cfg hp st doc
            (losses, final) = st' `deepseq` go st' rest
         in (loss : losses, final)

-- 6. Inference & Sampling

-- | random.choices: samples an index proportionally to the weights
sampleCategorical :: [Double] -> StdGen -> (Int, StdGen)
sampleCategorical weights gen =
  let (u, gen') = Random.randomR (0.0, 1.0) gen
      x = u * foldl' (+) 0 weights
      cumulative = drop 1 (List.scanl' (+) 0 weights)
   in (min (length weights - 1) (length (takeWhile (<= x) cumulative)), gen')

-- | Samples a document, starting from BOS, until BOS is predicted or block_size tokens are generated
generate :: Config -> Int -> Double -> Weights Double -> StdGen -> ([Int], StdGen)
generate cfg bosId temperature pureWeights gen0 =
  evalState (embedWeights pureWeights >>= \w -> go w 0 bosId (emptyCaches cfg) gen0) initialTape
  where
    go w posId tokenId caches gen =
      if posId >= cfg.blockSize
        then pure ([], gen)
        else do
          (logits, caches') <- forwardPassMonadic cfg w tokenId posId caches
          invTemperature <- constant (temperature ** (-1))
          probs <- mapM (`mul` invTemperature) logits >>= softmax
          let (nextToken, gen') = sampleCategorical (map (.node.val) probs) gen
          if nextToken == bosId
            then pure ([], gen')
            else do
              (rest, gen'') <- go w (posId + 1) nextToken caches' gen'
              pure (nextToken : rest, gen'')

-- 7. The whole microgpt.py script

shuffle :: [a] -> StdGen -> ([a], StdGen)
shuffle xs gen =
  let (ks, gen') = runState (replicateM (length xs) (state Random.random)) gen
   in (map snd (List.sortOn fst (zip (ks :: [Int]) xs)), gen')

microGPT :: IO ()
microGPT = do
  allDocs <- loadDocs datasetPath
  let (docs, gen1) = shuffle allDocs (Random.mkStdGen 42)
  putStrLn $ "num docs: " <> show (length docs)

  let tok = mkTokenizer docs
  putStrLn $ "vocab size: " <> show tok.vocabSize

  let cfg = microConfig tok.vocabSize
      hp = microHyperparams
      (w, gen2) = randomWeights cfg gen1
      padded n x = let s = show x in replicate (n - length s) ' ' <> s
  putStrLn $ "num params: " <> show (length w)

  final <-
    foldM
      ( \st (i, doc) -> do
          let (loss, st') = trainStep cfg hp st (encode tok doc)
          st' `deepseq` putStr ("step " <> padded 4 (i + 1) <> " / " <> padded 4 hp.numSteps <> " | loss " <> Numeric.showFFloat (Just 4) loss "\r")
          hFlush stdout
          pure st'
      )
      (initTrainState w)
      (zip [0 :: Int ..] (take hp.numSteps (cycle docs)))

  putStrLn "\n--- inference (new, hallucinated names) ---"
  let weights = fmap (.value) final.params
  foldM_
    ( \gen i -> do
        let (sample, gen') = generate cfg tok.bos 0.5 weights gen
        putStrLn $ "sample " <> padded 2 i <> ": " <> decode tok sample
        pure gen'
    )
    gen2
    [1 :: Int .. 20]