#[pyclass]
struct PySystem {
    system: RefCell<System>,
}

#[pymethods]
impl PySystem {
    #[staticmethod]
    fn from_pdb(path: &str) -> PyResult<Self> {
        let mut reader = PdbReader::new(path);
        let system = reader.read_system().map_err(to_py_err)?;
        Ok(Self {
            system: RefCell::new(system),
        })
    }

    #[staticmethod]
    fn from_gro(path: &str) -> PyResult<Self> {
        let mut reader = GroReader::new(path);
        let system = reader.read_system().map_err(to_py_err)?;
        Ok(Self {
            system: RefCell::new(system),
        })
    }

    fn select(&self, expr: &str) -> PyResult<PySelection> {
        let mut sys = self.system.borrow_mut();
        let selection = sys.select(expr).map_err(to_py_err)?;
        Ok(PySelection { selection })
    }

    fn select_indices(&self, indices: Vec<u32>) -> PyResult<PySelection> {
        let selection = Selection {
            expr: "__indices__".to_string(),
            indices: Arc::new(indices),
        };
        Ok(PySelection { selection })
    }

    fn n_atoms(&self) -> PyResult<usize> {
        Ok(self.system.borrow().n_atoms())
    }

    fn atom_table<'py>(&self, py: Python<'py>) -> PyResult<PyObject> {
        let sys = self.system.borrow();
        let atoms = &sys.atoms;
        let mut names = Vec::with_capacity(atoms.name_id.len());
        let mut resnames = Vec::with_capacity(atoms.resname_id.len());
        for &id in atoms.name_id.iter() {
            names.push(sys.interner.resolve(id).unwrap_or("").to_string());
        }
        for &id in atoms.resname_id.iter() {
            resnames.push(sys.interner.resolve(id).unwrap_or("").to_string());
        }
        let dict = PyDict::new_bound(py);
        dict.set_item("name", names)?;
        dict.set_item("resname", resnames)?;
        dict.set_item("resid", atoms.resid.clone())?;
        dict.set_item("chain_id", atoms.chain_id.clone())?;
        dict.set_item("mass", atoms.mass.clone())?;
        Ok(dict.into_py(py))
    }

    fn positions0<'py>(&self, py: Python<'py>) -> PyResult<PyObject> {
        let sys = self.system.borrow();
        let Some(pos) = &sys.positions0 else {
            return Ok(py.None());
        };
        let mut data = Vec::with_capacity(pos.len() * 3);
        for p in pos.iter() {
            data.extend_from_slice(&p[0..3]);
        }
        let arr = Array2::from_shape_vec((pos.len(), 3), data)
            .map_err(|_| PyRuntimeError::new_err("failed to build positions0 array"))?;
        Ok(arr.into_pyarray_bound(py).into_py(py))
    }
}


#[pyclass]
struct PySelection {
    selection: Selection,
}

#[pymethods]
impl PySelection {
    #[getter]
    fn indices<'py>(&self, py: Python<'py>) -> PyResult<PyObject> {
        let vec: Vec<u32> = self.selection.indices.as_ref().clone();
        Ok(vec.into_py(py))
    }
}

enum TrajKind {
    Dcd(DcdReader),
    Xtc(XtcReader),
}

#[pyclass(unsendable)]
struct PyTrajectory {
    inner: RefCell<TrajKind>,
}

#[pymethods]
impl PyTrajectory {
    #[staticmethod]
    #[pyo3(signature = (path, system, length_scale=None))]
    fn open_dcd(path: &str, system: &PySystem, length_scale: Option<f32>) -> PyResult<Self> {
        let reader = DcdReader::open(path, length_scale.unwrap_or(1.0)).map_err(to_py_err)?;
        let sys = system.system.borrow();
        if reader.n_atoms() != sys.n_atoms() {
            return Err(PyRuntimeError::new_err(
                "trajectory atom count does not match system",
            ));
        }
        Ok(Self {
            inner: RefCell::new(TrajKind::Dcd(reader)),
        })
    }

    #[staticmethod]
    fn open_xtc(path: &str, system: &PySystem) -> PyResult<Self> {
        let reader = XtcReader::open(path).map_err(to_py_err)?;
        let sys = system.system.borrow();
        if reader.n_atoms() != sys.n_atoms() {
            return Err(PyRuntimeError::new_err(
                "trajectory atom count does not match system",
            ));
        }
        Ok(Self {
            inner: RefCell::new(TrajKind::Xtc(reader)),
        })
    }

    fn n_atoms(&self) -> PyResult<usize> {
        let inner = self.inner.borrow();
        let n = match &*inner {
            TrajKind::Dcd(reader) => reader.n_atoms(),
            TrajKind::Xtc(reader) => reader.n_atoms(),
        };
        Ok(n)
    }

    #[pyo3(signature = (max_frames=128))]
    fn read_chunk<'py>(&self, py: Python<'py>, max_frames: usize) -> PyResult<PyObject> {
        let mut inner = self.inner.borrow_mut();
        let n_atoms = match &*inner {
            TrajKind::Dcd(reader) => reader.n_atoms(),
            TrajKind::Xtc(reader) => reader.n_atoms(),
        };
        let max_frames = max_frames.max(1);
        let mut builder = FrameChunkBuilder::new(n_atoms, max_frames);
        let frames = match &mut *inner {
            TrajKind::Dcd(reader) => reader.read_chunk(max_frames, &mut builder),
            TrajKind::Xtc(reader) => reader.read_chunk(max_frames, &mut builder),
        }
        .map_err(to_py_err)?;
        if frames == 0 {
            return Ok(py.None());
        }
        let chunk = builder.finish().map_err(to_py_err)?;
        let mut coord_data = Vec::with_capacity(chunk.n_frames * chunk.n_atoms * 3);
        for p in chunk.coords.iter() {
            coord_data.extend_from_slice(&p[0..3]);
        }
        let coords = Array3::from_shape_vec(
            (chunk.n_frames, chunk.n_atoms, 3),
            coord_data,
        )
        .map_err(|_| PyRuntimeError::new_err("failed to build coords array"))?;
        let coords_py = coords.into_pyarray_bound(py);

        let mut box_data = Vec::with_capacity(chunk.n_frames * 3);
        let mut box_ok = true;
        for b in chunk.box_.iter() {
            match *b {
                Box3::Orthorhombic { lx, ly, lz } => {
                    box_data.push(lx);
                    box_data.push(ly);
                    box_data.push(lz);
                }
                Box3::None => {
                    box_ok = false;
                    break;
                }
                Box3::Triclinic { .. } => {
                    box_ok = false;
                    break;
                }
            }
        }
        let box_py = if box_ok {
            let box_arr = Array2::from_shape_vec((chunk.n_frames, 3), box_data)
                .map_err(|_| PyRuntimeError::new_err("failed to build box array"))?;
            box_arr.into_pyarray_bound(py).into_py(py)
        } else {
            py.None()
        };

        let time_py = if let Some(times) = chunk.time_ps {
            PyArray1::from_vec_bound(py, times).into_py(py)
        } else {
            py.None()
        };

        let dict = PyDict::new_bound(py);
        dict.set_item("coords", coords_py)?;
        dict.set_item("box", box_py)?;
        dict.set_item("time_ps", time_py)?;
        dict.set_item("frames", chunk.n_frames)?;
        Ok(dict.into_py(py))
    }
}


#[pyclass]
struct PyRgPlan {
    plan: RefCell<RgPlan>,
}

#[pymethods]
impl PyRgPlan {
    #[new]
    fn new(selection: &PySelection, mass_weighted: Option<bool>) -> Self {
        let plan = RgPlan::new(selection.selection.clone(), mass_weighted.unwrap_or(false));
        Self {
            plan: RefCell::new(plan),
        }
    }

    #[pyo3(signature = (traj, system, chunk_frames=None, device="auto"))]
    fn run<'py>(
        &self,
        py: Python<'py>,
        traj: &PyTrajectory,
        system: &PySystem,
        chunk_frames: Option<usize>,
        device: &str,
    ) -> PyResult<&'py PyArray1<f32>> {
        let mut plan = self.plan.borrow_mut();
        let mut traj_ref = traj.inner.borrow_mut();
        let output = run_plan(&mut *plan, &mut traj_ref, &system.system.borrow(), chunk_frames, device)?;
        match output {
            PlanOutput::Series(values) => Ok(PyArray1::from_vec_bound(py, values).into_gil_ref()),
            _ => Err(PyRuntimeError::new_err("unexpected output")),
        }
    }
}


#[pyclass]
struct PyRadgyrTensorPlan {
    plan: RefCell<RadgyrTensorPlan>,
}

#[pymethods]
impl PyRadgyrTensorPlan {
    #[new]
    #[pyo3(signature = (selection, mass_weighted=false))]
    fn new(selection: &PySelection, mass_weighted: bool) -> Self {
        let plan = RadgyrTensorPlan::new(selection.selection.clone(), mass_weighted);
        Self {
            plan: RefCell::new(plan),
        }
    }

    #[pyo3(signature = (traj, system, chunk_frames=None, device="auto"))]
    fn run<'py>(
        &self,
        py: Python<'py>,
        traj: &PyTrajectory,
        system: &PySystem,
        chunk_frames: Option<usize>,
        device: &str,
    ) -> PyResult<&'py PyArray2<f32>> {
        let mut plan = self.plan.borrow_mut();
        let mut traj_ref = traj.inner.borrow_mut();
        let output = run_plan(&mut *plan, &mut traj_ref, &system.system.borrow(), chunk_frames, device)?;
        match output {
            PlanOutput::Matrix { data, rows, cols } => matrix_to_py(py, data, rows, cols),
            _ => Err(PyRuntimeError::new_err("unexpected output")),
        }
    }
}


#[pyclass]
struct PyRmsdPlan {
    plan: RefCell<RmsdPlan>,
}

#[pymethods]
impl PyRmsdPlan {
    #[new]
    #[pyo3(signature = (selection, reference="topology", align=true))]
    fn new(selection: &PySelection, reference: &str, align: bool) -> PyResult<Self> {
        let reference_mode = parse_reference(reference)?;
        let plan = RmsdPlan::new(selection.selection.clone(), reference_mode, align);
        Ok(Self {
            plan: RefCell::new(plan),
        })
    }

    #[pyo3(signature = (traj, system, chunk_frames=None, device="auto"))]
    fn run<'py>(
        &self,
        py: Python<'py>,
        traj: &PyTrajectory,
        system: &PySystem,
        chunk_frames: Option<usize>,
        device: &str,
    ) -> PyResult<&'py PyArray1<f32>> {
        let mut plan = self.plan.borrow_mut();
        let mut traj_ref = traj.inner.borrow_mut();
        let output = run_plan(&mut *plan, &mut traj_ref, &system.system.borrow(), chunk_frames, device)?;
        match output {
            PlanOutput::Series(values) => Ok(PyArray1::from_vec_bound(py, values).into_gil_ref()),
            _ => Err(PyRuntimeError::new_err("unexpected output")),
        }
    }
}


#[pyclass]
struct PySymmRmsdPlan {
    plan: RefCell<SymmRmsdPlan>,
}

#[pymethods]
impl PySymmRmsdPlan {
    #[new]
    #[pyo3(signature = (selection, reference="topology", align=true))]
    fn new(selection: &PySelection, reference: &str, align: bool) -> PyResult<Self> {
        let reference_mode = parse_reference(reference)?;
        let plan = SymmRmsdPlan::new(selection.selection.clone(), reference_mode, align);
        Ok(Self {
            plan: RefCell::new(plan),
        })
    }

    #[pyo3(signature = (traj, system, chunk_frames=None, device="auto"))]
    fn run<'py>(
        &self,
        py: Python<'py>,
        traj: &PyTrajectory,
        system: &PySystem,
        chunk_frames: Option<usize>,
        device: &str,
    ) -> PyResult<&'py PyArray1<f32>> {
        let mut plan = self.plan.borrow_mut();
        let mut traj_ref = traj.inner.borrow_mut();
        let output = run_plan(&mut *plan, &mut traj_ref, &system.system.borrow(), chunk_frames, device)?;
        match output {
            PlanOutput::Series(values) => Ok(PyArray1::from_vec_bound(py, values).into_gil_ref()),
            _ => Err(PyRuntimeError::new_err("unexpected output")),
        }
    }
}


#[pyclass]
struct PyDistanceRmsdPlan {
    plan: RefCell<DistanceRmsdPlan>,
}

#[pymethods]
impl PyDistanceRmsdPlan {
    #[new]
    #[pyo3(signature = (selection, reference="topology", pbc="none"))]
    fn new(selection: &PySelection, reference: &str, pbc: &str) -> PyResult<Self> {
        let reference_mode = parse_reference(reference)?;
        let pbc = parse_pbc(pbc)?;
        let plan = DistanceRmsdPlan::new(selection.selection.clone(), reference_mode, pbc);
        Ok(Self {
            plan: RefCell::new(plan),
        })
    }

    #[pyo3(signature = (traj, system, chunk_frames=None, device="auto"))]
    fn run<'py>(
        &self,
        py: Python<'py>,
        traj: &PyTrajectory,
        system: &PySystem,
        chunk_frames: Option<usize>,
        device: &str,
    ) -> PyResult<&'py PyArray1<f32>> {
        let mut plan = self.plan.borrow_mut();
        let mut traj_ref = traj.inner.borrow_mut();
        let output = run_plan(&mut *plan, &mut traj_ref, &system.system.borrow(), chunk_frames, device)?;
        match output {
            PlanOutput::Series(values) => Ok(PyArray1::from_vec_bound(py, values).into_gil_ref()),
            _ => Err(PyRuntimeError::new_err("unexpected output")),
        }
    }
}


#[pyclass]
struct PyPairwiseRmsdPlan {
    plan: RefCell<PairwiseRmsdPlan>,
}

#[pymethods]
impl PyPairwiseRmsdPlan {
    #[new]
    #[pyo3(signature = (selection, metric="rms", pbc="none"))]
    fn new(selection: &PySelection, metric: &str, pbc: &str) -> PyResult<Self> {
        let metric = parse_pairwise_metric(metric)?;
        let pbc = parse_pbc(pbc)?;
        let plan = PairwiseRmsdPlan::new(selection.selection.clone(), metric, pbc);
        Ok(Self {
            plan: RefCell::new(plan),
        })
    }

    #[pyo3(signature = (traj, system, chunk_frames=None, device="auto"))]
    fn run<'py>(
        &self,
        py: Python<'py>,
        traj: &PyTrajectory,
        system: &PySystem,
        chunk_frames: Option<usize>,
        device: &str,
    ) -> PyResult<&'py PyArray2<f32>> {
        let mut plan = self.plan.borrow_mut();
        let mut traj_ref = traj.inner.borrow_mut();
        let output = run_plan(&mut *plan, &mut traj_ref, &system.system.borrow(), chunk_frames, device)?;
        match output {
            PlanOutput::Matrix { data, rows, cols } => matrix_to_py(py, data, rows, cols),
            _ => Err(PyRuntimeError::new_err("unexpected output")),
        }
    }
}


#[pyclass]
struct PyMatrixPlan {
    plan: RefCell<MatrixPlan>,
}

#[pymethods]
impl PyMatrixPlan {
    #[new]
    #[pyo3(signature = (selection, mode, pbc="none"))]
    fn new(selection: &PySelection, mode: &str, pbc: &str) -> PyResult<Self> {
        let mode = parse_matrix_mode(mode)?;
        let pbc = parse_pbc(pbc)?;
        let plan = MatrixPlan::new(selection.selection.clone(), mode, pbc);
        Ok(Self {
            plan: RefCell::new(plan),
        })
    }

    #[pyo3(signature = (traj, system, chunk_frames=None, device="auto"))]
    fn run<'py>(
        &self,
        py: Python<'py>,
        traj: &PyTrajectory,
        system: &PySystem,
        chunk_frames: Option<usize>,
        device: &str,
    ) -> PyResult<&'py PyArray2<f32>> {
        let mut plan = self.plan.borrow_mut();
        let mut traj_ref = traj.inner.borrow_mut();
        let output = run_plan(&mut *plan, &mut traj_ref, &system.system.borrow(), chunk_frames, device)?;
        match output {
            PlanOutput::Matrix { data, rows, cols } => matrix_to_py(py, data, rows, cols),
            _ => Err(PyRuntimeError::new_err("unexpected output")),
        }
    }
}


#[pyclass]
struct PyPcaPlan {
    plan: RefCell<PcaPlan>,
}

#[pymethods]
impl PyPcaPlan {
    #[new]
    #[pyo3(signature = (selection, n_components, mass_weighted=false))]
    fn new(selection: &PySelection, n_components: usize, mass_weighted: bool) -> Self {
        let plan = PcaPlan::new(selection.selection.clone(), n_components, mass_weighted);
        Self {
            plan: RefCell::new(plan),
        }
    }

    #[pyo3(signature = (traj, system, chunk_frames=None, device="auto"))]
    fn run<'py>(
        &self,
        py: Python<'py>,
        traj: &PyTrajectory,
        system: &PySystem,
        chunk_frames: Option<usize>,
        device: &str,
    ) -> PyResult<PyObject> {
        let mut plan = self.plan.borrow_mut();
        let mut traj_ref = traj.inner.borrow_mut();
        let output = run_plan(&mut *plan, &mut traj_ref, &system.system.borrow(), chunk_frames, device)?;
        match output {
            PlanOutput::Pca(pca) => pca_to_py(py, pca),
            _ => Err(PyRuntimeError::new_err("unexpected output")),
        }
    }
}


#[pyclass]
struct PyAnalyzeModesPlan {
    plan: RefCell<AnalyzeModesPlan>,
}

#[pymethods]
impl PyAnalyzeModesPlan {
    #[new]
    #[pyo3(signature = (selection, n_components, mass_weighted=false))]
    fn new(selection: &PySelection, n_components: usize, mass_weighted: bool) -> Self {
        let plan = AnalyzeModesPlan::new(selection.selection.clone(), n_components, mass_weighted);
        Self {
            plan: RefCell::new(plan),
        }
    }

    #[pyo3(signature = (traj, system, chunk_frames=None, device="auto"))]
    fn run<'py>(
        &self,
        py: Python<'py>,
        traj: &PyTrajectory,
        system: &PySystem,
        chunk_frames: Option<usize>,
        device: &str,
    ) -> PyResult<PyObject> {
        let mut plan = self.plan.borrow_mut();
        let mut traj_ref = traj.inner.borrow_mut();
        let output = run_plan(&mut *plan, &mut traj_ref, &system.system.borrow(), chunk_frames, device)?;
        match output {
            PlanOutput::Pca(pca) => pca_to_py(py, pca),
            _ => Err(PyRuntimeError::new_err("unexpected output")),
        }
    }
}
