import bz2
import gzip
import numpy as np
import os
import pytest
from dynasor.trajectory.lammps_trajectory_reader import LammpsTrajectoryReader


@pytest.fixture
def first_frame():
    data = np.array([
        [1, 1, 1.91468, 3.02071, 0.528818, 0.00105876, 0.000389945, 0.000540445],
        [2, 2, 1.29395, 3.5652, 6.21471, 0.0116399, 0.0238521, 0.0115804],
        [2, 3, 1.39919, 2.29819, 0.989518, -0.0190113, 0.0234721, 0.0142206],
        [1, 4, 3.626, 4.14767, 2.28901, -0.00775121, -0.00332156, 0.00045625],
        [2, 5, 4.16502, 3.38633, 2.64931, -0.00916183, -0.00898758, -0.0093385],
        [2, 6, 3.00361, 3.81053, 1.58263, -0.00415819, 0.00792349, -0.00812868],
        [1, 7, 2.42769, 6.19665, 3.54307, 0.00516741, 0.00576708, -0.00187862],
        [2, 8, 3.06452, 0.52461, 4.05338, 0.0147389, -0.00707458, 0.000731629],
        [2, 9, 2.90869, 5.38709, 3.20654, -0.00964617, -0.00368993, -0.00034467],
        [1, 10, 3.80235, 1.58899, 5.41742, 0.000160311, 0.000400954, -0.00271396],
        [2, 11, 4.36344, 2.15982, 4.81798, 0.0317508, -0.00780526, 0.0190895],
        [2, 12, 3.13529, 2.16145, 5.8942, 0.00501066, 0.0133524, -0.0114753],
        [1, 13, 4.80923, 5.75237, 0.272529, -0.0041924, 0.00248113, 5.02226e-05],
        [2, 14, 4.24226, 0.143661, 6.00552, -0.00864071, 0.00125601, 0.00341038],
        [2, 15, 4.22911, 5.2088, 0.879155, 0.00156727, 0.00697033, 0.00957028],
        [1, 16, 0.417374, 1.05515, 1.90862, 0.00143456, -0.00280331, -0.00114098],
        [2, 17, 6.07448, 0.524278, 1.30311, 0.0218267, -0.00281436, -0.0211203],
        [2, 18, 0.957417, 0.433166, 2.47562, -0.0062517, -0.00294501, 0.00603581],
        [1, 19, 5.4205, 2.90285, 3.66719, 0.00245983, 0.000292166, 0.00455708],
        [2, 20, 5.97449, 2.21136, 3.20358, 0.00989121, 0.00880165, 0.000767886],
        [2, 21, 6.01398, 3.63229, 4.00734, -0.00868655, 0.0196349, -0.0174331],
        [1, 22, 0.616258, 4.77853, 5.1262, -0.000836257, -0.00691637, 0.000166677],
        [2, 23, 1.20745, 5.48874, 4.74398, 0.00678782, -0.0119932, 0.00252133],
        [2, 24, 6.16478, 5.20174, 5.69963, 0.00200832, -0.001067, -0.00067658],
        ])
    return data


@pytest.fixture
def filename_lammpstrj():
    this_dir = os.path.dirname(__file__)
    return os.path.join(this_dir, 'trajectory_files/positions_and_velocities.lammpstrj')


@pytest.fixture
def filename_lammpstrj_no_velocities():
    this_dir = os.path.dirname(__file__)
    return os.path.join(this_dir, 'trajectory_files/positions.lammpstrj')


def test_open_lammpstrj_no_velocities(filename_lammpstrj_no_velocities):
    reader = LammpsTrajectoryReader(filename_lammpstrj_no_velocities)
    reader.close()


def test_open_lammpstrj(filename_lammpstrj):
    reader = LammpsTrajectoryReader(filename_lammpstrj)
    reader.close()


def test_read_frames(filename_lammpstrj):
    reader = LammpsTrajectoryReader(filename_lammpstrj)
    frames = list(reader)
    assert len(frames) == 4
    reader.close()


def test_first_frame_contents(filename_lammpstrj, first_frame):
    positions_target = first_frame[:, 2:5]
    velocities_target = first_frame[:, 5:]
    reader = LammpsTrajectoryReader(filename_lammpstrj)
    frame = next(reader)
    assert frame.n_atoms == 24
    assert np.allclose(frame.positions, positions_target)
    assert np.allclose(frame.velocities, velocities_target)
    reader.close()


def test_first_frame_contents_no_velocities(filename_lammpstrj_no_velocities, first_frame):
    positions_target = first_frame[:, 2:5]
    reader = LammpsTrajectoryReader(filename_lammpstrj_no_velocities)
    frame = next(reader)
    assert frame.n_atoms == 24
    assert frame.velocities is None
    assert np.allclose(frame.positions, positions_target)
    reader.close()


def test_first_frame_contents_with_units(filename_lammpstrj, first_frame):
    positions_target = first_frame[:, 2:5]
    velocities_target = first_frame[:, 5:]
    reader = LammpsTrajectoryReader(filename_lammpstrj, length_unit='nm', time_unit='ps')
    frame = next(reader)
    assert frame.n_atoms == 24
    assert np.allclose(frame.positions, 10 * positions_target)
    assert np.allclose(frame.velocities, 1 / 100 * velocities_target)

    # cell
    cell_target = 10 * 6.25 * np.eye(3)
    assert np.allclose(frame.cell, cell_target)


@pytest.mark.parametrize(('length_unit', 'factor'), [('pm', 1e-2), ('fm', 1e-5)])
def test_small_length_unit_conversions(filename_lammpstrj, first_frame, length_unit, factor):
    reader = LammpsTrajectoryReader(filename_lammpstrj, length_unit=length_unit)
    frame = next(reader)
    reader.close()

    assert np.allclose(frame.positions, factor * first_frame[:, 2:5])
    assert np.allclose(frame.velocities, factor * first_frame[:, 5:])
    assert np.allclose(frame.cell, factor * 6.25 * np.eye(3))


# Edge-case and error-handling tests

def test_invalid_length_unit_raises(tmp_path):
    """An unrecognised length_unit raises ValueError."""
    f = tmp_path / 'dummy.lammpstrj'
    f.write_text('')
    with pytest.raises(ValueError, match='length unit'):
        LammpsTrajectoryReader(str(f), length_unit='furlongs')


def test_invalid_time_unit_raises(tmp_path):
    """An unrecognised time_unit raises ValueError."""
    f = tmp_path / 'dummy.lammpstrj'
    f.write_text('')
    with pytest.raises(ValueError, match='time unit'):
        LammpsTrajectoryReader(str(f), time_unit='fortnights')


def test_nonexistent_file_raises(tmp_path):
    """Opening a non-existent file raises FileNotFoundError."""
    with pytest.raises(FileNotFoundError):
        LammpsTrajectoryReader(str(tmp_path / 'no_such_file.lammpstrj'))


def test_empty_file_raises_stop_iteration(tmp_path):
    """An empty file yields StopIteration on the first next() call."""
    empty = tmp_path / 'empty.lammpstrj'
    empty.write_text('')
    reader = LammpsTrajectoryReader(str(empty))
    with pytest.raises(StopIteration):
        next(reader)


def test_malformed_header_raises(tmp_path):
    """A file whose content does not match the LAMMPS dump format raises IOError."""
    bad = tmp_path / 'bad.lammpstrj'
    bad.write_text('THIS IS NOT A LAMMPS DUMP FILE\ngarbage content\n')
    reader = LammpsTrajectoryReader(str(bad))
    with pytest.raises(IOError, match='TRJ_reader: Failed to parse TRJ frame header'):
        next(reader)


def test_stop_iteration_after_exhausting_frames(filename_lammpstrj):
    """next() raises StopIteration once all frames have been consumed."""
    reader = LammpsTrajectoryReader(filename_lammpstrj)
    list(reader)  # exhaust
    with pytest.raises(StopIteration):
        next(reader)


def test_minimal_single_atom_frame(tmp_path):
    """A minimal one-atom LAMMPS dump frame is parsed correctly."""
    content = (
        'ITEM: TIMESTEP\n0\n'
        'ITEM: NUMBER OF ATOMS\n1\n'
        'ITEM: BOX BOUNDS pp pp pp\n'
        '0.0 5.0\n0.0 5.0\n0.0 5.0\n'
        'ITEM: ATOMS id type x y z\n'
        '1 1 1.0 2.0 3.0\n'
    )
    f = tmp_path / 'single.lammpstrj'
    f.write_text(content)
    reader = LammpsTrajectoryReader(str(f))
    frame = next(reader)
    reader.close()
    assert frame.n_atoms == 1
    assert frame.frame_index == 0
    assert np.allclose(frame.positions[0], [1.0, 2.0, 3.0])
    assert np.allclose(frame.cell, 5.0 * np.eye(3))
    assert frame.velocities is None


_SINGLE_ATOM_CONTENT = (
    'ITEM: TIMESTEP\n0\n'
    'ITEM: NUMBER OF ATOMS\n1\n'
    'ITEM: BOX BOUNDS pp pp pp\n'
    '0.0 4.0\n0.0 4.0\n0.0 4.0\n'
    'ITEM: ATOMS id type x y z\n'
    '1 1 1.0 2.0 3.0\n'
)


def test_gzip_file(tmp_path):
    """Reader opens gzip-compressed .lammpstrj.gz files."""
    gz_path = tmp_path / 'traj.lammpstrj.gz'
    with gzip.open(str(gz_path), 'wt') as fh:
        fh.write(_SINGLE_ATOM_CONTENT)
    reader = LammpsTrajectoryReader(str(gz_path))
    frame = next(reader)
    reader.close()
    assert frame.n_atoms == 1
    assert np.allclose(frame.positions[0], [1.0, 2.0, 3.0])


def test_bz2_file(tmp_path):
    """Reader opens bz2-compressed .lammpstrj.bz2 files."""
    bz2_path = tmp_path / 'traj.lammpstrj.bz2'
    with bz2.open(str(bz2_path), 'wt') as fh:
        fh.write(_SINGLE_ATOM_CONTENT)
    reader = LammpsTrajectoryReader(str(bz2_path))
    frame = next(reader)
    reader.close()
    assert frame.n_atoms == 1
    assert np.allclose(frame.positions[0], [1.0, 2.0, 3.0])


def test_blank_lines_in_header_are_skipped(tmp_path):
    """Blank lines between ITEM sections are silently skipped."""
    content = (
        '\n'
        'ITEM: TIMESTEP\n0\n'
        '\n'
        'ITEM: NUMBER OF ATOMS\n1\n'
        'ITEM: BOX BOUNDS pp pp pp\n'
        '0.0 5.0\n0.0 5.0\n0.0 5.0\n'
        'ITEM: ATOMS id type x y z\n'
        '1 1 1.0 2.0 3.0\n'
    )
    f = tmp_path / 'blank.lammpstrj'
    f.write_text(content)
    reader = LammpsTrajectoryReader(str(f))
    frame = next(reader)
    reader.close()
    assert frame.n_atoms == 1


def test_triclinic_cell(tmp_path):
    """Triclinic box bounds (three values per line) are parsed into the correct cell."""
    # LAMMPS triclinic format: xlo xhi xy / ylo yhi xz / zlo zhi yz
    content = (
        'ITEM: TIMESTEP\n0\n'
        'ITEM: NUMBER OF ATOMS\n1\n'
        'ITEM: BOX BOUNDS xy xz yz pp pp pp\n'
        '0.0 5.0 1.0\n'
        '0.0 5.0 2.0\n'
        '0.0 5.0 3.0\n'
        'ITEM: ATOMS id type x y z\n'
        '1 1 1.0 2.0 3.0\n'
    )
    f = tmp_path / 'triclinic.lammpstrj'
    f.write_text(content)
    reader = LammpsTrajectoryReader(str(f))
    frame = next(reader)
    reader.close()
    # Diagonal: xhi-xlo = yhi-ylo = zhi-zlo = 5 for all three
    assert np.allclose(frame.cell.diagonal(), [5.0, 5.0, 5.0])
    # Off-diagonal tilt factors: cell[1,0]=xy=1, cell[2,0]=xz=2, cell[2,1]=yz=3
    assert np.isclose(frame.cell[1, 0], 1.0)
    assert np.isclose(frame.cell[2, 0], 2.0)
    assert np.isclose(frame.cell[2, 1], 3.0)


def test_malformed_cell_bounds_raises(tmp_path):
    """A BOX BOUNDS section with wrong column count raises IOError."""
    # Give 4 values per line — neither (3,2) nor (3,3)
    content = (
        'ITEM: TIMESTEP\n0\n'
        'ITEM: NUMBER OF ATOMS\n1\n'
        'ITEM: BOX BOUNDS pp pp pp\n'
        '0.0 5.0 1.0 9.0\n'
        '0.0 5.0 2.0 9.0\n'
        '0.0 5.0 3.0 9.0\n'
        'ITEM: ATOMS id type x y z\n'
        '1 1 1.0 2.0 3.0\n'
    )
    f = tmp_path / 'bad_cell.lammpstrj'
    f.write_text(content)
    reader = LammpsTrajectoryReader(str(f))
    with pytest.raises(IOError, match='Malformed cell bounds'):
        next(reader)


def test_scaled_positions(tmp_path):
    """Scaled positions (xs ys zs) are correctly unscaled by the cell diagonal."""
    cell_size = 10.0
    content = (
        'ITEM: TIMESTEP\n0\n'
        'ITEM: NUMBER OF ATOMS\n1\n'
        'ITEM: BOX BOUNDS pp pp pp\n'
        f'0.0 {cell_size}\n0.0 {cell_size}\n0.0 {cell_size}\n'
        'ITEM: ATOMS id type xs ys zs\n'
        '1 1 0.1 0.2 0.3\n'
    )
    f = tmp_path / 'scaled.lammpstrj'
    f.write_text(content)
    reader = LammpsTrajectoryReader(str(f))
    frame = next(reader)
    reader.close()
    assert np.allclose(frame.positions[0], [1.0, 2.0, 3.0])


def test_scaled_positions_carried_through_second_frame(tmp_path):
    """Scaled positions are also unscaled correctly in subsequent frames (_get_next)."""
    cell_size = 10.0
    content = (
        'ITEM: TIMESTEP\n0\n'
        'ITEM: NUMBER OF ATOMS\n1\n'
        'ITEM: BOX BOUNDS pp pp pp\n'
        f'0.0 {cell_size}\n0.0 {cell_size}\n0.0 {cell_size}\n'
        'ITEM: ATOMS id type xs ys zs\n'
        '1 1 0.1 0.2 0.3\n'
        'ITEM: TIMESTEP\n1\n'
        'ITEM: NUMBER OF ATOMS\n1\n'
        'ITEM: BOX BOUNDS pp pp pp\n'
        f'0.0 {cell_size}\n0.0 {cell_size}\n0.0 {cell_size}\n'
        'ITEM: ATOMS id type xs ys zs\n'
        '1 1 0.4 0.5 0.6\n'
    )
    f = tmp_path / 'scaled2.lammpstrj'
    f.write_text(content)
    reader = LammpsTrajectoryReader(str(f))
    next(reader)
    frame = next(reader)
    reader.close()
    assert np.allclose(frame.positions[0], [4.0, 5.0, 6.0])


def test_missing_position_columns_raises(tmp_path):
    """A dump file with no recognised position columns raises RuntimeError."""
    content = (
        'ITEM: TIMESTEP\n0\n'
        'ITEM: NUMBER OF ATOMS\n1\n'
        'ITEM: BOX BOUNDS pp pp pp\n'
        '0.0 5.0\n0.0 5.0\n0.0 5.0\n'
        'ITEM: ATOMS id type vx vy vz\n'
        '1 1 0.1 0.2 0.3\n'
    )
    f = tmp_path / 'no_pos.lammpstrj'
    f.write_text(content)
    reader = LammpsTrajectoryReader(str(f))
    with pytest.raises(RuntimeError):
        next(reader)


def test_unwrapped_positions(tmp_path):
    """Unwrapped positions (xu yu zu) outside the cell boundary are passed through as-is."""
    content = (
        'ITEM: TIMESTEP\n0\n'
        'ITEM: NUMBER OF ATOMS\n1\n'
        'ITEM: BOX BOUNDS pp pp pp\n'
        '0.0 5.0\n0.0 5.0\n0.0 5.0\n'
        'ITEM: ATOMS id type xu yu zu\n'
        '1 1 6.0 -1.0 12.0\n'
    )
    f = tmp_path / 'unwrapped.lammpstrj'
    f.write_text(content)
    reader = LammpsTrajectoryReader(str(f))
    frame = next(reader)
    reader.close()
    assert frame.n_atoms == 1
    assert np.allclose(frame.positions[0], [6.0, -1.0, 12.0])


def test_no_type_column(tmp_path):
    """Dump files without a 'type' column are accepted; frame.types is None/absent."""
    content = (
        'ITEM: TIMESTEP\n0\n'
        'ITEM: NUMBER OF ATOMS\n1\n'
        'ITEM: BOX BOUNDS pp pp pp\n'
        '0.0 5.0\n0.0 5.0\n0.0 5.0\n'
        'ITEM: ATOMS id x y z\n'
        '1 1.0 2.0 3.0\n'
    )
    f = tmp_path / 'no_type.lammpstrj'
    f.write_text(content)
    reader = LammpsTrajectoryReader(str(f))
    frame = next(reader)
    reader.close()
    assert frame.n_atoms == 1
    assert np.allclose(frame.positions[0], [1.0, 2.0, 3.0])
