ml.nim

import numericalnim, arraymancer, std/[math]
import std/strformat

# -----------------------------
# Helper: unpack θ into explicit matrices F, A, S
# -----------------------------

proc buildMatrices(theta: Tensor[float]): (Tensor[float], Tensor[float], Tensor[float]) =
  # F: 3x5 filter (rows=3, cols=5)
  var F = zeros[float]([3, 5])
  F[0, 1] = 1.0
  F[1, 2] = 1.0
  F[2, 3] = 1.0

  # A: 5x5 directed paths
  var A = zeros[float]([5, 5])
  let λ1 = theta[0]
  let λ2 = theta[1]
  let λ3 = theta[2]
  A[1, 0] = λ1   
  A[2, 0] = λ2   
  A[3, 0] = λ3   

  # S: 5x5 variances
  var S = zeros[float]([5, 5])
  S[0, 0] = theta[6]      # σ_s
  S[1, 1] = theta[3]      # e1
  S[2, 2] = theta[4]      # e2
  S[3, 3] = theta[5]      # e3
  S[4, 4] = 1e-4          # residual

  (F, A, S)

# -----------------------------
# Σ(θ) = F (I - A)^-1 S (I - A)^-T F^T
# -----------------------------
proc sigmaOf(theta: Tensor[float]): Tensor[float] =
  let (F, A, S) = buildMatrices(theta)
  # FIX: eye requires (rows, cols), so use (5, 5)
  let I  = eye[float](5, 5) 
  
  # Solve (I - A) * X = I  => X is (I - A)^-1
  let IAinv = solve(I - A, I)
  F * IAinv * S * IAinv.transpose * F.transpose

# -----------------------------
# ML fit function
# -----------------------------
proc logDetSPD(M: Tensor[float]): float =
  let (_, Svals, _) = M.svd()
  ln(Svals).sum()

proc mlFit(theta: Tensor[float], Sobs: Tensor[float]): float =
  let p = Sobs.shape[0] # automatically detects 3
  var Sigma = sigmaOf(theta)

  # tiny ridge if needed to ensure SPD
  let eps = 1e-8
  # FIX: eye requires (rows, cols), so use (p, p)
  Sigma = Sigma + eps * eye[float](p, p)

  # Σ^{-1} via solve(Σ, I)
  # FIX: eye requires (rows, cols), so use (p, p)
  let SigmaInv = solve(Sigma, eye[float](p, p))

  let logdetSigma = logDetSPD(Sigma)
  let logdetSobs  = logDetSPD(Sobs)
  let trTerm = (Sobs * SigmaInv).diagonal().sum()
   
  result = logdetSigma + trTerm - logdetSobs - float(p)

# -----------------------------
# Optimization Setup
# -----------------------------
let
  true_lambda1     = 0.6
  true_lambda2     = 0.6
  true_lambda3     = 0.6
  true_e           = 0.5
  true_factor_var  = 1.0

let thetaTrue = [true_lambda1, true_lambda2, true_lambda3,
                 true_e, true_e, true_e, true_factor_var].toTensor()

let Ctarget = sigmaOf(thetaTrue)

let theta0 = [0.3, 0.3, 0.3, 0.4, 0.4, 0.4, 0.8].toTensor()

# FIX: Add () to call the options constructor
let opts = lbfgsOptions[float]()

let thetaHat = lbfgs(
  proc (t: Tensor[float]): float {.closure.} =
    mlFit(t, Ctarget),
  theta0,
  options = opts
)

echo "\n--- ML RAM results ---"
echo "θ^ (λ1, λ2, λ3, e1, e2, e3, σ_s): ", thetaHat
echo "ML fit F_ML(θ^): ", mlFit(thetaHat, Ctarget)