use half::bf16;
use xserv_kernels::*;
use xserv_tensor::{Device, Tensor};

fn init() {
    xserv_cuda::device::set_device(0).unwrap();
}

// --- CPU reference implementations ---

fn cpu_rmsnorm(x: &[f32], gamma: &[f32], eps: f32, hidden: usize) -> Vec<f32> {
    let rows = x.len() / hidden;
    let mut out = vec![0.0f32; x.len()];
    for r in 0..rows {
        let row = &x[r * hidden..(r + 1) * hidden];
        let sum_sq: f32 = row.iter().map(|v| v * v).sum();
        let rms_inv = 1.0 / (sum_sq / hidden as f32 + eps).sqrt();
        for i in 0..hidden {
            out[r * hidden + i] = row[i] * rms_inv * gamma[i];
        }
    }
    out
}

fn cpu_layernorm(x: &[f32], gamma: &[f32], beta: &[f32], eps: f32, hidden: usize) -> Vec<f32> {
    let rows = x.len() / hidden;
    let mut out = vec![0.0f32; x.len()];
    for r in 0..rows {
        let row = &x[r * hidden..(r + 1) * hidden];
        let mean: f32 = row.iter().sum::<f32>() / hidden as f32;
        let var: f32 = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / hidden as f32;
        let inv_std = 1.0 / (var + eps).sqrt();
        for i in 0..hidden {
            out[r * hidden + i] = gamma[i] * (row[i] - mean) * inv_std + beta[i];
        }
    }
    out
}

fn cpu_gelu(x: &[f32]) -> Vec<f32> {
    let sqrt_2_over_pi = 0.7978845608f32;
    x.iter()
        .map(|&v| {
            let inner = sqrt_2_over_pi * (v + 0.044715 * v * v * v);
            0.5 * v * (1.0 + inner.tanh())
        })
        .collect()
}

fn cpu_silu(x: &[f32]) -> Vec<f32> {
    x.iter().map(|&v| v / (1.0 + (-v).exp())).collect()
}

fn cpu_softmax(x: &[f32], cols: usize) -> Vec<f32> {
    let rows = x.len() / cols;
    let mut out = vec![0.0f32; x.len()];
    for r in 0..rows {
        let row = &x[r * cols..(r + 1) * cols];
        let max = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let exps: Vec<f32> = row.iter().map(|v| (v - max).exp()).collect();
        let sum: f32 = exps.iter().sum();
        for i in 0..cols {
            out[r * cols + i] = exps[i] / sum;
        }
    }
    out
}

fn cpu_rope(x: &mut [f32], positions: &[u32], num_heads: usize, head_dim: usize, theta: f32) {
    let half_dim = head_dim / 2;
    let num_tokens = positions.len();
    for t in 0..num_tokens {
        let pos = positions[t] as f32;
        for h in 0..num_heads {
            for i in 0..half_dim {
                let freq = 1.0 / theta.powf(2.0 * i as f32 / head_dim as f32);
                let angle = pos * freq;
                let cos_val = angle.cos();
                let sin_val = angle.sin();
                let base = (t * num_heads + h) * head_dim;
                let x0 = x[base + i];
                let x1 = x[base + i + half_dim];
                x[base + i] = x0 * cos_val - x1 * sin_val;
                x[base + i + half_dim] = x1 * cos_val + x0 * sin_val;
            }
        }
    }
}

fn check_close(result: &[f32], expected: &[f32], atol: f32, name: &str) {
    assert_eq!(result.len(), expected.len(), "{name}: length mismatch");
    let mut max_err = 0.0f32;
    for (i, (r, e)) in result.iter().zip(expected).enumerate() {
        let err = (r - e).abs();
        if err > max_err {
            max_err = err;
        }
        assert!(
            err <= atol,
            "{name}: mismatch at [{i}]: got {r}, expected {e}, err {err}"
        );
    }
    println!("{name}: max_err = {max_err:.6e}");
}

fn make_data(n: usize) -> Vec<f32> {
    (0..n).map(|i| ((i % 17) as f32 - 8.0) * 0.1).collect()
}

// === RMSNorm ===

#[test]
fn test_rmsnorm_f32() {
    init();
    let hidden = 768;
    let rows = 4;
    let x_data = make_data(rows * hidden);
    let gamma_data: Vec<f32> = (0..hidden).map(|i| 0.5 + (i % 3) as f32 * 0.2).collect();
    let expected = cpu_rmsnorm(&x_data, &gamma_data, 1e-5, hidden);

    let x = Tensor::from_slice(&x_data, &[rows, hidden]).to_device(Device::Cuda(0));
    let gamma = Tensor::from_slice(&gamma_data, &[hidden]).to_device(Device::Cuda(0));
    let out = rmsnorm(&x, &gamma, 1e-5).to_device(Device::Cpu);
    check_close(out.as_slice::<f32>(), &expected, 1e-4, "rmsnorm_f32");
}

#[test]
fn test_rmsnorm_bf16() {
    init();
    let hidden = 768;
    let rows = 4;
    let x_f32 = make_data(rows * hidden);
    let gamma_f32: Vec<f32> = (0..hidden).map(|i| 0.5 + (i % 3) as f32 * 0.2).collect();
    let expected = cpu_rmsnorm(&x_f32, &gamma_f32, 1e-5, hidden);

    let x_bf16: Vec<bf16> = x_f32.iter().map(|&v| bf16::from_f32(v)).collect();
    let gamma_bf16: Vec<bf16> = gamma_f32.iter().map(|&v| bf16::from_f32(v)).collect();
    let x = Tensor::from_slice(&x_bf16, &[rows, hidden]).to_device(Device::Cuda(0));
    let gamma = Tensor::from_slice(&gamma_bf16, &[hidden]).to_device(Device::Cuda(0));
    let out = rmsnorm(&x, &gamma, 1e-5).to_device(Device::Cpu);

    let result: Vec<f32> = out.as_slice::<bf16>().iter().map(|v| v.to_f32()).collect();
    check_close(&result, &expected, 0.05, "rmsnorm_bf16");
}

// === LayerNorm ===

#[test]
fn test_layernorm_f32() {
    init();
    let hidden = 768;
    let rows = 4;
    let x_data = make_data(rows * hidden);
    let gamma_data: Vec<f32> = (0..hidden).map(|i| 0.8 + (i % 5) as f32 * 0.1).collect();
    let beta_data: Vec<f32> = (0..hidden).map(|i| ((i % 7) as f32 - 3.0) * 0.01).collect();
    let expected = cpu_layernorm(&x_data, &gamma_data, &beta_data, 1e-5, hidden);

    let x = Tensor::from_slice(&x_data, &[rows, hidden]).to_device(Device::Cuda(0));
    let gamma = Tensor::from_slice(&gamma_data, &[hidden]).to_device(Device::Cuda(0));
    let beta = Tensor::from_slice(&beta_data, &[hidden]).to_device(Device::Cuda(0));
    let out = layernorm(&x, &gamma, &beta, 1e-5).to_device(Device::Cpu);
    check_close(out.as_slice::<f32>(), &expected, 1e-4, "layernorm_f32");
}

// === GELU ===

#[test]
fn test_gelu_f32() {
    init();
    let data = make_data(10000);
    let expected = cpu_gelu(&data);
    let x = Tensor::from_slice(&data, &[10000]).to_device(Device::Cuda(0));
    let out = gelu(&x).to_device(Device::Cpu);
    check_close(out.as_slice::<f32>(), &expected, 1e-5, "gelu_f32");
}

#[test]
fn test_gelu_bf16() {
    init();
    let data_f32 = make_data(10000);
    let expected = cpu_gelu(&data_f32);
    let data_bf16: Vec<bf16> = data_f32.iter().map(|&v| bf16::from_f32(v)).collect();
    let x = Tensor::from_slice(&data_bf16, &[10000]).to_device(Device::Cuda(0));
    let out = gelu(&x).to_device(Device::Cpu);
    let result: Vec<f32> = out.as_slice::<bf16>().iter().map(|v| v.to_f32()).collect();
    check_close(&result, &expected, 0.02, "gelu_bf16");
}

// === SiLU ===

#[test]
fn test_silu_f32() {
    init();
    let data = make_data(10000);
    let expected = cpu_silu(&data);
    let x = Tensor::from_slice(&data, &[10000]).to_device(Device::Cuda(0));
    let out = silu(&x).to_device(Device::Cpu);
    check_close(out.as_slice::<f32>(), &expected, 1e-5, "silu_f32");
}

// === Softmax ===

#[test]
fn test_softmax_f32() {
    init();
    let rows = 8;
    let cols = 256;
    let data = make_data(rows * cols);
    let expected = cpu_softmax(&data, cols);
    let x = Tensor::from_slice(&data, &[rows, cols]).to_device(Device::Cuda(0));
    let out = softmax(&x).to_device(Device::Cpu);
    check_close(out.as_slice::<f32>(), &expected, 1e-5, "softmax_f32");
}

#[test]
fn test_softmax_sum_to_one() {
    init();
    let rows = 4;
    let cols = 2048;
    let data: Vec<f32> = (0..rows * cols)
        .map(|i| ((i % 31) as f32 - 15.0) * 0.5)
        .collect();
    let x = Tensor::from_slice(&data, &[rows, cols]).to_device(Device::Cuda(0));
    let out = softmax(&x).to_device(Device::Cpu);
    let result = out.as_slice::<f32>();
    for r in 0..rows {
        let row_sum: f32 = result[r * cols..(r + 1) * cols].iter().sum();
        assert!(
            (row_sum - 1.0).abs() < 1e-5,
            "softmax row {r} sum = {row_sum}"
        );
    }
}

#[test]
fn test_softmax_large_values() {
    init();
    let data = vec![1000.0f32, 1001.0, 999.0, 1000.5];
    let expected = cpu_softmax(&data, 4);
    let x = Tensor::from_slice(&data, &[1, 4]).to_device(Device::Cuda(0));
    let out = softmax(&x).to_device(Device::Cpu);
    check_close(out.as_slice::<f32>(), &expected, 1e-5, "softmax_large");
}

// === Embedding ===

#[test]
fn test_embedding_f32() {
    init();
    let vocab_size = 100;
    let hidden = 64;
    let table_data: Vec<f32> = (0..vocab_size * hidden).map(|i| i as f32 * 0.01).collect();
    let token_ids: Vec<u32> = vec![0, 5, 99, 42, 1];

    let table = Tensor::from_slice(&table_data, &[vocab_size, hidden]).to_device(Device::Cuda(0));
    let out = embedding(&table, &token_ids).to_device(Device::Cpu);

    assert_eq!(out.shape(), &[5, hidden]);
    let result = out.as_slice::<f32>();
    for (seq_idx, &tid) in token_ids.iter().enumerate() {
        for i in 0..hidden {
            let expected = table_data[tid as usize * hidden + i];
            let got = result[seq_idx * hidden + i];
            assert!(
                (got - expected).abs() < 1e-6,
                "embedding mismatch at [{seq_idx},{i}]: got {got}, expected {expected}"
            );
        }
    }
}

// === RoPE ===

#[test]
fn test_rope_f32() {
    init();
    let num_tokens = 4;
    let num_heads = 2;
    let head_dim = 8;
    let theta = 10000.0f32;
    let positions: Vec<u32> = vec![0, 1, 2, 3];

    let x_data: Vec<f32> = (0..num_tokens * num_heads * head_dim)
        .map(|i| ((i % 13) as f32 - 6.0) * 0.1)
        .collect();
    let mut expected = x_data.clone();
    cpu_rope(&mut expected, &positions, num_heads, head_dim, theta);

    let x =
        Tensor::from_slice(&x_data, &[num_tokens, num_heads, head_dim]).to_device(Device::Cuda(0));
    let cache = RopeCache::new(64, head_dim, theta);
    rope_inplace(&x, &cache, &positions);

    let out = x.to_device(Device::Cpu);
    check_close(out.as_slice::<f32>(), &expected, 1e-4, "rope_f32");
}

#[test]
fn test_rope_position_0_identity() {
    init();
    // At position 0, all angles are 0, so cos=1, sin=0 → identity transform
    let num_tokens = 1;
    let num_heads = 2;
    let head_dim = 8;
    let positions: Vec<u32> = vec![0];

    let x_data: Vec<f32> = (0..num_tokens * num_heads * head_dim)
        .map(|i| (i as f32 + 1.0) * 0.1)
        .collect();

    let x =
        Tensor::from_slice(&x_data, &[num_tokens, num_heads, head_dim]).to_device(Device::Cuda(0));
    let cache = RopeCache::new(64, head_dim, 10000.0);
    rope_inplace(&x, &cache, &positions);

    let out = x.to_device(Device::Cpu);
    check_close(out.as_slice::<f32>(), &x_data, 1e-6, "rope_pos0");
}