import numpy as np
import pytest
import taco_format


def test_center_of_mass_equal_masses():
    """Test center of mass calculation with equal masses"""
    # Create a simple system with 3 atoms
    positions = np.array(
        [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32
    )

    # Calculate center of mass without masses (default: equal masses)
    com = taco_format.center_of_mass(positions)

    # Expected result: average of positions
    expected = np.array([[1 / 3, 1 / 3, 0.0]])
    np.testing.assert_allclose(expected, com)


def test_center_of_mass_with_masses():
    """Test center of mass calculation with different masses"""
    # Create a simple system with 3 atoms
    positions = np.array(
        [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32
    )

    # Define different masses (e.g., H=1, C=12, O=16)
    masses = np.array([[1.0], [12.0], [16.0]], dtype=np.float32)

    # Calculate center of mass with masses
    com = taco_format.center_of_mass(positions, masses)

    # Expected result: weighted average of positions
    total_mass = 1.0 + 12.0 + 16.0
    expected_x = (0.0 * 1.0 + 1.0 * 12.0 + 0.0 * 16.0) / total_mass
    expected_y = (0.0 * 1.0 + 0.0 * 12.0 + 1.0 * 16.0) / total_mass
    expected = np.array([[expected_x, expected_y, 0.0]])
    np.testing.assert_allclose(expected, com)


def test_extract_subset():
    """Test extraction of a subset of atoms"""
    # Create a system with 5 atoms
    positions = np.array(
        [
            [0.0, 0.0, 0.0],  # atom 0
            [1.0, 0.0, 0.0],  # atom 1
            [0.0, 1.0, 0.0],  # atom 2
            [1.0, 1.0, 0.0],  # atom 3
            [0.5, 0.5, 1.0],  # atom 4
        ],
        dtype=np.float32,
    )

    # Extract atoms 0, 2, 4
    indices = [0, 2, 4]
    subset = taco_format.extract_subset(positions, indices)

    # Check shape
    assert subset.shape == (3, 3)

    # Check values
    expected = np.array(
        [
            [0.0, 0.0, 0.0],  # atom 0
            [0.0, 1.0, 0.0],  # atom 2
            [0.5, 0.5, 1.0],  # atom 4
        ],
        dtype=np.float32,
    )
    np.testing.assert_allclose(expected, subset)


def test_calc_rmsd():
    """Test RMSD calculation"""
    # Create two sets of coordinates
    coords1 = np.array(
        [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32
    )

    # Second set with small displacement
    coords2 = np.array(
        [[0.1, 0.1, 0.1], [1.1, 0.1, 0.1], [0.1, 1.1, 0.1]], dtype=np.float32
    )

    # Calculate RMSD
    rmsd = taco_format.calc_rmsd(coords1, coords2)

    # Each atom is displaced by [0.1, 0.1, 0.1] which is a distance of sqrt(0.03)
    # The mean squared displacement is 0.03, and the RMSD is sqrt(0.03) ≈ 0.173
    expected_rmsd = np.sqrt(0.03)
    assert abs(rmsd - expected_rmsd) < 1e-6


def test_calc_rmsd_different_shapes():
    """Test RMSD calculation with different shaped arrays"""
    coords1 = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float32)

    coords2 = np.array(
        [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32
    )

    # This should raise an error
    with pytest.raises(Exception):
        taco_format.calc_rmsd(coords1, coords2)
