// grad_correctness.sy -- Gradient correctness test suite.
// Tests: multi-branch, reused weights, deep chains, numerical stability.
// Aislinn McCubbins, April 2026

fn main() {
    println("=== GRADIENT CORRECTNESS SUITE ===")
    println("")

    // TEST 1: Multi-branch graph
    // Two branches share the same input, gradients must accumulate
    println("TEST 1: Multi-branch (gradient accumulation)")
    begin_tape()
    let w = rand_tensor(2, 2)
    let w = register_param("w", w)
    let x = ones(1, 2)
    let branch_a = t_matmul(x, w)
    let branch_b = t_matmul(x, w)
    let merged = t_add(branch_a, branch_b)
    end_tape()
    backward(merged)
    let gw = get_grad("w")
    println("  grad_w (should be 2x single branch):")
    println(gw)
    println("")

    // TEST 2: Reused weights across layers
    // w used in BOTH layers: grad must accumulate from both paths
    println("TEST 2: Weight reuse across layers")
    begin_tape()
    let w2 = rand_tensor(2, 2)
    let w2 = register_param("w2", w2)
    let x2 = ones(1, 2)
    let h1 = t_matmul(x2, w2)
    let h1 = t_relu(h1)
    let h2 = t_matmul(h1, w2)
    end_tape()
    backward(h2)
    let gw2 = get_grad("w2")
    println("  grad_w2 (accumulated from 2 matmuls):")
    println(gw2)
    println("")

    // TEST 3: Deep non-linear chain
    // relu -> matmul -> relu -> add -> softmax
    println("TEST 3: Deep chain (relu -> matmul -> relu -> add -> softmax)")
    begin_tape()
    let w3 = rand_tensor(3, 3)
    let w3 = register_param("w3", w3)
    let b3 = rand_tensor(1, 3)
    let b3 = register_param("b3", b3)
    let x3 = ones(1, 3)
    let h = t_matmul(x3, w3)
    let h = t_relu(h)
    let h = t_matmul(h, w3)
    let h = t_relu(h)
    let h = t_add(h, b3)
    let out = t_softmax(h)
    end_tape()
    backward(out)
    let gw3 = get_grad("w3")
    let gb3 = get_grad("b3")
    println("  grad_w3:")
    println(gw3)
    println("  grad_b3:")
    println(gb3)
    println("")

    // TEST 4: Numerical stability - softmax with large values
    println("TEST 4: Softmax stability (large inputs)")
    begin_tape()
    let big_w = rand_tensor(3, 3)
    let big_w = register_param("big_w", big_w)
    let big_x = ones(1, 3)
    let h4 = t_matmul(big_x, big_w)
    // Stack additions to amplify values
    let h4 = t_add(h4, h4)
    let h4 = t_add(h4, h4)
    let h4 = t_add(h4, h4)
    let h4 = t_add(h4, h4)
    let h4 = t_add(h4, h4)
    let out4 = t_softmax(h4)
    end_tape()
    backward(out4)
    let g_big = get_grad("big_w")
    println("  softmax output (should sum to 1, no NaN/Inf):")
    println(out4)
    println("  grad (should be finite, no NaN):")
    println(g_big)
    println("")

    // TEST 5: Element-wise multiply gradient
    println("TEST 5: Element-wise multiply")
    begin_tape()
    let a5 = rand_tensor(2, 2)
    let a5 = register_param("a5", a5)
    let b5 = rand_tensor(2, 2)
    let b5 = register_param("b5", b5)
    let c5 = t_mul(a5, b5)
    end_tape()
    backward(c5)
    let ga5 = get_grad("a5")
    let gb5 = get_grad("b5")
    println("  a5:")
    println(a5)
    println("  b5:")
    println(b5)
    println("  grad_a5 (should equal b5):")
    println(ga5)
    println("  grad_b5 (should equal a5):")
    println(gb5)
    println("")

    // TEST 6: Subtraction gradient
    println("TEST 6: Subtraction (a - b)")
    begin_tape()
    let s_a = ones(2, 2)
    let s_a = register_param("s_a", s_a)
    let s_b = ones(2, 2)
    let s_b = register_param("s_b", s_b)
    let s_c = t_sub(s_a, s_b)
    end_tape()
    backward(s_c)
    let gsa = get_grad("s_a")
    let gsb = get_grad("s_b")
    println("  grad_a (should be all 1.0):")
    println(gsa)
    println("  grad_b (should be all -1.0):")
    println(gsb)
    println("")

    println("=== ALL GRADIENT TESTS COMPLETE ===")
}
