// tape_stress.sy -- Autodiff pressure test.
// 1000+ ops, deep reuse, long training loops.
// If this holds, the autodiff is production-grade.
// Aislinn McCubbins, April 2026

fn main() {
    println("=== TAPE STRESS TEST ===")

    // TEST A: 1000+ ops in one forward pass
    // Chain 200 matmul+relu pairs = 400 tape ops minimum
    println("TEST A: Deep graph (200 layers, 400+ tape ops)")
    begin_tape()
    let w_deep = rand_tensor(4, 4)
    let w_deep = register_param("w_deep", w_deep)
    let x = rand_tensor(1, 4)
    let h = t_matmul(x, w_deep)
    let h = t_relu(h)
    for i in 0..199 {
        let h = t_matmul(h, w_deep)
        let h = t_relu(h)
    }
    end_tape()
    backward(h)
    let g = get_grad("w_deep")
    println("  grad shape: " + tensor_shape(g))
    println("  grad (should be finite):")
    println(g)
    println("")

    // TEST B: Weight reuse stress — same weight used 50 times
    println("TEST B: Heavy weight reuse (50x same matrix)")
    begin_tape()
    let w_reuse = rand_tensor(3, 3)
    let w_reuse = register_param("w_reuse", w_reuse)
    let h2 = ones(1, 3)
    for i in 0..50 {
        let h2 = t_matmul(h2, w_reuse)
    }
    end_tape()
    backward(h2)
    let g2 = get_grad("w_reuse")
    println("  grad (accumulated from 50 matmuls):")
    println(g2)
    println("")

    // TEST C: Long training loop — 100 steps of SGD
    println("TEST C: 100-step training loop")
    let w_train = rand_tensor(2, 2)
    let x_train = ones(1, 2)
    let lr = 0.01

    for step in 0..100 {
        begin_tape()
        let w_train = register_param("w_t", w_train)
        let h3 = t_matmul(x_train, w_train)
        let h3 = t_relu(h3)
        end_tape()
        backward(h3)
        let grad = get_grad("w_t")

        // SGD: w = w - lr * grad (manual scale)
        let update = grad
        let w_train = w_train - update
    }
    println("  Final weights after 100 steps:")
    println(w_train)
    println("  (should be finite, not NaN/Inf)")
    println("")

    // TEST D: Multi-branch fan-out (gradient accumulation stress)
    println("TEST D: 20-branch fan-out")
    begin_tape()
    let w_fan = rand_tensor(3, 3)
    let w_fan = register_param("w_fan", w_fan)
    let x_fan = ones(1, 3)
    let sum_out = t_matmul(x_fan, w_fan)
    for i in 0..19 {
        let branch = t_matmul(x_fan, w_fan)
        let sum_out = t_add(sum_out, branch)
    }
    end_tape()
    backward(sum_out)
    let g_fan = get_grad("w_fan")
    println("  grad (should be 20x single matmul grad):")
    println(g_fan)
    println("")

    println("=== STRESS TEST COMPLETE ===")
}
