// ops_phase2_test.sy -- Test Phase 2 ops: comparison, sampling, FFT

fn main() {
    println("=== PHASE 2 OPS TEST ===")

    let a = rand_tensor(1, 5)
    println("a:")
    println(a)

    // Comparison
    let gt = t_gt(a, 0.5)
    println("a > 0.5:")
    println(gt)

    let lt = t_lt(a, 0.3)
    println("a < 0.3:")
    println(lt)

    // Masked fill
    let filled = masked_fill(a, gt, -999.0)
    println("masked_fill(a, a>0.5, -999):")
    println(filled)

    // Top-k
    let logits = rand_tensor(1, 10)
    println("logits:")
    println(logits)
    let masked = topk_mask(logits, 3)
    println("top-3 mask:")
    println(masked)

    // Cumsum
    let probs = ones(1, 5)
    let cs = t_cumsum(probs)
    println("cumsum([1,1,1,1,1]):")
    println(cs)

    // Argmax
    println("argmax:")
    println(t_argmax(logits))

    // Multinomial
    let probs2 = softmax(logits)
    let sample = t_multinomial(probs2)
    println("multinomial sample:")
    println(sample)

    // Index select
    let table = rand_tensor(5, 3)
    let idx = tensor_from_ints(0, 3, 1)
    let selected = index_select(table, idx)
    println("index_select rows [0,3,1] from 5x3:")
    println(selected)

    // FFT
    let signal = rand_tensor(1, 8)
    println("signal:")
    println(signal)
    let spectrum = rfft(signal)
    println("rfft magnitudes:")
    println(spectrum)

    // Variance
    let data = rand_tensor(2, 4)
    let v = t_var(data)
    println("variance per row:")
    println(v)

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