import pytest
import torch
from neural_optimiser import test_dir
from torch_geometric.data import Data


def test_calculate_returns_expected_shapes(mace_calculator, batch):
    """Test that calculate returns tensors of expected shapes and dtypes."""
    e, f = mace_calculator(batch)

    assert isinstance(e, torch.Tensor)
    assert isinstance(f, torch.Tensor)
    assert e.shape == (batch.n_conformers,)
    assert f.shape == batch.pos.shape
    assert f.dtype == batch.pos.dtype


def test_to_atomic_data_builds_expected_fields_and_dtypes(mace_calculator, batch):
    """Test conversion to atomic Data object with expected fields and dtypes."""
    atomic = mace_calculator.to_atomic_data(batch)

    # positions concatenated
    assert hasattr(atomic, "positions")
    assert atomic.positions.shape == batch.pos.shape
    assert torch.allclose(atomic.positions, batch.pos)

    # head per-graph index tensor
    assert hasattr(atomic, "head")
    assert atomic.head.dtype == torch.long
    assert atomic.head.numel() == batch.n_conformers

    # edge_index present (we patched to empty)
    assert hasattr(atomic, "edge_index")
    assert atomic.edge_index.shape[0] == 2


def test_atomic_numbers_mapping_and_one_hot(mace_calculator):
    """Test atomic number to index mapping and one-hot encoding."""
    z = torch.tensor([1, 8, 6, 1], dtype=torch.long)
    idx = mace_calculator.atomic_numbers_to_indices(z, z_table=mace_calculator._z_table)
    assert idx.tolist() == [0, 3, 1, 0]

    one_hot = mace_calculator.to_one_hot(idx, num_classes=len(mace_calculator._z_table))
    assert one_hot.shape == (4, len(mace_calculator._z_table))
    assert torch.all(one_hot.sum(dim=-1) == 1)


def test_validate_batch_missing_fields_raises(mace_calculator):
    """Test that _validate_batch raises errors for bad inputs."""
    bad = Data(pos=torch.zeros((1, 3)), batch=torch.zeros(1, dtype=torch.long))
    with pytest.raises(AttributeError, match="Batch must have attributes"):
        mace_calculator._validate_batch(bad)


def test_get_energies_matches_ase(mace_calculator, batch, atoms, device):
    """Compare energies to ASE MACE for the same model."""
    pytest.importorskip("mace", reason="MACE not installed")
    from mace.calculators.mace import MACECalculator as MACECalc

    model_path = test_dir / "models" / "MACE_SPICE2_NEUTRAL.model"
    e = mace_calculator.get_energies(batch)

    mace_calc = MACECalc(model_paths=str(model_path), device=device)
    atoms.calc = mace_calc
    _e = torch.tensor(atoms.get_potential_energy(), dtype=e.dtype, device=device)

    assert torch.isclose(e.squeeze(), _e, atol=1e-4)


def test_mace_calculator(mace_calculator, batch, atoms, device):
    """Compare MACECalculator results to ASE MACE calculator."""
    pytest.importorskip("mace", reason="MACE not installed")
    from mace.calculators.mace import MACECalculator as MACECalc

    model_paths = test_dir / "models" / "MACE_SPICE2_NEUTRAL.model"
    e, f = mace_calculator(batch)

    mace_calc = MACECalc(model_paths=str(model_paths), device=device)
    atoms.calc = mace_calc
    _e = torch.tensor(atoms.get_potential_energy(), dtype=e.dtype, device=device)
    _f = torch.tensor(atoms.get_forces(), dtype=f.dtype, device=device)

    # Ensure comparable shapes/dtypes
    assert torch.isclose(e.squeeze(), _e, atol=1e-4)
    assert torch.allclose(f, _f, atol=1e-4)


def test_mace_calculator2(mace_calculator, batch, atoms, device):
    """Compare MACECalculator results to ASE MACE calculator."""
    pytest.importorskip("mace", reason="MACE not installed")
    from mace.calculators.mace import MACECalculator as MACECalc

    model_paths = test_dir / "models" / "MACE_SPICE2_NEUTRAL.model"
    e = mace_calculator.get_energies(batch)

    mace_calc = MACECalc(model_paths=str(model_paths), device=device)
    atoms.calc = mace_calc
    _e = torch.tensor(atoms.get_potential_energy(), dtype=e.dtype, device=device)

    # Ensure comparable shapes/dtypes
    assert torch.isclose(e.squeeze(), _e, atol=1e-4)
