//! Benchmark all new model-agnostic kernels added for architecture alignment.
//!
//! Profiles: standalone activations (SiLU, GELU-tanh), LayerNorm vs RMSNorm,
//! V-norm, scale_vector, partial RoPE, and sliding window attention.
//!
//! Run: cargo run --release --features metal -p infer-compute --example profile_new_kernels

#[cfg(not(feature = "metal"))]
fn main() {
    eprintln!("This example requires --features metal");
}

#[cfg(feature = "metal")]
fn main() {
    use std::time::Instant;
    let metal = infer_compute::metal::MetalBackend::new().expect("Metal required");
    let bufs = metal.bufs();
    let queue = metal.queue();

    println!("=== New Kernel Benchmarks (model-agnostic alignment) ===\n");

    let hidden = 2560;
    let inter = 10240;
    let head_dim = 256;
    let iters = 100;

    // ── Standalone Activations ──
    println!("--- Standalone Activations (inter={inter}) ---\n");
    {
        let input: Vec<f32> = (0..inter)
            .map(|i| (i as f32 - inter as f32 / 2.0) * 0.001)
            .collect();
        let input_buf = bufs.transient_from_f32(&input);
        let out_buf = bufs.output((inter * 4) as u64);
        let n_val = inter as u32;

        // Warm up
        for _ in 0..5 {
            let cmd = queue.new_command_buffer();
            let enc = cmd.new_compute_command_encoder();
            enc.set_compute_pipeline_state(&metal.silu_pipeline);
            enc.set_buffer(0, Some(&input_buf), 0);
            enc.set_buffer(1, Some(&out_buf), 0);
            enc.set_bytes(2, 4, &n_val as *const u32 as *const std::ffi::c_void);
            enc.dispatch_threads(
                metal::MTLSize::new(inter as u64, 1, 1),
                metal::MTLSize::new(256, 1, 1),
            );
            enc.end_encoding();
            cmd.commit();
            cmd.wait_until_completed();
        }

        // SiLU standalone
        let t = Instant::now();
        for _ in 0..iters {
            let cmd = queue.new_command_buffer();
            let enc = cmd.new_compute_command_encoder();
            enc.set_compute_pipeline_state(&metal.silu_pipeline);
            enc.set_buffer(0, Some(&input_buf), 0);
            enc.set_buffer(1, Some(&out_buf), 0);
            enc.set_bytes(2, 4, &n_val as *const u32 as *const std::ffi::c_void);
            enc.dispatch_threads(
                metal::MTLSize::new(inter as u64, 1, 1),
                metal::MTLSize::new(256, 1, 1),
            );
            enc.end_encoding();
            cmd.commit();
            cmd.wait_until_completed();
        }
        let silu_us = t.elapsed().as_micros() as f64 / iters as f64;

        // GELU-tanh standalone
        let t = Instant::now();
        for _ in 0..iters {
            let cmd = queue.new_command_buffer();
            let enc = cmd.new_compute_command_encoder();
            enc.set_compute_pipeline_state(&metal.gelu_tanh_pipeline);
            enc.set_buffer(0, Some(&input_buf), 0);
            enc.set_buffer(1, Some(&out_buf), 0);
            enc.set_bytes(2, 4, &n_val as *const u32 as *const std::ffi::c_void);
            enc.dispatch_threads(
                metal::MTLSize::new(inter as u64, 1, 1),
                metal::MTLSize::new(256, 1, 1),
            );
            enc.end_encoding();
            cmd.commit();
            cmd.wait_until_completed();
        }
        let gelu_us = t.elapsed().as_micros() as f64 / iters as f64;

        // GEGLU SiLU (gated, for comparison)
        let gate_buf = bufs.transient_from_f32(&input);
        let up_buf = bufs.transient_from_f32(&input);
        let t = Instant::now();
        for _ in 0..iters {
            let cmd = queue.new_command_buffer();
            let enc = cmd.new_compute_command_encoder();
            enc.set_compute_pipeline_state(&metal.geglu_pipeline);
            enc.set_buffer(0, Some(&gate_buf), 0);
            enc.set_buffer(1, Some(&up_buf), 0);
            enc.set_buffer(2, Some(&out_buf), 0);
            enc.set_bytes(3, 4, &n_val as *const u32 as *const std::ffi::c_void);
            enc.dispatch_threads(
                metal::MTLSize::new(inter as u64, 1, 1),
                metal::MTLSize::new(256, 1, 1),
            );
            enc.end_encoding();
            cmd.commit();
            cmd.wait_until_completed();
        }
        let geglu_us = t.elapsed().as_micros() as f64 / iters as f64;

        println!("  SiLU standalone:     {silu_us:7.1}µs");
        println!("  GELU-tanh standalone:{gelu_us:7.1}µs");
        println!("  GEGLU SiLU (gated):  {geglu_us:7.1}µs  (reads 2 buffers)");
        println!();
    }

    // ── LayerNorm vs RMSNorm ──
    println!("--- LayerNorm vs RMSNorm (hidden={hidden}) ---\n");
    {
        let x: Vec<f32> = (0..hidden)
            .map(|i| (i as f32 - hidden as f32 / 2.0) * 0.01)
            .collect();
        let weight: Vec<f32> = vec![1.0; hidden];
        let bias: Vec<f32> = vec![0.0; hidden];
        let x_buf = bufs.transient_from_f32(&x);
        let w_buf = bufs.transient_from_f32(&weight);
        let b_buf = bufs.transient_from_f32(&bias);
        let out_buf = bufs.output((hidden * 4) as u64);
        let n_val = hidden as u32;
        let eps = 1e-6f32;
        let offset = 0.0f32;

        // RMSNorm
        let t = Instant::now();
        for _ in 0..iters {
            let cmd = queue.new_command_buffer();
            let enc = cmd.new_compute_command_encoder();
            enc.set_compute_pipeline_state(&metal.rms_norm_pipeline);
            enc.set_buffer(0, Some(&x_buf), 0);
            enc.set_buffer(1, Some(&w_buf), 0);
            enc.set_buffer(2, Some(&out_buf), 0);
            enc.set_bytes(3, 4, &n_val as *const u32 as *const std::ffi::c_void);
            enc.set_bytes(4, 4, &eps as *const f32 as *const std::ffi::c_void);
            enc.set_bytes(5, 4, &offset as *const f32 as *const std::ffi::c_void);
            enc.dispatch_threads(
                metal::MTLSize::new(hidden as u64, 1, 1),
                metal::MTLSize::new(256, 1, 1),
            );
            enc.end_encoding();
            cmd.commit();
            cmd.wait_until_completed();
        }
        let rms_us = t.elapsed().as_micros() as f64 / iters as f64;

        // LayerNorm (with bias)
        let t = Instant::now();
        for _ in 0..iters {
            let cmd = queue.new_command_buffer();
            let enc = cmd.new_compute_command_encoder();
            enc.set_compute_pipeline_state(&metal.layer_norm_pipeline);
            enc.set_buffer(0, Some(&x_buf), 0);
            enc.set_buffer(1, Some(&w_buf), 0);
            enc.set_buffer(2, Some(&b_buf), 0);
            enc.set_buffer(3, Some(&out_buf), 0);
            enc.set_bytes(4, 4, &n_val as *const u32 as *const std::ffi::c_void);
            enc.set_bytes(5, 4, &eps as *const f32 as *const std::ffi::c_void);
            enc.set_bytes(6, 4, &offset as *const f32 as *const std::ffi::c_void);
            enc.dispatch_threads(
                metal::MTLSize::new(hidden as u64, 1, 1),
                metal::MTLSize::new(256, 1, 1),
            );
            enc.end_encoding();
            cmd.commit();
            cmd.wait_until_completed();
        }
        let ln_us = t.elapsed().as_micros() as f64 / iters as f64;

        // LayerNorm (no bias)
        let t = Instant::now();
        for _ in 0..iters {
            let cmd = queue.new_command_buffer();
            let enc = cmd.new_compute_command_encoder();
            enc.set_compute_pipeline_state(&metal.layer_norm_no_bias_pipeline);
            enc.set_buffer(0, Some(&x_buf), 0);
            enc.set_buffer(1, Some(&w_buf), 0);
            enc.set_buffer(2, Some(&out_buf), 0);
            enc.set_bytes(3, 4, &n_val as *const u32 as *const std::ffi::c_void);
            enc.set_bytes(4, 4, &eps as *const f32 as *const std::ffi::c_void);
            enc.set_bytes(5, 4, &offset as *const f32 as *const std::ffi::c_void);
            enc.dispatch_threads(
                metal::MTLSize::new(hidden as u64, 1, 1),
                metal::MTLSize::new(256, 1, 1),
            );
            enc.end_encoding();
            cmd.commit();
            cmd.wait_until_completed();
        }
        let ln_nb_us = t.elapsed().as_micros() as f64 / iters as f64;

        println!("  RMSNorm:             {rms_us:7.1}µs");
        println!(
            "  LayerNorm (bias):    {ln_us:7.1}µs  ({:.2}x RMSNorm)",
            ln_us / rms_us
        );
        println!(
            "  LayerNorm (no bias): {ln_nb_us:7.1}µs  ({:.2}x RMSNorm)",
            ln_nb_us / rms_us
        );
        println!();
    }

    // ── V-norm ──
    println!("--- V-norm (head_dim={head_dim}, per-head) ---\n");
    {
        let v: Vec<f32> = (0..head_dim).map(|i| (i as f32) * 0.01).collect();
        let v_buf = bufs.transient_from_f32(&v);
        let out_buf = bufs.output((head_dim * 4) as u64);
        let n_val = head_dim as u32;
        let eps = 1e-6f32;

        let t = Instant::now();
        for _ in 0..iters {
            let cmd = queue.new_command_buffer();
            let enc = cmd.new_compute_command_encoder();
            enc.set_compute_pipeline_state(&metal.v_norm_pipeline);
            enc.set_buffer(0, Some(&v_buf), 0);
            enc.set_buffer(1, Some(&out_buf), 0);
            enc.set_bytes(2, 4, &n_val as *const u32 as *const std::ffi::c_void);
            enc.set_bytes(3, 4, &eps as *const f32 as *const std::ffi::c_void);
            enc.dispatch_threads(
                metal::MTLSize::new(head_dim as u64, 1, 1),
                metal::MTLSize::new(256, 1, 1),
            );
            enc.end_encoding();
            cmd.commit();
            cmd.wait_until_completed();
        }
        let vnorm_us = t.elapsed().as_micros() as f64 / iters as f64;

        // Cost for 4 KV heads (typical Gemma)
        let per_layer_4heads = vnorm_us * 4.0;
        println!("  V-norm (1 head):     {vnorm_us:7.1}µs");
        println!("  V-norm (4 KV heads): {per_layer_4heads:7.1}µs/layer");
        println!();
    }

    // ── Scale vector ──
    println!("--- Scale vector (hidden={hidden}) ---\n");
    {
        let x: Vec<f32> = (0..hidden).map(|i| i as f32 * 0.001).collect();
        let x_buf = bufs.transient_from_f32(&x);
        let out_buf = bufs.output((hidden * 4) as u64);
        let n_val = hidden as u32;
        let scalar = 0.73f32;

        let t = Instant::now();
        for _ in 0..iters {
            let cmd = queue.new_command_buffer();
            let enc = cmd.new_compute_command_encoder();
            enc.set_compute_pipeline_state(&metal.scale_vector_pipeline);
            enc.set_buffer(0, Some(&x_buf), 0);
            enc.set_buffer(1, Some(&out_buf), 0);
            enc.set_bytes(2, 4, &n_val as *const u32 as *const std::ffi::c_void);
            enc.set_bytes(3, 4, &scalar as *const f32 as *const std::ffi::c_void);
            enc.dispatch_threads(
                metal::MTLSize::new(hidden as u64, 1, 1),
                metal::MTLSize::new(256, 1, 1),
            );
            enc.end_encoding();
            cmd.commit();
            cmd.wait_until_completed();
        }
        let scale_us = t.elapsed().as_micros() as f64 / iters as f64;
        println!("  scale_vector:        {scale_us:7.1}µs");
        println!();
    }

    // ── Partial RoPE ──
    println!("--- Partial RoPE (head_dim={head_dim}) ---\n");
    {
        let q: Vec<f32> = (0..head_dim).map(|i| (i as f32) * 0.01).collect();
        let q_buf = bufs.transient_from_f32(&q);
        let hd = head_dim as u32;
        let pos = 42u32;
        let base = 1_000_000.0f32;

        // Full rotation (rotary_dim=0 means full)
        let rdim_full = 0u32;
        let t = Instant::now();
        for _ in 0..iters {
            let cmd = queue.new_command_buffer();
            let enc = cmd.new_compute_command_encoder();
            enc.set_compute_pipeline_state(&metal.rope_at_pos_pipeline);
            enc.set_buffer(0, Some(&q_buf), 0);
            enc.set_bytes(1, 4, &hd as *const u32 as *const std::ffi::c_void);
            enc.set_bytes(2, 4, &base as *const f32 as *const std::ffi::c_void);
            enc.set_bytes(3, 4, &pos as *const u32 as *const std::ffi::c_void);
            enc.set_bytes(4, 4, &rdim_full as *const u32 as *const std::ffi::c_void);
            enc.dispatch_threads(
                metal::MTLSize::new((head_dim / 2) as u64, 1, 1),
                metal::MTLSize::new(128, 1, 1),
            );
            enc.end_encoding();
            cmd.commit();
            cmd.wait_until_completed();
        }
        let full_us = t.elapsed().as_micros() as f64 / iters as f64;

        // 25% rotation (Gemma 4 global: rotary_dim = head_dim/4)
        let rdim_25 = (head_dim / 4) as u32;
        let t = Instant::now();
        for _ in 0..iters {
            let cmd = queue.new_command_buffer();
            let enc = cmd.new_compute_command_encoder();
            enc.set_compute_pipeline_state(&metal.rope_at_pos_pipeline);
            enc.set_buffer(0, Some(&q_buf), 0);
            enc.set_bytes(1, 4, &hd as *const u32 as *const std::ffi::c_void);
            enc.set_bytes(2, 4, &base as *const f32 as *const std::ffi::c_void);
            enc.set_bytes(3, 4, &pos as *const u32 as *const std::ffi::c_void);
            enc.set_bytes(4, 4, &rdim_25 as *const u32 as *const std::ffi::c_void);
            enc.dispatch_threads(
                metal::MTLSize::new((head_dim / 8) as u64, 1, 1),
                metal::MTLSize::new(32, 1, 1),
            );
            enc.end_encoding();
            cmd.commit();
            cmd.wait_until_completed();
        }
        let partial_us = t.elapsed().as_micros() as f64 / iters as f64;

        println!("  Full RoPE (256 dims):    {full_us:7.1}µs");
        println!(
            "  Partial RoPE (64 dims):  {partial_us:7.1}µs  ({:.1}x speedup)",
            full_us / partial_us
        );
        println!();
    }

    // ── Summary: per-layer overhead of new features ──
    println!("--- Per-Layer Overhead Summary (Gemma 4 style) ---\n");
    println!("  These are the costs added by new model-agnostic features.");
    println!("  Baseline decode layer: ~0.8ms (from profile_components)\n");
    println!("  Feature                 Cost/layer    % of baseline");
    println!("  ─────────────────────── ──────────── ─────────────");
    // Note: actual numbers computed above, just reference the concept
    println!("  V-norm (4 KV heads)     ~dispatch     <0.1%");
    println!("  Layer scalar            ~dispatch     <0.1%");
    println!("  Partial RoPE (25%)      saves ~75%    net gain");
    println!("  LayerNorm vs RMSNorm    ~same         neutral");
    println!("  Standard FFN (no gate)  saves 1 proj  net gain");
    println!();
    println!("=== Done ===");
}