#![allow(
    clippy::all,
    unused_imports,
    unused_variables,
    dead_code,
    unused_mut,
    unused_assignments,
    non_camel_case_types,
    clippy::approx_constant,
    unexpected_cfgs,
    unused_must_use,
    unused_parens
)]
//! End-to-End LLM Inference Benchmarks for M4 Pro
//!
//! Comprehensive benchmarks for complete inference pipeline:
//! - Time to first token (TTFT)
//! - Tokens per second (throughput)
//! - Memory usage tracking
//! - Full transformer layer forward pass
//!
//! Performance targets for M4 Pro:
//! - TTFT: <100ms for 7B model
//! - Throughput: 100+ tokens/sec for 7B model
//! - Memory: <16GB for 7B model inference

use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use rand::Rng;
use std::time::Instant;

// Simulated model configuration
#[derive(Clone, Copy)]
struct ModelConfig {
    hidden_size: usize,
    intermediate_size: usize,
    num_attention_heads: usize,
    num_kv_heads: usize,
    head_dim: usize,
    num_layers: usize,
    vocab_size: usize,
    max_seq_len: usize,
}

impl ModelConfig {
    fn llama2_7b() -> Self {
        Self {
            hidden_size: 4096,
            intermediate_size: 11008,
            num_attention_heads: 32,
            num_kv_heads: 32,
            head_dim: 128,
            num_layers: 32,
            vocab_size: 32000,
            max_seq_len: 4096,
        }
    }

    fn llama2_13b() -> Self {
        Self {
            hidden_size: 5120,
            intermediate_size: 13824,
            num_attention_heads: 40,
            num_kv_heads: 40,
            head_dim: 128,
            num_layers: 40,
            vocab_size: 32000,
            max_seq_len: 4096,
        }
    }

    fn llama3_8b() -> Self {
        Self {
            hidden_size: 4096,
            intermediate_size: 14336,
            num_attention_heads: 32,
            num_kv_heads: 8, // GQA
            head_dim: 128,
            num_layers: 32,
            vocab_size: 128256,
            max_seq_len: 8192,
        }
    }

    fn mistral_7b() -> Self {
        Self {
            hidden_size: 4096,
            intermediate_size: 14336,
            num_attention_heads: 32,
            num_kv_heads: 8, // GQA
            head_dim: 128,
            num_layers: 32,
            vocab_size: 32000,
            max_seq_len: 32768,
        }
    }

    fn params_per_layer(&self) -> usize {
        // Attention: Q, K, V, O projections
        let attn_params = self.hidden_size * self.hidden_size * 4;

        // MLP: gate, up, down projections
        let mlp_params = self.hidden_size * self.intermediate_size * 3;

        // Norms (2 per layer)
        let norm_params = self.hidden_size * 2;

        attn_params + mlp_params + norm_params
    }

    fn total_params(&self) -> usize {
        // Embedding
        let embed_params = self.vocab_size * self.hidden_size;

        // All layers
        let layer_params = self.params_per_layer() * self.num_layers;

        // Final norm + LM head
        let final_params = self.hidden_size + self.vocab_size * self.hidden_size;

        embed_params + layer_params + final_params
    }

    fn memory_bytes_fp16(&self) -> usize {
        self.total_params() * 2 // FP16
    }

    fn memory_bytes_int4(&self) -> usize {
        self.total_params() / 2 // INT4
    }
}

// Simulated transformer layer operations
struct TransformerLayer {
    // Weights (simulated)
    q_proj: Vec<f32>,
    k_proj: Vec<f32>,
    v_proj: Vec<f32>,
    o_proj: Vec<f32>,
    gate_proj: Vec<f32>,
    up_proj: Vec<f32>,
    down_proj: Vec<f32>,
    input_norm_weight: Vec<f32>,
    post_attn_norm_weight: Vec<f32>,
    config: ModelConfig,
}

impl TransformerLayer {
    fn new(config: ModelConfig) -> Self {
        let hidden = config.hidden_size;
        let intermediate = config.intermediate_size;

        Self {
            q_proj: random_tensor(hidden * hidden),
            k_proj: random_tensor(
                hidden * (hidden / config.num_attention_heads * config.num_kv_heads),
            ),
            v_proj: random_tensor(
                hidden * (hidden / config.num_attention_heads * config.num_kv_heads),
            ),
            o_proj: random_tensor(hidden * hidden),
            gate_proj: random_tensor(hidden * intermediate),
            up_proj: random_tensor(hidden * intermediate),
            down_proj: random_tensor(intermediate * hidden),
            input_norm_weight: random_tensor(hidden),
            post_attn_norm_weight: random_tensor(hidden),
            config,
        }
    }

    // Simulated forward pass for a single token
    fn forward_single_token(&self, hidden_state: &mut [f32], kv_cache_len: usize) {
        let hidden = self.config.hidden_size;

        // 1. Input LayerNorm/RMSNorm
        rms_norm_inplace(hidden_state, &self.input_norm_weight, 1e-6);

        // 2. Attention projections (Q, K, V)
        let mut q = gemv(&self.q_proj, hidden_state, hidden, hidden);
        let k = gemv(
            &self.k_proj,
            hidden_state,
            hidden,
            hidden / self.config.num_attention_heads * self.config.num_kv_heads,
        );
        let v = gemv(
            &self.v_proj,
            hidden_state,
            hidden,
            hidden / self.config.num_attention_heads * self.config.num_kv_heads,
        );

        // 3. Apply RoPE (simplified)
        apply_rope_simple(&mut q, self.config.head_dim, kv_cache_len);

        // 4. Attention (simplified - would use flash attention in practice)
        // For single token decode, this is essentially a dot product with cached KV
        let attn_output = attention_decode(
            &q,
            &k,
            &v,
            self.config.num_attention_heads,
            self.config.head_dim,
        );

        // 5. Output projection
        let attn_projected = gemv(&self.o_proj, &attn_output, hidden, hidden);

        // 6. Residual connection
        for i in 0..hidden {
            hidden_state[i] += attn_projected[i];
        }

        // 7. Post-attention LayerNorm
        rms_norm_inplace(hidden_state, &self.post_attn_norm_weight, 1e-6);

        // 8. MLP forward
        let gate_out = gemv(
            &self.gate_proj,
            hidden_state,
            hidden,
            self.config.intermediate_size,
        );
        let up_out = gemv(
            &self.up_proj,
            hidden_state,
            hidden,
            self.config.intermediate_size,
        );

        // SiLU activation and element-wise multiply
        let mut mlp_intermediate = Vec::with_capacity(self.config.intermediate_size);
        for i in 0..self.config.intermediate_size {
            let silu = gate_out[i] / (1.0 + (-gate_out[i]).exp());
            mlp_intermediate.push(silu * up_out[i]);
        }

        // Down projection
        let mlp_output = gemv(
            &self.down_proj,
            &mlp_intermediate,
            self.config.intermediate_size,
            hidden,
        );

        // 9. Residual connection
        for i in 0..hidden {
            hidden_state[i] += mlp_output[i];
        }
    }
}

// Helper functions
fn random_tensor(size: usize) -> Vec<f32> {
    let mut rng = rand::thread_rng();
    (0..size).map(|_| rng.gen_range(-0.1..0.1)).collect()
}

fn rms_norm_inplace(x: &mut [f32], weight: &[f32], eps: f32) {
    let sum_sq: f32 = x.iter().map(|v| v * v).sum();
    let inv_rms = 1.0 / (sum_sq / x.len() as f32 + eps).sqrt();
    for (i, w) in weight.iter().enumerate() {
        x[i] = x[i] * inv_rms * w;
    }
}

fn gemv(matrix: &[f32], vector: &[f32], m: usize, n: usize) -> Vec<f32> {
    let mut output = vec![0.0f32; n];
    for j in 0..n {
        let mut sum = 0.0f32;
        for i in 0..m {
            sum += matrix[i * n + j] * vector[i];
        }
        output[j] = sum;
    }
    output
}

fn apply_rope_simple(x: &mut [f32], head_dim: usize, position: usize) {
    let half_dim = head_dim / 2;
    for i in 0..half_dim {
        let freq = 1.0 / 10000.0f32.powf((2 * i) as f32 / head_dim as f32);
        let theta = position as f32 * freq;
        let cos_theta = theta.cos();
        let sin_theta = theta.sin();

        let x0 = x[i * 2];
        let x1 = x[i * 2 + 1];
        x[i * 2] = x0 * cos_theta - x1 * sin_theta;
        x[i * 2 + 1] = x1 * cos_theta + x0 * sin_theta;
    }
}

fn attention_decode(
    q: &[f32],
    k: &[f32],
    v: &[f32],
    num_heads: usize,
    head_dim: usize,
) -> Vec<f32> {
    // Simplified single-token attention decode
    let mut output = vec![0.0f32; num_heads * head_dim];

    for h in 0..num_heads {
        let q_offset = h * head_dim;
        let q_slice = &q[q_offset..q_offset + head_dim];

        // Dot product with single K (simplified - in practice would use KV cache)
        let k_offset = (h % (k.len() / head_dim)) * head_dim;
        let k_slice = &k[k_offset..k_offset + head_dim];

        let score: f32 = q_slice.iter().zip(k_slice).map(|(q, k)| q * k).sum();
        let scale = 1.0 / (head_dim as f32).sqrt();
        let weight = (score * scale).exp(); // Simplified softmax for single token

        let v_offset = (h % (v.len() / head_dim)) * head_dim;
        let v_slice = &v[v_offset..v_offset + head_dim];

        for i in 0..head_dim {
            output[q_offset + i] = v_slice[i] * weight;
        }
    }

    output
}

// KV Cache simulation
struct KvCache {
    keys: Vec<f32>,
    values: Vec<f32>,
    num_tokens: usize,
    num_kv_heads: usize,
    head_dim: usize,
    max_seq_len: usize,
}

impl KvCache {
    fn new(config: &ModelConfig) -> Self {
        let capacity = config.max_seq_len * config.num_kv_heads * config.head_dim;
        Self {
            keys: vec![0.0; capacity],
            values: vec![0.0; capacity],
            num_tokens: 0,
            num_kv_heads: config.num_kv_heads,
            head_dim: config.head_dim,
            max_seq_len: config.max_seq_len,
        }
    }

    fn append(&mut self, k: &[f32], v: &[f32]) {
        if self.num_tokens >= self.max_seq_len {
            return;
        }

        let stride = self.num_kv_heads * self.head_dim;
        let offset = self.num_tokens * stride;

        self.keys[offset..offset + stride].copy_from_slice(&k[..stride.min(k.len())]);
        self.values[offset..offset + stride].copy_from_slice(&v[..stride.min(v.len())]);
        self.num_tokens += 1;
    }

    fn memory_bytes(&self) -> usize {
        (self.keys.len() + self.values.len()) * std::mem::size_of::<f32>()
    }
}

// === Benchmark Functions ===

fn bench_single_layer_forward(c: &mut Criterion) {
    let mut group = c.benchmark_group("single_layer_forward");
    group.sample_size(30);

    let configs = [
        ("llama2_7b", ModelConfig::llama2_7b()),
        ("llama3_8b", ModelConfig::llama3_8b()),
        ("mistral_7b", ModelConfig::mistral_7b()),
    ];

    for (name, config) in configs {
        let layer = TransformerLayer::new(config);
        let mut hidden_state = random_tensor(config.hidden_size);

        let id = BenchmarkId::new(name, config.hidden_size);

        group.throughput(Throughput::Elements(config.params_per_layer() as u64));
        group.bench_function(id, |b| {
            b.iter(|| {
                let mut h = hidden_state.clone();
                layer.forward_single_token(black_box(&mut h), 100);
                h
            })
        });
    }

    group.finish();
}

fn bench_multi_layer_forward(c: &mut Criterion) {
    let mut group = c.benchmark_group("multi_layer_forward");
    group.sample_size(20);

    let config = ModelConfig::llama2_7b();

    for num_layers in [1, 4, 8, 16, 32] {
        let layers: Vec<TransformerLayer> = (0..num_layers)
            .map(|_| TransformerLayer::new(config))
            .collect();
        let mut hidden_state = random_tensor(config.hidden_size);

        let id = BenchmarkId::new(format!("{}_layers", num_layers), num_layers);

        group.throughput(Throughput::Elements(
            (config.params_per_layer() * num_layers) as u64,
        ));
        group.bench_function(id, |b| {
            b.iter(|| {
                let mut h = hidden_state.clone();
                for layer in &layers {
                    layer.forward_single_token(black_box(&mut h), 100);
                }
                h
            })
        });
    }

    group.finish();
}

fn bench_kv_cache_operations(c: &mut Criterion) {
    let mut group = c.benchmark_group("kv_cache");
    group.sample_size(50);

    let configs = [
        ("llama2_7b", ModelConfig::llama2_7b()),
        ("llama3_8b", ModelConfig::llama3_8b()),
    ];

    for (name, config) in configs {
        // Append operation
        let mut cache = KvCache::new(&config);
        let k = random_tensor(config.num_kv_heads * config.head_dim);
        let v = random_tensor(config.num_kv_heads * config.head_dim);

        group.bench_function(
            BenchmarkId::new(format!("{}_append", name), config.num_kv_heads),
            |b| {
                b.iter_batched(
                    || KvCache::new(&config),
                    |mut cache| {
                        cache.append(black_box(&k), black_box(&v));
                        cache
                    },
                    criterion::BatchSize::SmallInput,
                )
            },
        );

        // Memory footprint at various sequence lengths
        for seq_len in [256, 512, 1024, 2048] {
            let mut cache = KvCache::new(&config);
            for _ in 0..seq_len {
                cache.append(&k, &v);
            }

            let memory_mb = cache.memory_bytes() / (1024 * 1024);
            let id = BenchmarkId::new(format!("{}_seq_{}_{}MB", name, seq_len, memory_mb), seq_len);

            group.throughput(Throughput::Bytes(cache.memory_bytes() as u64));
            group.bench_function(id, |b| {
                b.iter(|| {
                    let mut c = KvCache::new(&config);
                    for _ in 0..seq_len {
                        c.append(black_box(&k), black_box(&v));
                    }
                    c
                })
            });
        }
    }

    group.finish();
}

fn bench_decode_throughput(c: &mut Criterion) {
    let mut group = c.benchmark_group("decode_throughput");
    group.sample_size(20);

    // Measure tokens per second for decode phase
    let config = ModelConfig::llama2_7b();
    let layers: Vec<TransformerLayer> = (0..config.num_layers)
        .map(|_| TransformerLayer::new(config))
        .collect();

    // Simulate decoding multiple tokens
    for num_tokens in [1, 10, 50, 100] {
        let id = BenchmarkId::new(format!("{}_tokens", num_tokens), num_tokens);

        group.throughput(Throughput::Elements(num_tokens as u64));
        group.bench_function(id, |b| {
            b.iter(|| {
                let mut hidden_state = random_tensor(config.hidden_size);
                for token_idx in 0..num_tokens {
                    for layer in &layers {
                        layer.forward_single_token(black_box(&mut hidden_state), token_idx);
                    }
                }
                hidden_state
            })
        });
    }

    group.finish();
}

fn bench_prefill_latency(c: &mut Criterion) {
    let mut group = c.benchmark_group("prefill_latency");
    group.sample_size(10);

    // Simulate prefill phase (processing prompt)
    let config = ModelConfig::llama2_7b();
    let layer = TransformerLayer::new(config);

    for seq_len in [32, 64, 128, 256] {
        // Process multiple tokens (simplified - in practice would batch)
        let id = BenchmarkId::new(format!("seq_{}", seq_len), seq_len);

        group.throughput(Throughput::Elements(seq_len as u64));
        group.bench_function(id, |b| {
            b.iter(|| {
                let mut total_output = vec![0.0f32; config.hidden_size];
                for pos in 0..seq_len {
                    let mut hidden_state = random_tensor(config.hidden_size);
                    layer.forward_single_token(black_box(&mut hidden_state), pos);
                    // Accumulate (simplified)
                    for i in 0..config.hidden_size {
                        total_output[i] += hidden_state[i] / seq_len as f32;
                    }
                }
                total_output
            })
        });
    }

    group.finish();
}

fn bench_model_memory(c: &mut Criterion) {
    let mut group = c.benchmark_group("model_memory_estimate");
    group.sample_size(20);

    let configs = [
        ("llama2_7b", ModelConfig::llama2_7b()),
        ("llama2_13b", ModelConfig::llama2_13b()),
        ("llama3_8b", ModelConfig::llama3_8b()),
        ("mistral_7b", ModelConfig::mistral_7b()),
    ];

    for (name, config) in configs {
        let fp16_gb = config.memory_bytes_fp16() as f64 / (1024.0 * 1024.0 * 1024.0);
        let int4_gb = config.memory_bytes_int4() as f64 / (1024.0 * 1024.0 * 1024.0);

        println!(
            "{}: FP16={:.2}GB, INT4={:.2}GB, params={}M",
            name,
            fp16_gb,
            int4_gb,
            config.total_params() / 1_000_000
        );

        // Benchmark single layer to estimate per-layer latency
        let layer = TransformerLayer::new(config);
        let mut hidden_state = random_tensor(config.hidden_size);

        let id = BenchmarkId::new(
            format!("{}_fp16_{:.1}GB", name, fp16_gb),
            config.total_params(),
        );

        group.throughput(Throughput::Elements(config.params_per_layer() as u64));
        group.bench_function(id, |b| {
            b.iter(|| {
                let mut h = hidden_state.clone();
                layer.forward_single_token(black_box(&mut h), 100);
                h
            })
        });
    }

    group.finish();
}

fn bench_inference_components(c: &mut Criterion) {
    let mut group = c.benchmark_group("inference_components");
    group.sample_size(50);

    let config = ModelConfig::llama2_7b();
    let hidden = config.hidden_size;
    let intermediate = config.intermediate_size;

    // Individual component benchmarks
    let input = random_tensor(hidden);
    let weight = random_tensor(hidden);

    // RMSNorm
    group.bench_function("rmsnorm_4096", |b| {
        b.iter_batched(
            || input.clone(),
            |mut x| {
                rms_norm_inplace(black_box(&mut x), black_box(&weight), 1e-6);
                x
            },
            criterion::BatchSize::SmallInput,
        )
    });

    // Linear projection (hidden -> hidden)
    let proj_matrix = random_tensor(hidden * hidden);
    group.bench_function("linear_4096x4096", |b| {
        b.iter(|| gemv(black_box(&proj_matrix), black_box(&input), hidden, hidden))
    });

    // Linear projection (hidden -> intermediate)
    let mlp_up_matrix = random_tensor(hidden * intermediate);
    group.bench_function("linear_4096x11008", |b| {
        b.iter(|| {
            gemv(
                black_box(&mlp_up_matrix),
                black_box(&input),
                hidden,
                intermediate,
            )
        })
    });

    // RoPE
    let mut rope_input = random_tensor(config.num_attention_heads * config.head_dim);
    group.bench_function("rope_32heads", |b| {
        b.iter_batched(
            || rope_input.clone(),
            |mut x| {
                for h in 0..config.num_attention_heads {
                    let offset = h * config.head_dim;
                    apply_rope_simple(
                        black_box(&mut x[offset..offset + config.head_dim]),
                        config.head_dim,
                        100,
                    );
                }
                x
            },
            criterion::BatchSize::SmallInput,
        )
    });

    group.finish();
}

fn bench_tokens_per_second_estimation(c: &mut Criterion) {
    let mut group = c.benchmark_group("tokens_per_second");
    group.sample_size(10);

    // Full model throughput estimation
    let config = ModelConfig::llama2_7b();

    // Create a simplified full model
    let layers: Vec<TransformerLayer> = (0..4) // Use 4 layers for faster benchmarking
        .map(|_| TransformerLayer::new(config))
        .collect();

    let id = BenchmarkId::new("llama2_7b_4layers", 4);

    // Time how long it takes to process tokens
    group.bench_function(id, |b| {
        b.iter_custom(|iters| {
            let mut total_time = std::time::Duration::ZERO;

            for _ in 0..iters {
                let mut hidden_state = random_tensor(config.hidden_size);
                let start = Instant::now();

                for layer in &layers {
                    layer.forward_single_token(black_box(&mut hidden_state), 100);
                }

                total_time += start.elapsed();
            }

            total_time
        })
    });

    group.finish();
}

criterion_group!(
    benches,
    bench_single_layer_forward,
    bench_multi_layer_forward,
    bench_kv_cache_operations,
    bench_decode_throughput,
    bench_prefill_latency,
    bench_model_memory,
    bench_inference_components,
    bench_tokens_per_second_estimation,
);

criterion_main!(benches);