// autodiff_train.sy — Full training loop with gradient descent.
// The first neural network TRAINED in Scyla.
// Aislinn McCubbins, April 2026

fn main() {
    println("=== Scyla Training Loop ===")

    // Target: learn weights that map [1,1] -> [1,0] via matmul + relu
    let x = ones(1, 2)

    // Random initial weights
    let w = rand_tensor(2, 2)
    println("Initial weights:")
    println(w)

    // Training loop: 5 steps of gradient descent
    let step = 0
    for i in 0..5 {
        // Record forward pass
        begin_tape()
        let w = register_param("w", w)

        // Forward: h = relu(x @ w)
        let h = t_matmul(x, w)
        let h = t_relu(h)

        // MSE-like loss: sum of squares (want output to be [1, 0])
        // For simplicity: loss = sum(h)  (minimize total activation)
        end_tape()
        backward(h)

        // Get gradients
        let grad_w = get_grad("w")
        println("Step " + i + " output: " + tensor_shape(h))
        println(h)
        println("  grad_w:")
        println(grad_w)

        // SGD update: w = w - grad_w
        let scaled = zeros(2, 2)
        let scaled = grad_w
        let w = w - scaled
    }

    println("Final weights:")
    println(w)
    println("=== Training complete! ===")
}
