"""Tests for ciffy.nn module."""

import glob
import os
import pytest
import numpy as np

import ciffy
from ciffy import Scale, operations

from tests.utils import get_test_cif, TORCH_AVAILABLE, DATA_DIR


# =============================================================================
# Test KNN
# =============================================================================

class TestKNN:
    """Tests for operations.knn() function."""

    def test_knn_shape(self, backend):
        """Test that knn returns correct shape."""
        p = ciffy.load(get_test_cif("3SKW"), backend=backend)
        k = 5
        neighbors = operations.knn(p, k=k, scale=Scale.ATOM)

        assert neighbors.shape == (k, p.size())

    def test_knn_residue_scale(self, backend):
        """Test KNN at residue scale."""
        p = ciffy.load(get_test_cif("3SKW"), backend=backend)
        k = 3
        neighbors = operations.knn(p, k=k, scale=Scale.RESIDUE)

        assert neighbors.shape == (k, p.size(Scale.RESIDUE))

    def test_knn_excludes_self(self, backend):
        """Test that knn excludes self (no point is its own neighbor)."""
        p = ciffy.load(get_test_cif("3SKW"), backend=backend)
        neighbors = operations.knn(p, k=3, scale=Scale.ATOM)

        # Check that no atom is its own neighbor
        n_atoms = p.size()
        for i in range(n_atoms):
            neighbor_list = neighbors[:, i]
            if backend == "torch":
                neighbor_list = neighbor_list.numpy()
            assert i not in neighbor_list

    def test_knn_k_too_large(self):
        """Test that knn raises error when k >= n."""
        p = ciffy.load(get_test_cif("3SKW"), backend="numpy")
        with pytest.raises(ValueError, match="k=.* must be less than"):
            operations.knn(p, k=p.size(), scale=Scale.ATOM)


# =============================================================================
# Test PolymerDataset
# =============================================================================

@pytest.mark.skipif(not TORCH_AVAILABLE, reason="PyTorch not available")
class TestPolymerDataset:
    """Tests for PolymerDataset class."""

    def test_dataset_molecule_scale(self):
        """Test dataset at molecule scale."""
        from ciffy.nn import PolymerDataset

        dataset = PolymerDataset(DATA_DIR, scale=Scale.MOLECULE)
        assert len(dataset) > 0

        # Check that items are Polymers
        p = dataset[0]
        assert isinstance(p, ciffy.Polymer)

    def test_dataset_chain_scale(self):
        """Test dataset at chain scale."""
        from ciffy.nn import PolymerDataset

        dataset = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)
        # Chain scale should have more items than molecule scale
        mol_dataset = PolymerDataset(DATA_DIR, scale=Scale.MOLECULE)
        assert len(dataset) >= len(mol_dataset)

    def test_dataset_max_atoms_filter(self):
        """Test that max_atoms filters correctly."""
        from ciffy.nn import PolymerDataset

        # Very small max_atoms should filter out most/all
        dataset_small = PolymerDataset(DATA_DIR, max_atoms=10)
        dataset_large = PolymerDataset(DATA_DIR, max_atoms=100000)

        assert len(dataset_small) <= len(dataset_large)

    def test_dataset_invalid_scale(self):
        """Test that invalid scale raises error."""
        from ciffy.nn import PolymerDataset

        with pytest.raises(ValueError, match="scale must be MOLECULE or CHAIN"):
            PolymerDataset(DATA_DIR, scale=Scale.ATOM)

    def test_dataset_invalid_directory(self):
        """Test that invalid directory raises error."""
        from ciffy.nn import PolymerDataset

        with pytest.raises(FileNotFoundError):
            PolymerDataset("/nonexistent/path/")


# =============================================================================
# Test PolymerEmbedding
# =============================================================================

@pytest.mark.skipif(not TORCH_AVAILABLE, reason="PyTorch not available")
class TestPolymerEmbedding:
    """Tests for PolymerEmbedding class."""

    def test_embedding_atom_scale(self):
        """Test embedding at atom scale."""
        from ciffy.nn import PolymerEmbedding

        embed = PolymerEmbedding(
            scale=Scale.ATOM,
            atom_dim=64,
            residue_dim=32,
            element_dim=16,
        )

        p = ciffy.load(get_test_cif("3SKW"), backend="torch")
        features = embed(p)

        assert features.shape == (p.size(), embed.output_dim)
        assert embed.output_dim == 64 + 32 + 16

    def test_embedding_residue_scale(self):
        """Test embedding at residue scale."""
        from ciffy.nn import PolymerEmbedding

        embed = PolymerEmbedding(
            scale=Scale.RESIDUE,
            residue_dim=64,
        )

        p = ciffy.load(get_test_cif("3SKW"), backend="torch")
        features = embed(p)

        assert features.shape == (p.size(Scale.RESIDUE), 64)

    def test_embedding_invalid_scale_residue_with_atom(self):
        """Test that atom_dim with RESIDUE scale raises error."""
        from ciffy.nn import PolymerEmbedding

        with pytest.raises(ValueError, match="atom_dim cannot be used"):
            PolymerEmbedding(scale=Scale.RESIDUE, atom_dim=64)

    def test_embedding_invalid_scale_residue_with_element(self):
        """Test that element_dim with RESIDUE scale raises error."""
        from ciffy.nn import PolymerEmbedding

        with pytest.raises(ValueError, match="element_dim cannot be used"):
            PolymerEmbedding(scale=Scale.RESIDUE, element_dim=64)

    def test_embedding_no_dims_raises(self):
        """Test that no embedding dims raises error."""
        from ciffy.nn import PolymerEmbedding

        with pytest.raises(ValueError, match="At least one embedding"):
            PolymerEmbedding(scale=Scale.ATOM)

    def test_embedding_gradients(self):
        """Test that embeddings have gradients."""
        from ciffy.nn import PolymerEmbedding

        embed = PolymerEmbedding(scale=Scale.ATOM, atom_dim=32)
        p = ciffy.load(get_test_cif("3SKW"), backend="torch")

        features = embed(p)
        loss = features.sum()
        loss.backward()

        # Check that embedding weights have gradients
        assert embed.atom_embedding.weight.grad is not None


# =============================================================================
# Edge Case Tests
# =============================================================================

@pytest.mark.skipif(not TORCH_AVAILABLE, reason="PyTorch not available")
class TestPolymerDatasetEdgeCases:
    """Edge case tests for PolymerDataset."""

    def test_dataset_empty_directory(self, tmp_path):
        """Dataset with empty directory has length 0."""
        from ciffy.nn import PolymerDataset

        empty_dir = tmp_path / "empty"
        empty_dir.mkdir()

        dataset = PolymerDataset(str(empty_dir))
        assert len(dataset) == 0

    def test_dataset_filters_match_nothing(self):
        """Dataset with impossible filters is empty."""
        from ciffy.nn import PolymerDataset

        # max_atoms=1 should filter out everything
        dataset = PolymerDataset(DATA_DIR, max_atoms=1)
        assert len(dataset) == 0

    def test_dataset_min_atoms_filter(self):
        """Dataset min_atoms filters correctly."""
        from ciffy.nn import PolymerDataset

        # Very large min_atoms should filter out everything
        dataset = PolymerDataset(DATA_DIR, min_atoms=1000000)
        assert len(dataset) == 0

    def test_dataset_all_excluded(self):
        """Dataset with all IDs excluded is empty."""
        from ciffy.nn import PolymerDataset

        # Get all CIF files and exclude them
        import glob
        cif_files = glob.glob(os.path.join(DATA_DIR, "*.cif"))
        exclude_ids = [os.path.basename(f).replace(".cif", "") for f in cif_files]

        dataset = PolymerDataset(DATA_DIR, exclude_ids=exclude_ids)
        assert len(dataset) == 0

    def test_dataset_molecule_types_filter_no_match(self):
        """Dataset with non-matching molecule_types is empty or smaller."""
        from ciffy.nn import PolymerDataset
        from ciffy import Molecule

        # DNA type likely not present in test data
        dataset = PolymerDataset(DATA_DIR, molecule_types=[Molecule.DNA])
        # May be empty or have some entries
        assert len(dataset) >= 0  # Shouldn't crash

    def test_dataset_directory_with_non_cif_files(self, tmp_path):
        """Dataset ignores non-CIF files."""
        from ciffy.nn import PolymerDataset
        import shutil

        # Create directory with mixed files
        mixed_dir = tmp_path / "mixed"
        mixed_dir.mkdir()

        # Add a CIF file
        shutil.copy(get_test_cif("3SKW"), mixed_dir / "3SKW.cif")

        # Add non-CIF files
        (mixed_dir / "readme.txt").write_text("test")
        (mixed_dir / "data.json").write_text("{}")

        dataset = PolymerDataset(str(mixed_dir))
        # Should only find the CIF file
        assert len(dataset) >= 1

    def test_dataset_limit(self):
        """Dataset respects limit parameter."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale

        # Create dataset without limit
        full_dataset = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)
        full_count = len(full_dataset)

        # Skip if dataset too small
        if full_count < 3:
            pytest.skip("Need at least 3 chains to test limit")

        # Create dataset with limit
        limited_dataset = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, limit=2)
        assert len(limited_dataset) == 2

        # Limit larger than dataset should return all
        large_limit = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, limit=10000)
        assert len(large_limit) == full_count

    def test_dataset_min_residues_filter(self):
        """Dataset min_residues filters correctly."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale

        # Very large min_residues should filter out everything
        dataset = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, min_residues=1000000)
        assert len(dataset) == 0

        # Reasonable min_residues should keep some
        dataset_some = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, min_residues=10)
        dataset_all = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)
        assert len(dataset_some) <= len(dataset_all)

    def test_dataset_max_residues_filter(self):
        """Dataset max_residues filters correctly."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale

        # max_residues=1 filters out chains with >1 residue
        # (0-residue chains like ions/water pass this filter)
        dataset_small = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, max_residues=1)
        dataset_medium = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, max_residues=50)
        dataset_all = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)

        # Stricter filter should have fewer or equal entries
        assert len(dataset_small) <= len(dataset_medium)
        assert len(dataset_medium) <= len(dataset_all)

        # Large max_residues should keep all
        dataset_large = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, max_residues=100000)
        assert len(dataset_large) == len(dataset_all)

    def test_dataset_residue_range_filter(self):
        """Dataset with both min and max residues filters correctly."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale

        # Impossible range (min > max effectively) should be empty
        dataset = PolymerDataset(
            DATA_DIR, scale=Scale.CHAIN, min_residues=1000, max_residues=10
        )
        assert len(dataset) == 0

        # Reasonable range
        dataset_range = PolymerDataset(
            DATA_DIR, scale=Scale.CHAIN, min_residues=10, max_residues=1000
        )
        dataset_all = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)
        assert len(dataset_range) <= len(dataset_all)

    def test_dataset_cache(self):
        """Dataset cache returns same chain object on repeated access."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale

        # With cache enabled, repeated access returns same object
        ds_cache = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, cache=True)

        p1 = ds_cache[0]
        p2 = ds_cache[0]
        assert p1 is p2  # Same object returned from cache

    def test_dataset_cache_disabled(self):
        """Dataset without cache returns different objects."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, cache=False)

        p1 = ds[0]
        p2 = ds[0]
        assert p1 is not p2  # Different objects without cache

    def test_dataset_load_preloads_all(self):
        """Dataset load() preloads all items into cache."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, limit=5)
        assert ds._cache is not None

        # Before load, cache is empty
        assert len(ds._cache._items) == 0

        # After load, all items cached
        ds.load()
        assert len(ds._cache._items) == len(ds)

        # Items are accessible
        for i in range(len(ds)):
            assert ds[i] is not None

    def test_dataset_load_returns_self(self):
        """Dataset load() returns self for chaining."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, limit=3).load()
        assert isinstance(ds, PolymerDataset)
        assert len(ds._cache._items) == len(ds)

    def test_dataset_load_raises_without_cache(self):
        """Dataset load() raises ValueError when cache disabled."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, cache=False)

        with pytest.raises(ValueError, match="cache is disabled"):
            ds.load()

    def test_dataset_chain_correctness(self):
        """Dataset returns correct chain for chain_idx."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)
        if len(ds) == 0:
            pytest.skip("No chains in test data")

        # Get first item and verify it's a single chain
        chain = ds[0]
        assert chain.size(Scale.CHAIN) == 1

    def test_dataset_backend_numpy(self):
        """Dataset respects backend='numpy'."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, backend="numpy")
        if len(ds) == 0:
            pytest.skip("No chains in test data")

        p = ds[0]
        assert isinstance(p.coordinates, np.ndarray)

    def test_dataset_backend_torch(self):
        """Dataset respects backend='torch'."""
        import torch
        from ciffy.nn import PolymerDataset
        from ciffy import Scale

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, backend="torch")
        if len(ds) == 0:
            pytest.skip("No chains in test data")

        p = ds[0]
        assert isinstance(p.coordinates, torch.Tensor)


@pytest.mark.skipif(not TORCH_AVAILABLE, reason="PyTorch not available")
class TestPolymerDatasetPathSequence:
    """Tests for PolymerDataset with path sequence input."""

    def test_dataset_from_path_list(self):
        """Dataset accepts list of file paths."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale
        from pathlib import Path

        cif_files = list(Path(DATA_DIR).glob("*.cif"))
        dataset = PolymerDataset(cif_files, scale=Scale.CHAIN)

        # Should have same chains as directory input
        dir_dataset = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)
        assert len(dataset) == len(dir_dataset)

    def test_dataset_from_path_subset(self):
        """Dataset works with subset of files."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale
        from pathlib import Path

        cif_files = sorted(Path(DATA_DIR).glob("*.cif"))
        if len(cif_files) < 2:
            pytest.skip("Need at least 2 CIF files")

        # Use only first file
        subset = cif_files[:1]
        dataset = PolymerDataset(subset, scale=Scale.CHAIN)

        full_dataset = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)
        assert len(dataset) < len(full_dataset)

    def test_dataset_datasplit_integration(self):
        """Dataset integrates with split() method."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale
        from pathlib import Path

        cif_files = list(Path(DATA_DIR).glob("*.cif"))
        if len(cif_files) < 4:
            pytest.skip("Need at least 4 CIF files")

        ds = PolymerDataset(cif_files, scale=Scale.CHAIN)
        train_ds, _, test_ds = ds.split(train=0.5, val=0.25, test=0.25, seed=42)

        # No file overlap between train and test
        train_files = set(p for p, _ in train_ds._index)
        test_files = set(p for p, _ in test_ds._index)
        assert len(train_files & test_files) == 0

    def test_dataset_empty_list(self):
        """Dataset handles empty path list."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale

        dataset = PolymerDataset([], scale=Scale.CHAIN)
        assert len(dataset) == 0

    def test_dataset_path_list_with_filters(self):
        """Filters work with path list input."""
        from ciffy.nn import PolymerDataset
        from ciffy import Scale, Molecule
        from pathlib import Path

        cif_files = list(Path(DATA_DIR).glob("*.cif"))
        dataset = PolymerDataset(cif_files, scale=Scale.CHAIN, molecule_types=Molecule.RNA)

        # Should match directory input with same filter
        dir_dataset = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, molecule_types=Molecule.RNA)
        assert len(dataset) == len(dir_dataset)


@pytest.mark.skipif(not TORCH_AVAILABLE, reason="PyTorch not available")
class TestPolymerEmbeddingEdgeCases:
    """Edge case tests for PolymerEmbedding."""

    def test_embedding_single_dim(self):
        """Embedding with dimension 1 works."""
        from ciffy.nn import PolymerEmbedding

        embed = PolymerEmbedding(scale=Scale.ATOM, atom_dim=1)
        p = ciffy.load(get_test_cif("3SKW"), backend="torch")

        features = embed(p)
        assert features.shape == (p.size(), 1)

    def test_embedding_large_dim(self):
        """Embedding with large dimension works."""
        from ciffy.nn import PolymerEmbedding

        embed = PolymerEmbedding(scale=Scale.ATOM, atom_dim=1024)
        p = ciffy.load(get_test_cif("3SKW"), backend="torch")

        features = embed(p)
        assert features.shape == (p.size(), 1024)

    def test_embedding_only_residue_dim(self):
        """Embedding with only residue_dim at atom scale."""
        from ciffy.nn import PolymerEmbedding

        embed = PolymerEmbedding(scale=Scale.ATOM, residue_dim=64)
        p = ciffy.load(get_test_cif("3SKW"), backend="torch")

        features = embed(p)
        # Should expand residue embeddings to atom level
        assert features.shape == (p.size(), 64)

    def test_embedding_only_element_dim(self):
        """Embedding with only element_dim at atom scale."""
        from ciffy.nn import PolymerEmbedding

        embed = PolymerEmbedding(scale=Scale.ATOM, element_dim=32)
        p = ciffy.load(get_test_cif("3SKW"), backend="torch")

        features = embed(p)
        assert features.shape == (p.size(), 32)

    def test_embedding_all_dims(self):
        """Embedding with all dims at atom scale."""
        from ciffy.nn import PolymerEmbedding

        embed = PolymerEmbedding(
            scale=Scale.ATOM,
            atom_dim=64,
            residue_dim=32,
            element_dim=16,
        )
        p = ciffy.load(get_test_cif("3SKW"), backend="torch")

        features = embed(p)
        assert features.shape == (p.size(), 64 + 32 + 16)

    def test_embedding_on__template_polymer(self):
        """Embedding works on _template polymer."""
        from ciffy.nn import PolymerEmbedding

        embed = PolymerEmbedding(scale=Scale.ATOM, atom_dim=32)
        p = ciffy.template("acgu", backend="torch")

        features = embed(p)
        assert features.shape == (p.size(), 32)

    def test_embedding_on_single_residue(self):
        """Embedding works on single-residue polymer."""
        from ciffy.nn import PolymerEmbedding

        embed = PolymerEmbedding(scale=Scale.RESIDUE, residue_dim=32)
        p = ciffy.template("a", backend="torch")

        features = embed(p)
        assert features.shape == (1, 32)


# =============================================================================
# Test PolymerDataset.split()
# =============================================================================

@pytest.mark.skipif(not TORCH_AVAILABLE, reason="PyTorch not available")
class TestPolymerDatasetSplit:
    """Tests for PolymerDataset.split() method."""

    def test_split_basic_three_way(self):
        """Basic 3-way split works correctly."""
        from ciffy.nn import PolymerDataset
        from pathlib import Path

        cif_files = sorted(Path(DATA_DIR).glob("*.cif"))
        if len(cif_files) < 3:
            pytest.skip("Need at least 3 CIF files for 3-way split test")

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)
        train, val, test = ds.split(0.6, 0.2, 0.2, seed=42)

        # Each split should be a PolymerDataset
        assert isinstance(train, PolymerDataset)
        assert isinstance(val, PolymerDataset)
        assert isinstance(test, PolymerDataset)

        # Total items should match original
        total = len(train) + len(val) + len(test)
        assert total == len(ds)

    def test_split_two_way(self):
        """2-way split (train/test only) works correctly."""
        from ciffy.nn import PolymerDataset
        from pathlib import Path

        cif_files = sorted(Path(DATA_DIR).glob("*.cif"))
        if len(cif_files) < 2:
            pytest.skip("Need at least 2 CIF files for 2-way split test")

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)
        result = ds.split(0.8, 0.0, 0.2, seed=42)

        # Should return tuple of 2
        assert len(result) == 2
        train, test = result

        # Total items should match original
        assert len(train) + len(test) == len(ds)

    def test_split_no_file_overlap(self):
        """Files don't overlap between splits."""
        from ciffy.nn import PolymerDataset
        from pathlib import Path

        cif_files = sorted(Path(DATA_DIR).glob("*.cif"))
        if len(cif_files) < 3:
            pytest.skip("Need at least 3 CIF files")

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)
        train, val, test = ds.split(0.5, 0.25, 0.25, seed=42)

        # Extract files from each split
        train_files = set(p for p, _ in train._index)
        val_files = set(p for p, _ in val._index)
        test_files = set(p for p, _ in test._index)

        # No overlap between splits
        assert len(train_files & val_files) == 0
        assert len(train_files & test_files) == 0
        assert len(val_files & test_files) == 0

    def test_split_preserves_config(self):
        """Split child datasets preserve parent configuration."""
        from ciffy.nn import PolymerDataset
        from ciffy import Molecule

        ds = PolymerDataset(
            DATA_DIR,
            scale=Scale.CHAIN,
            molecule_types=Molecule.RNA,
            backend="numpy",
        )
        if len(ds) < 2:
            pytest.skip("Need at least 2 items")

        train, test = ds.split(0.5, 0.0, 0.5, seed=42)

        # Check configuration is preserved
        assert train.scale == ds.scale
        assert train.molecule_types == ds.molecule_types
        assert train.backend == ds.backend

    def test_split_reproducible_with_seed(self):
        """Same seed produces same split."""
        from ciffy.nn import PolymerDataset
        from pathlib import Path

        cif_files = sorted(Path(DATA_DIR).glob("*.cif"))
        if len(cif_files) < 2:
            pytest.skip("Need at least 2 CIF files")

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)

        train1, test1 = ds.split(0.5, 0.0, 0.5, seed=123)
        train2, test2 = ds.split(0.5, 0.0, 0.5, seed=123)

        train1_files = set(p for p, _ in train1._index)
        train2_files = set(p for p, _ in train2._index)

        assert train1_files == train2_files

    def test_split_invalid_ratios_raises(self):
        """Invalid split ratios raise ValueError."""
        from ciffy.nn import PolymerDataset

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)

        # Negative ratios
        with pytest.raises(ValueError, match="non-negative"):
            ds.split(-0.1, 0.5, 0.6)

        # Sum != 1.0
        with pytest.raises(ValueError, match="sum to 1.0"):
            ds.split(0.5, 0.5, 0.5)

    def test_split_empty_dataset(self):
        """Splitting empty dataset returns empty splits."""
        from ciffy.nn import PolymerDataset

        ds = PolymerDataset(DATA_DIR, max_atoms=1)  # Should filter out everything
        assert len(ds) == 0

        train, test = ds.split(0.8, 0.0, 0.2)
        assert len(train) == 0
        assert len(test) == 0

    def test_split_child_dataset_works(self):
        """Child datasets can load items correctly."""
        from ciffy.nn import PolymerDataset

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)
        if len(ds) < 2:
            pytest.skip("Need at least 2 items")

        train, test = ds.split(0.5, 0.0, 0.5, seed=42)

        # Load an item from each split
        if len(train) > 0:
            p = train[0]
            assert isinstance(p, ciffy.Polymer)
            assert p.size(Scale.CHAIN) == 1

        if len(test) > 0:
            p = test[0]
            assert isinstance(p, ciffy.Polymer)
            assert p.size(Scale.CHAIN) == 1

    def test_split_molecule_scale(self):
        """Split works at molecule scale."""
        from ciffy.nn import PolymerDataset

        ds = PolymerDataset(DATA_DIR, scale=Scale.MOLECULE)
        if len(ds) < 2:
            pytest.skip("Need at least 2 molecules")

        train, test = ds.split(0.5, 0.0, 0.5, seed=42)

        assert len(train) + len(test) == len(ds)

        # Check items are full molecules
        if len(train) > 0:
            p = train[0]
            assert isinstance(p, ciffy.Polymer)

    def test_split_with_cache(self):
        """Split datasets with cache work correctly."""
        from ciffy.nn import PolymerDataset

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, cache=True)
        if len(ds) < 2:
            pytest.skip("Need at least 2 items")

        train, _ = ds.split(0.5, 0.0, 0.5, seed=42)

        # Child dataset should have its own cache
        assert train._cache is not None

        # Load item twice - should return same object
        if len(train) > 0:
            p1 = train[0]
            p2 = train[0]
            assert p1 is p2

    def test_split_without_cache(self):
        """Split datasets without cache work correctly."""
        from ciffy.nn import PolymerDataset

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN, cache=False)
        if len(ds) < 2:
            pytest.skip("Need at least 2 items")

        train, _ = ds.split(0.5, 0.0, 0.5, seed=42)

        # Child dataset should not have cache
        assert train._cache is None

    def test_split_train_only(self):
        """Split with only train returns single dataset."""
        from ciffy.nn import PolymerDataset

        ds = PolymerDataset(DATA_DIR, scale=Scale.CHAIN)

        result = ds.split(1.0, 0.0, 0.0, seed=42)

        # Should return tuple of 1
        assert len(result) == 1
        train = result[0]
        assert len(train) == len(ds)
