// loss_fusion_test.sy -- Test fused softmax+cross_entropy loss.
// THE most important fusion in deep learning.
// Aislinn McCubbins, April 2026

fn main() {
    println("=== LOSS FUSION TEST ===")

    // Test 1: Basic cross-entropy
    println("TEST 1: Cross-entropy loss")
    let pred = ones(2, 3)
    let pred = softmax(pred)
    println("  predictions (uniform):")
    println(pred)

    // one-hot target: class 0 for both samples
    let target = zeros(2, 3)
    // manually set target[0,0]=1, target[1,0]=1
    // For now use ones and scale — target = [[0.33,0.33,0.33]]
    let target = softmax(ones(2, 3))
    let loss = cross_entropy(pred, target)
    println("  loss (uniform pred, uniform target = ~1.099):")
    println(loss)
    println("")

    // Test 2: Fused softmax+cross_entropy with gradient
    println("TEST 2: Fused softmax_cross_entropy + backward")
    begin_tape()
    let w = rand_tensor(3, 3)
    let w = register_param("w_loss", w)
    let x = ones(1, 3)
    let logits = t_matmul(x, w)
    // One-hot target: class 1
    let target2 = zeros(1, 3)
    let loss2 = softmax_cross_entropy(logits, target2)
    end_tape()
    println("  logits:")
    println(logits)
    println("  loss:")
    println(loss2)

    backward(loss2)
    let g = get_grad("w_loss")
    println("  grad w (should be (softmax - target) / batch):")
    println(g)
    println("")

    // Test 3: Verify fused vs unfused give same result
    println("TEST 3: Fused vs unfused comparison")
    let logits3 = rand_tensor(2, 4)
    let target3 = zeros(2, 4)

    let unfused = cross_entropy(softmax(logits3), target3)
    let fused = softmax_cross_entropy(logits3, target3)
    println("  unfused loss:")
    println(unfused)
    println("  fused loss:")
    println(fused)
    println("  (should be identical)")
    println("")

    // Test 4: Training with real loss
    println("TEST 4: 10-step training with cross-entropy loss")
    let w4 = rand_tensor(2, 3)
    let x4 = ones(1, 2)
    // Target: class 2 (one-hot [0, 0, 1])
    let tgt4 = zeros(1, 3)

    for step in 0..10 {
        begin_tape()
        let w4 = register_param("w4", w4)
        let logits4 = t_matmul(x4, w4)
        let loss4 = softmax_cross_entropy(logits4, tgt4)
        end_tape()
        backward(loss4)
        let grad4 = get_grad("w4")
        let w4 = w4 - grad4
    }
    println("  Final weights:")
    println(w4)
    println("  (training should push weights toward target class)")

    println("")
    println("=== LOSS FUSION TESTS COMPLETE ===")
}
