//! Benchmarks for Metal-accelerated Mamba (State Space Model)
//!
//! Run with: cargo bench --bench metal_mamba_bench --features metal

use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main};
use rtx_tensor::{DType, Device, Tensor};

/// Benchmark configuration parameters for Mamba
struct MambaBenchConfig {
    batch_size: usize,
    seq_len: usize,
    d_model: usize,
    d_state: usize,
    d_conv: usize,
    expand: usize,
}

impl MambaBenchConfig {
    fn small() -> Self {
        Self {
            batch_size: 1,
            seq_len: 128,
            d_model: 256,
            d_state: 8,
            d_conv: 4,
            expand: 2,
        }
    }

    fn medium() -> Self {
        Self {
            batch_size: 4,
            seq_len: 512,
            d_model: 768,
            d_state: 16,
            d_conv: 4,
            expand: 2,
        }
    }

    fn large() -> Self {
        Self {
            batch_size: 8,
            seq_len: 2048,
            d_model: 1024,
            d_state: 16,
            d_conv: 4,
            expand: 2,
        }
    }

    fn d_inner(&self) -> usize {
        self.d_model * self.expand
    }

    fn total_tokens(&self) -> usize {
        self.batch_size * self.seq_len
    }

    fn name(&self) -> String {
        format!(
            "b{}s{}d{}n{}",
            self.batch_size, self.seq_len, self.d_model, self.d_state
        )
    }
}

/// Benchmark input projection (d_model -> 2 * d_inner)
fn bench_input_projection(c: &mut Criterion) {
    let mut group = c.benchmark_group("mamba_input_projection");

    for config in [MambaBenchConfig::small(), MambaBenchConfig::medium()] {
        let device = Device::cpu();
        let d_inner = config.d_inner();
        let total_tokens = config.total_tokens();

        let input = Tensor::randn(&[total_tokens, config.d_model], DType::F32, &device).unwrap();

        let in_proj = Tensor::randn(&[2 * d_inner, config.d_model], DType::F32, &device).unwrap();

        let flops = 2 * total_tokens * config.d_model * 2 * d_inner;
        group.throughput(Throughput::Elements(flops as u64));

        group.bench_with_input(
            BenchmarkId::new("cpu", config.name()),
            &(&input, &in_proj),
            |b, (inp, proj)| {
                b.iter(|| {
                    let proj_t = proj.t().unwrap();
                    black_box(inp.matmul(&proj_t).unwrap())
                });
            },
        );
    }

    group.finish();
}

/// Benchmark causal 1D convolution (critical for Mamba)
fn bench_causal_conv1d(c: &mut Criterion) {
    let mut group = c.benchmark_group("mamba_causal_conv1d");
    group.sample_size(30);

    for config in [MambaBenchConfig::small(), MambaBenchConfig::medium()] {
        let device = Device::cpu();
        let d_inner = config.d_inner();

        // Input: [batch, seq, d_inner]
        let input = Tensor::randn(
            &[config.batch_size, config.seq_len, d_inner],
            DType::F32,
            &device,
        )
        .unwrap();

        // Convolution weight: [d_inner, 1, kernel_size]
        let conv_weight = Tensor::randn(&[d_inner, 1, config.d_conv], DType::F32, &device).unwrap();

        // FLOPs: batch * seq * d_inner * kernel_size
        let flops = config.batch_size * config.seq_len * d_inner * config.d_conv;
        group.throughput(Throughput::Elements(flops as u64));

        group.bench_with_input(
            BenchmarkId::new("cpu_naive", config.name()),
            &(&input, &conv_weight),
            |b, (inp, weight)| {
                b.iter(|| {
                    // Naive implementation for benchmarking
                    let inp_data = inp.to_cpu().unwrap();
                    let weight_data = weight.to_cpu().unwrap();
                    let mut output = vec![0.0f32; config.batch_size * config.seq_len * d_inner];

                    for batch in 0..config.batch_size {
                        for t in 0..config.seq_len {
                            for d in 0..d_inner {
                                let mut sum = 0.0f32;
                                for k in 0..config.d_conv {
                                    let t_src = t as i64 - k as i64;
                                    if t_src >= 0 {
                                        let idx = batch * config.seq_len * d_inner
                                            + t_src as usize * d_inner
                                            + d;
                                        let w_idx = d * config.d_conv + (config.d_conv - 1 - k);
                                        sum += inp_data[idx] * weight_data[w_idx];
                                    }
                                }
                                output[batch * config.seq_len * d_inner + t * d_inner + d] = sum;
                            }
                        }
                    }
                    black_box(output)
                });
            },
        );
    }

    group.finish();
}

/// Benchmark selective scan (the core SSM operation)
fn bench_selective_scan(c: &mut Criterion) {
    let mut group = c.benchmark_group("mamba_selective_scan");
    group.sample_size(20);

    for config in [MambaBenchConfig::small()] {
        let device = Device::cpu();
        let d_inner = config.d_inner();

        // Input u: [batch, seq, d_inner]
        let u = Tensor::randn(
            &[config.batch_size, config.seq_len, d_inner],
            DType::F32,
            &device,
        )
        .unwrap();

        // Delta (time step): [batch, seq, d_inner]
        let delta = Tensor::randn(
            &[config.batch_size, config.seq_len, d_inner],
            DType::F32,
            &device,
        )
        .unwrap()
        .abs()
        .unwrap();

        // A: [d_inner, d_state]
        let A = Tensor::randn(&[d_inner, config.d_state], DType::F32, &device)
            .unwrap()
            .neg()
            .unwrap()
            .exp()
            .unwrap();

        // B: [batch, seq, d_state]
        let B = Tensor::randn(
            &[config.batch_size, config.seq_len, config.d_state],
            DType::F32,
            &device,
        )
        .unwrap();

        // C: [batch, seq, d_state]
        let C = Tensor::randn(
            &[config.batch_size, config.seq_len, config.d_state],
            DType::F32,
            &device,
        )
        .unwrap();

        // Operations: batch * seq * d_inner * d_state * (state update + output)
        let ops = config.batch_size * config.seq_len * d_inner * config.d_state * 4;
        group.throughput(Throughput::Elements(ops as u64));

        group.bench_function(BenchmarkId::new("cpu_sequential", config.name()), |b| {
            b.iter(|| {
                let u_data = u.to_cpu().unwrap();
                let delta_data = delta.to_cpu().unwrap();
                let A_data = A.to_cpu().unwrap();
                let B_data = B.to_cpu().unwrap();
                let C_data = C.to_cpu().unwrap();

                let mut output = vec![0.0f32; config.batch_size * config.seq_len * d_inner];
                let mut state = vec![0.0f32; config.batch_size * d_inner * config.d_state];

                for batch in 0..config.batch_size {
                    for t in 0..config.seq_len {
                        for d in 0..d_inner {
                            let u_val = u_data[batch * config.seq_len * d_inner + t * d_inner + d];
                            let delta_val =
                                delta_data[batch * config.seq_len * d_inner + t * d_inner + d];

                            let mut y = 0.0f32;
                            for n in 0..config.d_state {
                                let state_idx =
                                    batch * d_inner * config.d_state + d * config.d_state + n;
                                let A_val = A_data[d * config.d_state + n];
                                let B_val = B_data[batch * config.seq_len * config.d_state
                                    + t * config.d_state
                                    + n];
                                let C_val = C_data[batch * config.seq_len * config.d_state
                                    + t * config.d_state
                                    + n];

                                let deltaA = (delta_val * A_val).exp();
                                let deltaB = delta_val * B_val;

                                state[state_idx] = deltaA * state[state_idx] + deltaB * u_val;
                                y += C_val * state[state_idx];
                            }

                            output[batch * config.seq_len * d_inner + t * d_inner + d] = y;
                        }
                    }
                }

                black_box(output)
            });
        });
    }

    group.finish();
}

/// Benchmark SiLU activation (used throughout Mamba)
fn bench_silu_activation(c: &mut Criterion) {
    let mut group = c.benchmark_group("mamba_silu");

    for size in [65536, 262144, 1048576] {
        let device = Device::cpu();
        let tensor = Tensor::randn(&[size], DType::F32, &device).unwrap();

        group.throughput(Throughput::Elements(size as u64));
        group.bench_with_input(BenchmarkId::new("cpu", size), &tensor, |b, t| {
            b.iter(|| black_box(t.swish().unwrap()));
        });
    }

    group.finish();
}

/// Benchmark output projection (d_inner -> d_model)
fn bench_output_projection(c: &mut Criterion) {
    let mut group = c.benchmark_group("mamba_output_projection");

    for config in [MambaBenchConfig::small(), MambaBenchConfig::medium()] {
        let device = Device::cpu();
        let d_inner = config.d_inner();
        let total_tokens = config.total_tokens();

        let gated = Tensor::randn(&[total_tokens, d_inner], DType::F32, &device).unwrap();

        let out_proj = Tensor::randn(&[config.d_model, d_inner], DType::F32, &device).unwrap();

        let flops = 2 * total_tokens * d_inner * config.d_model;
        group.throughput(Throughput::Elements(flops as u64));

        group.bench_with_input(
            BenchmarkId::new("cpu", config.name()),
            &(&gated, &out_proj),
            |b, (g, proj)| {
                b.iter(|| {
                    let proj_t = proj.t().unwrap();
                    black_box(g.matmul(&proj_t).unwrap())
                });
            },
        );
    }

    group.finish();
}

/// Benchmark state initialization
fn bench_state_init(c: &mut Criterion) {
    let mut group = c.benchmark_group("mamba_state_init");

    for config in [
        MambaBenchConfig::small(),
        MambaBenchConfig::medium(),
        MambaBenchConfig::large(),
    ] {
        let device = Device::cpu();
        let d_inner = config.d_inner();
        let state_size = config.batch_size * d_inner * config.d_state;

        group.throughput(Throughput::Elements(state_size as u64));
        group.bench_function(BenchmarkId::new("zeros", config.name()), |b| {
            b.iter(|| {
                black_box(
                    Tensor::zeros(&[config.batch_size, d_inner, config.d_state], &device).unwrap(),
                )
            });
        });
    }

    group.finish();
}

/// Benchmark full Mamba forward simulation
fn bench_mamba_forward_simulation(c: &mut Criterion) {
    let mut group = c.benchmark_group("mamba_forward_simulation");
    group.sample_size(10);

    let config = MambaBenchConfig::small();
    let device = Device::cpu();
    let d_inner = config.d_inner();
    let total_tokens = config.total_tokens();

    // Input
    let input = Tensor::randn(
        &[config.batch_size, config.seq_len, config.d_model],
        DType::F32,
        &device,
    )
    .unwrap();

    // Weights
    let in_proj = Tensor::randn(&[2 * d_inner, config.d_model], DType::F32, &device).unwrap();
    let conv_weight = Tensor::randn(&[d_inner, 1, config.d_conv], DType::F32, &device).unwrap();
    let out_proj = Tensor::randn(&[config.d_model, d_inner], DType::F32, &device).unwrap();

    group.throughput(Throughput::Elements(total_tokens as u64));

    group.bench_function(BenchmarkId::new("cpu_simulation", config.name()), |b| {
        b.iter(|| {
            // Input projection
            let flat = input.view(&[total_tokens, config.d_model]).unwrap();
            let in_proj_t = in_proj.t().unwrap();
            let projected = flat.matmul(&in_proj_t).unwrap();

            // Split x and z
            let x = projected.narrow(1, 0, d_inner).unwrap();
            let z = projected.narrow(1, d_inner, d_inner).unwrap();

            // Apply SiLU to both
            let x_act = x.swish().unwrap();
            let z_act = z.swish().unwrap();

            // Gating
            let gated = x_act.mul(&z_act).unwrap();

            // Output projection
            let out_proj_t = out_proj.t().unwrap();
            let output = gated.matmul(&out_proj_t).unwrap();

            // Reshape
            let output = output
                .view(&[config.batch_size, config.seq_len, config.d_model])
                .unwrap();

            // Residual
            black_box(output.add(&input).unwrap())
        });
    });

    group.finish();
}

/// Compare Mamba vs Attention complexity scaling
fn bench_complexity_scaling(c: &mut Criterion) {
    let mut group = c.benchmark_group("mamba_vs_attention_scaling");
    group.sample_size(10);

    let d_model = 256;
    let device = Device::cpu();

    for seq_len in [128, 256, 512] {
        let batch_size = 2;
        let total_tokens = batch_size * seq_len;

        // Mamba-like: O(n) - just two projections and element-wise ops
        let input = Tensor::randn(&[total_tokens, d_model], DType::F32, &device).unwrap();
        let weight = Tensor::randn(&[d_model, d_model], DType::F32, &device).unwrap();

        group.throughput(Throughput::Elements(seq_len as u64));

        group.bench_with_input(
            BenchmarkId::new("mamba_linear", seq_len),
            &(&input, &weight),
            |b, (inp, w)| {
                b.iter(|| {
                    let w_t = w.t().unwrap();
                    let proj = inp.matmul(&w_t).unwrap();
                    black_box(proj.swish().unwrap())
                });
            },
        );

        // Attention-like: O(n^2) - Q * K^T
        let q = Tensor::randn(&[batch_size, seq_len, d_model], DType::F32, &device).unwrap();
        let k = Tensor::randn(&[batch_size, seq_len, d_model], DType::F32, &device).unwrap();

        group.bench_with_input(
            BenchmarkId::new("attention_quadratic", seq_len),
            &(&q, &k),
            |b, (q_tensor, k_tensor)| {
                b.iter(|| {
                    // Q * K^T: [batch, seq, d] * [batch, d, seq] -> [batch, seq, seq]
                    let k_t = k_tensor.transpose(1, 2).unwrap();
                    black_box(q_tensor.matmul(&k_t).unwrap())
                });
            },
        );
    }

    group.finish();
}

criterion_group!(
    benches,
    bench_input_projection,
    bench_causal_conv1d,
    bench_selective_scan,
    bench_silu_activation,
    bench_output_projection,
    bench_state_init,
    bench_mamba_forward_simulation,
    bench_complexity_scaling,
);

criterion_main!(benches);