## network.nim — MLP and ActorCritic forward pass (inference only, no autograd).

import arraymancer
import std/[math, random]

const
  STATE_DIM*  = 44
  ACTION_DIM* = 6

type
  MLP* = object
    w1*, b1*: Tensor[float32]  # [hidden, input], [hidden]
    w2*, b2*: Tensor[float32]  # [hidden, hidden], [hidden]
    w3*, b3*: Tensor[float32]  # [output, hidden], [output]

  ActorCritic* = object
    actor*:  MLP
    critic*: MLP
    logStd*: Tensor[float32]  # [ACTION_DIM] — one per action dim

proc initMLP*(inputDim, hiddenDim, outputDim: int): MLP =
  # Xavier/He-style init: scale weights by sqrt(2/fan_in)
  result.w1 = randomNormalTensor[float32]([hiddenDim, inputDim])  *. sqrt(2.0'f32 / inputDim.float32)
  result.b1 = zeros[float32](hiddenDim)
  result.w2 = randomNormalTensor[float32]([hiddenDim, hiddenDim]) *. sqrt(2.0'f32 / hiddenDim.float32)
  result.b2 = zeros[float32](hiddenDim)
  result.w3 = randomNormalTensor[float32]([outputDim, hiddenDim]) *. sqrt(1.0'f32 / hiddenDim.float32)
  result.b3 = zeros[float32](outputDim)

proc initActorCritic*(): ActorCritic =
  result.actor  = initMLP(STATE_DIM, 64, ACTION_DIM)
  result.critic = initMLP(STATE_DIM, 64, 1)
  result.logStd = zeros[float32](ACTION_DIM)  # init to 0 → std=1

proc forward*(mlp: MLP, x: Tensor[float32]): Tensor[float32] =
  ## x shape: [inputDim] (1D vector)
  let h1 = tanh(mlp.w1 * x + mlp.b1)
  let h2 = tanh(mlp.w2 * h1 + mlp.b2)
  result  = mlp.w3 * h2 + mlp.b3

proc actorForward*(ac: ActorCritic, state: Tensor[float32]): tuple[actions: Tensor[float32], logProb: float32] =
  ## state: [STATE_DIM]. Returns sampled actions [ACTION_DIM] and sum log-prob.
  let mean = ac.actor.forward(state)
  # Floor logStd at -3 before exp → min std ≈ 0.05
  var clampedLogStd = ac.logStd.map(proc(v: float32): float32 = max(v, -3.0'f32))
  let std = clampedLogStd.map(proc(v: float32): float32 = exp(v))

  var actions = newTensor[float32](ACTION_DIM)
  var logP    = 0.0'f32
  for i in 0..<ACTION_DIM:
    let mu = mean[i]
    let s  = std[i]
    let z  = gauss(0.0'f64, 1.0'f64).float32
    actions[i] = mu + s * z
    # log N(a; mu, s) = -0.5*((a-mu)/s)^2 - log(s) - 0.5*log(2π)
    let diff = (actions[i] - mu) / s
    logP += -0.5'f32 * diff * diff - ln(s) - 0.5'f32 * ln(2.0'f32 * PI.float32)

  result = (actions: actions, logProb: logP)

proc criticForward*(ac: ActorCritic, state: Tensor[float32]): float32 =
  ## state: [STATE_DIM]. Returns scalar value estimate.
  let val = ac.critic.forward(state)
  result = val[0]

proc computeLogProb*(ac: ActorCritic, state, action: Tensor[float32]): float32 =
  ## Log-probability of action under current policy (no sampling).
  let mean = ac.actor.forward(state)
  let logStdClamped = ac.logStd.map(proc(v: float32): float32 = max(v, -3.0'f32))
  let std = logStdClamped.map(proc(v: float32): float32 = exp(v))
  var logP = 0.0'f32
  for i in 0..<ACTION_DIM:
    let mu = mean[i]
    let s  = std[i]
    let diff = (action[i] - mu) / s
    logP += -0.5'f32 * diff * diff - ln(s) - 0.5'f32 * ln(2.0'f32 * PI.float32)
  result = logP