// Custom OOP model: explicit forward, loss, backward, and parameter updates.
// Four synthetic XOR cases; no automatic parameter discovery or inheritance.
import aner.nn
import aner.dev
import aner.viz

class XorNetwork {
    var hidden_weights: Tensor
    var hidden_bias: Tensor
    var output_weights: Tensor
    var output_bias: Tensor

    fn forward(input: Tensor) {
        let hidden = nn.tanh(nn.add_bias(nn.matmul(input, hidden_weights), hidden_bias))
        return nn.sigmoid(nn.add_bias(nn.matmul(hidden, output_weights), output_bias))
    }
    fn update(rate: Float64) {
        hidden_weights = nn.sgd(hidden_weights, rate)
        hidden_bias = nn.sgd(hidden_bias, rate)
        output_weights = nn.sgd(output_weights, rate)
        output_bias = nn.sgd(output_bias, rate)
    }
}
fn objective(prediction: Tensor, target: Tensor) { return nn.mse(prediction, target) }
let input: Tensor = [[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]]
let target: Tensor = [[0.0], [1.0], [1.0], [0.0]]
let network = XorNetwork(
    hidden_weights: nn.parameter(nn.random(2, 8, seed: 42, scale: 1.0)),
    hidden_bias: nn.parameter(nn.zeros(1, 8)),
    output_weights: nn.parameter(nn.random(8, 1, seed: 43, scale: 1.0)),
    output_bias: nn.parameter(nn.zeros(1, 1)))
print("Initial MSE")
print(nn.item(objective(network.forward(input), target)))
var epoch = 0
while epoch < 3000 {
    let loss = objective(network.forward(input), target)
    nn.backward(loss)
    // Capture after backward, before replacing the parameter leaves.
    if epoch % 500 == 0 { viz.capture(epoch, loss) }
    network.update(rate: 0.5)
    epoch = epoch + 1
}
let prediction = network.forward(input)
dev.capture("Custom forward predictions", prediction)
print("Final MSE")
print(nn.item(objective(prediction, target)))
print("XOR predictions")
print(nn.value(prediction, 0, 0))
print(nn.value(prediction, 1, 0))
print(nn.value(prediction, 2, 0))
print(nn.value(prediction, 3, 0))