import os
import numpy as np
import pytest
from ase import Atoms
from ase.calculators.singlepoint import SinglePointCalculator
import taco_format


# Create a temporary directory for test files
@pytest.fixture
def temp_dir(tmpdir):
    return tmpdir


def create_test_atom():
    """Create a simple test atom with forces and energy"""
    atoms = Atoms("H", positions=[[1.0, 0.0, 0.0]])

    # Set velocities
    atoms.set_velocities([[1.0, 0.0, 0.0]])

    # Set forces and energy with a calculator
    forces = np.array([[0.5, 0.5, 0.5]])
    energy = -1.5
    calc = SinglePointCalculator(atoms, energy=energy, forces=forces)
    atoms.calc = calc

    return atoms


def create_test_molecule():
    """Create a water molecule with forces and energy"""
    atoms = Atoms(
        "H2O", positions=[[0.0, 0.0, 0.0], [0.95, 0.0, 0.0], [0.45, 0.85, 0.0]]
    )

    # Set cell and PBC
    atoms.set_cell([10.0, 10.0, 10.0])
    atoms.set_pbc([True, True, True])

    # Set velocities
    atoms.set_velocities([[0.0, 0.0, 0.1], [0.0, 0.1, 0.0], [0.1, 0.0, 0.0]])

    # Set forces and energy with a calculator
    forces = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]])
    energy = -10.5
    calc = SinglePointCalculator(atoms, energy=energy, forces=forces)
    atoms.calc = calc

    return atoms


def test_write_read_single_atom(temp_dir):
    """Test writing and reading a single atom"""
    filename = os.path.join(temp_dir, "single_atom.taco")

    # Create a test atom
    atoms = create_test_atom()
    print(atoms.calc)

    # Write to TACO file
    taco_format.write(filename, atoms)

    # Read from TACO file
    atoms_list = taco_format.read(filename)
    assert len(atoms_list) == 1
    atoms_read = atoms_list[0]
    print(atoms_read.calc)

    # Check that everything was preserved
    np.testing.assert_allclose(atoms_read.get_positions(), atoms.get_positions())
    np.testing.assert_allclose(atoms_read.get_velocities(), atoms.get_velocities())
    np.testing.assert_allclose(atoms_read.get_pbc(), atoms.get_pbc())
    np.testing.assert_allclose(atoms_read.get_forces(), atoms.get_forces())
    assert abs(atoms_read.get_potential_energy() - atoms.get_potential_energy()) < 1e-6


def test_write_read_molecule(temp_dir):
    """Test writing and reading a molecule"""
    filename = os.path.join(temp_dir, "water.taco")

    # Create a test molecule
    atoms = create_test_molecule()

    # Write to TACO file
    taco_format.write(filename, atoms)

    # Read from TACO file
    atoms_list = taco_format.read(filename)
    assert len(atoms_list) == 1
    atoms_read = atoms_list[0]

    # Check that everything was preserved
    np.testing.assert_allclose(atoms_read.get_positions(), atoms.get_positions())
    np.testing.assert_allclose(atoms_read.get_velocities(), atoms.get_velocities())
    np.testing.assert_allclose(atoms_read.get_pbc(), atoms.get_pbc())
    np.testing.assert_allclose(atoms_read.get_forces(), atoms.get_forces())
    assert abs(atoms_read.get_potential_energy() - atoms.get_potential_energy()) < 1e-6

    # Check cell and PBC
    np.testing.assert_allclose(atoms_read.get_cell(), atoms.get_cell())
    np.testing.assert_array_equal(atoms_read.get_pbc(), atoms.get_pbc())


def test_file_info(temp_dir):
    """Test get_file_info function"""
    filename = os.path.join(temp_dir, "info_test.taco")

    # Create and save a test molecule
    atoms = create_test_molecule()
    taco_format.write(filename, atoms)

    # Get file info
    info = taco_format.get_file_info(filename)

    # Check that the info contains the correct information
    assert "TACO file:" in info
    assert "Number of atoms: 3" in info
    assert "Number of frames: 1" in info


def test_is_taco_file(temp_dir):
    """Test is_taco_file function"""
    taco_file = os.path.join(temp_dir, "is_taco.taco")
    non_taco_file = os.path.join(temp_dir, "not_taco.txt")

    # Create a TACO file
    atoms = create_test_atom()
    taco_format.write(taco_file, atoms)

    # Create a non-TACO file
    with open(non_taco_file, "w") as f:
        f.write("This is not a TACO file")

    # Test the function
    assert taco_format.is_taco_file(taco_file) == True
    assert taco_format.is_taco_file(non_taco_file) == False


def test_compression_options(temp_dir):
    """Test different compression options"""
    filename_default = os.path.join(temp_dir, "default.taco")
    filename_high = os.path.join(temp_dir, "high.taco")
    filename_lossless = os.path.join(temp_dir, "lossless.taco")

    # Create a test molecule
    atoms = create_test_molecule()

    # Write with different compression settings
    taco_format.write(filename_default, atoms)  # Default settings
    taco_format.write(filename_high, atoms, compression_level=19)  # High compression
    taco_format.write(filename_lossless, atoms, lossless=True)  # Lossless

    # Check that the files exist and have different sizes
    assert os.path.exists(filename_default)
    assert os.path.exists(filename_high)
    assert os.path.exists(filename_lossless)

    # Read the files to ensure they're valid
    atoms_default = taco_format.read(filename_default)[0]
    atoms_high = taco_format.read(filename_high)[0]
    atoms_lossless = taco_format.read(filename_lossless)[0]

    # Check that the content is the same
    np.testing.assert_allclose(atoms_default.get_positions(), atoms.get_positions())
    np.testing.assert_allclose(atoms_high.get_positions(), atoms.get_positions())
    np.testing.assert_allclose(atoms_lossless.get_positions(), atoms.get_positions())

    # The lossless file should have the most accurate representation
    pos_default = atoms_default.get_positions()
    pos_lossless = atoms_lossless.get_positions()
    pos_original = atoms.get_positions()

    # Default might have tiny differences due to lossy compression
    default_error = np.sum(np.abs(pos_default - pos_original))
    lossless_error = np.sum(np.abs(pos_lossless - pos_original))

    # The lossless error should be very close to zero
    assert lossless_error < 1e-6
