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

#[repr(C)]
#[derive(Copy, Clone, Pod, Zeroable)]
struct ConvParams {
    n: u32, cin: u32, cout: u32,
    din: u32, hin: u32, win: u32,
    kd: u32, kh: u32, kw: u32,
    pad: u32, stride: u32,
    dout: u32, hout: u32, wout: u32,
    _p0: u32, _p1: u32,
}
#[repr(C)]
#[derive(Copy, Clone, Pod, Zeroable)]
struct GNParams { n: u32, c: u32, d: u32, h: u32, w: u32, groups: u32, eps: f32, _p: u32 }
#[repr(C)]
#[derive(Copy, Clone, Pod, Zeroable)]
struct UpParams {
    n: u32, c: u32, din: u32, hin: u32, win: u32,
    scale: u32, dout: u32, hout: u32, wout: u32,
    _p0: u32, _p1: u32, _p2: u32,
}

const CONV_WGSL: &str = r#"
struct P { n: u32, cin: u32, cout: u32, din: u32, hin: u32, win: u32, kd: u32, kh: u32, kw: u32,
           pad: u32, stride: u32, dout: u32, hout: u32, wout: u32, p0: u32, p1: u32 };
@group(0) @binding(0) var<storage, read> x: array<f32>;
@group(0) @binding(1) var<storage, read> w: array<f32>;
@group(0) @binding(2) var<storage, read> b: array<f32>;
@group(0) @binding(3) var<storage, read_write> y: array<f32>;
@group(0) @binding(4) var<uniform> p: P;
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x;
    let total = p.n * p.cout * p.dout * p.hout * p.wout;
    if (idx >= total) { return; }
    var t = idx;
    let ow = t % p.wout; t = t / p.wout;
    let oh = t % p.hout; t = t / p.hout;
    let od = t % p.dout; t = t / p.dout;
    let co = t % p.cout;
    let bn = t / p.cout;
    var acc = b[co];
    for (var ci = 0u; ci < p.cin; ci = ci + 1u) {
        for (var kd = 0u; kd < p.kd; kd = kd + 1u) {
            let id = i32(od * p.stride + kd) - i32(p.pad);
            if (id < 0 || id >= i32(p.din)) { continue; }
            for (var kh = 0u; kh < p.kh; kh = kh + 1u) {
                let ih = i32(oh * p.stride + kh) - i32(p.pad);
                if (ih < 0 || ih >= i32(p.hin)) { continue; }
                for (var kw = 0u; kw < p.kw; kw = kw + 1u) {
                    let iw = i32(ow * p.stride + kw) - i32(p.pad);
                    if (iw < 0 || iw >= i32(p.win)) { continue; }
                    let xidx = (((bn * p.cin + ci) * p.din + u32(id)) * p.hin + u32(ih)) * p.win + u32(iw);
                    let widx = (((co * p.cin + ci) * p.kd + kd) * p.kh + kh) * p.kw + kw;
                    acc = acc + x[xidx] * w[widx];
                }
            }
        }
    }
    y[idx] = acc;
}
"#;

const GN_WGSL: &str = r#"
struct P { n: u32, c: u32, d: u32, h: u32, w: u32, groups: u32, eps: f32, p: u32 };
@group(0) @binding(0) var<storage, read> x: array<f32>;
@group(0) @binding(1) var<storage, read> gamma: array<f32>;
@group(0) @binding(2) var<storage, read> beta: array<f32>;
@group(0) @binding(3) var<storage, read_write> y: array<f32>;
@group(0) @binding(4) var<uniform> p: P;
@compute @workgroup_size(1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let tg = p.n * p.groups;
    if (gid.x >= tg) { return; }
    let bn = gid.x / p.groups;
    let g = gid.x % p.groups;
    let cg = p.c / p.groups;
    let c0 = g * cg;
    let spatial = p.d * p.h * p.w;
    let count = cg * spatial;
    var sum = 0.0;
    var sumsq = 0.0;
    for (var c = c0; c < c0 + cg; c = c + 1u) {
        for (var s = 0u; s < spatial; s = s + 1u) {
            let v = x[(bn * p.c + c) * spatial + s];
            sum = sum + v; sumsq = sumsq + v * v;
        }
    }
    let mean = sum / f32(count);
    let variance = sumsq / f32(count) - mean * mean;
    let inv = inverseSqrt(variance + p.eps);
    for (var c = c0; c < c0 + cg; c = c + 1u) {
        for (var s = 0u; s < spatial; s = s + 1u) {
            let idx = (bn * p.c + c) * spatial + s;
            let norm = (x[idx] - mean) * inv;
            let v = norm * gamma[c] + beta[c];
            y[idx] = v / (1.0 + exp(-v));
        }
    }
}
"#;

const UP_WGSL: &str = r#"
struct P { n: u32, c: u32, din: u32, hin: u32, win: u32, scale: u32, dout: u32, hout: u32, wout: u32,
           p0: u32, p1: u32, p2: u32 };
@group(0) @binding(0) var<storage, read> x: array<f32>;
@group(0) @binding(1) var<storage, read_write> y: array<f32>;
@group(0) @binding(2) var<uniform> p: P;
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x;
    let total = p.n * p.c * p.dout * p.hout * p.wout;
    if (idx >= total) { return; }
    var t = idx;
    let ow = t % p.wout; t = t / p.wout;
    let oh = t % p.hout; t = t / p.hout;
    let od = t % p.dout; t = t / p.dout;
    let co = t % p.c;
    let bn = t / p.c;
    let iw = ow / p.scale;
    let ih = oh / p.scale;
    let id = od / p.scale;
    let src = (((bn * p.c + co) * p.din + id) * p.hin + ih) * p.win + iw;
    y[idx] = x[src];
}
"#;

fn cpu_conv(x: &[f32], w: &[f32], b: &[f32], p: &ConvParams) -> Vec<f32> {
    let total = (p.n * p.cout * p.dout * p.hout * p.wout) as usize;
    let mut y = vec![0.0f32; total];
    for idx in 0..total {
        let mut t = idx;
        let ow = t % p.wout as usize; t /= p.wout as usize;
        let oh = t % p.hout as usize; t /= p.hout as usize;
        let od = t % p.dout as usize; t /= p.dout as usize;
        let co = t % p.cout as usize;
        let bn = t / p.cout as usize;
        let mut acc = b[co];
        for ci in 0..p.cin as usize {
            for kd in 0..p.kd as usize {
                let id = od as i32 * p.stride as i32 + kd as i32 - p.pad as i32;
                if id < 0 || id >= p.din as i32 { continue; }
                for kh in 0..p.kh as usize {
                    let ih = oh as i32 * p.stride as i32 + kh as i32 - p.pad as i32;
                    if ih < 0 || ih >= p.hin as i32 { continue; }
                    for kw in 0..p.kw as usize {
                        let iw = ow as i32 * p.stride as i32 + kw as i32 - p.pad as i32;
                        if iw < 0 || iw >= p.win as i32 { continue; }
                        let xidx = (((bn * p.cin as usize + ci) * p.din as usize + id as usize) * p.hin as usize + ih as usize) * p.win as usize + iw as usize;
                        let widx = (((co * p.cin as usize + ci) * p.kd as usize + kd) * p.kh as usize + kh) * p.kw as usize + kw;
                        acc += x[xidx] * w[widx];
                    }
                }
            }
        }
        y[idx] = acc;
    }
    y
}
fn cpu_gn(x: &[f32], gamma: &[f32], beta: &[f32], p: &GNParams) -> Vec<f32> {
    let spatial = (p.d * p.h * p.w) as usize;
    let cg = (p.c / p.groups) as usize;
    let mut y = vec![0.0f32; x.len()];
    for bn in 0..p.n as usize {
        for g in 0..p.groups as usize {
            let c0 = g * cg;
            let count = cg * spatial;
            let mut sum = 0.0f32;
            let mut sumsq = 0.0f32;
            for c in c0..c0 + cg {
                for s in 0..spatial {
                    let v = x[(bn * p.c as usize + c) * spatial + s];
                    sum += v; sumsq += v * v;
                }
            }
            let mean = sum / count as f32;
            let variance = sumsq / count as f32 - mean * mean;
            let inv = (variance + p.eps).sqrt().recip();
            for c in c0..c0 + cg {
                for s in 0..spatial {
                    let idx = (bn * p.c as usize + c) * spatial + s;
                    let norm = (x[idx] - mean) * inv;
                    let v = norm * gamma[c] + beta[c];
                    y[idx] = v / (1.0 + (-v).exp());
                }
            }
        }
    }
    y
}
fn cpu_up(x: &[f32], p: &UpParams) -> Vec<f32> {
    let mut y = vec![0.0f32; (p.n * p.c * p.dout * p.hout * p.wout) as usize];
    for idx in 0..y.len() {
        let mut t = idx;
        let ow = t % p.wout as usize; t /= p.wout as usize;
        let oh = t % p.hout as usize; t /= p.hout as usize;
        let od = t % p.dout as usize; t /= p.dout as usize;
        let co = t % p.c as usize;
        let bn = t / p.c as usize;
        let iw = ow / p.scale as usize;
        let ih = oh / p.scale as usize;
        let id = od / p.scale as usize;
        y[idx] = x[(((bn * p.c as usize + co) * p.din as usize + id) * p.hin as usize + ih) * p.win as usize + iw];
    }
    y
}

fn main() {
    let cp = ConvParams { n: 1, cin: 2, cout: 2, din: 2, hin: 2, win: 2,
                         kd: 3, kh: 3, kw: 3, pad: 1, stride: 1,
                         dout: 2, hout: 2, wout: 2, _p0: 0, _p1: 0 };
    let gp = GNParams { n: 1, c: 2, d: 2, h: 2, w: 2, groups: 1, eps: 1e-5, _p: 0 };
    let up = UpParams { n: 1, c: 2, din: 2, hin: 2, win: 2, scale: 2, dout: 4, hout: 4, wout: 4,
                        _p0: 0, _p1: 0, _p2: 0 };
    let xn = (cp.n * cp.cin * cp.din * cp.hin * cp.win) as usize;
    let wn = (cp.cout * cp.cin * cp.kd * cp.kh * cp.kw) as usize;
    let yn = (cp.n * cp.cout * cp.dout * cp.hout * cp.wout) as usize;
    let un = (up.n * up.c * up.dout * up.hout * up.wout) as usize;
    let mut x = vec![0.0f32; xn];
    let mut w = vec![0.0f32; wn];
    let mut b = vec![0.0f32; cp.cout as usize];
    let mut gamma = vec![0.0f32; gp.c as usize];
    let mut beta = vec![0.0f32; gp.c as usize];
    for i in 0..xn { x[i] = ((i * 13 % 71) as f32 - 35.0) / 19.0; }
    for i in 0..wn { w[i] = ((i * 17 % 53) as f32 - 26.0) / 31.0; }
    for i in 0..b.len() { b[i] = ((i * 7 % 11) as f32 - 5.0) * 0.01; }
    for i in 0..gamma.len() { gamma[i] = 1.0 + ((i * 5 % 9) as f32 - 4.0) * 0.02; }
    for i in 0..beta.len() { beta[i] = ((i * 3 % 7) as f32 - 3.0) * 0.01; }
    let conv_cpu = cpu_conv(&x, &w, &b, &cp);
    let gn_cpu = cpu_gn(&conv_cpu, &gamma, &beta, &gp);
    let expect = cpu_up(&gn_cpu, &up);

    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 (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
        label: Some("vae-mini"), required_features: wgpu::Features::empty(),
        required_limits: wgpu::Limits::downlevel_defaults(), ..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 w_buf = mk(&w); let b_buf = mk(&b);
    let conv_buf = device.create_buffer(&wgpu::BufferDescriptor {
        label: None, size: (yn * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE, mapped_at_creation: false,
    });
    let gamma_buf = mk(&gamma); let beta_buf = mk(&beta);
    let gn_buf = device.create_buffer(&wgpu::BufferDescriptor {
        label: None, size: (yn * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE, mapped_at_creation: false,
    });
    let up_buf = device.create_buffer(&wgpu::BufferDescriptor {
        label: None, size: (un * 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: (un * 4) as u64,
        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, mapped_at_creation: false,
    });
    let conv_p = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: None, contents: bytemuck::bytes_of(&cp), usage: wgpu::BufferUsages::UNIFORM,
    });
    let gn_p = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: None, contents: bytemuck::bytes_of(&gp), usage: wgpu::BufferUsages::UNIFORM,
    });
    let up_p = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: None, contents: bytemuck::bytes_of(&up), usage: wgpu::BufferUsages::UNIFORM,
    });

    let conv_mod = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: None, source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(CONV_WGSL)) });
    let gn_mod = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: None, source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(GN_WGSL)) });
    let up_mod = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: None, source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(UP_WGSL)) });
    let conv_pipe = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { label: None, layout: None, module: &conv_mod, entry_point: Some("main"), compilation_options: Default::default(), cache: None });
    let gn_pipe = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { label: None, layout: None, module: &gn_mod, entry_point: Some("main"), compilation_options: Default::default(), cache: None });
    let up_pipe = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { label: None, layout: None, module: &up_mod, entry_point: Some("main"), compilation_options: Default::default(), cache: None });

    let conv_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None, layout: &conv_pipe.get_bind_group_layout(0),
        entries: &[
            wgpu::BindGroupEntry { binding: 0, resource: x_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 1, resource: w_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 2, resource: b_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 3, resource: conv_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 4, resource: conv_p.as_entire_binding() },
        ],
    });
    let gn_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None, layout: &gn_pipe.get_bind_group_layout(0),
        entries: &[
            wgpu::BindGroupEntry { binding: 0, resource: conv_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 1, resource: gamma_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 2, resource: beta_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 3, resource: gn_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 4, resource: gn_p.as_entire_binding() },
        ],
    });
    let up_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None, layout: &up_pipe.get_bind_group_layout(0),
        entries: &[
            wgpu::BindGroupEntry { binding: 0, resource: gn_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 1, resource: up_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 2, resource: up_p.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(&conv_pipe); pass.set_bind_group(0, &conv_bg, &[]);
        pass.dispatch_workgroups(((yn + 63) / 64) as u32, 1, 1);
    }
    {
        let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None, timestamp_writes: None });
        pass.set_pipeline(&gn_pipe); pass.set_bind_group(0, &gn_bg, &[]);
        pass.dispatch_workgroups(gp.n * gp.groups, 1, 1);
    }
    {
        let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None, timestamp_writes: None });
        pass.set_pipeline(&up_pipe); pass.set_bind_group(0, &up_bg, &[]);
        pass.dispatch_workgroups(((un + 63) / 64) as u32, 1, 1);
    }
    enc.copy_buffer_to_buffer(&up_buf, 0, &readback, 0, (un * 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!("vae-mini conv3d->gn+silu->upsample3d out={} max_abs_err={max_err:.3e}", un);
    assert!(max_err < 1e-5, "vae mini mismatch");
    println!("VULKAN VAE MINI OK");
}