"""
Tests for delta encoding functionality with stress and extra data.

This module tests the delta encoding and decoding of stress tensors and
extra arrays in TACO format trajectories.
"""
import numpy as np
import pytest
from ase import Atoms
from ase.calculators.singlepoint import SinglePointCalculator
import taco_format


def test_delta_encoding_stress_only(tmp_path):
    """Test delta encoding with stress tensors only."""
    frames = []
    
    # Create 5 frames with changing stress values
    for i in range(5):
        atoms = Atoms("H2", positions=[[0, 0, 0], [1.0 + i * 0.1, 0, 0]])
        atoms.set_velocities([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
        
        # Create stress tensor that changes over time
        stress = np.array([
            1.0 + i * 0.1,  # xx
            2.0 + i * 0.2,  # yy
            3.0 + i * 0.3,  # zz
            0.1 + i * 0.01, # xy
            0.2 + i * 0.02, # xz
            0.3 + i * 0.03  # yz
        ])
        
        calc = SinglePointCalculator(
            atoms, 
            energy=100.0 + i * 10.0,
            forces=np.ones((2, 3)) * (1.0 + i * 0.1),
            stress=stress
        )
        atoms.calc = calc
        frames.append(atoms)
    
    # Write trajectory
    filename = tmp_path / "delta_stress.taco"
    taco_format.write(str(filename), frames)
    
    # Read back and verify
    read_frames = taco_format.read(str(filename))
    
    assert len(read_frames) == 5
    
    for i, atoms in enumerate(read_frames):
        # Verify stress values
        expected_stress = np.array([
            1.0 + i * 0.1,
            2.0 + i * 0.2,
            3.0 + i * 0.3,
            0.1 + i * 0.01,
            0.2 + i * 0.02,
            0.3 + i * 0.03
        ])
        np.testing.assert_allclose(atoms.get_stress(), expected_stress, rtol=1e-4)
        
        # Verify other properties still work
        assert abs(atoms.get_potential_energy() - (100.0 + i * 10.0)) < 1e-4
        np.testing.assert_allclose(
            atoms.get_forces(), 
            np.ones((2, 3)) * (1.0 + i * 0.1), 
            rtol=1e-4
        )


def test_delta_encoding_extra_arrays(tmp_path):
    """Test delta encoding with extra arrays of different dimensions."""
    frames = []
    
    # Create 4 frames with changing extra arrays
    for i in range(4):
        atoms = Atoms("CO", positions=[[0, 0, 0], [1.0 + i * 0.05, 0, 0]])
        
        # Create extra arrays that change over time
        extra1d = np.array([10.0 + i, 20.0 + i * 2])  # 1D array
        extra2d = np.array([
            [1.0 + i * 0.1, 2.0 + i * 0.2],
            [3.0 + i * 0.3, 4.0 + i * 0.4]
        ])  # 2D array
        extra3d = np.array([
            [[1.0 + i, 2.0 + i], [3.0 + i, 4.0 + i]],
            [[5.0 + i, 6.0 + i], [7.0 + i, 8.0 + i]]
        ])  # 3D array
        
        # Attach extra arrays
        atoms.arrays["temperature_profile"] = extra1d
        atoms.arrays["pressure_tensor"] = extra2d
        atoms.arrays["local_density"] = extra3d
        
        calc = SinglePointCalculator(
            atoms,
            energy=50.0 + i * 5.0,
            forces=np.random.rand(2, 3) * 0.1  # Small random forces
        )
        atoms.calc = calc
        frames.append(atoms)
    
    # Write trajectory
    filename = tmp_path / "delta_extra.taco"
    taco_format.write(str(filename), frames)
    
    # Read back and verify
    read_frames = taco_format.read(str(filename))
    
    assert len(read_frames) == 4
    
    for i, atoms in enumerate(read_frames):
        # Verify extra arrays
        expected_1d = np.array([10.0 + i, 20.0 + i * 2])
        expected_2d = np.array([
            [1.0 + i * 0.1, 2.0 + i * 0.2],
            [3.0 + i * 0.3, 4.0 + i * 0.4]
        ])
        expected_3d = np.array([
            [[1.0 + i, 2.0 + i], [3.0 + i, 4.0 + i]],
            [[5.0 + i, 6.0 + i], [7.0 + i, 8.0 + i]]
        ])
        
        np.testing.assert_allclose(
            atoms.arrays["temperature_profile"], 
            expected_1d, 
            rtol=1e-4
        )
        np.testing.assert_allclose(
            atoms.arrays["pressure_tensor"], 
            expected_2d, 
            rtol=1e-4
        )
        np.testing.assert_allclose(
            atoms.arrays["local_density"], 
            expected_3d, 
            rtol=1e-4
        )


def test_delta_encoding_stress_and_extra_combined(tmp_path):
    """Test delta encoding with both stress and extra data combined."""
    frames = []
    
    # Create 6 frames with both stress and extra data
    for i in range(6):
        atoms = Atoms("H2O", positions=[
            [0, 0, 0], 
            [1.0 + i * 0.02, 0, 0], 
            [0.5, 0.8 + i * 0.01, 0]
        ])
        atoms.set_velocities(np.random.rand(3, 3) * 0.1)
        
        # Stress tensor with gradual changes
        stress = np.array([
            0.5 + i * 0.05,  # xx
            0.6 + i * 0.06,  # yy
            0.7 + i * 0.07,  # zz
            0.01 + i * 0.001, # xy
            0.02 + i * 0.002, # xz
            0.03 + i * 0.003  # yz
        ])
        
        # Extra data that varies smoothly
        bond_lengths = np.array([1.0 + i * 0.01, 0.96 + i * 0.005])  # 1D
        virial_tensor = np.array([
            [10.0 + i, 1.0 + i * 0.1],
            [1.0 + i * 0.1, 11.0 + i]
        ])  # 2D
        
        atoms.arrays["bond_lengths"] = bond_lengths
        atoms.arrays["virial"] = virial_tensor
        
        calc = SinglePointCalculator(
            atoms,
            energy=-150.0 + i * 2.0,
            forces=np.random.rand(3, 3) * 0.2,
            stress=stress
        )
        atoms.calc = calc
        frames.append(atoms)
    
    # Write trajectory
    filename = tmp_path / "delta_combined.taco"
    taco_format.write(str(filename), frames)
    
    # Read back and verify
    read_frames = taco_format.read(str(filename))
    
    assert len(read_frames) == 6
    
    for i, atoms in enumerate(read_frames):
        # Verify stress
        expected_stress = np.array([
            0.5 + i * 0.05,
            0.6 + i * 0.06,
            0.7 + i * 0.07,
            0.01 + i * 0.001,
            0.02 + i * 0.002,
            0.03 + i * 0.003
        ])
        np.testing.assert_allclose(atoms.get_stress(), expected_stress, rtol=1e-3)
        
        # Verify extra data
        expected_bonds = np.array([1.0 + i * 0.01, 0.96 + i * 0.005])
        expected_virial = np.array([
            [10.0 + i, 1.0 + i * 0.1],
            [1.0 + i * 0.1, 11.0 + i]
        ])
        
        np.testing.assert_allclose(
            atoms.arrays["bond_lengths"], 
            expected_bonds, 
            rtol=1e-4
        )
        np.testing.assert_allclose(
            atoms.arrays["virial"], 
            expected_virial, 
            rtol=1e-4
        )
        
        # Verify energy
        assert abs(atoms.get_potential_energy() - (-150.0 + i * 2.0)) < 1e-3


def test_delta_encoding_new_extra_data_mid_trajectory(tmp_path):
    """Test delta encoding when new extra data appears mid-trajectory."""
    frames = []
    
    # Create frames where extra data is introduced in the middle
    for i in range(5):
        atoms = Atoms("Ne", positions=[[i * 0.1, 0, 0]])
        
        # Always present extra data
        atoms.arrays["always_present"] = np.array([100.0 + i])
        
        # Extra data that appears only from frame 2 onwards
        if i >= 2:
            atoms.arrays["appears_later"] = np.array([
                [200.0 + i, 300.0 + i],
                [400.0 + i, 500.0 + i]
            ])
        
        # Extra data that disappears after frame 2
        if i <= 2:
            atoms.arrays["disappears_early"] = np.array([50.0 + i * 10])
        
        calc = SinglePointCalculator(atoms, energy=i * 10.0)
        atoms.calc = calc
        frames.append(atoms)
    
    # Write trajectory
    filename = tmp_path / "delta_new_extra.taco"
    taco_format.write(str(filename), frames)
    
    # Read back and verify
    read_frames = taco_format.read(str(filename))
    
    assert len(read_frames) == 5
    
    for i, atoms in enumerate(read_frames):
        # Always present data should be there
        np.testing.assert_allclose(
            atoms.arrays["always_present"],
            np.array([100.0 + i]),
            rtol=1e-4
        )
        
        # Check conditional extra data
        if i >= 2:
            expected = np.array([
                [200.0 + i, 300.0 + i],
                [400.0 + i, 500.0 + i]
            ])
            np.testing.assert_allclose(
                atoms.arrays["appears_later"],
                expected,
                rtol=1e-4
            )
        else:
            assert "appears_later" not in atoms.arrays
        
        if i <= 2:
            np.testing.assert_allclose(
                atoms.arrays["disappears_early"],
                np.array([50.0 + i * 10]),
                rtol=1e-4
            )
        else:
            # Extra data that wasn't present in the original frame should not appear
            assert "disappears_early" not in atoms.arrays


def test_delta_encoding_random_access_with_extra_data(tmp_path):
    """Test random access to frames with delta-encoded extra data."""
    frames = []
    
    # Create many frames for random access testing
    for i in range(20):
        atoms = Atoms("Ar", positions=[[i * 0.05, 0, 0]])
        
        # Create extra data that's dependent on frame index
        extra_1d = np.array([i * 1.0, i * 2.0, i * 3.0])
        extra_2d = np.array([
            [i * 0.1, i * 0.2],
            [i * 0.3, i * 0.4]
        ])
        
        atoms.arrays["frame_dependent_1d"] = extra_1d
        atoms.arrays["frame_dependent_2d"] = extra_2d
        
        # Add stress that depends on frame index
        stress = np.array([i * 0.01] * 6)  # Simple stress
        
        calc = SinglePointCalculator(
            atoms,
            energy=i * 1.5,
            stress=stress
        )
        atoms.calc = calc
        frames.append(atoms)
    
    # Write trajectory
    filename = tmp_path / "delta_random_access.taco"
    taco_format.write(str(filename), frames)
    
    # Test random access by reading specific frames
    random_indices = [19, 5, 12, 0, 15, 8, 3]
    
    for idx in random_indices:
        # Read single frame using frame_index parameter
        atoms_list = taco_format.read(str(filename), frame_index=idx)
        atoms = atoms_list[0] if isinstance(atoms_list, list) else atoms_list
        
        # Verify frame-dependent data
        expected_1d = np.array([idx * 1.0, idx * 2.0, idx * 3.0])
        expected_2d = np.array([
            [idx * 0.1, idx * 0.2],
            [idx * 0.3, idx * 0.4]
        ])
        expected_stress = np.array([idx * 0.01] * 6)
        
        np.testing.assert_allclose(
            atoms.arrays["frame_dependent_1d"],
            expected_1d,
            rtol=1e-3,  # Relaxed tolerance for compression
            err_msg=f"Failed for frame {idx}"
        )
        np.testing.assert_allclose(
            atoms.arrays["frame_dependent_2d"],
            expected_2d,
            rtol=1e-3,  # Relaxed tolerance for compression
            err_msg=f"Failed for frame {idx}"
        )
        np.testing.assert_allclose(
            atoms.get_stress(),
            expected_stress,
            rtol=1e-3,  # Relaxed tolerance for compression
            err_msg=f"Failed for frame {idx}"
        )
        assert abs(atoms.get_potential_energy() - idx * 1.5) < 1e-4


def test_delta_encoding_precision_preservation(tmp_path):
    """Test that delta encoding preserves precision for small changes."""
    frames = []
    
    # Create frames with very small changes in extra data
    base_value = 1.23456789
    
    for i in range(10):
        atoms = Atoms("Li", positions=[[0, 0, 0]])
        
        # Very small incremental changes
        precise_data = np.array([
            base_value + i * 1e-6,  # Very small increments
            base_value * 2 + i * 1e-7,
            base_value * 3 + i * 1e-8
        ])
        
        atoms.arrays["precise_data"] = precise_data
        
        calc = SinglePointCalculator(atoms, energy=0.0)
        atoms.calc = calc
        frames.append(atoms)
    
    # Write trajectory
    filename = tmp_path / "delta_precision.taco"
    taco_format.write(str(filename), frames)
    
    # Read back and verify precision is maintained
    read_frames = taco_format.read(str(filename))
    
    for i, atoms in enumerate(read_frames):
        expected = np.array([
            base_value + i * 1e-6,
            base_value * 2 + i * 1e-7,
            base_value * 3 + i * 1e-8
        ])
        
        # Use a tighter tolerance to ensure precision is preserved
        np.testing.assert_allclose(
            atoms.arrays["precise_data"],
            expected,
            rtol=1e-6,  # Tight tolerance
            atol=1e-8,  # Small absolute tolerance
            err_msg=f"Precision lost in frame {i}"
        )


if __name__ == "__main__":
    pytest.main([__file__])
