import torch
from metatomic.torch import ModelEvaluationOptions, ModelMetadata, System

from metatrain.utils.data import DatasetInfo
from metatrain.utils.neighbor_lists import (
    get_requested_neighbor_lists,
    get_system_with_neighbor_lists,
)

from .architectures import ArchitectureTests


class ExportedTests(ArchitectureTests):
    """Test suite to test exported models."""

    avoid_consistency_check: list[str] = []
    """List of output names for which to avoid consistency checks."""

    def test_to(
        self,
        device: torch.device,
        dtype: torch.dtype,
        model_hypers: dict,
        dataset_info: DatasetInfo,
    ) -> None:
        """Tests that the `.to()` method of the exported model works.

        In other words, it tests that the exported model can be moved to
        different devices and dtypes.

        :param device: The device to move the exported model to.
        :param dtype: The dtype to move the exported model to.
        :param model_hypers: Hyperparameters to initialize the model.
        :param dataset_info: Dataset information to initialize the model.
        """

        model = self.model_cls(model_hypers, dataset_info).to(dtype=dtype)

        exported = model.export(metadata=ModelMetadata(name="test"))

        # test correct metadata
        assert "This is the test model" in str(exported.metadata())

        exported.to(device=device)

        system = System(
            types=torch.tensor([6, 6]),
            positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
            cell=torch.zeros(3, 3),
            pbc=torch.tensor([False, False, False]),
        )
        requested_neighbor_lists = get_requested_neighbor_lists(exported)
        system = get_system_with_neighbor_lists(system, requested_neighbor_lists)
        system = system.to(device=device, dtype=dtype)

        strict_evaluation_options = ModelEvaluationOptions(
            length_unit=dataset_info.length_unit,
            outputs={
                k: v
                for k, v in model.outputs.items()
                if k not in self.avoid_consistency_check
            },
        )

        exported([system], strict_evaluation_options, check_consistency=True)

        if len(self.avoid_consistency_check) > 0:
            evaluation_options = ModelEvaluationOptions(
                length_unit=dataset_info.length_unit,
                outputs={k: model.outputs[k] for k in self.avoid_consistency_check},
            )
            exported([system], evaluation_options, check_consistency=False)
