//! Tensor operations for TACO format
//!
//! This module provides utilities for handling tensors in the TACO format,
//! including transformations and operations specific to molecular dynamics data.

use crate::Result;
use ndarray::{Array1, Array2};

/// Calculate the center of mass of a set of positions
pub fn center_of_mass(positions: &Array2<f32>, masses: &Option<Array1<f32>>) -> Array1<f32> {
    let n_atoms = positions.shape()[0];

    match masses {
        Some(mass_array) => {
            let mut com = Array1::<f32>::zeros(3);
            let mut total_mass = 0.0;

            for i in 0..n_atoms {
                let mass = mass_array[i];
                for j in 0..3 {
                    com[j] += positions[[i, j]] * mass;
                }
                total_mass += mass;
            }

            if total_mass > 0.0 {
                com /= total_mass;
            }

            com
        }
        None => {
            // If no masses are provided, use geometric center (equal masses)
            let mut center = Array1::<f32>::zeros(3);

            for i in 0..n_atoms {
                for j in 0..3 {
                    center[j] += positions[[i, j]];
                }
            }

            if n_atoms > 0 {
                center /= n_atoms as f32;
            }
            center
        }
    }
}

/// Extract a subset of atoms from positions
pub fn extract_subset(positions: &Array2<f32>, indices: &[usize]) -> Array2<f32> {
    let mut result = Array2::<f32>::zeros((indices.len(), 3));

    for (i, &idx) in indices.iter().enumerate() {
        for j in 0..3 {
            result[[i, j]] = positions[[idx, j]];
        }
    }

    result
}

/// Calculate root-mean-square deviation (RMSD) between two coordinate sets
pub fn calc_rmsd(coords1: &Array2<f32>, coords2: &Array2<f32>) -> Result<f32> {
    if coords1.shape() != coords2.shape() {
        return Err(crate::Error::ShapeMismatch(format!(
            "Coordinate shapes don't match: {:?} vs {:?}",
            coords1.shape(),
            coords2.shape()
        )));
    }

    let n_atoms = coords1.shape()[0];
    let mut sum_sq_diff = 0.0;

    for i in 0..n_atoms {
        for j in 0..3 {
            let diff = coords1[[i, j]] - coords2[[i, j]];
            sum_sq_diff += diff * diff;
        }
    }

    Ok((sum_sq_diff / (n_atoms as f32)).sqrt())
}
