use traj_core::error::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, GpuSelection};

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

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

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

    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 selection = ctx.selection(&self.selection.indices, None)?;
                self.gpu = Some(selection);
            }
        }
        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) {
            let coords = convert_coords(&chunk.coords);
            let out = ctx.bbox_area(&coords, chunk.n_atoms, chunk.n_frames, gpu)?;
            self.results.extend(out);
            return Ok(());
        }
        let n_atoms = chunk.n_atoms;
        let sel = &self.selection.indices;
        for frame in 0..chunk.n_frames {
            if sel.is_empty() {
                self.results.push(0.0);
                continue;
            }
            let mut min = [f64::INFINITY; 3];
            let mut max = [f64::NEG_INFINITY; 3];
            for &idx in sel.iter() {
                let p = chunk.coords[frame * n_atoms + idx as usize];
                let x = p[0] as f64;
                let y = p[1] as f64;
                let z = p[2] as f64;
                if x < min[0] {
                    min[0] = x;
                }
                if y < min[1] {
                    min[1] = y;
                }
                if z < min[2] {
                    min[2] = z;
                }
                if x > max[0] {
                    max[0] = x;
                }
                if y > max[1] {
                    max[1] = y;
                }
                if z > max[2] {
                    max[2] = z;
                }
            }
            let dx = (max[0] - min[0]).max(0.0);
            let dy = (max[1] - min[1]).max(0.0);
            let dz = (max[2] - min[2]).max(0.0);
            let area = 2.0 * (dx * dy + dy * dz + dx * dz);
            self.results.push(area as f32);
        }
        Ok(())
    }

    fn finalize(&mut self) -> TrajResult<PlanOutput> {
        Ok(PlanOutput::Series(std::mem::take(&mut self.results)))
    }
}

pub struct MolSurfPlan {
    inner: SurfPlan,
}

impl MolSurfPlan {
    pub fn new(selection: Selection) -> Self {
        Self {
            inner: SurfPlan::new(selection),
        }
    }
}

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

    fn init(&mut self, system: &System, device: &Device) -> TrajResult<()> {
        self.inner.init(system, device)
    }

    fn process_chunk(
        &mut self,
        chunk: &FrameChunk,
        system: &System,
        device: &Device,
    ) -> TrajResult<()> {
        self.inner.process_chunk(chunk, system, device)
    }

    fn finalize(&mut self) -> TrajResult<PlanOutput> {
        self.inner.finalize()
    }
}
