//! Python interface for taco_format using PyO3.
//!
//! This module exposes functions to save and load ASE Atoms objects from Python.
//!
//! Example::
//!
//!     >>> import ase
//!     >>> import taco_format
//!     >>> atoms = ase.Atoms('H2O')
//!     >>> taco_format.save_ase('test.taco', atoms)
//!     >>> atoms2 = taco_format.load_ase('test.taco')
//!     >>> assert atoms == atoms2
//!
//! For trajectories::
//!
//!     >>> atoms_list = [ase.Atoms('H2O') for _ in range(10)]
//!     >>> taco_format.save_trajectory('traj.taco', atoms_list)
//!     >>> atoms_list2 = taco_format.load_trajectory('traj.taco')

use std::ffi::CString;

use crate::Header;
use crate::compression::CompressionSettings;
use crate::frame::{Frame, FrameData};
use crate::io::{Reader, Writer};
use crate::metadata::{AtomMetadata, SimulationMetadata};
use numpy::{PyArray1, PyArray2, PyArray3, PyArrayMethods, PyReadonlyArray1};
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyAny, PyDict, PyList, PyModule};

/// Convert ASE Atoms object to a TACO Frame
fn atoms_to_frame(
    py: Python,
    atoms: &Bound<PyAny>,
    frame_number: u64,
    time: f64,
) -> PyResult<Frame> {
    // Extract positions from atoms
    let positions = atoms.getattr("positions")?;
    let positions_array = positions.extract::<Bound<PyArray2<f64>>>()?;
    let positions_f32 = positions_array.readonly().as_array().mapv(|x| x as f32);
    let mut frame_data = FrameData::new(positions_f32);

    // Extract velocities if available
    if atoms.hasattr("get_velocities")? {
        let get_velocities_method = atoms.getattr("get_velocities")?;
        let velocities = get_velocities_method.call0()?;
        if !velocities.is_none() {
            let velocities_array = velocities.extract::<Bound<PyArray2<f64>>>()?;
            let velocities_f32 = velocities_array.readonly().as_array().mapv(|x| x as f32);
            frame_data = frame_data.with_velocities(velocities_f32);
        }
    }

    // Extract forces if available
    if atoms.hasattr("get_forces")? {
        let get_forces_method = atoms.getattr("get_forces")?;
        match get_forces_method.call0() {
            Ok(f) if !f.is_none() => {
                let forces_array = f.extract::<Bound<PyArray2<f64>>>()?;
                let forces_f32 = forces_array.readonly().as_array().mapv(|x| x as f32);
                frame_data = frame_data.with_forces(forces_f32);
            }
            Err(e) => {
                eprintln!("Warning: failed to extract forces: {}", e);
            }
            _ => {}
        }
    }

    // Extract stress if available using get_stress (not from info/arrays)
    if atoms.hasattr("get_stress")? {
        let get_stress_method = atoms.getattr("get_stress")?;
        match get_stress_method.call1((false,)) {
            Ok(s) if !s.is_none() => {
                if let Ok(stress_array) = s.extract::<Bound<PyArray1<f64>>>() {
                    let stress_vec = stress_array.readonly().as_array().mapv(|x| x as f32);
                    let stress_mat = stress_vec.insert_axis(ndarray::Axis(0));
                    frame_data = frame_data.with_stress(stress_mat);
                } else if let Ok(stress_array) = s.extract::<Bound<PyArray1<f32>>>() {
                    let stress_vec = stress_array.readonly().as_array().to_owned();
                    let stress_mat = stress_vec.insert_axis(ndarray::Axis(0));
                    frame_data = frame_data.with_stress(stress_mat);
                } else if let Ok(stress_array) = s.extract::<Bound<PyArray2<f64>>>() {
                    let stress_f32 = stress_array.readonly().as_array().mapv(|x| x as f32);
                    frame_data = frame_data.with_stress(stress_f32);
                } else if let Ok(stress_array) = s.extract::<Bound<PyArray2<f32>>>() {
                    let stress_f32 = stress_array.readonly().as_array().to_owned();
                    frame_data = frame_data.with_stress(stress_f32);
                }
            }
            Err(e) => {
                eprintln!("Warning: failed to extract stress: {}", e);
            }
            _ => {}
        }
    }

    // Extract all extra arrays (Array1, Array2, Array3) from atoms.arrays (excluding positions, velocities, forces, stress)
    if let Ok(arrays) = atoms.getattr("arrays") {
        let keys = arrays.call_method0("keys")?;
        for key in keys.try_iter()? {
            let key_str = key?.extract::<String>()?;
            if key_str != "positions"
                && key_str != "velocities"
                && key_str != "forces"
                && key_str != "stress"
                && key_str != "numbers"
            {
                let arr = arrays.get_item(&key_str)?;
                if !arr.is_none() {
                    macro_rules! try_extract_array {
                        ($ty:ty, $with_fn:ident) => {
                            if let Ok(array) = arr.extract::<Bound<$ty>>() {
                                let arr_f32 = array.readonly().as_array().mapv(|x| x as f32);
                                frame_data = frame_data.$with_fn(&key_str, arr_f32);
                                continue;
                            }
                        };
                    }
                    // 1D
                    try_extract_array!(numpy::PyArray1<f32>, with_extra_1d);
                    try_extract_array!(numpy::PyArray1<f64>, with_extra_1d);
                    try_extract_array!(numpy::PyArray1<i64>, with_extra_1d);
                    try_extract_array!(numpy::PyArray1<i32>, with_extra_1d);
                    try_extract_array!(numpy::PyArray1<u64>, with_extra_1d);
                    try_extract_array!(numpy::PyArray1<u32>, with_extra_1d);
                    // 2D
                    try_extract_array!(PyArray2<f32>, with_extra_2d);
                    try_extract_array!(PyArray2<f64>, with_extra_2d);
                    try_extract_array!(PyArray2<i64>, with_extra_2d);
                    try_extract_array!(PyArray2<i32>, with_extra_2d);
                    try_extract_array!(PyArray2<u64>, with_extra_2d);
                    try_extract_array!(PyArray2<u32>, with_extra_2d);
                    // 3D
                    try_extract_array!(PyArray3<f32>, with_extra_3d);
                    try_extract_array!(PyArray3<f64>, with_extra_3d);
                    try_extract_array!(PyArray3<i64>, with_extra_3d);
                    try_extract_array!(PyArray3<i32>, with_extra_3d);
                    try_extract_array!(PyArray3<u64>, with_extra_3d);
                    try_extract_array!(PyArray3<u32>, with_extra_3d);
                }
            }
        }
    }

    // Create the frame
    let mut frame = Frame::new(frame_number, time, frame_data);

    // Extract cell/box information if available
    let locals = PyDict::new(py);
    locals.set_item("fas_atoms", atoms)?;
    let cell: Bound<PyAny> = py.eval(
        &CString::new("fas_atoms.cell.array").unwrap(),
        None,
        Some(&locals),
    )?;
    let cell: &Bound<PyArray2<f64>> = cell.as_unbound().cast_bound::<PyArray2<f64>>(py)?;

    // Set full cell vectors to preserve triclinic cells
    let cell_guard = cell.readonly();
    let cell_array = cell_guard.as_array();
    let vectors = [
        [
            cell_array[(0, 0)] as f32,
            cell_array[(0, 1)] as f32,
            cell_array[(0, 2)] as f32,
        ],
        [
            cell_array[(1, 0)] as f32,
            cell_array[(1, 1)] as f32,
            cell_array[(1, 2)] as f32,
        ],
        [
            cell_array[(2, 0)] as f32,
            cell_array[(2, 1)] as f32,
            cell_array[(2, 2)] as f32,
        ],
    ];
    frame.box_vectors = Some(vectors);

    // Also store diagonal dimensions for compatibility and simple readers
    let dimensions = [
        cell_array[(0, 0)]
            .hypot(cell_array[(0, 1)])
            .hypot(cell_array[(0, 2)]) as f32,
        cell_array[(1, 0)]
            .hypot(cell_array[(1, 1)])
            .hypot(cell_array[(1, 2)]) as f32,
        cell_array[(2, 0)]
            .hypot(cell_array[(2, 1)])
            .hypot(cell_array[(2, 2)]) as f32,
    ];
    frame.box_dimensions = Some(dimensions);

    // Try to extract energy information if available
    if atoms.hasattr("get_potential_energy")? {
        let potential_energy_method = atoms.getattr("get_potential_energy")?;
        if let Ok(potential_energy) = potential_energy_method.call0() {
            let kinetic_energy = match atoms.getattr("get_kinetic_energy")?.call0() {
                Ok(ke) => ke.extract::<f64>()? as f32,
                Err(_) => 0.0f32,
            };
            frame = frame.with_energy(potential_energy.extract::<f64>()? as f32, kinetic_energy);
        }
    }

    // Try to extract temperature and pressure if available
    if atoms.hasattr("get_temperature")? {
        let temp_method = atoms.getattr("get_temperature")?;
        if let Ok(temperature) = temp_method.call0() {
            // let pressure = match atoms.getattr("get_pressure")?.call0() {
            //     Ok(p) => p.extract::<f64>()? as f32,
            //     Err(_) => 0.0f32,
            // };
            let pressure = 0.0f32; // Placeholder for pressure extraction
            frame = frame.with_thermo(temperature.extract::<f64>()? as f32, pressure);
        }
    }

    Ok(frame)
}

/// Convert a TACO Frame to ASE Atoms object
fn frame_to_atoms<'py>(
    py: Python<'py>,
    frame: &Frame,
    pbc: Option<[bool; 3]>,
    atomic_numbers: Option<&Vec<u8>>,
) -> PyResult<Bound<'py, PyAny>> {
    // Get ASE module
    let ase = PyModule::import(py, "ase")?;
    let atoms_class = ase.getattr("Atoms")?;

    // Get positions from frame
    let positions = frame.data.positions.as_ref().ok_or_else(|| {
        PyValueError::new_err("Frame is missing positions; cannot create ASE Atoms")
    })?;

    // Convert positions to a Python array
    let py_positions = PyArray2::from_array(py, positions);

    // Get atomic numbers for this frame
    let atomic_numbers = if let Some(numbers) = atomic_numbers {
        numbers.clone()
    } else {
        // Fallback to zeros if not available
        vec![0; positions.shape()[0]]
    };

    // Create atoms with positions and atomic numbers
    let py_atomic_numbers = PyList::new(py, &atomic_numbers)?;

    // Create atoms object with positions and atomic numbers
    let atoms = atoms_class.call1((py_atomic_numbers, py_positions))?;

    // Set cell/box if available, using box_vectors (full triclinic) when present
    if let Some(box_vectors) = frame.box_vectors {
        let cell = py.eval(
            &CString::new(format!(
                "[[{}, {}, {}], [{}, {}, {}], [{}, {}, {}]]",
                box_vectors[0][0],
                box_vectors[0][1],
                box_vectors[0][2],
                box_vectors[1][0],
                box_vectors[1][1],
                box_vectors[1][2],
                box_vectors[2][0],
                box_vectors[2][1],
                box_vectors[2][2],
            ))
            .unwrap(),
            None,
            None,
        )?;
        atoms.setattr("cell", cell)?;
    } else if let Some(box_dims) = frame.box_dimensions {
        if let Some(box_angles) = frame.box_angles {
            // Reconstruct full cell matrix from dimensions and angles
            let (a, b, c) = (box_dims[0], box_dims[1], box_dims[2]);
            let alpha = box_angles[0].to_radians();
            let beta = box_angles[1].to_radians();
            let gamma = box_angles[2].to_radians();
            let sinb = beta.sin();
            let cosb = beta.cos();
            let cosg = gamma.cos();
            let sing = gamma.sin();

            let cos_alpha = alpha.cos();
            let num = (cos_alpha - cosb * cosg) / (sinb * sing);
            let calpha = num.max(-1.0f32).min(1.0f32);
            let alpha_prime = calpha.acos();
            let ca = alpha_prime.cos();
            let sa = alpha_prime.sin();

            let cell = py.eval(
                &CString::new(format!(
                    "[[{}, 0.0, 0.0], [{}, {}, 0.0], [{}, {}, {}]]",
                    a,
                    b * cosg,
                    b * sing,
                    c * cosb,
                    c * (ca - cosb * cosg) / sing,
                    c * sa * sinb,
                ))
                .unwrap(),
                None,
                None,
            )?;
            atoms.setattr("cell", cell)?;
        } else {
            let cell = py.eval(
                &CString::new(format!(
                    "[[{}, 0.0, 0.0], [0.0, {}, 0.0], [0.0, 0.0, {}]]",
                    box_dims[0], box_dims[1], box_dims[2]
                ))
                .unwrap(),
                None,
                None,
            )?;
            atoms.setattr("cell", cell)?;
        }
    }

    // Set PBC from direct pbc value
    let pbc = pbc.unwrap_or([false, false, false]); // Default to non-periodic if not specified

    // Set PBC on the atoms object
    let py_pbc = PyList::new(py, [pbc[0], pbc[1], pbc[2]])?;
    atoms.setattr("pbc", py_pbc)?;

    // Handle velocities if available and set them correctly
    if let Some(velocities) = &frame.data.velocities {
        let py_velocities = PyArray2::from_array(py, velocities);
        atoms.call_method1("set_velocities", (py_velocities,))?;
    }

    // Create calculator even if only forces OR energy is present, not just when both are present
    let needs_calculator = frame.data.forces.is_some()
        || frame.potential_energy.is_some()
        || frame.data.stress.is_some();

    if needs_calculator {
        // Import ASE's SinglePointCalculator
        let ase_calculators = PyModule::import(py, "ase.calculators.singlepoint")?;
        let spc = ase_calculators.getattr("SinglePointCalculator")?;

        // Create arguments for the calculator
        let kwargs = PyDict::new(py);

        // Add energy if available
        if let Some(potential_energy) = frame.potential_energy {
            kwargs.set_item("energy", potential_energy)?;
        }

        // Add forces if available
        if let Some(forces) = &frame.data.forces {
            let py_forces = PyArray2::from_array(py, forces);
            kwargs.set_item("forces", py_forces)?;
        }

        // Add stress if available
        if let Some(stress) = &frame.data.stress {
            let shape = stress.shape();
            if shape.len() == 2 && shape[0] == 1 && shape[1] > 1 {
                let flat = stress.row(0).to_owned();
                let py_stress = PyArray1::from_array(py, &flat);
                kwargs.set_item("stress", py_stress)?;
            } else if shape.len() == 2 && shape[1] == 1 && shape[0] > 1 {
                let flat = stress.column(0).to_owned();
                let py_stress = PyArray1::from_array(py, &flat);
                kwargs.set_item("stress", py_stress)?;
            } else {
                let py_stress = PyArray2::from_array(py, stress);
                kwargs.set_item("stress", py_stress)?;
            }
        }

        // Instantiate the calculator and attach to atoms
        let calc = spc.call((&atoms,), Some(&kwargs))?;
        atoms.setattr("calc", calc)?;
    }

    // Set extra arrays (excluding stress) from FrameData.extra
    let arrays = atoms.getattr("arrays")?;
    for (key, extra) in &frame.data.extra {
        match extra {
            crate::frame::ExtraArray::Array1(arr) => {
                let py_arr = numpy::PyArray1::from_array(py, arr);
                arrays.set_item(key, py_arr)?;
            }
            crate::frame::ExtraArray::Array2(arr) => {
                let py_arr = PyArray2::from_array(py, arr);
                arrays.set_item(key, py_arr)?;
            }
            crate::frame::ExtraArray::Array3(arr) => {
                let py_arr = numpy::PyArray3::from_array(py, arr);
                arrays.set_item(key, py_arr)?;
            }
            crate::frame::ExtraArray::Array1F64(arr) => {
                let py_arr = numpy::PyArray1::<f64>::from_array(py, arr);
                arrays.set_item(key, py_arr)?;
            }
            crate::frame::ExtraArray::Array2F64(arr) => {
                let py_arr = numpy::PyArray2::<f64>::from_array(py, arr);
                arrays.set_item(key, py_arr)?;
            }
            crate::frame::ExtraArray::Array3F64(arr) => {
                let py_arr = numpy::PyArray3::<f64>::from_array(py, arr);
                arrays.set_item(key, py_arr)?;
            }
        }
    }

    Ok(atoms)
}

/// Write ASE Atoms object(s) to a TACO file.
///
/// This function automatically detects whether it's dealing with a single Atoms object
/// or a list of Atoms objects and handles them accordingly.
///
/// :param path: Path to the output file.
/// :type path: str
/// :param atoms: ASE Atoms object or list of ASE Atoms objects.
/// :type atoms: Union[ase.Atoms, list[ase.Atoms]]
/// :param time_step: Time step between frames in picoseconds. Default is 0.001 (1 fs).
/// :type time_step: float, optional
/// :param full_frame_interval: Interval for storing full frames (non-delta encoded) to prevent error accumulation. Default is 10.
/// :type full_frame_interval: int, optional
/// :param compression_level: Compression level (1-22). Default is 3.
/// :type compression_level: int, optional
/// :param lossless: Use lossless compression for all data. Default is False.
/// :type lossless: bool, optional
///
/// **Example with a single Atoms object**::
///
///     >>> import ase
///     >>> import taco_format
///     >>> atoms = ase.Atoms('H2O')
///     >>> taco_format.write('test.taco', atoms, compression_level=5)
///
/// **Example with a trajectory**::
///
///     >>> import ase
///     >>> import taco_format
///     >>> atoms_list = [ase.Atoms('H2O') for _ in range(10)]
///     >>> taco_format.write('traj.taco', atoms_list, time_step=0.002,
///                          full_frame_interval=50,
///                          compression_level=5)
#[allow(clippy::too_many_arguments)]
#[pyfunction(signature = (
    path,
    atoms,
    time_step = 0.001,
    full_frame_interval = 10,
    compression_level = 3,
    lossless = false
))]
pub fn write(
    py: Python,
    path: &str,
    atoms: &Bound<PyAny>,
    time_step: f64,
    full_frame_interval: u32,
    compression_level: i32,
    lossless: bool,
) -> PyResult<()> {
    // Check if atoms is a list by trying to convert it to a PyList
    if let Ok(atoms_list) = atoms.cast::<PyList>() {
        // It's a list, save as trajectory
        if atoms_list.is_empty() {
            return Err(PyValueError::new_err("Atoms list cannot be empty"));
        }

        // Handle trajectory - multiple frames
        write_trajectory(
            py,
            path,
            atoms_list,
            time_step,
            full_frame_interval,
            compression_level,
            lossless,
        )
    } else {
        // It's a single Atoms object, save as snapshot
        write_single(py, path, atoms, time_step, compression_level, lossless)
    }
}

/// Internal function to write a single ASE Atoms object to a TACO file.
fn write_single(
    py: Python,
    path: &str,
    atoms: &Bound<PyAny>,
    time_step: f64,
    compression_level: i32,
    lossless: bool,
) -> PyResult<()> {
    // Convert atoms to frame
    let frame = atoms_to_frame(py, atoms, 0, 0.0)?;

    // Get number of atoms
    let num_atoms = atoms
        .getattr("positions")?
        .getattr("shape")?
        .extract::<(usize, usize)>()?
        .0 as u32;

    // Create default sim metadata
    let simulation_metadata = SimulationMetadata::default();

    // Create atom metadata with atomic numbers
    let mut atom_metadata = AtomMetadata::new();

    // Extract atomic numbers if available
    if atoms.hasattr("numbers")? {
        let numbers = atoms.getattr("numbers")?;
        let numbers_list = numbers.extract::<Vec<u8>>()?;
        atom_metadata.atomic_numbers = Some(numbers_list);
    } else if atoms.hasattr("get_atomic_numbers")? {
        // Try the method to get atomic numbers
        let numbers_method = atoms.getattr("get_atomic_numbers")?;
        let numbers = numbers_method.call0()?;
        let numbers_list = numbers.extract::<Vec<u8>>()?;
        atom_metadata.atomic_numbers = Some(numbers_list);
    }

    // Create compression settings based on parameters
    let compression_settings = if lossless {
        CompressionSettings::lossless()
    } else {
        CompressionSettings {
            zstd_level: compression_level,
            ..Default::default()
        }
    };

    // Get periodic boundary conditions from atoms
    let mut pbc = [false; 3];
    if atoms.hasattr("pbc")? {
        let py_pbc = atoms.getattr("pbc")?;

        // ASE may provide pbc as a single boolean or an array of 3 booleans
        // Try to extract as a single boolean
        if let Ok(single_pbc) = py_pbc.extract::<bool>() {
            pbc = [single_pbc, single_pbc, single_pbc];
        }
        // Try to extract as a Vec<bool>
        else if let Ok(pbc_array) = py_pbc.extract::<Vec<bool>>() {
            if pbc_array.len() >= 3 {
                pbc.copy_from_slice(&pbc_array[0..3]);
            }
        }
        // Try to handle NumPy array case
        else if let Ok(array) = py_pbc.extract::<PyReadonlyArray1<bool>>() {
            let ndarray = array.as_array();
            if ndarray.shape()[0] >= 3 {
                pbc = [ndarray[0], ndarray[1], ndarray[2]];
            }
        } else if let Ok(array) = py_pbc.extract::<PyReadonlyArray1<i32>>() {
            // Handle case where booleans are stored as integers
            let ndarray = array.as_array();
            if ndarray.shape()[0] >= 3 {
                pbc = [ndarray[0] != 0, ndarray[1] != 0, ndarray[2] != 0];
            }
        }
    }

    // Create header with appropriate settings
    let mut header = Header::new(
        num_atoms,
        time_step,
        simulation_metadata,
        atom_metadata,
        compression_settings,
    );

    // Set the PBC information from the atoms object
    header.pbc = pbc;

    // Create and write trajectory
    let mut writer = Writer::create_with_header(path, header)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;

    writer
        .write_frame(frame)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;
    writer
        .finish()
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;

    Ok(())
}

/// Internal function to write a list of ASE Atoms objects as a trajectory to a TACO file.
fn write_trajectory(
    py: Python,
    path: &str,
    atoms_list: &Bound<PyList>,
    time_step: f64,
    full_frame_interval: u32,
    compression_level: i32,
    lossless: bool,
) -> PyResult<()> {
    if atoms_list.is_empty() {
        return Err(PyValueError::new_err("Atoms list cannot be empty"));
    }

    // Get the first atoms object to determine metadata
    let first_atoms = &atoms_list.get_item(0)?;
    let num_atoms = first_atoms
        .getattr("positions")?
        .getattr("shape")?
        .extract::<(usize, usize)>()?
        .0 as u32;

    // Create default simulation metadata
    let simulation_metadata = SimulationMetadata::default();

    // Create atom metadata with atomic numbers from the first frame
    let mut atom_metadata = AtomMetadata::new();

    // Extract atomic numbers if available
    if first_atoms.hasattr("numbers")? {
        let numbers = first_atoms.getattr("numbers")?;
        let numbers_list = numbers.extract::<Vec<u8>>()?;
        atom_metadata.atomic_numbers = Some(numbers_list);
    } else if first_atoms.hasattr("get_atomic_numbers")? {
        // Try the method to get atomic numbers
        let numbers_method = first_atoms.getattr("get_atomic_numbers")?;
        let numbers = numbers_method.call0()?;
        let numbers_list = numbers.extract::<Vec<u8>>()?;
        atom_metadata.atomic_numbers = Some(numbers_list);
    }

    // Get periodic boundary conditions from the first frame
    let mut pbc = [false; 3];
    if first_atoms.hasattr("pbc")? {
        let py_pbc = first_atoms.getattr("pbc")?;

        // ASE may provide pbc as a single boolean or an array of 3 booleans
        // Try to extract as a single boolean
        if let Ok(single_pbc) = py_pbc.extract::<bool>() {
            pbc = [single_pbc, single_pbc, single_pbc];
        }
        // Try to extract as a Vec<bool>
        else if let Ok(pbc_array) = py_pbc.extract::<Vec<bool>>() {
            if pbc_array.len() >= 3 {
                pbc.copy_from_slice(&pbc_array[0..3]);
            }
        }
        // Try to handle NumPy array case
        else if let Ok(array) = py_pbc.extract::<PyReadonlyArray1<bool>>() {
            let ndarray = array.as_array();
            if ndarray.shape()[0] >= 3 {
                pbc = [ndarray[0], ndarray[1], ndarray[2]];
            }
        } else if let Ok(array) = py_pbc.extract::<PyReadonlyArray1<i32>>() {
            // Handle case where booleans are stored as integers
            let ndarray = array.as_array();
            if ndarray.shape()[0] >= 3 {
                pbc = [ndarray[0] != 0, ndarray[1] != 0, ndarray[2] != 0];
            }
        }
    }

    // Create compression settings based on parameters
    let compression_settings = if lossless {
        CompressionSettings::lossless()
    } else {
        CompressionSettings {
            zstd_level: compression_level,
            ..Default::default()
        }
    };

    // Create header with appropriate settings and PBC information
    let mut header = Header::new(
        num_atoms,
        time_step,
        simulation_metadata,
        atom_metadata,
        compression_settings,
    );

    // Set the PBC information
    header.pbc = pbc;

    // Set the full frame interval (keyframe interval)
    header.full_frame_interval = full_frame_interval;

    // Create writer with our custom header
    let mut writer = Writer::create_with_header(path, header)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;

    // Write each frame
    for (idx, atoms) in atoms_list.iter().enumerate() {
        // Convert atoms to frame, with proper frame number and time
        let frame_number = idx as u64;
        let time = frame_number as f64 * time_step;
        let frame = atoms_to_frame(py, &atoms, frame_number, time)?;

        // Write the frame
        writer
            .write_frame(frame)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;
    }

    // Finish writing
    writer
        .finish()
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;

    Ok(())
}

/// Read ASE Atoms object(s) from a TACO file.
///
/// Always returns a list of ASE Atoms objects. For single-frame files the list
/// contains exactly one Atoms object.
///
/// :param path: Path to the input file.
/// :type path: str
/// :param frame_index: Index of the frame to read (ignored when reading multiple frames). Default is 0.
/// :type frame_index: int, optional
/// :param read_as_trajectory: Deprecated. The function always returns a list.
/// :type read_as_trajectory: bool, optional
/// :param start_frame: Start index when reading a range of frames (inclusive). Default is 0.
/// :type start_frame: int, optional
/// :param end_frame: End index when reading a range of frames (exclusive). Default is None (read all frames).
/// :type end_frame: int, optional
/// :returns: List of ASE Atoms objects.
/// :rtype: list[ase.Atoms]
///
/// **Example reading a single structure**::
///
///     >>> import taco_format
///     >>> atoms_list = taco_format.read('test.taco')
///     >>> atoms = atoms_list[0]
///
/// **Example reading a specific frame from a trajectory**::
///
///     >>> # Read the 10th frame (index 9) from a trajectory
///     >>> frames = taco_format.read('traj.taco', frame_index=9)
///     >>> atoms = frames[0]
///
/// **Example reading a range of frames**::
///
///     >>> # Read frames 100 through 999
///     >>> frames = taco_format.read('traj.taco', start_frame=100, end_frame=1000)
#[allow(clippy::too_many_arguments)]
#[pyfunction(signature = (
    path,
    frame_index = 0,
    read_as_trajectory = false,
    start_frame = 0,
    end_frame = None,
))]
pub fn read<'py>(
    py: Python<'py>,
    path: &str,
    frame_index: u64,
    read_as_trajectory: bool,
    start_frame: u64,
    end_frame: Option<u64>,
) -> PyResult<Bound<'py, PyList>> {
    // Open the file
    let mut reader = Reader::open(path)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;

    // Get frame count to determine if it's a single structure or trajectory
    let num_frames = reader.header().num_frames;

    // Extract necessary metadata from reader
    let atomic_numbers = reader.header().atom_metadata.atomic_numbers.clone();
    let pbc = reader.header().pbc;

    // Case 2: A range of frames is specified
    if let Some(end) = end_frame
        && end > start_frame
    {
        // Read range sequentially
        let atoms_list = PyList::empty(py);

        for idx in start_frame..end {
            if idx >= num_frames {
                break; // Stop if we reach the end of the file
            }

            let frame = reader
                .read_frame(idx)
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;

            let atoms = frame_to_atoms(py, &frame, Some(pbc), atomic_numbers.as_ref())?;
            atoms_list.append(atoms)?;
        }

        return Ok(atoms_list);
    }

    // Case 3: Single structure or default behavior
    if num_frames == 1 && !read_as_trajectory {
        // Single structure - read the only frame
        let frame = reader
            .read_frame(0)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;

        let atoms_list = PyList::empty(py);
        atoms_list.append(frame_to_atoms(
            py,
            &frame,
            Some(pbc),
            atomic_numbers.as_ref(),
        )?)?;
        Ok(atoms_list)
    } else if read_as_trajectory || num_frames > 1 {
        // Multiple frames or forced trajectory mode

        // If frame_index is specified and non-zero, create a trajectory with just that frame
        if frame_index > 0 && frame_index < num_frames && !read_as_trajectory {
            let atoms_list = PyList::empty(py);

            // Find the previous keyframe if we're not starting at frame 0
            let full_frame_interval = reader.header().full_frame_interval as u64;
            let mut read_start = frame_index;

            if full_frame_interval > 0 {
                // Calculate the previous keyframe
                let prev_keyframe = (frame_index / full_frame_interval) * full_frame_interval;
                if prev_keyframe < frame_index {
                    // Start reading from the previous keyframe to properly apply deltas
                    read_start = prev_keyframe;
                }
            }

            // Read sequentially from keyframe to the target frame to handle delta encoding
            for i in read_start..=frame_index {
                let frame = reader
                    .read_frame(i)
                    .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;

                // Only return the requested frame
                if i == frame_index {
                    atoms_list.append(frame_to_atoms(
                        py,
                        &frame,
                        Some(pbc),
                        atomic_numbers.as_ref(),
                    )?)?;
                    return Ok(atoms_list);
                }
            }
        }

        // Create a Python list to store atoms objects
        let atoms_list = PyList::empty(py);

        // Read frames sequentially to properly handle delta encoding
        for frame_index in 0..num_frames {
            let frame = reader
                .read_frame(frame_index)
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;

            // Convert frame to atoms, passing the PBC data directly
            let atoms = frame_to_atoms(py, &frame, Some(pbc), atomic_numbers.as_ref())?;
            atoms_list.append(atoms)?;
        }

        Ok(atoms_list)
    } else {
        // Empty file or other issue
        Err(PyValueError::new_err(format!(
            "Invalid TACO file: contains {num_frames} frames"
        )))
    }
}

/// Save an ASE Atoms object to a TACO file.
///
/// This is a convenience alias for write() with a single Atoms object.
///
/// :param path: Path to the output file.
/// :type path: str
/// :param atoms: ASE Atoms object.
/// :type atoms: ase.Atoms
/// :param compression_level: Compression level (1-22). Default is 3.
/// :type compression_level: int, optional
/// :param lossless: Use lossless compression for all data. Default is False.
/// :type lossless: bool, optional
///
/// **Example**::
///
///     >>> import taco_format
///     >>> atoms = ase.Atoms('H2O')
///     >>> taco_format.save_ase('test.taco', atoms)
#[pyfunction(signature = (path, atoms, compression_level = 3, lossless = false))]
pub fn save_ase(
    py: Python,
    path: &str,
    atoms: &Bound<PyAny>,
    compression_level: i32,
    lossless: bool,
) -> PyResult<()> {
    write_single(py, path, atoms, 0.001, compression_level, lossless)
}

/// Read a range of frames from a TACO file.
///
/// :param path: Path to the input file.
/// :type path: str
/// :param start_frame: Index of the first frame to read.
/// :type start_frame: int
/// :param end_frame: Index of the frame after the last one to read (exclusive).
/// :param end_frame: int
/// :returns: List of ASE Atoms objects.
/// :rtype: list[ase.Atoms]
///
/// **Example reading a range of frames**::
///
///     >>> import taco_format
///     >>> # Read frames 10 through 19
///     >>> atoms_list = taco_format.read_frame_range('traj.taco', 10, 20)
#[pyfunction]
pub fn read_frame_range<'py>(
    py: Python<'py>,
    path: &str,
    start_frame: u64,
    end_frame: u64,
) -> PyResult<Bound<'py, PyList>> {
    // Open the file
    let mut reader = Reader::open(path)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;

    // Validate frame range
    if end_frame <= start_frame {
        return Err(PyValueError::new_err(format!(
            "Invalid frame range: end_frame ({end_frame}) must be greater than start_frame ({start_frame})"
        )));
    }

    if start_frame >= reader.num_frames() {
        return Err(PyValueError::new_err(format!(
            "Start frame {} out of bounds (max: {})",
            start_frame,
            reader.num_frames() - 1
        )));
    }

    // Extract necessary metadata from reader
    let atomic_numbers = reader.header().atom_metadata.atomic_numbers.clone();
    let pbc = reader.header().pbc;

    // Create a Python list to store atoms objects
    let atoms_list = PyList::empty(py);

    // If we're not starting from the first frame and it's not a keyframe,
    // we need to start reading from the most recent keyframe to get correct deltas
    let full_frame_interval = reader.header().full_frame_interval as u64;

    // Find the previous keyframe if we're not starting at frame 0
    let mut read_start = start_frame;
    if start_frame > 0 && full_frame_interval > 0 {
        // Calculate the previous keyframe
        let prev_keyframe = (start_frame / full_frame_interval) * full_frame_interval;
        if prev_keyframe < start_frame {
            // Start reading from the previous keyframe to properly apply deltas
            read_start = prev_keyframe;
        }
    }

    // Read frames sequentially
    let max_frame = std::cmp::min(end_frame, reader.num_frames());

    // Read all frames from the starting keyframe up to the end
    for frame_index in read_start..max_frame {
        let frame = reader
            .read_frame(frame_index)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;

        // Only add frames that are in the requested range
        if frame_index >= start_frame {
            let atoms = frame_to_atoms(py, &frame, Some(pbc), atomic_numbers.as_ref())?;
            atoms_list.append(atoms)?;
        }
    }

    Ok(atoms_list)
}

/// Read specific frames from a TACO file.
///
/// :param path: Path to the input file
/// :type path: str
/// :param frame_indices: List of frame indices to read
/// :type frame_indices: list[int]
/// :returns: List of ASE Atoms objects
/// :rtype: list[ase.Atoms]
///
/// **Example reading specific frames**::
///
///     >>> import taco_format
///     >>> # Read frames 0, 10, 20, 30, 40
///     >>> frames = taco_format.read_frames('traj.taco', [0, 10, 20, 30, 40])
#[pyfunction]
pub fn read_frames<'py>(
    py: Python<'py>,
    path: &str,
    frame_indices: Vec<u64>,
) -> PyResult<Bound<'py, PyList>> {
    if frame_indices.is_empty() {
        return Ok(PyList::empty(py));
    }

    // Open the file
    let mut reader = Reader::open(path)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;

    // Extract necessary metadata from reader
    let atomic_numbers = reader.header().atom_metadata.atomic_numbers.clone();
    let pbc = reader.header().pbc;
    let full_frame_interval = reader.header().full_frame_interval as u64;

    // Validate frame indices
    for &idx in &frame_indices {
        if idx >= reader.num_frames() {
            return Err(PyValueError::new_err(format!(
                "Frame index {} out of bounds (max: {})",
                idx,
                reader.num_frames() - 1
            )));
        }
    }

    // Create a Python list to store atoms objects with the same order as input indices
    let atoms_list = PyList::empty(py);
    let mut frame_map: std::collections::HashMap<u64, Bound<'py, PyAny>> =
        std::collections::HashMap::new();

    // Find all required keyframes
    let mut keyframes_needed: std::collections::HashSet<u64> = std::collections::HashSet::new();

    if full_frame_interval > 0 {
        for &idx in &frame_indices {
            // Calculate the keyframe that this frame depends on
            let keyframe = (idx / full_frame_interval) * full_frame_interval;
            keyframes_needed.insert(keyframe);
        }
    } else {
        // If no delta encoding, each frame is its own keyframe
        for &idx in &frame_indices {
            keyframes_needed.insert(idx);
        }
    }

    // Now process each keyframe and the frames that depend on it
    for &keyframe in &keyframes_needed {
        // Reset the reader's state for each new keyframe
        reader
            .reset()
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;

        // Find all frames that depend on this keyframe
        let end_of_keyframe_block = if full_frame_interval > 0 {
            std::cmp::min(keyframe + full_frame_interval, reader.num_frames())
        } else {
            keyframe + 1
        };

        // Find the maximum frame we need in this keyframe block
        let max_needed_frame = frame_indices
            .iter()
            .filter(|&&idx| idx >= keyframe && idx < end_of_keyframe_block)
            .max()
            .cloned()
            .unwrap_or(keyframe);

        // Read from keyframe up to the maximum needed frame in this block
        for frame_idx in keyframe..=std::cmp::min(max_needed_frame, reader.num_frames() - 1) {
            let frame = reader
                .read_frame(frame_idx)
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;

            // If this frame is in our target list, store it
            if frame_indices.contains(&frame_idx) {
                let atoms = frame_to_atoms(py, &frame, Some(pbc), atomic_numbers.as_ref())?;
                frame_map.insert(frame_idx, atoms);
            }
        }
    }

    // Fill the output list in the original order requested
    for &idx in &frame_indices {
        if let Some(atoms) = frame_map.get(&idx) {
            atoms_list.append(atoms.clone())?;
        }
    }

    Ok(atoms_list)
}

/// Write a list of ASE Atoms objects to a TACO file.
///
/// :param path: Path to the output file.
/// :type path: str
/// :param atoms_list: List of ASE Atoms objects.
/// :type atoms_list: list[ase.Atoms]
/// :param time_step: Time step between frames in picoseconds. Default is 0.001 (1 fs).
/// :type time_step: float, optional
/// :param full_frame_interval: Interval for storing full frames. Default is 10.
/// :type full_frame_interval: int, optional
/// :param compression_level: Compression level (1-22). Default is 3.
/// :type compression_level: int, optional
/// :param lossless: Use lossless compression for all data. Default is False.
/// :type lossless: bool, optional
///
/// **Example**::
///
///     >>> import taco_format
///     >>> atoms_list = [...] # List of ASE Atoms objects
///     >>> taco_format.write_frames('traj.taco', atoms_list)
#[allow(clippy::too_many_arguments)]
#[pyfunction(signature = (
    path,
    atoms_list,
    time_step = 0.001,
    full_frame_interval = 10,
    compression_level = 3,
    lossless = false,
))]
pub fn write_frames(
    py: Python,
    path: &str,
    atoms_list: &Bound<PyList>,
    time_step: f64,
    full_frame_interval: u32,
    compression_level: i32,
    lossless: bool,
) -> PyResult<()> {
    write_trajectory(
        py,
        path,
        atoms_list,
        time_step,
        full_frame_interval,
        compression_level,
        lossless,
    )
}

// Utility and tensor operations for Python

/// Check if a file is a TACO file
///
/// :param path: Path to the file
/// :type path: str
/// :returns: True if the file is a TACO file, False otherwise
/// :rtype: bool
#[pyfunction]
pub fn is_taco_file(_py: Python, path: &str) -> PyResult<bool> {
    crate::utils::is_taco_file(path)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))
}

/// Get file information as a string
///
/// :param path: Path to the TACO file
/// :type path: str
/// :returns: Formatted string with file information
/// :rtype: str
#[pyfunction]
pub fn get_file_info(_py: Python, path: &str) -> PyResult<String> {
    crate::utils::get_file_info(path)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))
}

/// Copy frames from one TACO file to another
///
/// :param src_path: Source file path
/// :type src_path: str
/// :param dst_path: Destination file path
/// :type dst_path: str
/// :param start_frame: Index of the first frame to copy
/// :type start_frame: int
/// :param num_frames: Number of frames to copy (None to copy all remaining frames)
/// :type num_frames: Optional[int]
#[pyfunction(signature = (src_path, dst_path, start_frame, num_frames = None))]
pub fn copy_frames(
    _py: Python,
    src_path: &str,
    dst_path: &str,
    start_frame: u64,
    num_frames: Option<u64>,
) -> PyResult<()> {
    crate::utils::copy_frames(src_path, dst_path, start_frame, num_frames)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))
}

/// Extract a subset of atoms from a TACO file
///
/// :param src_path: Source file path
/// :type src_path: str
/// :param dst_path: Destination file path
/// :type dst_path: str
/// :param atom_indices: List of atom indices to extract
/// :type atom_indices: list[int]
#[pyfunction]
pub fn extract_atoms(
    _py: Python,
    src_path: &str,
    dst_path: &str,
    atom_indices: Vec<usize>,
) -> PyResult<()> {
    crate::utils::extract_atoms(src_path, dst_path, &atom_indices)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))
}

// Tensor operations for Python

/// Calculate the center of mass of a set of positions
///
/// :param positions: Atomic positions
/// :type positions: numpy.ndarray
/// :param masses: Atomic masses (optional)
/// :type masses: Optional[numpy.ndarray]
/// :returns: Center of mass coordinates [x, y, z]
/// :rtype: numpy.ndarray
#[pyfunction(signature = (positions, masses = None))]
pub fn center_of_mass(
    py: Python,
    positions: &Bound<PyAny>,
    masses: Option<&Bound<PyAny>>,
) -> PyResult<Py<PyAny>> {
    // Convert positions to ndarray
    let positions_array = positions.extract::<Bound<PyArray2<f32>>>()?;
    let positions_view = positions_array.readonly();

    // Convert masses to ndarray if provided
    let masses_view = if let Some(masses_any) = masses {
        if let Ok(masses_array) = masses_any.extract::<Bound<numpy::PyArray1<f32>>>() {
            let masses = masses_array.readonly().as_array().to_owned();
            Some(masses)
        } else {
            let masses_array = masses_any.extract::<Bound<PyArray2<f32>>>()?;
            let masses = masses_array
                .readonly()
                .as_array()
                .index_axis(ndarray::Axis(1), 0)
                .to_owned();
            Some(masses)
        }
    } else {
        None
    };

    // Calculate center of mass
    let com = crate::tensor::center_of_mass(&positions_view.as_array().to_owned(), &masses_view);

    // Convert result back to Python
    let py_com = PyArray2::from_array(py, &com.into_shape_with_order((1, 3)).unwrap());
    Ok(py_com.into())
}

/// Extract a subset of atoms from positions
///
/// :param positions: Atomic positions
/// :type positions: numpy.ndarray
/// :param indices: Atom indices to extract
/// :type indices: list[int]
/// :returns: Positions of the selected atoms
/// :rtype: numpy.ndarray
#[pyfunction]
pub fn extract_subset(
    py: Python,
    positions: &Bound<PyAny>,
    indices: Vec<usize>,
) -> PyResult<Py<PyAny>> {
    // Convert positions to ndarray
    let positions_array = positions.extract::<Bound<PyArray2<f32>>>()?;
    let positions_view = positions_array.readonly();

    // Extract subset
    let result = crate::tensor::extract_subset(&positions_view.as_array().to_owned(), &indices);

    // Convert result back to Python
    let py_result = PyArray2::from_array(py, &result);
    Ok(py_result.into())
}

/// Calculate root-mean-square deviation (RMSD) between two coordinate sets
///
/// :param coords1: First coordinate set
/// :type coords1: numpy.ndarray
/// :param coords2: Second coordinate set
/// :type coords2: numpy.ndarray
/// :returns: RMSD value
/// :rtype: float
#[pyfunction]
pub fn calc_rmsd(_py: Python, coords1: &Bound<PyAny>, coords2: &Bound<PyAny>) -> PyResult<f32> {
    // Convert coordinates to ndarrays
    let coords1_array = coords1.extract::<Bound<PyArray2<f32>>>()?;
    let coords2_array = coords2.extract::<Bound<PyArray2<f32>>>()?;

    let coords1_view = coords1_array.readonly();
    let coords2_view = coords2_array.readonly();

    // Calculate RMSD
    crate::tensor::calc_rmsd(
        &coords1_view.as_array().to_owned(),
        &coords2_view.as_array().to_owned(),
    )
    .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))
}

/// Run the TACO CLI with the given arguments
///
/// This function provides access to the same CLI functionality as the standalone
/// `taco` binary, but callable from Python.
///
/// Args:
///     args: List of command line arguments (same as you would pass to the taco binary)
///           The first argument should be the command name (e.g., "info", "extract", etc.)
///
/// Example:
///     >>> import taco_format
///     >>> taco_format.run_cli(["info", "trajectory.taco"])
///     >>> taco_format.run_cli(["extract", "input.taco", "output.taco", "--start", "0", "--end", "100"])
#[pyfunction]
pub fn run_cli(args: Vec<String>) -> PyResult<()> {
    // Prepend "taco" as the program name (clap expects this)
    let mut full_args = vec!["taco".to_string()];
    full_args.extend(args);

    match crate::cli::run_cli_with_args(full_args) {
        Ok(()) => Ok(()),
        Err(e) => Err(PyRuntimeError::new_err(format!("CLI command failed: {e}"))),
    }
}

/// CLI entry point for the installed 'taco' command
///
/// This function is called when the user runs 'taco' from the command line.
/// It automatically gets the command line arguments from Python's sys.argv.
#[pyfunction]
pub fn cli_main(py: Python) -> PyResult<()> {
    // Get sys.argv from Python
    let sys = py.import("sys")?;
    let argv = sys.getattr("argv")?;
    let args: Vec<String> = argv.extract()?;

    // Replace the script name with "taco" for proper clap parsing
    let mut cli_args = args;
    if !cli_args.is_empty() {
        cli_args[0] = "taco".to_string();
    }

    // Pass all args to the CLI parser
    match crate::cli::run_cli_with_args(cli_args) {
        Ok(()) => Ok(()),
        Err(e) => Err(PyRuntimeError::new_err(format!("CLI command failed: {e}"))),
    }
}

/// Register all `taco_format` Python functions on a module.
///
/// The `#[pymodule]` entry point itself lives in the `taco_format_python`
/// wrapper crate (which builds a `cdylib`); this helper does the actual
/// registration so the wrapper stays minimal.
pub fn register_module(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(write, m)?)?;
    m.add_function(wrap_pyfunction!(read, m)?)?;

    // Add sequential reading functions
    m.add_function(wrap_pyfunction!(read_frame_range, m)?)?;
    m.add_function(wrap_pyfunction!(read_frames, m)?)?;
    m.add_function(wrap_pyfunction!(write_frames, m)?)?;

    // Add utility functions
    m.add_function(wrap_pyfunction!(is_taco_file, m)?)?;
    m.add_function(wrap_pyfunction!(get_file_info, m)?)?;
    m.add_function(wrap_pyfunction!(copy_frames, m)?)?;
    m.add_function(wrap_pyfunction!(extract_atoms, m)?)?;

    // Add tensor functions
    m.add_function(wrap_pyfunction!(center_of_mass, m)?)?;
    m.add_function(wrap_pyfunction!(extract_subset, m)?)?;
    m.add_function(wrap_pyfunction!(calc_rmsd, m)?)?;

    // Add CLI function
    m.add_function(wrap_pyfunction!(run_cli, m)?)?;
    m.add_function(wrap_pyfunction!(cli_main, m)?)?;

    Ok(())
}
