import arraymancer
import std/[strformat, times]

# XOR dataset — [4,2] inputs, [4,1] targets
let xData = [[0.0'f32, 0.0'f32],
             [0.0'f32, 1.0'f32],
             [1.0'f32, 0.0'f32],
             [1.0'f32, 1.0'f32]].toTensor()

let yData = [[0.0'f32], [1.0'f32], [1.0'f32], [0.0'f32]].toTensor()

var ctx = newContext Tensor[float32]

let l1 = ctx.init(Linear[float32], 2, 4)
let l2 = ctx.init(Linear[float32], 4, 1)

type XorNet = object
  l1: Linear[float32]
  l2: Linear[float32]

var net = XorNet(l1: l1, l2: l2)
var adam = optimizer(net, Adam, learning_rate = 0.01'f32)

let t0 = cpuTime()

for epoch in 1..1000:
  let xVar = ctx.variable(xData)
  let yVar = ctx.variable(yData)

  # Forward
  let h    = l1.forward(xVar).sigmoid()
  let pred = l2.forward(h).sigmoid()

  # MSE loss: mean((pred - y)^2)
  let diff = pred - yVar
  let loss = mean(diff *. diff)

  if epoch mod 200 == 0:
    echo &"Epoch {epoch:4d}  loss = {loss.value[0]:.6f}"

  loss.backprop()
  adam.update()

let trainTime = cpuTime() - t0

# Inference (no grad)
echo ""
echo "Final predictions:"
let xVar  = ctx.variable(xData)
let hInf  = l1.forward(xVar).sigmoid()
let preds = l2.forward(hInf).sigmoid()

var allCorrect = true
for i in 0..<4:
  let p       = preds.value[i, 0]
  let rounded = if p >= 0.5'f32: 1 else: 0
  let expected = [0,1,1,0][i]
  if rounded != expected: allCorrect = false
  echo &"  [{xData[i,0].int},{xData[i,1].int}] -> {p:.4f}  (rounded: {rounded}, expected: {expected})"

echo ""
echo &"All correct: {allCorrect}"
echo &"Training time: {trainTime:.2f}s"