// ops_test.sy -- Test all new tensor ops.
// Every op tested in isolation before using in the real model.
// Aislinn McCubbins, April 2026

fn main() {
    println("=== NEW OPS TEST SUITE ===")

    // Element-wise math
    let x = ones(2, 3)
    println("exp(ones):")
    println(t_exp(x))

    println("tanh(ones):")
    println(t_tanh(x))

    println("sigmoid(ones):")
    println(t_sigmoid(x))

    println("gelu(ones):")
    println(gelu(x))

    println("softplus(ones):")
    println(t_softplus(x))

    println("clamp(ones*5, -2, 2):")
    let big = ones(1, 4)
    let big = big + big
    let big = big + big
    println(t_clamp(big, -2.0, 2.0))

    // Shape ops
    let a = rand_tensor(2, 6)
    println("")
    println("Original [2,6]:")
    println(a)

    println("reshape to [3,4]:")
    println(t_reshape(a, 3, 4))

    println("reshape to [6,2]:")
    println(t_reshape(a, 6, 2))

    println("slice last dim [0:3]:")
    println(t_slice(a, -1, 0, 3))

    println("slice last dim [3:6]:")
    println(t_slice(a, -1, 3, 6))

    // Cat
    let left = ones(2, 2)
    let right = zeros(2, 3)
    println("")
    println("cat([2,2] ones, [2,3] zeros):")
    println(t_cat(left, right))

    // Mean
    let m = ones(2, 4)
    let m = m + rand_tensor(2, 4)
    println("")
    println("tensor:")
    println(m)
    println("mean over last dim:")
    println(t_mean(m))

    // Mean over seq dim (3D)
    let m3 = ones(1, 4, 3)
    println("mean_seq of [1,4,3] ones:")
    println(t_mean_seq(m3))

    // Unsqueeze + expand
    let v = ones(1, 3)
    println("")
    println("unsqueeze [1,3] at dim 1 -> [1,1,3]:")
    println(t_unsqueeze(v, 1))

    let expanded = t_expand(t_unsqueeze(v, 1), 1, 4)
    println("expand dim 1 to 4 -> [1,4,3]:")
    println(expanded)

    // Multi-head attention!
    println("")
    println("=== MULTI-HEAD ATTENTION ===")
    let q = rand_tensor(1, 4, 8)
    let k = rand_tensor(1, 4, 8)
    let v2 = rand_tensor(1, 4, 8)
    let attn_out = multi_head_attention(q, k, v2, 2, true)
    println("MHA output [1,4,8] (2 heads, causal):")
    println(attn_out)

    // Gather
    println("")
    let logits = rand_tensor(1, 3)
    let indices = tensor_from_ints(2, 0, 1)
    let indices = t_reshape(indices, 1, 3)
    println("logits:")
    println(logits)
    println("gather indices [2,0,1]:")
    println(t_gather(logits, indices))

    println("")
    println("=== ALL OPS PASS ===")
}
