//! Metal Operations Benchmark - Compare Metal vs CPU for common tensor ops
//! Run with: cargo run -p rtx-flash-metal-attention --example metal_ops_bench --release

use rtx_tensor::{Device, Tensor};
use std::time::Instant;

/// CPU GELU activation
fn cpu_gelu(data: &[f32]) -> Vec<f32> {
    data.iter()
        .map(|x| {
            let sqrt_2_over_pi = 0.7978845608028654f32;
            let coeff = 0.044715f32;
            let inner = sqrt_2_over_pi * (x + coeff * x * x * x);
            0.5 * x * (1.0 + inner.tanh())
        })
        .collect()
}

/// CPU ReLU activation
fn cpu_relu(data: &[f32]) -> Vec<f32> {
    data.iter().map(|x| x.max(0.0)).collect()
}

/// CPU Softmax (over last dim)
fn cpu_softmax(data: &[f32], last_dim: usize) -> Vec<f32> {
    let mut output = vec![0.0f32; data.len()];
    for chunk_start in (0..data.len()).step_by(last_dim) {
        let chunk = &data[chunk_start..chunk_start + last_dim];
        let max_val = chunk.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let exp_sum: f32 = chunk.iter().map(|x| (x - max_val).exp()).sum();
        for (i, x) in chunk.iter().enumerate() {
            output[chunk_start + i] = (x - max_val).exp() / exp_sum;
        }
    }
    output
}

/// CPU Matrix Multiply (C = A @ B)
fn cpu_matmul(a: &[f32], b: &[f32], m: usize, n: usize, k: usize) -> Vec<f32> {
    let mut c = vec![0.0f32; m * n];
    for i in 0..m {
        for j in 0..n {
            let mut sum = 0.0f32;
            for l in 0..k {
                sum += a[i * k + l] * b[l * n + j];
            }
            c[i * n + j] = sum;
        }
    }
    c
}

fn bench_op<F>(name: &str, iterations: usize, warmup: usize, mut f: F) -> f64
where
    F: FnMut(),
{
    // Warmup
    for _ in 0..warmup {
        f();
    }

    // Benchmark
    let start = Instant::now();
    for _ in 0..iterations {
        f();
    }
    let elapsed = start.elapsed();
    elapsed.as_secs_f64() * 1000.0 / iterations as f64
}

fn main() {
    println!("=== Metal vs CPU Operations Benchmark ===\n");

    // Get Metal device
    let metal_device = match Device::metal(0) {
        Ok(d) => {
            println!("Metal device: Metal(0)");
            d
        }
        Err(e) => {
            eprintln!("Metal device not available: {:?}", e);
            return;
        }
    };

    let sizes = vec![
        ("Small (1K)", 1024usize),
        ("Medium (64K)", 65536),
        ("Large (1M)", 1048576),
    ];

    println!("\n{:=^80}", " Element-wise Operations ");
    println!(
        "{:<15} {:>12} {:>12} {:>12} {:>12}",
        "Size", "Metal (ms)", "CPU (ms)", "Speedup", "Elements"
    );
    println!("{:-<80}", "");

    // GELU Benchmark
    println!("\n--- GELU Activation ---");
    for (name, size) in &sizes {
        let shape = &[*size];
        let cpu_data = vec![0.5f32; *size];

        // Metal
        let t = Tensor::ones(shape, &metal_device).expect("tensor");
        let metal_ms = bench_op("metal_gelu", 100, 10, || {
            let _ = t.gelu();
        });

        // CPU
        let cpu_ms = bench_op("cpu_gelu", 100, 10, || {
            let _ = cpu_gelu(&cpu_data);
        });

        let speedup = cpu_ms / metal_ms;
        println!(
            "{:<15} {:>12.4} {:>12.4} {:>11.1}x {:>12}",
            name, metal_ms, cpu_ms, speedup, size
        );
    }

    // ReLU Benchmark
    println!("\n--- ReLU Activation ---");
    for (name, size) in &sizes {
        let shape = &[*size];
        let cpu_data = vec![0.5f32; *size];

        let t = Tensor::ones(shape, &metal_device).expect("tensor");
        let metal_ms = bench_op("metal_relu", 100, 10, || {
            let _ = t.relu();
        });

        let cpu_ms = bench_op("cpu_relu", 100, 10, || {
            let _ = cpu_relu(&cpu_data);
        });

        let speedup = cpu_ms / metal_ms;
        println!(
            "{:<15} {:>12.4} {:>12.4} {:>11.1}x {:>12}",
            name, metal_ms, cpu_ms, speedup, size
        );
    }

    // Softmax Benchmark
    println!("\n--- Softmax ---");
    let softmax_sizes = vec![
        ("128x64", 128, 64),
        ("512x128", 512, 128),
        ("2048x256", 2048, 256),
    ];
    for (name, rows, cols) in &softmax_sizes {
        let shape = &[*rows, *cols];
        let size = rows * cols;
        let cpu_data = vec![0.5f32; size];

        let t = Tensor::ones(shape, &metal_device).expect("tensor");
        let metal_ms = bench_op("metal_softmax", 100, 10, || {
            let _ = t.softmax(-1);
        });

        let cpu_ms = bench_op("cpu_softmax", 100, 10, || {
            let _ = cpu_softmax(&cpu_data, *cols);
        });

        let speedup = cpu_ms / metal_ms;
        println!(
            "{:<15} {:>12.4} {:>12.4} {:>11.1}x {:>12}",
            name, metal_ms, cpu_ms, speedup, size
        );
    }

    // Matrix Multiplication Benchmark
    println!("\n{:=^80}", " Matrix Multiplication ");
    println!(
        "{:<15} {:>12} {:>12} {:>12} {:>12}",
        "Size", "Metal (ms)", "CPU (ms)", "Speedup", "GFLOPS"
    );
    println!("{:-<80}", "");

    let matmul_sizes = vec![
        ("64x64", 64, 64, 64),
        ("256x256", 256, 256, 256),
        ("512x512", 512, 512, 512),
        ("1024x1024", 1024, 1024, 1024),
    ];

    for (name, m, n, k) in &matmul_sizes {
        let cpu_a = vec![1.0f32; m * k];
        let cpu_b = vec![1.0f32; k * n];

        let a = Tensor::ones(&[*m, *k], &metal_device).expect("A");
        let b = Tensor::ones(&[*k, *n], &metal_device).expect("B");

        let metal_ms = bench_op("metal_matmul", 50, 10, || {
            let _ = a.matmul(&b);
        });

        let cpu_iters = if *m >= 512 { 5 } else { 50 };
        let cpu_ms = bench_op("cpu_matmul", cpu_iters, 2, || {
            let _ = cpu_matmul(&cpu_a, &cpu_b, *m, *n, *k);
        });

        let flops = 2.0 * (*m as f64) * (*n as f64) * (*k as f64);
        let metal_gflops = flops / (metal_ms / 1000.0) / 1e9;
        let speedup = cpu_ms / metal_ms;

        println!(
            "{:<15} {:>12.4} {:>12.4} {:>11.1}x {:>12.2}",
            name, metal_ms, cpu_ms, speedup, metal_gflops
        );
    }

    // Reduction Operations
    println!("\n{:=^80}", " Reduction Operations ");
    println!(
        "{:<15} {:>12} {:>12} {:>12} {:>12}",
        "Size", "Metal (ms)", "CPU (ms)", "Speedup", "Elements"
    );
    println!("{:-<80}", "");

    println!("\n--- Sum Reduction ---");
    for (name, size) in &sizes {
        let shape = &[*size];
        let cpu_data = vec![1.0f32; *size];

        let t = Tensor::ones(shape, &metal_device).expect("tensor");
        let metal_ms = bench_op("metal_sum", 100, 10, || {
            let _ = t.sum(None);
        });

        let cpu_ms = bench_op("cpu_sum", 100, 10, || {
            let _: f32 = cpu_data.iter().sum();
        });

        let speedup = cpu_ms / metal_ms;
        println!(
            "{:<15} {:>12.4} {:>12.4} {:>11.1}x {:>12}",
            name, metal_ms, cpu_ms, speedup, size
        );
    }

    println!("\n{:-<80}", "");
    println!("\nDone!");
}