{-# 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 ((>=>))
import Control.Monad.State.Strict (State, evalState, evalStateT, get, lift, put, runState, runStateT, 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 Data.Vector (Vector)
import qualified Data.Vector as Vector
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 (Vector.singleton "abcdefghijklmnopqrstuvwxyz")
          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 (Vector.singleton "abcdefghijklmnopqrstuvwxyz")
          cfg = microConfig tok.vocabSize
          w = deterministicWeights cfg
          (emptyDocLoss, _) = trainingStep cfg w (encode tok "")
          (singleTokLoss, _) = trainingStep cfg w (Vector.singleton tok.bos)
      emptyDocLoss `shouldSatisfy` (> 0.0)
      singleTokLoss `shouldBe` 0.0
  describe "Equivalence with microgpt.py (values from microgpt/reference.py)" $ do
    let tok = mkTokenizer (Vector.singleton "abcdefghijklmnopqrstuvwxyz")
        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 = Vector.fromList (Foldable.toList (parameterGradients grads w))
      Vector.forM_ referenceGradients $ \(i, expected) ->
        (paramGrads Vector.!? i)
          `shouldSatisfy` maybe False (approxEq expected)
    it "follows the same training trajectory (Adam with linear decay)" $ do
      let docs = Vector.map (encode tok) (Vector.fromList ["emma", "olivia", "ava", "isabella", "sophia"])
          (losses, final) = train cfg microHyperparams (initTrainState w) docs
          finalParams = Vector.fromList (Foldable.toList final.params)
      Vector.length losses `shouldBe` 5
      Vector.zipWithM_
        shouldBeApprox
        losses
        (Vector.fromList [3.506449207715624, 3.2498071277112257, 2.7546668273691433, 3.496728826446422, 3.056047453583676])
      fmap (.value) (finalParams Vector.!? 0)
        `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 (Vector.singleton "abcdefghijklmnopqrstuvwxyz")
          cfg = microConfig tok.vocabSize
          (sample, _) = generate cfg tok.bos 0.5 (deterministicWeights cfg) (Random.mkStdGen 42)
      Vector.length sample `shouldSatisfy` (<= cfg.blockSize)
      sample `shouldSatisfy` Vector.all (\t -> t >= 0 && t < tok.bos)

-- | (parameter index, gradient) for "emma", from microgpt/reference.py
referenceGradients :: Vector (Int, Double)
referenceGradients =
  Vector.fromList
    [ (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 (Vector String)
loadDocs path = do
  contents <- readFile path
  let strip = List.dropWhileEnd Char.isSpace . dropWhile Char.isSpace
  pure $ Vector.fromList (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 (Vector.map (encode tok) (Vector.take n docs)))

  Criterion.defaultMain
    [ Criterion.bgroup
        "microGPT"
        [ Criterion.bench ("forward (" <> show n <> " docs)") $
            Criterion.nf (Vector.map (forwardLoss cfg w)) batch,
          Criterion.bench ("forward+backward (" <> show n <> " docs)") $
            Criterion.nf (Vector.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 :: Vector (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 = Vector.fromListN 2 [(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 = Vector.fromListN 2 [(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 = Vector.singleton (x.idx, LocalPartialDerivative $ p * (x.node.val ** (p - 1)))
    }

nexp :: (Floating a) => NodeIdx a -> Node a
nexp x =
  Node
    { deps = Vector.singleton (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 = Vector.singleton (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 = Vector.singleton (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 = Vector.empty}

-- | 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 =
  Vector.foldl' step initialGrads (Vector.enumFromStepN finalNode.idx (-1) (startId + 1))
  where
    NodeId startId = finalNode.idx

    -- \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
              Vector.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

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

-- | Sums a vector of nodes into a single node (like Python's sum, the empty sum is 0)
sumNodes :: Vector (NodeIdx Double) -> GraphBuilder (NodeIdx Double)
sumNodes xs = case Vector.uncons xs of
  Nothing -> constant 0
  Just (x, rest) -> Vector.foldM add x rest

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

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

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

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

-- | Element-wise vector addition
addVectors :: Vector (NodeIdx Double) -> Vector (NodeIdx Double) -> GraphBuilder (Vector (NodeIdx Double))
addVectors = Vector.zipWithM add

transpose :: Matrix a -> Matrix a
transpose rows =
  let width = Vector.foldl' (\acc r -> max acc (Vector.length r)) 0 rows
   in Vector.generate width $ \col -> Vector.mapMaybe (Vector.!? col) rows

-- 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 :: Matrix a,
    attnWk :: Matrix a,
    attnWv :: Matrix a,
    attnWo :: Matrix a,
    mlpFc1 :: Matrix a,
    mlpFc2 :: Matrix 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 :: Matrix a,
    wpe :: Matrix a,
    lmHead :: Matrix a,
    layers :: Vector (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 = Vector.replicate cfg.nLayer layer
    }
  where
    matrix nout nin = Vector.replicate nout (Vector.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 :: Matrix (NodeIdx Double),
    values :: Matrix (NodeIdx Double)
  }

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

-- | Safe row lookup into an embedding matrix
lookupRow :: Matrix a -> Int -> Vector a
lookupRow rows idx = Maybe.fromMaybe Vector.empty (rows Vector.!? idx)

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

  let keys = Vector.snoc cache.keys k
      values = Vector.snoc cache.values v
      hd = headDim cfg
      slice h = Vector.take hd . Vector.drop (h * hd)

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

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

  out <- linear (Vector.concatMap id 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 (NodeIdx Double) -> GraphBuilder (Vector (NodeIdx Double), 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 >>= Vector.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 -> Vector KVCache -> GraphBuilder (Vector (NodeIdx Double), Vector KVCache)
forwardPassMonadic cfg w tokenId posId caches = do
  x0 <- addVectors (lookupRow w.wte tokenId) (lookupRow w.wpe posId)
  x1 <- rmsnorm x0

  (caches', x) <-
    runStateT
      ( Vector.zipWithM
          ( \layer cache -> do
              h <- get
              (h', cache') <- lift $ transformerLayer cfg layer cache h
              put h'
              pure cache'
          )
          w.layers
          caches
      )
      x1

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

-- 4. Tokenizer

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

mkTokenizer :: Vector String -> Tokenizer
mkTokenizer docs =
  Tokenizer
    { uchars = uchars,
      tokenIds = Vector.ifoldl' (\m i c -> Map.insert c i m) Map.empty uchars,
      bos = Vector.length uchars,
      vocabSize = Vector.length uchars + 1
    }
  where
    uchars = Vector.fromList (Set.toAscList (Foldable.foldMap Set.fromList docs))

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

decode :: Tokenizer -> Vector Int -> String
decode tok = Vector.toList . Vector.mapMaybe (tok.uchars Vector.!?)

-- 5. Training

-- | Mean cross-entropy loss over a document (at most block_size predictions)
sequenceLoss :: Config -> Weights (NodeIdx Double) -> Vector Int -> GraphBuilder (NodeIdx Double)
sequenceLoss cfg w tokens =
  let pairs = Vector.take cfg.blockSize (Vector.zip tokens (Vector.drop 1 tokens))
   in if Vector.null pairs
        then constant 0
        else do
          losses <- evalStateT (Vector.imapM step pairs) (emptyCaches cfg)
          totalLoss <- sumNodes losses
          mul totalLoss =<< constant (1 / fromIntegral (Vector.length pairs))
  where
    step posId (tokenId, targetId) = do
      caches <- get
      (logits, caches') <- lift $ forwardPassMonadic cfg w tokenId posId caches
      put caches'
      probs <- lift $ softmax logits
      lift $ case probs Vector.!? targetId of
        Nothing -> constant 0
        Just targetProb -> log' targetProb >>= (\lp -> mul lp =<< constant (-1.0))

-- | Loss of a document, forward pass only
forwardLoss :: Config -> Weights Double -> Vector 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 -> Vector 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 -> Vector 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 -> Vector (Vector Int) -> (Vector Double, TrainState)
train cfg hp st0 docs =
  let (final, losses) =
        Traversable.mapAccumL
          ( \st doc ->
              let (loss, st') = trainStep cfg hp st doc
               in st' `deepseq` (st', loss)
          )
          st0
          docs
   in (losses, final)

-- 6. Inference & Sampling

-- | random.choices: samples an index proportionally to the weights
sampleCategorical :: Vector Double -> StdGen -> (Int, StdGen)
sampleCategorical weights gen =
  let (u, gen') = Random.randomR (0.0, 1.0) gen
      x = u * Vector.foldl' (+) 0 weights
      cumulative = Vector.postscanl' (+) 0 weights
   in (min (Vector.length weights - 1) (Vector.length (Vector.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 -> (Vector Int, StdGen)
generate cfg bosId temperature pureWeights gen0 =
  evalState (embedWeights pureWeights >>= \w -> go w 0 bosId (emptyCaches cfg) Vector.empty gen0) initialTape
  where
    go w posId tokenId caches acc gen =
      if posId >= cfg.blockSize
        then pure (acc, gen)
        else do
          (logits, caches') <- forwardPassMonadic cfg w tokenId posId caches
          invTemperature <- constant (temperature ** (-1))
          probs <- Vector.mapM (`mul` invTemperature) logits >>= softmax
          let (nextToken, gen') = sampleCategorical (Vector.map (.node.val) probs) gen
          if nextToken == bosId
            then pure (acc, gen')
            else go w (posId + 1) nextToken caches' (Vector.snoc acc nextToken) gen'

-- 7. The whole microgpt.py script

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

microGPT :: IO ()
microGPT = do
  allDocs <- loadDocs datasetPath
  let (docs, gen1) = shuffle allDocs (Random.mkStdGen 42)
  putStrLn $ "num docs: " <> show (Vector.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)

  let stepDocs
        | Vector.null docs = Vector.empty
        | otherwise =
            Vector.generate hp.numSteps $ \i ->
              Maybe.fromMaybe "" (docs Vector.!? (i `mod` Vector.length docs))

  final <-
    Vector.ifoldM
      ( \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)
      stepDocs

  putStrLn "\n--- inference (new, hallucinated names) ---"
  let weights = fmap (.value) final.params
  Vector.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
    (Vector.enumFromN (1 :: Int) 20)