use bytemuck::{Pod, Zeroable};
use std::borrow::Cow;
use wgpu::util::DeviceExt;

#[repr(C)]
#[derive(Copy, Clone, Pod, Zeroable)]
struct Params { scale: f32, eps: f32, _p0: u32, _p1: u32 }

const D: usize = 32;
const L: usize = 4;
const F: usize = 64;

const WGSL: &str = r#"
const D: u32 = 32u;
const L: u32 = 4u;
const F: u32 = 64u;
struct Params { scale: f32, eps: f32, p0: u32, p1: u32 };
@group(0) @binding(0) var<storage, read> x: array<f32>;      // [L,D]
@group(0) @binding(1) var<storage, read> g1: array<f32>;     // [D]
@group(0) @binding(2) var<storage, read> g2: array<f32>;     // [D]
@group(0) @binding(3) var<storage, read> wq: array<f32>;     // [D,D]
@group(0) @binding(4) var<storage, read> wk: array<f32>;
@group(0) @binding(5) var<storage, read> wv: array<f32>;
@group(0) @binding(6) var<storage, read> wo: array<f32>;
@group(0) @binding(7) var<storage, read> w1: array<f32>;     // [F,D]
@group(0) @binding(8) var<storage, read> w2: array<f32>;     // [D,F]
@group(0) @binding(9) var<storage, read_write> out: array<f32>; // [L,D]
@group(0) @binding(10) var<storage, read> e: array<f32>; // [6,D]
@group(0) @binding(11) var<uniform> p: Params;

@compute @workgroup_size(1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i >= L) { return; }

    var h: array<f32, 32>;
    var q: array<f32, 32>;
    var kk: array<f32, 128>;
    var vv: array<f32, 128>;
    var scores: array<f32, 4>;
    var attn: array<f32, 32>;
    var x1: array<f32, 32>;
    var h2: array<f32, 32>;
    var ffn: array<f32, 64>;

    // h = RMSNorm(x_i) * g1
    var ss = 0.0;
    for (var d = 0u; d < D; d = d + 1u) { ss = ss + x[i * D + d] * x[i * D + d]; }
    let inv = inverseSqrt(ss / f32(D) + p.eps);
    for (var d = 0u; d < D; d = d + 1u) { h[d] = (x[i * D + d] * inv * g1[d]) * (1.0 + e[D + d]) + e[d]; }

    // q
    for (var d = 0u; d < D; d = d + 1u) {
        var acc = 0.0;
        for (var e = 0u; e < D; e = e + 1u) { acc = acc + h[e] * wq[d * D + e]; }
        q[d] = acc;
    }

    // K/V for causal keys, and scores
    var maxs = -1e30;
    for (var j = 0u; j < L; j = j + 1u) {
        if (j > i) { break; }
        var hj: array<f32, 32>;
        var sj = 0.0;
        for (var d = 0u; d < D; d = d + 1u) { sj = sj + x[j * D + d] * x[j * D + d]; }
        let ij = inverseSqrt(sj / f32(D) + p.eps);
        for (var d = 0u; d < D; d = d + 1u) { hj[d] = (x[j * D + d] * ij * g1[d]) * (1.0 + e[D + d]) + e[d]; }
        for (var d = 0u; d < D; d = d + 1u) {
            var ks = 0.0;
            var vs = 0.0;
            for (var e = 0u; e < D; e = e + 1u) {
                ks = ks + hj[e] * wk[d * D + e];
                vs = vs + hj[e] * wv[d * D + e];
            }
            kk[j * D + d] = ks;
            vv[j * D + d] = vs;
        }
        var dot = 0.0;
        for (var d = 0u; d < D; d = d + 1u) { dot = dot + q[d] * kk[j * D + d]; }
        scores[j] = dot * p.scale;
        maxs = max(maxs, scores[j]);
    }

    var sum = 0.0;
    for (var j = 0u; j <= i; j = j + 1u) { sum = sum + exp(scores[j] - maxs); }
    for (var d = 0u; d < D; d = d + 1u) {
        var acc = 0.0;
        for (var j = 0u; j <= i; j = j + 1u) {
            acc = acc + (exp(scores[j] - maxs) / sum) * vv[j * D + d];
        }
        attn[d] = acc;
    }

    // out projection + residual
    for (var d = 0u; d < D; d = d + 1u) {
        var acc = 0.0;
        for (var e = 0u; e < D; e = e + 1u) { acc = acc + attn[e] * wo[d * D + e]; }
        x1[d] = x[i * D + d] + acc * e[2u * D + d];
    }

    // h2 = RMSNorm(x1) * g2
    var s2 = 0.0;
    for (var d = 0u; d < D; d = d + 1u) { s2 = s2 + x1[d] * x1[d]; }
    let i2 = inverseSqrt(s2 / f32(D) + p.eps);
    for (var d = 0u; d < D; d = d + 1u) { h2[d] = (x1[d] * i2 * g2[d]) * (1.0 + e[4u * D + d]) + e[3u * D + d]; }

    // FFN: W1 -> SiLU -> W2 + residual
    for (var f = 0u; f < F; f = f + 1u) {
        var acc = 0.0;
        for (var d = 0u; d < D; d = d + 1u) { acc = acc + h2[d] * w1[f * D + d]; }
        ffn[f] = acc / (1.0 + exp(-acc));
    }
    for (var d = 0u; d < D; d = d + 1u) {
        var acc = 0.0;
        for (var f = 0u; f < F; f = f + 1u) { acc = acc + ffn[f] * w2[d * F + f]; }
        out[i * D + d] = x1[d] + acc * e[5u * D + d];
    }
}
"#;

fn rms(v: &[f32], g: &[f32], eps: f32) -> Vec<f32> {
    let mut out = vec![0.0f32; D];
    let ss: f32 = v.iter().map(|x| x * x).sum();
    let inv = (ss / D as f32 + eps).sqrt().recip();
    for d in 0..D { out[d] = v[d] * inv * g[d]; }
    out
}
fn matvec(w: &[f32], n: usize, k: usize, v: &[f32]) -> Vec<f32> {
    let mut out = vec![0.0f32; n];
    for i in 0..n {
        let mut acc = 0.0f32;
        for j in 0..k { acc += w[i * k + j] * v[j]; }
        out[i] = acc;
    }
    out
}
fn cpu_block(x: &[f32], g1: &[f32], g2: &[f32], wq: &[f32], wk: &[f32], wv: &[f32], wo: &[f32],
             w1: &[f32], w2: &[f32], e: &[f32], eps: f32, scale: f32) -> Vec<f32> {
    let mut out = vec![0.0f32; L * D];
    for i in 0..L {
        let h0 = rms(&x[i * D..(i + 1) * D], g1, eps);
        let mut h = vec![0.0f32; D];
        for d in 0..D { h[d] = h0[d] * (1.0 + e[D + d]) + e[d]; }
        let q = matvec(wq, D, D, &h);
        let mut kk = vec![0.0f32; L * D];
        let mut vv = vec![0.0f32; L * D];
        let mut scores = vec![0.0f32; L];
        let mut maxs = f32::NEG_INFINITY;
        for j in 0..=i {
            let hj0 = rms(&x[j * D..(j + 1) * D], g1, eps);
            let mut hj = vec![0.0f32; D];
            for d in 0..D { hj[d] = hj0[d] * (1.0 + e[D + d]) + e[d]; }
            let k = matvec(wk, D, D, &hj);
            let v = matvec(wv, D, D, &hj);
            kk[j * D..(j + 1) * D].copy_from_slice(&k);
            vv[j * D..(j + 1) * D].copy_from_slice(&v);
            let dot: f32 = (0..D).map(|d| q[d] * k[d]).sum();
            scores[j] = dot * scale;
            maxs = maxs.max(scores[j]);
        }
        let sum: f32 = scores[..=i].iter().map(|s| (s - maxs).exp()).sum();
        let mut attn = vec![0.0f32; D];
        for d in 0..D {
            for j in 0..=i { attn[d] += ((scores[j] - maxs).exp() / sum) * vv[j * D + d]; }
        }
        let o = matvec(wo, D, D, &attn);
        let mut x1 = vec![0.0f32; D];
        for d in 0..D { x1[d] = x[i * D + d] + o[d] * e[2 * D + d]; }
        let h20 = rms(&x1, g2, eps);
        let mut h2 = vec![0.0f32; D];
        for d in 0..D { h2[d] = h20[d] * (1.0 + e[4 * D + d]) + e[3 * D + d]; }
        let mut ffn = matvec(w1, F, D, &h2);
        for f in 0..F { ffn[f] = ffn[f] / (1.0 + (-ffn[f]).exp()); }
        let o2 = matvec(w2, D, F, &ffn);
        for d in 0..D { out[i * D + d] = x1[d] + o2[d] * e[5 * D + d]; }
    }
    out
}

fn main() {
    let eps = 1e-6f32;
    let scale = 1.0 / (D as f32).sqrt();
    let mut x = vec![0.0f32; L * D];
    let mut g1 = vec![0.0f32; D];
    let mut g2 = vec![0.0f32; D];
    let mut wq = vec![0.0f32; D * D];
    let mut wk = vec![0.0f32; D * D];
    let mut wv = vec![0.0f32; D * D];
    let mut wo = vec![0.0f32; D * D];
    let mut w1 = vec![0.0f32; F * D];
    let mut w2 = vec![0.0f32; D * F];
    let mut e = vec![0.0f32; 6 * D];
    for i in 0..x.len() { x[i] = ((i * 13 % 71) as f32 - 35.0) / 19.0; }
    for i in 0..D {
        g1[i] = 1.0 + ((i * 7 % 11) as f32 - 5.0) * 0.01;
        g2[i] = 1.0 + ((i * 5 % 13) as f32 - 6.0) * 0.01;
    }
    let fill = |dst: &mut [f32], k: usize| {
        for i in 0..dst.len() { dst[i] = ((i * 17 % 97) as f32 - 48.0) / (k as f32 * 8.0); }
    };
    fill(&mut wq, D); fill(&mut wk, D); fill(&mut wv, D); fill(&mut wo, D);
    fill(&mut w1, D); fill(&mut w2, F);
    for cc in 0..6 {
        for d in 0..D {
            let gate = cc == 2 || cc == 5;
            e[cc * D + d] = if gate { 1.0 + ((d * 3 + cc) as f32 % 11.0 - 5.0) * 0.03 }
                            else { ((d * 7 + cc * 13) as f32 % 9.0 - 4.0) * 0.02 };
        }
    }
    let expect = cpu_block(&x, &g1, &g2, &wq, &wk, &wv, &wo, &w1, &w2, &e, eps, scale);

    let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor { backends: wgpu::Backends::VULKAN, ..Default::default() });
    let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
        power_preference: wgpu::PowerPreference::HighPerformance, compatible_surface: None, force_fallback_adapter: false,
    })).expect("no Vulkan adapter");
    let info = adapter.get_info();
    println!("adapter: {} | backend: {:?}", info.name, info.backend);
    assert_eq!(info.backend, wgpu::Backend::Vulkan);

    let mut limits = wgpu::Limits::downlevel_defaults();
    limits.max_storage_buffers_per_shader_stage = 16;
    let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
        label: Some("dit-block"), required_features: wgpu::Features::empty(),
        required_limits: limits, ..Default::default()
    })).expect("device");

    let mk = |data: &[f32]| device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: None, contents: bytemuck::cast_slice(data), usage: wgpu::BufferUsages::STORAGE,
    });
    let x_buf = mk(&x); let g1_buf = mk(&g1); let g2_buf = mk(&g2);
    let wq_buf = mk(&wq); let wk_buf = mk(&wk); let wv_buf = mk(&wv); let wo_buf = mk(&wo);
    let w1_buf = mk(&w1); let w2_buf = mk(&w2); let e_buf = mk(&e);
    let out_buf = device.create_buffer(&wgpu::BufferDescriptor {
        label: None, size: (L * D * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, mapped_at_creation: false,
    });
    let readback = device.create_buffer(&wgpu::BufferDescriptor {
        label: None, size: (L * D * 4) as u64,
        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, mapped_at_creation: false,
    });
    let params = Params { scale, eps, _p0: 0, _p1: 0 };
    let p_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: None, contents: bytemuck::bytes_of(&params), usage: wgpu::BufferUsages::UNIFORM,
    });
    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some("dit-block"), source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(WGSL)),
    });
    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
        label: Some("dit-block"), layout: None, module: &module, entry_point: Some("main"),
        compilation_options: Default::default(), cache: None,
    });
    let bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None, layout: &pipeline.get_bind_group_layout(0),
        entries: &[
            wgpu::BindGroupEntry { binding: 0, resource: x_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 1, resource: g1_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 2, resource: g2_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 3, resource: wq_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 4, resource: wk_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 5, resource: wv_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 6, resource: wo_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 7, resource: w1_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 8, resource: w2_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 9, resource: out_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 10, resource: e_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 11, resource: p_buf.as_entire_binding() },
        ],
    });
    let mut enc = device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
    {
        let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None, timestamp_writes: None });
        pass.set_pipeline(&pipeline); pass.set_bind_group(0, &bg, &[]);
        pass.dispatch_workgroups(L as u32, 1, 1);
    }
    enc.copy_buffer_to_buffer(&out_buf, 0, &readback, 0, (L * D * 4) as u64);
    queue.submit(Some(enc.finish()));
    let slice = readback.slice(..);
    slice.map_async(wgpu::MapMode::Read, |r| r.expect("map"));
    device.poll(wgpu::PollType::wait_indefinitely()).expect("poll");
    let data = slice.get_mapped_range();
    let got: &[f32] = bytemuck::cast_slice(&data);
    let mut max_err = 0.0f32;
    for (a, b) in got.iter().zip(expect.iter()) { max_err = max_err.max((a - b).abs()); }
    println!("L={L} D={D} F={F} max_abs_err={max_err:.3e}");
    assert!(max_err < 1e-4, "DiT block mismatch");
    println!("VULKAN MICRO DIT BLOCK OK");
}