//! Utility functions for TACO format
//!
//! This module contains various utilities and helper functions for working with TACO files.

use std::io::{self, Read};
use std::path::Path;

use crate::{Error, Result};

/// Check if a file is a TACO file by examining its magic number
pub fn is_taco_file<P: AsRef<Path>>(path: P) -> Result<bool> {
    let mut file = std::fs::File::open(path)?;
    let mut magic = [0u8; 4];

    match file.read_exact(&mut magic) {
        Ok(()) => Ok(magic == crate::MAGIC),
        Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => Ok(false),
        Err(e) => Err(Error::Io(e)),
    }
}

/// Get file info as a printable string
pub fn get_file_info<P: AsRef<Path>>(path: P) -> Result<String> {
    let reader = crate::io::Reader::open(&path)?;
    let header = reader.header();

    let mut info = String::new();
    info.push_str(&format!("TACO file: {}\n", path.as_ref().display()));
    info.push_str(&format!("Version: {}\n", header.version));
    info.push_str(&format!("Number of atoms: {}\n", header.num_atoms));
    info.push_str(&format!("Number of frames: {}\n", header.num_frames));
    info.push_str(&format!("Time step: {} ps\n", header.time_step));

    // Compression settings
    info.push_str("\nCompression Settings:\n");
    info.push_str(&format!(
        "  Precision: {:?}\n",
        header.compression_settings.precision
    ));

    // Simulation metadata
    if let Some(name) = &header.simulation_metadata.name {
        info.push_str(&format!("\nSimulation name: {name}\n"));
    }
    if let Some(temp) = header.simulation_metadata.temperature {
        info.push_str(&format!("Temperature: {temp} K\n"));
    }
    if let Some(press) = header.simulation_metadata.pressure {
        info.push_str(&format!("Pressure: {press} bar\n"));
    }
    if let Some(ensemble) = &header.simulation_metadata.ensemble {
        info.push_str(&format!("Ensemble: {ensemble}\n"));
    }
    if let Some(software) = &header.simulation_metadata.software {
        info.push_str(&format!("Software: {software}"));
        if let Some(version) = &header.simulation_metadata.software_version {
            info.push_str(&format!(" (version {version})"));
        }
        info.push('\n');
    }
    if let Some(force_field) = &header.simulation_metadata.force_field {
        info.push_str(&format!("Force field: {force_field}\n"));
    }

    // Custom parameters
    if !header.simulation_metadata.custom_parameters.is_empty() {
        info.push_str("\nCustom parameters:\n");
        for (key, value) in &header.simulation_metadata.custom_parameters {
            info.push_str(&format!("  {key}: {value}\n"));
        }
    }

    Ok(info)
}

/// Copy frames from one TACO file to another
pub fn copy_frames<P1: AsRef<Path>, P2: AsRef<Path>>(
    src_path: P1,
    dst_path: P2,
    start_frame: u64,
    num_frames: Option<u64>,
) -> Result<()> {
    // Open source file
    let mut reader = crate::io::Reader::open(src_path)?;
    let header = reader.header().clone();

    // Create destination file
    let mut new_header = header.clone();
    new_header.num_frames = 0;
    new_header.frame_index_offset = 0;
    let mut writer = crate::io::Writer::create_with_header(dst_path, new_header)?;

    // Determine end frame
    let end_frame = match num_frames {
        Some(n) => std::cmp::min(start_frame + n, header.num_frames),
        None => header.num_frames,
    };

    // Copy the frames
    for i in start_frame..end_frame {
        let frame = reader.read_frame(i)?;
        writer.write_frame(frame)?;
    }

    // Finish writing
    writer.finish()?;

    Ok(())
}

/// Extract a subset of atoms from a TACO file
pub fn extract_atoms<P1: AsRef<Path>, P2: AsRef<Path>>(
    src_path: P1,
    dst_path: P2,
    atom_indices: &[usize],
) -> Result<()> {
    use crate::frame::{Frame, FrameData};
    use crate::tensor::extract_subset;

    // Open source file
    let mut reader = crate::io::Reader::open(src_path)?;
    let header = reader.header().clone();

    // Validate indices
    for &idx in atom_indices {
        if idx >= header.num_atoms as usize {
            return Err(Error::InvalidFormat(format!(
                "Atom index {} out of bounds (max: {})",
                idx,
                header.num_atoms - 1
            )));
        }
    }

    // Extract subset of atom metadata
    let atom_metadata = {
        let mut subset = crate::metadata::AtomMetadata::new();

        // Extract atom names
        if let Some(names) = &header.atom_metadata.names {
            subset.names = Some(atom_indices.iter().map(|&i| names[i].clone()).collect());
        }

        // Extract atom types
        if let Some(types) = &header.atom_metadata.types {
            subset.types = Some(atom_indices.iter().map(|&i| types[i].clone()).collect());
        }

        // Extract atom elements
        if let Some(elements) = &header.atom_metadata.elements {
            subset.elements = Some(atom_indices.iter().map(|&i| elements[i].clone()).collect());
        }

        // Extract atom masses
        if let Some(masses) = &header.atom_metadata.masses {
            subset.masses = Some(atom_indices.iter().map(|&i| masses[i]).collect());
        }

        // Extract atom charges
        if let Some(charges) = &header.atom_metadata.charges {
            subset.charges = Some(atom_indices.iter().map(|&i| charges[i]).collect());
        }

        // Extract residue numbers
        if let Some(res_nums) = &header.atom_metadata.residue_numbers {
            subset.residue_numbers = Some(atom_indices.iter().map(|&i| res_nums[i]).collect());
        }

        // Extract residue names
        if let Some(res_names) = &header.atom_metadata.residue_names {
            subset.residue_names =
                Some(atom_indices.iter().map(|&i| res_names[i].clone()).collect());
        }

        // Extract chain IDs
        if let Some(chains) = &header.atom_metadata.chain_ids {
            subset.chain_ids = Some(atom_indices.iter().map(|&i| chains[i].clone()).collect());
        }

        // Extract segment IDs
        if let Some(segments) = &header.atom_metadata.segment_ids {
            subset.segment_ids = Some(atom_indices.iter().map(|&i| segments[i].clone()).collect());
        }

        // Extract custom properties
        for (name, values) in &header.atom_metadata.custom_properties {
            subset.custom_properties.insert(
                name.clone(),
                atom_indices.iter().map(|&i| values[i].clone()).collect(),
            );
        }

        subset
    };

    // Remap constraints to new atom indices
    let index_map: std::collections::HashMap<usize, usize> = atom_indices
        .iter()
        .copied()
        .enumerate()
        .map(|(new, old)| (old, new))
        .collect();

    let new_constraints: Vec<crate::metadata::Constraint> = header
        .constraints
        .iter()
        .filter_map(|c| match c {
            crate::metadata::Constraint::FixAtoms { indices } => {
                let new_indices: Vec<usize> = indices
                    .iter()
                    .filter_map(|&old| index_map.get(&old).copied())
                    .collect();
                if new_indices.is_empty() {
                    None
                } else {
                    Some(crate::metadata::Constraint::FixAtoms {
                        indices: new_indices,
                    })
                }
            }
            other => Some(other.clone()),
        })
        .collect();

    let mut new_header = header.clone();
    new_header.num_atoms = atom_indices.len() as u32;
    new_header.num_frames = 0;
    new_header.frame_index_offset = 0;
    new_header.atom_metadata = atom_metadata;
    new_header.constraints = new_constraints;
    for spec in new_header.extra_datasets.values_mut() {
        if let Some(shape) = &mut spec.shape {
            if !shape.is_empty() && shape[0] == header.num_atoms {
                shape[0] = atom_indices.len() as u32;
            }
        }
    }

    let mut writer = crate::io::Writer::create_with_header(dst_path, new_header)?;

    // Process frames
    for i in 0..header.num_frames {
        let frame = reader.read_frame(i)?;

        // Extract subset of coordinates
        let mut new_frame_data = FrameData {
            positions: None,
            velocities: None,
            forces: None,
            stress: None,
            extra: std::collections::HashMap::new(),
        };

        if let Some(pos) = &frame.data.positions {
            new_frame_data.positions = Some(extract_subset(pos, atom_indices));
        }

        if let Some(vel) = &frame.data.velocities {
            new_frame_data.velocities = Some(extract_subset(vel, atom_indices));
        }

        if let Some(force) = &frame.data.forces {
            new_frame_data.forces = Some(extract_subset(force, atom_indices));
        }

        if let Some(stress) = &frame.data.stress {
            if stress.nrows() == header.num_atoms as usize {
                new_frame_data.stress = Some(extract_subset(stress, atom_indices));
            } else {
                new_frame_data.stress = Some(stress.clone());
            }
        }

        for (key, extra) in &frame.data.extra {
            match extra {
                crate::frame::ExtraArray::Array1(arr) => {
                    if arr.len() == header.num_atoms as usize {
                        let subset: Vec<f32> = atom_indices.iter().map(|&i| arr[i]).collect();
                        new_frame_data
                            .extra
                            .insert(key.clone(), crate::frame::ExtraArray::Array1(subset.into()));
                    } else {
                        new_frame_data.extra.insert(key.clone(), extra.clone());
                    }
                }
                crate::frame::ExtraArray::Array2(arr) => {
                    if arr.nrows() == header.num_atoms as usize {
                        let mut subset =
                            ndarray::Array2::<f32>::zeros((atom_indices.len(), arr.ncols()));
                        for (row, &idx) in atom_indices.iter().enumerate() {
                            for col in 0..arr.ncols() {
                                subset[[row, col]] = arr[[idx, col]];
                            }
                        }
                        new_frame_data
                            .extra
                            .insert(key.clone(), crate::frame::ExtraArray::Array2(subset));
                    } else {
                        new_frame_data.extra.insert(key.clone(), extra.clone());
                    }
                }
                crate::frame::ExtraArray::Array3(arr) => {
                    if arr.shape()[0] == header.num_atoms as usize {
                        let dim1 = arr.shape()[1];
                        let dim2 = arr.shape()[2];
                        let mut subset =
                            ndarray::Array3::<f32>::zeros((atom_indices.len(), dim1, dim2));
                        for (row, &idx) in atom_indices.iter().enumerate() {
                            for j in 0..dim1 {
                                for k in 0..dim2 {
                                    subset[[row, j, k]] = arr[[idx, j, k]];
                                }
                            }
                        }
                        new_frame_data
                            .extra
                            .insert(key.clone(), crate::frame::ExtraArray::Array3(subset));
                    } else {
                        new_frame_data.extra.insert(key.clone(), extra.clone());
                    }
                }
                crate::frame::ExtraArray::Array1F64(arr) => {
                    if arr.len() == header.num_atoms as usize {
                        let subset: Vec<f64> = atom_indices.iter().map(|&i| arr[i]).collect();
                        new_frame_data.extra.insert(
                            key.clone(),
                            crate::frame::ExtraArray::Array1F64(subset.into()),
                        );
                    } else {
                        new_frame_data.extra.insert(key.clone(), extra.clone());
                    }
                }
                crate::frame::ExtraArray::Array2F64(arr) => {
                    if arr.nrows() == header.num_atoms as usize {
                        let mut subset =
                            ndarray::Array2::<f64>::zeros((atom_indices.len(), arr.ncols()));
                        for (row, &idx) in atom_indices.iter().enumerate() {
                            for col in 0..arr.ncols() {
                                subset[[row, col]] = arr[[idx, col]];
                            }
                        }
                        new_frame_data
                            .extra
                            .insert(key.clone(), crate::frame::ExtraArray::Array2F64(subset));
                    } else {
                        new_frame_data.extra.insert(key.clone(), extra.clone());
                    }
                }
                crate::frame::ExtraArray::Array3F64(arr) => {
                    if arr.shape()[0] == header.num_atoms as usize {
                        let dim1 = arr.shape()[1];
                        let dim2 = arr.shape()[2];
                        let mut subset =
                            ndarray::Array3::<f64>::zeros((atom_indices.len(), dim1, dim2));
                        for (row, &idx) in atom_indices.iter().enumerate() {
                            for j in 0..dim1 {
                                for k in 0..dim2 {
                                    subset[[row, j, k]] = arr[[idx, j, k]];
                                }
                            }
                        }
                        new_frame_data
                            .extra
                            .insert(key.clone(), crate::frame::ExtraArray::Array3F64(subset));
                    } else {
                        new_frame_data.extra.insert(key.clone(), extra.clone());
                    }
                }
            }
        }

        // Create new frame with extracted data
        let new_frame = Frame {
            data: new_frame_data,
            ..frame
        };

        writer.write_frame(new_frame)?;
    }

    // Finish writing
    writer.finish()?;

    Ok(())
}
