"""
Tests for loading and saving molecular structures.

Tests are parameterized to run with both numpy and torch backends.
"""

import glob
import os
import tempfile
import pytest

TESTS_DIR = os.path.dirname(__file__)
DATA_DIR = os.path.join(TESTS_DIR, "data")
CIF_FILES = sorted(glob.glob(os.path.join(DATA_DIR, "*.cif")))
BACKENDS = ["numpy", "torch"]


class TestLoad:
    """Test CIF file loading with both numpy and torch backends."""

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    @pytest.mark.parametrize("backend", BACKENDS)
    def test_load_file(self, cif_file, backend):
        from ciffy import load

        polymer = load(cif_file, backend=backend)
        assert polymer is not None
        assert not polymer.empty()
        assert polymer.size() > 0

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    @pytest.mark.parametrize("backend", BACKENDS)
    def test_load_has_coordinates(self, cif_file, backend):
        from ciffy import load

        polymer = load(cif_file, backend=backend)
        assert polymer.coordinates.shape[0] == polymer.size()
        assert polymer.coordinates.shape[1] == 3

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    @pytest.mark.parametrize("backend", BACKENDS)
    def test_load_has_atoms(self, cif_file, backend):
        from ciffy import load

        polymer = load(cif_file, backend=backend)
        assert polymer.atoms.shape[0] == polymer.size()
        assert polymer.elements.shape[0] == polymer.size()

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    @pytest.mark.parametrize("backend", BACKENDS)
    def test_load_has_chains(self, cif_file, backend):
        from ciffy import load, Scale

        polymer = load(cif_file, backend=backend)
        assert polymer.size(Scale.CHAIN) > 0
        assert len(polymer.names) == polymer.size(Scale.CHAIN)
        assert len(polymer.strands) == polymer.size(Scale.CHAIN)

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    @pytest.mark.parametrize("backend", BACKENDS)
    def test_load_has_residues(self, cif_file, backend):
        from ciffy import load, Scale

        polymer = load(cif_file, backend=backend)
        assert polymer.size(Scale.RESIDUE) > 0
        assert polymer.sequence.shape[0] == polymer.size(Scale.RESIDUE)

    @pytest.mark.parametrize("backend", BACKENDS)
    def test_load_nonexistent_file(self, backend):
        from ciffy import load

        with pytest.raises(OSError):
            load("nonexistent_file.cif", backend=backend)

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    def test_backend_conversion(self, cif_file):
        """Test that polymers can be converted between backends."""
        from ciffy import load
        import numpy as np

        # Load as numpy
        np_polymer = load(cif_file, backend="numpy")
        assert isinstance(np_polymer.coordinates, np.ndarray)

        # Convert to torch
        torch_polymer = np_polymer.torch()
        import torch
        assert isinstance(torch_polymer.coordinates, torch.Tensor)

        # Convert back to numpy
        np_polymer2 = torch_polymer.numpy()
        assert isinstance(np_polymer2.coordinates, np.ndarray)

        # Verify coordinates are equivalent
        assert np.allclose(np_polymer.coordinates, np_polymer2.coordinates)


class TestSave:
    """Test PDB file saving with both backends."""

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    @pytest.mark.parametrize("backend", BACKENDS)
    def test_save_file(self, cif_file, backend):
        from ciffy import load, RNA

        polymer = load(cif_file, backend=backend)

        # Get RNA chains only (PDB writer currently supports RNA)
        rna = polymer.subset(RNA)
        if rna.empty():
            pytest.skip("No RNA chains in structure")
        if rna.polymer_count == 0:
            pytest.skip("RNA subset has no polymer atoms (all HETATM)")

        with tempfile.NamedTemporaryFile(suffix=".pdb", delete=False) as f:
            output_path = f.name

        try:
            rna.write_pdb(output_path)
            assert os.path.exists(output_path)
            assert os.path.getsize(output_path) > 0
        finally:
            if os.path.exists(output_path):
                os.unlink(output_path)

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    @pytest.mark.parametrize("backend", BACKENDS)
    def test_save_and_check_content(self, cif_file, backend):
        from ciffy import load, RNA

        polymer = load(cif_file, backend=backend)

        rna = polymer.subset(RNA)
        if rna.empty():
            pytest.skip("No RNA chains in structure")
        if rna.polymer_count == 0:
            pytest.skip("RNA subset has no polymer atoms (all HETATM)")

        with tempfile.NamedTemporaryFile(suffix=".pdb", delete=False) as f:
            output_path = f.name

        try:
            rna.write_pdb(output_path)

            with open(output_path, 'r') as f:
                content = f.read()

            # Check PDB format
            lines = content.strip().split('\n')
            assert len(lines) > 0

            for line in lines:
                assert line.startswith("ATOM")
                # Check line has expected PDB format length
                assert len(line) >= 54  # Minimum PDB ATOM record length

        finally:
            if os.path.exists(output_path):
                os.unlink(output_path)


class TestRoundTrip:
    """Test loading, manipulating, and saving structures with both backends."""

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    @pytest.mark.parametrize("backend", BACKENDS)
    def test_load_center_save(self, cif_file, backend):
        from ciffy import load, RNA, MOLECULE

        polymer = load(cif_file, backend=backend)

        rna = polymer.subset(RNA)
        if rna.empty():
            pytest.skip("No RNA chains in structure")
        if rna.polymer_count == 0:
            pytest.skip("RNA subset has no polymer atoms (all HETATM)")

        # Center the structure
        centered, means = rna.center(MOLECULE)
        assert centered.size() == rna.size()

        with tempfile.NamedTemporaryFile(suffix=".pdb", delete=False) as f:
            output_path = f.name

        try:
            centered.write_pdb(output_path)
            assert os.path.exists(output_path)
            assert os.path.getsize(output_path) > 0
        finally:
            if os.path.exists(output_path):
                os.unlink(output_path)

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    @pytest.mark.parametrize("backend", BACKENDS)
    def test_load_select_chain_save(self, cif_file, backend):
        from ciffy import load, RNA

        polymer = load(cif_file, backend=backend)

        rna = polymer.subset(RNA)
        if rna.empty():
            pytest.skip("No RNA chains in structure")
        if rna.polymer_count == 0:
            pytest.skip("RNA subset has no polymer atoms (all HETATM)")

        # Select first chain
        first_chain = rna.select(0)
        assert first_chain.size() > 0

        with tempfile.NamedTemporaryFile(suffix=".pdb", delete=False) as f:
            output_path = f.name

        try:
            first_chain.write_pdb(output_path)
            assert os.path.exists(output_path)
        finally:
            if os.path.exists(output_path):
                os.unlink(output_path)


class TestCifSave:
    """Test CIF file saving with both backends."""

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    @pytest.mark.parametrize("backend", BACKENDS)
    def test_save_cif(self, cif_file, backend):
        """Test basic CIF writing."""
        from ciffy import load

        polymer = load(cif_file, backend=backend)

        with tempfile.NamedTemporaryFile(suffix=".cif", delete=False) as f:
            output_path = f.name

        try:
            polymer.write(output_path)
            assert os.path.exists(output_path)
            assert os.path.getsize(output_path) > 0
        finally:
            if os.path.exists(output_path):
                os.unlink(output_path)

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    @pytest.mark.parametrize("backend", BACKENDS)
    def test_cif_has_header(self, cif_file, backend):
        """Test that saved CIF has proper header."""
        from ciffy import load

        polymer = load(cif_file, backend=backend)

        with tempfile.NamedTemporaryFile(suffix=".cif", delete=False) as f:
            output_path = f.name

        try:
            polymer.write(output_path)

            with open(output_path, 'r') as f:
                content = f.read()

            # Check CIF format has data_ header
            assert content.startswith("data_")
            # Check has loop_ sections
            assert "loop_" in content
            # Check has _atom_site block
            assert "_atom_site." in content

        finally:
            if os.path.exists(output_path):
                os.unlink(output_path)

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    @pytest.mark.parametrize("backend", BACKENDS)
    def test_round_trip_cif(self, cif_file, backend):
        """Test load -> save -> load preserves data."""
        from ciffy import load
        import numpy as np

        original = load(cif_file, backend=backend)

        with tempfile.NamedTemporaryFile(suffix=".cif", delete=False) as f:
            output_path = f.name

        try:
            original.write(output_path)
            reloaded = load(output_path, backend=backend)

            # Note: CIF writer currently only writes polymer atoms,
            # so we compare polymer counts instead of total counts
            assert reloaded.size() == original.polymer_count

            # Verify polymer coordinates are close (allow small float precision loss)
            orig_coords = np.asarray(original.coordinates[:original.polymer_count])
            reload_coords = np.asarray(reloaded.coordinates)
            assert np.allclose(orig_coords, reload_coords, atol=0.001)

            # Verify chain count matches
            assert len(reloaded.names) == len(original.names)

        finally:
            if os.path.exists(output_path):
                os.unlink(output_path)

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    @pytest.mark.parametrize("backend", BACKENDS)
    def test_round_trip_preserves_sequence(self, cif_file, backend):
        """Test that round-trip preserves residue sequence for polymer residues.

        Note: The CIF writer only writes polymer atoms (group_PDB="ATOM").
        Non-polymer residues (HETATM) are excluded, so the reloaded sequence
        may be shorter than the original.
        """
        from ciffy import load
        import numpy as np

        original = load(cif_file, backend=backend)

        with tempfile.NamedTemporaryFile(suffix=".cif", delete=False) as f:
            output_path = f.name

        try:
            original.write(output_path)
            reloaded = load(output_path, backend=backend)

            # Verify reloaded polymer has valid sequence
            reload_seq = np.asarray(reloaded.sequence)
            assert len(reload_seq) > 0, "Reloaded structure should have residues"

            # Verify residue count matches what was written (polymer atoms only)
            # The reloaded polymer_count should match what we wrote
            assert reloaded.polymer_count == reloaded.size(), \
                "Reloaded should be all polymer (no HETATM in round-trip)"

        finally:
            if os.path.exists(output_path):
                os.unlink(output_path)

    @pytest.mark.parametrize("cif_file", CIF_FILES)
    def test_cif_save_from_subset(self, cif_file):
        """Test saving a subset of the structure to CIF."""
        from ciffy import load, RNA

        polymer = load(cif_file, backend="numpy")

        rna = polymer.subset(RNA)
        if rna.empty():
            pytest.skip("No RNA chains in structure")
        if rna.polymer_count == 0:
            pytest.skip("RNA subset has no polymer atoms (all HETATM)")

        with tempfile.NamedTemporaryFile(suffix=".cif", delete=False) as f:
            output_path = f.name

        try:
            rna.write(output_path)
            assert os.path.exists(output_path)
            assert os.path.getsize(output_path) > 0

            # Reload and verify polymer size
            # Note: CIF writer only writes polymer atoms
            reloaded = load(output_path, backend="numpy")
            assert reloaded.size() == rna.polymer_count

        finally:
            if os.path.exists(output_path):
                os.unlink(output_path)
