//! GPU-accelerated power iteration PCA via tiled matrix multiplication.
//!
//! The expensive part (matmul) runs on GPU; the cheap part (QR/normalize on
//! the small D×k matrix) stays on CPU via nalgebra.

use std::sync::Arc;

use super::buffers::DeviceExt;
use super::{GpuContext, GpuError};


/// Uniform parameters for the matmul shader.
#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct Dims {
    m: u32,
    n: u32,
    k: u32,
    _pad: u32,
}

/// GPU power iteration PCA pipeline.
pub struct GpuPowerIteration {
    pipeline: wgpu::ComputePipeline,
    dims_bind_group_layout: wgpu::BindGroupLayout,
    data_bind_group_layout: wgpu::BindGroupLayout,
    ctx: Arc<GpuContext>,
}

impl GpuPowerIteration {
    /// Create the matmul pipeline.
    pub fn new(ctx: Arc<GpuContext>) -> Result<Self, GpuError> {
        let shader = ctx.device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("pca_matmul"),
            source: wgpu::ShaderSource::Wgsl(include_str!("pca.wgsl").into()),
        });

        let data_bind_group_layout =
            ctx.device
                .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                    label: Some("matmul_data"),
                    entries: &[
                        wgpu::BindGroupLayoutEntry {
                            binding: 0,
                            visibility: wgpu::ShaderStages::COMPUTE,
                            ty: wgpu::BindingType::Buffer {
                                ty: wgpu::BufferBindingType::Storage { read_only: true },
                                has_dynamic_offset: false,
                                min_binding_size: None,
                            },
                            count: None,
                        },
                        wgpu::BindGroupLayoutEntry {
                            binding: 1,
                            visibility: wgpu::ShaderStages::COMPUTE,
                            ty: wgpu::BindingType::Buffer {
                                ty: wgpu::BufferBindingType::Storage { read_only: true },
                                has_dynamic_offset: false,
                                min_binding_size: None,
                            },
                            count: None,
                        },
                        wgpu::BindGroupLayoutEntry {
                            binding: 2,
                            visibility: wgpu::ShaderStages::COMPUTE,
                            ty: wgpu::BindingType::Buffer {
                                ty: wgpu::BufferBindingType::Storage { read_only: false },
                                has_dynamic_offset: false,
                                min_binding_size: None,
                            },
                            count: None,
                        },
                    ],
                });

        let dims_bind_group_layout =
            ctx.device
                .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                    label: Some("matmul_dims"),
                    entries: &[wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Uniform,
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        count: None,
                    }],
                });

        let pipeline_layout =
            ctx.device
                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                    label: Some("pca_layout"),
                    bind_group_layouts: &[&data_bind_group_layout, &dims_bind_group_layout],
                    immediate_size: 0,
                });

        let pipeline = ctx
            .device
            .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("pca_matmul"),
                layout: Some(&pipeline_layout),
                module: &shader,
                entry_point: Some("matmul"),
                compilation_options: Default::default(),
                cache: None,
            });

        Ok(GpuPowerIteration {
            pipeline,
            dims_bind_group_layout,
            data_bind_group_layout,
            ctx,
        })
    }

    /// GPU matrix multiply: C = A @ B
    /// A: (m × k), B: (k × n), C: (m × n) — all row-major f32.
    fn gpu_matmul(
        &self,
        a: &[f32],
        b: &[f32],
        m: usize,
        k: usize,
        n: usize,
    ) -> Result<Vec<f32>, GpuError> {
        let device = &self.ctx.device;
        let queue = &self.ctx.queue;

        let buf_a = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("mat_a"),
            contents: bytemuck::cast_slice(a),
            usage: wgpu::BufferUsages::STORAGE,
        });
        let buf_b = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("mat_b"),
            contents: bytemuck::cast_slice(b),
            usage: wgpu::BufferUsages::STORAGE,
        });
        let buf_c = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("mat_c"),
            size: (m * n * std::mem::size_of::<f32>()) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        let staging = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("staging"),
            size: (m * n * std::mem::size_of::<f32>()) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let data_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("matmul_data_bg"),
            layout: &self.data_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry { binding: 0, resource: buf_a.as_entire_binding() },
                wgpu::BindGroupEntry { binding: 1, resource: buf_b.as_entire_binding() },
                wgpu::BindGroupEntry { binding: 2, resource: buf_c.as_entire_binding() },
            ],
        });

        let dims = Dims {
            m: m as u32,
            n: n as u32,
            k: k as u32,
            _pad: 0,
        };
        let dims_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("dims"),
            contents: bytemuck::bytes_of(&dims),
            usage: wgpu::BufferUsages::UNIFORM,
        });
        let dims_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("dims_bg"),
            layout: &self.dims_bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: dims_buf.as_entire_binding(),
            }],
        });

        let tile = 16u32;
        let wg_x = (m as u32 + tile - 1) / tile;
        let wg_y = (n as u32 + tile - 1) / tile;

        let mut encoder = device.create_command_encoder(&Default::default());
        {
            let mut pass = encoder.begin_compute_pass(&Default::default());
            pass.set_pipeline(&self.pipeline);
            pass.set_bind_group(0, &data_bg, &[]);
            pass.set_bind_group(1, &dims_bg, &[]);
            pass.dispatch_workgroups(wg_x, wg_y, 1);
        }
        encoder.copy_buffer_to_buffer(&buf_c, 0, &staging, 0, (m * n * 4) as u64);
        queue.submit(std::iter::once(encoder.finish()));

        let slice = staging.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |r| { tx.send(r).unwrap(); });
        let _ = device.poll(wgpu::PollType::wait_indefinitely());
        rx.recv()
            .unwrap()
            .map_err(|e| GpuError::ValidationError(format!("Map failed: {e:?}")))?;

        let data = slice.get_mapped_range();
        let result: Vec<f32> = bytemuck::cast_slice(&data).to_vec();
        drop(data);
        staging.unmap();

        Ok(result)
    }

    /// Compute top-k PCA components via power iteration.
    ///
    /// `xtx`: precomputed covariance matrix X^T X, row-major (d × d)
    /// `d`: embedding dimension
    /// `k`: number of components
    /// `max_iter`: power iteration steps (20 is typical)
    /// `seed`: for deterministic initialization
    ///
    /// Returns (components as k×d row-major, eigenvalues).
    pub fn power_iteration(
        &self,
        xtx: &[f32],
        d: usize,
        k: usize,
        max_iter: usize,
        seed: u64,
    ) -> Result<(Vec<f32>, Vec<f32>), GpuError> {
        use rand::prelude::*;
        use rand_chacha::ChaCha8Rng;

        assert_eq!(xtx.len(), d * d, "xtx must be d×d");

        let mut components: Vec<Vec<f32>> = Vec::with_capacity(k);
        let mut eigenvalues = Vec::with_capacity(k);
        let mut rng = ChaCha8Rng::seed_from_u64(seed + 999);

        for _comp_idx in 0..k {
            // Initialize random vector (d×1)
            let mut v: Vec<f32> = (0..d).map(|_| rng.gen_range(-1.0f32..1.0f32)).collect();
            let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
            v.iter_mut().for_each(|x| *x /= norm);

            for _iter in 0..max_iter {
                // GPU: xtxv = xtx @ v  (d×d) @ (d×1) = (d×1)
                let mut xtxv = self.gpu_matmul(xtx, &v, d, d, 1)?;

                // CPU: deflation — orthogonalize against previous components
                for prev in &components {
                    let dot: f32 = xtxv.iter().zip(prev.iter()).map(|(a, b)| a * b).sum();
                    for (x, p) in xtxv.iter_mut().zip(prev.iter()) {
                        *x -= dot * p;
                    }
                }

                // CPU: normalize
                let norm: f32 = xtxv.iter().map(|x| x * x).sum::<f32>().sqrt();
                if norm > 1e-10 {
                    v = xtxv.iter().map(|x| x / norm).collect();
                }
            }

            // Eigenvalue: λ = v^T (X^T X) v
            let xtxv = self.gpu_matmul(xtx, &v, d, d, 1)?;
            let eigenvalue: f32 = v.iter().zip(xtxv.iter()).map(|(a, b)| a * b).sum();
            eigenvalues.push(eigenvalue);

            components.push(v);
        }

        // Pack into k×d row-major
        let mut result = Vec::with_capacity(k * d);
        for comp in &components {
            result.extend_from_slice(comp);
        }

        Ok((result, eigenvalues))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn gpu_matmul_identity() {
        // Multiply by identity matrix — result should be input
        let ctx = match crate::gpu::GpuContext::new() {
            Ok(ctx) => Arc::new(ctx),
            Err(_) => {
                eprintln!("No GPU available, skipping test");
                return;
            }
        };
        let pca = GpuPowerIteration::new(ctx).unwrap();

        let d = 4;
        // Identity matrix
        let mut eye = vec![0.0f32; d * d];
        for i in 0..d {
            eye[i * d + i] = 1.0;
        }
        // Input vector
        let v = vec![1.0f32, 2.0, 3.0, 4.0];

        let result = pca.gpu_matmul(&eye, &v, d, d, 1).unwrap();
        for i in 0..d {
            assert!(
                (result[i] - v[i]).abs() < 1e-5,
                "matmul identity: result[{i}]={} != v[{i}]={}",
                result[i],
                v[i]
            );
        }
    }

    #[test]
    fn gpu_pca_convergence() {
        // Create a simple covariance matrix with known eigenstructure
        let ctx = match crate::gpu::GpuContext::new() {
            Ok(ctx) => Arc::new(ctx),
            Err(_) => {
                eprintln!("No GPU available, skipping test");
                return;
            }
        };
        let pca = GpuPowerIteration::new(ctx).unwrap();

        let d = 8;
        // Diagonal covariance with well-separated eigenvalues (powers of 2)
        let mut xtx = vec![0.0f32; d * d];
        let eigen_vals = [100.0, 50.0, 10.0, 5.0, 2.0, 1.0, 0.5, 0.1];
        for i in 0..d {
            xtx[i * d + i] = eigen_vals[i];
        }

        let k = 3;
        let (components, eigenvalues) = pca.power_iteration(&xtx, d, k, 30, 42).unwrap();

        // Eigenvalues should be approximately 100, 50, 10
        assert!(
            (eigenvalues[0] - 100.0).abs() < 1.0,
            "First eigenvalue: {} (expected ~100)",
            eigenvalues[0]
        );
        assert!(
            (eigenvalues[1] - 50.0).abs() < 1.0,
            "Second eigenvalue: {} (expected ~50)",
            eigenvalues[1]
        );
        assert!(
            (eigenvalues[2] - 10.0).abs() < 1.0,
            "Third eigenvalue: {} (expected ~10)",
            eigenvalues[2]
        );

        // Each component should be close to a standard basis vector
        for comp_idx in 0..k {
            let comp = &components[comp_idx * d..(comp_idx + 1) * d];
            let max_idx = comp
                .iter()
                .enumerate()
                .max_by(|(_, a), (_, b)| a.abs().partial_cmp(&b.abs()).unwrap())
                .unwrap()
                .0;
            // The dominant component should align with the corresponding axis
            assert!(
                comp[max_idx].abs() > 0.99,
                "Component {comp_idx} max={:.3} at idx {max_idx}",
                comp[max_idx].abs()
            );
        }
    }
}
