use std::collections::HashMap;

use traj_core::error::{TrajError, TrajResult};
use traj_core::frame::FrameChunk;
use traj_core::selection::Selection;
use traj_core::system::System;

use crate::executor::{Device, Plan, PlanOutput};

#[cfg(feature = "cuda")]
use traj_gpu::{convert_coords, GpuBufferF32, GpuBufferU32, GpuGroups, GpuSelection};

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct TorsionStat {
    pub trans: f32,
    pub cis: f32,
    pub g_plus: f32,
    pub g_minus: f32,
}

pub struct TorsionDiffusionPlan {
    selection: Selection,
    results: Vec<TorsionStat>,
    #[cfg(feature = "cuda")]
    gpu: Option<TorsionGpuState>,
}

pub struct ToroidalDiffusionPlan {
    selection: Selection,
    results: Vec<TorsionStat>,
    #[cfg(feature = "cuda")]
    gpu: Option<TorsionGpuState>,
}

pub struct MultiPuckerPlan {
    selection: Selection,
    results: Vec<f32>,
    frames: usize,
    bins: usize,
    #[cfg(feature = "cuda")]
    gpu: Option<MultiPuckerGpuState>,
}

#[cfg(feature = "cuda")]
struct MultiPuckerGpuState {
    sel: GpuSelection,
    groups: GpuGroups,
    masses: GpuBufferF32,
}

#[cfg(feature = "cuda")]
struct TorsionGpuState {
    torsions: GpuBufferU32,
    n_torsions: usize,
}

impl TorsionDiffusionPlan {
    pub fn new(selection: Selection) -> Self {
        Self {
            selection,
            results: Vec::new(),
            #[cfg(feature = "cuda")]
            gpu: None,
        }
    }
}

impl ToroidalDiffusionPlan {
    pub fn new(selection: Selection) -> Self {
        Self {
            selection,
            results: Vec::new(),
            #[cfg(feature = "cuda")]
            gpu: None,
        }
    }
}

impl MultiPuckerPlan {
    pub fn new(selection: Selection, bins: usize) -> Self {
        Self {
            selection,
            results: Vec::new(),
            frames: 0,
            bins: bins.max(1),
            #[cfg(feature = "cuda")]
            gpu: None,
        }
    }
}

impl Plan for TorsionDiffusionPlan {
    fn name(&self) -> &'static str {
        "tordiff"
    }

    fn init(&mut self, _system: &System, _device: &Device) -> TrajResult<()> {
        self.results.clear();
        #[cfg(feature = "cuda")]
        {
            self.gpu = None;
            if let Device::Cuda(ctx) = _device {
                let n_torsions = self.selection.indices.len() / 4;
                if n_torsions > 0 {
                    let indices = self.selection.indices[..n_torsions * 4].to_vec();
                    let torsions = ctx.upload_u32(&indices)?;
                    self.gpu = Some(TorsionGpuState { torsions, n_torsions });
                }
            }
        }
        Ok(())
    }

    fn process_chunk(
        &mut self,
        chunk: &FrameChunk,
        system: &System,
        _device: &Device,
    ) -> TrajResult<()> {
        #[cfg(feature = "cuda")]
        if let (Device::Cuda(ctx), Some(gpu)) = (_device, &self.gpu) {
            if gpu.n_torsions == 0 {
                self.results.extend(
                    std::iter::repeat(TorsionStat {
                        trans: 0.0,
                        cis: 0.0,
                        g_plus: 0.0,
                        g_minus: 0.0,
                    })
                    .take(chunk.n_frames),
                );
                return Ok(());
            }
            let coords = convert_coords(&chunk.coords);
            let counts = ctx.torsion_diffusion_counts(
                &coords,
                chunk.n_atoms,
                chunk.n_frames,
                &gpu.torsions,
                gpu.n_torsions,
            )?;
            let denom = gpu.n_torsions.max(1) as f32;
            for frame in 0..chunk.n_frames {
                let base = frame * 4;
                self.results.push(TorsionStat {
                    trans: counts[base] as f32 / denom,
                    cis: counts[base + 1] as f32 / denom,
                    g_plus: counts[base + 2] as f32 / denom,
                    g_minus: counts[base + 3] as f32 / denom,
                });
            }
            return Ok(());
        }
        let stats = torsion_stats(chunk, system, &self.selection)?;
        self.results.extend(stats);
        Ok(())
    }

    fn finalize(&mut self) -> TrajResult<PlanOutput> {
        let mut data = Vec::with_capacity(self.results.len() * 4);
        for stat in self.results.drain(..) {
            data.push(stat.trans);
            data.push(stat.cis);
            data.push(stat.g_plus);
            data.push(stat.g_minus);
        }
        let rows = data.len() / 4;
        Ok(PlanOutput::Matrix {
            data,
            rows,
            cols: 4,
        })
    }
}

impl Plan for ToroidalDiffusionPlan {
    fn name(&self) -> &'static str {
        "toroidal_diffusion"
    }

    fn init(&mut self, _system: &System, _device: &Device) -> TrajResult<()> {
        self.results.clear();
        #[cfg(feature = "cuda")]
        {
            self.gpu = None;
            if let Device::Cuda(ctx) = _device {
                let n_torsions = self.selection.indices.len() / 4;
                if n_torsions > 0 {
                    let indices = self.selection.indices[..n_torsions * 4].to_vec();
                    let torsions = ctx.upload_u32(&indices)?;
                    self.gpu = Some(TorsionGpuState { torsions, n_torsions });
                }
            }
        }
        Ok(())
    }

    fn process_chunk(
        &mut self,
        chunk: &FrameChunk,
        system: &System,
        _device: &Device,
    ) -> TrajResult<()> {
        #[cfg(feature = "cuda")]
        if let (Device::Cuda(ctx), Some(gpu)) = (_device, &self.gpu) {
            if gpu.n_torsions == 0 {
                self.results.extend(
                    std::iter::repeat(TorsionStat {
                        trans: 0.0,
                        cis: 0.0,
                        g_plus: 0.0,
                        g_minus: 0.0,
                    })
                    .take(chunk.n_frames),
                );
                return Ok(());
            }
            let coords = convert_coords(&chunk.coords);
            let counts = ctx.torsion_diffusion_counts(
                &coords,
                chunk.n_atoms,
                chunk.n_frames,
                &gpu.torsions,
                gpu.n_torsions,
            )?;
            let denom = gpu.n_torsions.max(1) as f32;
            for frame in 0..chunk.n_frames {
                let base = frame * 4;
                self.results.push(TorsionStat {
                    trans: counts[base] as f32 / denom,
                    cis: counts[base + 1] as f32 / denom,
                    g_plus: counts[base + 2] as f32 / denom,
                    g_minus: counts[base + 3] as f32 / denom,
                });
            }
            return Ok(());
        }
        let stats = torsion_stats(chunk, system, &self.selection)?;
        self.results.extend(stats);
        Ok(())
    }

    fn finalize(&mut self) -> TrajResult<PlanOutput> {
        let mut data = Vec::with_capacity(self.results.len() * 4);
        for stat in self.results.drain(..) {
            data.push(stat.trans);
            data.push(stat.cis);
            data.push(stat.g_plus);
            data.push(stat.g_minus);
        }
        let rows = data.len() / 4;
        Ok(PlanOutput::Matrix {
            data,
            rows,
            cols: 4,
        })
    }
}

impl Plan for MultiPuckerPlan {
    fn name(&self) -> &'static str {
        "multipucker"
    }

    fn init(&mut self, _system: &System, _device: &Device) -> TrajResult<()> {
        self.results.clear();
        self.frames = 0;
        #[cfg(feature = "cuda")]
        {
            self.gpu = None;
            if let Device::Cuda(ctx) = _device {
                if !self.selection.indices.is_empty() {
                    let offsets = vec![0u32, self.selection.indices.len() as u32];
                    let indices = self.selection.indices.clone();
                    let groups = ctx.groups(&offsets, &indices, indices.len())?;
                    let masses = vec![1.0f32; _system.n_atoms()];
                    let masses = ctx.upload_f32(&masses)?;
                    let sel = ctx.selection(&self.selection.indices, None)?;
                    self.gpu = Some(MultiPuckerGpuState { sel, groups, masses });
                }
            }
        }
        Ok(())
    }

    fn process_chunk(
        &mut self,
        chunk: &FrameChunk,
        _system: &System,
        _device: &Device,
    ) -> TrajResult<()> {
        #[cfg(feature = "cuda")]
        if let (Device::Cuda(ctx), Some(gpu)) = (_device, &self.gpu) {
            if self.selection.indices.len() < 5 {
                self.frames += chunk.n_frames;
                return Ok(());
            }
            let coords = convert_coords(&chunk.coords);
            let coms = ctx.group_com(
                &coords,
                chunk.n_atoms,
                chunk.n_frames,
                &gpu.groups,
                &gpu.masses,
                1.0,
            )?;
            let maxes = ctx.max_dist_points(
                &coords,
                chunk.n_atoms,
                chunk.n_frames,
                &gpu.sel,
                &coms,
            )?;
            for max_r in maxes {
                let bin = if max_r == 0.0 { 0 } else { self.bins - 1 };
                for i in 0..self.bins {
                    self.results.push(if i == bin { 1.0 } else { 0.0 });
                }
                self.frames += 1;
            }
            return Ok(());
        }
        let n_atoms = chunk.n_atoms;
        let sel = &self.selection.indices;
        if sel.len() < 5 {
            self.frames += chunk.n_frames;
            return Ok(());
        }
        for frame in 0..chunk.n_frames {
            let mut values = Vec::with_capacity(sel.len());
            for &idx in sel.iter() {
                let p = chunk.coords[frame * n_atoms + idx as usize];
                values.push([p[0] as f64, p[1] as f64, p[2] as f64]);
            }
            let mut sum = [0.0f64; 3];
            for v in values.iter() {
                sum[0] += v[0];
                sum[1] += v[1];
                sum[2] += v[2];
            }
            let n = values.len() as f64;
            let center = [sum[0] / n, sum[1] / n, sum[2] / n];
            let mut max_r = 0.0f64;
            for v in values.iter() {
                let dx = v[0] - center[0];
                let dy = v[1] - center[1];
                let dz = v[2] - center[2];
                let r = (dx * dx + dy * dy + dz * dz).sqrt();
                if r > max_r {
                    max_r = r;
                }
            }
            let bin = if max_r == 0.0 { 0 } else { self.bins - 1 };
            for i in 0..self.bins {
                self.results.push(if i == bin { 1.0 } else { 0.0 });
            }
            self.frames += 1;
        }
        Ok(())
    }

    fn finalize(&mut self) -> TrajResult<PlanOutput> {
        Ok(PlanOutput::Matrix {
            data: std::mem::take(&mut self.results),
            rows: self.frames,
            cols: self.bins,
        })
    }
}

fn torsion_stats(
    chunk: &FrameChunk,
    system: &System,
    selection: &Selection,
) -> TrajResult<Vec<TorsionStat>> {
    let _ = chunk.n_atoms;
    let sel = &selection.indices;
    if sel.len() < 4 {
        return Ok(vec![TorsionStat {
            trans: 0.0,
            cis: 0.0,
            g_plus: 0.0,
            g_minus: 0.0,
        }; chunk.n_frames]);
    }
    let masses = &system.atoms.mass;
    let mut out = Vec::with_capacity(chunk.n_frames);
    for frame in 0..chunk.n_frames {
        let mut counts = HashMap::from([
            ("trans", 0usize),
            ("cis", 0usize),
            ("g_plus", 0usize),
            ("g_minus", 0usize),
        ]);
        let mut total = 0usize;
        for group in sel.chunks(4) {
            if group.len() < 4 {
                break;
            }
            let a = center_of_selection(chunk, frame, &group[0..1], masses, false);
            let b = center_of_selection(chunk, frame, &group[1..2], masses, false);
            let c = center_of_selection(chunk, frame, &group[2..3], masses, false);
            let d = center_of_selection(chunk, frame, &group[3..4], masses, false);
            let angle = dihedral_value(a, b, c, d)?;
            let deg = angle.to_degrees();
            if deg > 150.0 || deg < -150.0 {
                *counts.get_mut("trans").unwrap() += 1;
            } else if deg.abs() < 30.0 {
                *counts.get_mut("cis").unwrap() += 1;
            } else if deg > 0.0 {
                *counts.get_mut("g_plus").unwrap() += 1;
            } else {
                *counts.get_mut("g_minus").unwrap() += 1;
            }
            total += 1;
        }
        let denom = total.max(1) as f32;
        out.push(TorsionStat {
            trans: *counts.get("trans").unwrap() as f32 / denom,
            cis: *counts.get("cis").unwrap() as f32 / denom,
            g_plus: *counts.get("g_plus").unwrap() as f32 / denom,
            g_minus: *counts.get("g_minus").unwrap() as f32 / denom,
        });
    }
    Ok(out)
}

fn center_of_selection(
    chunk: &FrameChunk,
    frame: usize,
    indices: &[u32],
    masses: &[f32],
    mass_weighted: bool,
) -> [f64; 3] {
    let n_atoms = chunk.n_atoms;
    let mut sum = [0.0f64; 3];
    let mut mass_sum = 0.0f64;
    for &idx in indices.iter() {
        let atom_idx = idx as usize;
        let p = chunk.coords[frame * n_atoms + atom_idx];
        let m = if mass_weighted {
            masses[atom_idx] as f64
        } else {
            1.0
        };
        sum[0] += p[0] as f64 * m;
        sum[1] += p[1] as f64 * m;
        sum[2] += p[2] as f64 * m;
        mass_sum += m;
    }
    if mass_sum == 0.0 {
        return [0.0, 0.0, 0.0];
    }
    [sum[0] / mass_sum, sum[1] / mass_sum, sum[2] / mass_sum]
}

fn dihedral_value(a: [f64; 3], b: [f64; 3], c: [f64; 3], d: [f64; 3]) -> TrajResult<f64> {
    let b0 = [a[0] - b[0], a[1] - b[1], a[2] - b[2]];
    let b1 = [c[0] - b[0], c[1] - b[1], c[2] - b[2]];
    let b2 = [d[0] - c[0], d[1] - c[1], d[2] - c[2]];
    let b1_norm = {
        let n = (b1[0] * b1[0] + b1[1] * b1[1] + b1[2] * b1[2]).sqrt();
        if n == 0.0 {
            return Err(TrajError::Mismatch("dihedral axis has zero length".into()));
        }
        [b1[0] / n, b1[1] / n, b1[2] / n]
    };
    let dot_b0 = b0[0] * b1_norm[0] + b0[1] * b1_norm[1] + b0[2] * b1_norm[2];
    let dot_b2 = b2[0] * b1_norm[0] + b2[1] * b1_norm[1] + b2[2] * b1_norm[2];
    let v = [
        b0[0] - dot_b0 * b1_norm[0],
        b0[1] - dot_b0 * b1_norm[1],
        b0[2] - dot_b0 * b1_norm[2],
    ];
    let w = [
        b2[0] - dot_b2 * b1_norm[0],
        b2[1] - dot_b2 * b1_norm[1],
        b2[2] - dot_b2 * b1_norm[2],
    ];
    let x = v[0] * w[0] + v[1] * w[1] + v[2] * w[2];
    let m = [
        b1_norm[1] * v[2] - b1_norm[2] * v[1],
        b1_norm[2] * v[0] - b1_norm[0] * v[2],
        b1_norm[0] * v[1] - b1_norm[1] * v[0],
    ];
    let y = m[0] * w[0] + m[1] * w[1] + m[2] * w[2];
    Ok(y.atan2(x))
}
