"""
Copyright (c) Meta Platforms, Inc. and affiliates.

This source code is licensed under the MIT license found in the
LICENSE file in the root directory of this source tree.
"""

from __future__ import annotations

import copy
import logging
from abc import ABCMeta, abstractmethod
from collections import defaultdict
from typing import TYPE_CHECKING, Sequence

import hydra
import torch
from torch import nn

from fairchem.core.common.registry import registry
from fairchem.core.common.utils import (
    load_model_and_weights_from_checkpoint,
)

if TYPE_CHECKING:
    from ase import Atoms

    from fairchem.core.datasets.atomic_data import AtomicData
    from fairchem.core.units.mlip_unit.api.inference import InferenceSettings
    from fairchem.core.units.mlip_unit.mlip_unit import Task


def _get_dataset_to_tasks_map(tasks: Sequence[Task]) -> dict[str, list[Task]]:
    """Create a mapping from dataset names to their associated tasks."""
    dset_to_tasks_map = defaultdict(list)
    for task in tasks:
        for dataset_name in task.datasets:
            dset_to_tasks_map[dataset_name].append(task)
    return dict(dset_to_tasks_map)


class HeadInterface(metaclass=ABCMeta):
    @property
    def use_amp(self):
        return False

    @abstractmethod
    def forward(
        self, data: AtomicData, emb: dict[str, torch.Tensor]
    ) -> dict[str, torch.Tensor]:
        """Head forward.

        Arguments
        ---------
        data: AtomicData
            Atomic systems as input
        emb: dict[str->torch.Tensor]
            Embeddings of the input as generated by the backbone

        Returns
        -------
        outputs: dict[str->torch.Tensor]
            Return one or more targets generated by this head
        """
        return


class BackboneInterface(metaclass=ABCMeta):
    @abstractmethod
    def forward(self, data: AtomicData) -> dict[str, torch.Tensor]:
        """Backbone forward.

        Arguments
        ---------
        data: AtomicData
            Atomic systems as input

        Returns
        -------
        embedding: dict[str->torch.Tensor]
            Return backbone embeddings for the given input
        """
        return

    @classmethod
    @abstractmethod
    def build_inference_settings(cls, settings: InferenceSettings) -> dict:
        """
        Build config overrides from inference settings.

        Called BEFORE model instantiation. Override in subclasses to handle
        model-specific settings like activation_checkpointing.

        Can also perform validation of settings and raise errors if incompatible.

        Returns:
            Dict of backbone config key-value pairs to override.
        """
        return {}

    @abstractmethod
    def validate_tasks(self, dataset_to_tasks: dict[str, list]) -> None:
        """Validate that task datasets are compatible with this backbone."""
        pass  # noqa

    @abstractmethod
    def prepare_for_inference(self, data: AtomicData, settings: InferenceSettings):
        """Prepare backbone for inference. Called once on first prediction.

        Returns:
            self or a new backbone if replacement occurred (e.g., after MOLE merge).
        """
        return self

    @abstractmethod
    def on_predict_check(self, data: AtomicData) -> None:
        """Called before each prediction for any per-prediction checks."""
        pass  # noqa

    @abstractmethod
    def validate_atoms_data(self, atoms: Atoms, task_name: str) -> None:
        """Validate and set defaults for calculator input data.

        Models should set appropriate defaults for atoms.info (e.g., charge, spin)
        and validate that required fields are present and within valid ranges.
        """
        pass  # noqa


@registry.register_model("hydra")
class HydraModel(nn.Module):
    def __init__(
        self,
        backbone: dict | None = None,
        heads: dict | None = None,
        finetune_config: dict | None = None,
        otf_graph: bool = True,
        pass_through_head_outputs: bool = False,
        freeze_backbone: bool = False,
    ):
        super().__init__()
        self.device = None
        self.otf_graph = otf_graph
        # This is required for hydras with models that have multiple outputs per head, since we will deprecate
        # the old config system at some point, this will prevent the need to make major modifications to the trainer
        # because they all expect the name of the outputs directly instead of the head_name.property_name
        self.pass_through_head_outputs = pass_through_head_outputs

        # if finetune_config is provided, then attempt to load the model from the given finetune checkpoint
        starting_model = None
        if finetune_config is not None:
            # Make it hard to sneak more fields into finetuneconfig
            assert (
                len(set(finetune_config.keys()) - {"starting_checkpoint", "override"})
                == 0
            )
            starting_model: HydraModel = load_model_and_weights_from_checkpoint(
                finetune_config["starting_checkpoint"]
            )
            logging.info(
                f"Found and loaded fine-tuning checkpoint: {finetune_config['starting_checkpoint']} (Note we are NOT loading the training state from this checkpoint, only parts of the model and weights)"
            )
            assert isinstance(
                starting_model, HydraModel
            ), "Can only finetune starting from other hydra models!"
            # TODO this is a bit hacky to overrride attrs in the backbone
            if "override" in finetune_config:
                for key, value in finetune_config["override"].items():
                    setattr(starting_model.backbone, key, value)

        if backbone is not None:
            backbone = copy.deepcopy(backbone)
            backbone_model_name = backbone.pop("model")
            self.backbone: BackboneInterface = registry.get_model_class(
                backbone_model_name
            )(
                **backbone,
            )
        elif starting_model is not None:
            self.backbone = starting_model.backbone
            logging.info(
                f"User did not specify a backbone, using the backbone from the starting checkpoint {self.backbone}"
            )
        else:
            raise RuntimeError(
                "Backbone not specified and not found in the starting checkpoint"
            )

        if freeze_backbone:
            for param in self.backbone.parameters():
                param.requires_grad = False

        if heads is not None:
            heads = copy.deepcopy(heads)
            # Iterate through outputs_cfg and create heads
            self.output_heads: dict[str, HeadInterface] = {}

            head_names_sorted = sorted(heads.keys())
            assert len(set(head_names_sorted)) == len(
                head_names_sorted
            ), "Head names must be unique!"
            for head_name in head_names_sorted:
                head_config = heads[head_name]
                if "module" not in head_config:
                    raise ValueError(
                        f"{head_name} head does not specify module to use for the head"
                    )

                module_name = head_config.pop("module")
                self.output_heads[head_name] = registry.get_model_class(module_name)(
                    self.backbone,
                    **head_config,
                )

            self.output_heads = torch.nn.ModuleDict(self.output_heads)
        elif starting_model is not None:
            self.output_heads = starting_model.output_heads
            logging.info(
                f"User did not specify heads, using the output heads from the starting checkpoint {self.output_heads}"
            )
        else:
            raise RuntimeError(
                "Heads not specified and not found in the starting checkpoint"
            )

    def forward(self, data: AtomicData):
        # lazily get device from input to use with amp, at least one input must be a tensor to figure out it's device
        if not self.device:
            device_from_tensors = {
                x.device.type for x in data.values() if isinstance(x, torch.Tensor)
            }
            assert (
                len(device_from_tensors) == 1
            ), f"all inputs must be on the same device, found the following devices {device_from_tensors}"
            self.device = device_from_tensors.pop()

        emb = self.backbone(data)
        # Predict all output properties for all structures in the batch for now.
        out = {}
        for k in self.output_heads:
            with torch.autocast(
                device_type=self.device, enabled=self.output_heads[k].use_amp
            ):
                if self.pass_through_head_outputs:
                    out.update(self.output_heads[k](data, emb))
                else:
                    out[k] = self.output_heads[k](data, emb)

        return out

    @property
    def direct_forces(self) -> bool:
        """
        Whether this model uses direct force prediction.
        """
        return getattr(self.backbone, "direct_forces", False)

    def prepare_for_inference(
        self, data: AtomicData, settings: InferenceSettings
    ) -> None:
        """
        Prepare model for inference. Called once on first prediction.

        Delegates to backbone and each head that implements prepare_for_inference.
        """
        need_eval = False
        # Backbone may return a new backbone (e.g., after MOLE merge)
        new_backbone = self.backbone.prepare_for_inference(data, settings)
        if new_backbone is not self.backbone:
            self.backbone = new_backbone
            # Re-set eval mode after backbone replacement
            need_eval = True

        # Let each head prepare itself if needed (returns new head or self)
        new_output_heads = torch.nn.ModuleDict()
        heads_changed = False
        for head_name, head in self.output_heads.items():
            if hasattr(head, "prepare_for_inference"):
                new_head = head.prepare_for_inference(data, settings)
                new_output_heads[head_name] = new_head
                if new_head is not head:
                    heads_changed = True
            else:
                new_output_heads[head_name] = head

        if heads_changed:
            self.output_heads = new_output_heads
            need_eval = True

        if need_eval:
            self.eval()

    def on_predict_check(self, data: AtomicData) -> None:
        """
        Called before each prediction for any per-prediction checks.
        """
        self.backbone.on_predict_check(data)

    def validate_atoms_data(self, atoms: Atoms, task_name: str) -> None:
        """
        Validate and set defaults for calculator input data.
        """
        self.backbone.validate_atoms_data(atoms, task_name)

    def setup_tasks(self, tasks_config: list) -> None:
        """
        Setup tasks from checkpoint config.

        Args:
            tasks_config: List of task configurations from checkpoint
        """
        tasks = [hydra.utils.instantiate(task_config) for task_config in tasks_config]
        self.tasks = {t.name: t for t in tasks}
        self._dataset_to_tasks = _get_dataset_to_tasks_map(tasks)

        # Let backbone validate tasks
        self.backbone.validate_tasks(self._dataset_to_tasks)

    @property
    def dataset_to_tasks(self) -> dict[str, list]:
        """
        Mapping from dataset names to their associated tasks.
        """
        if not hasattr(self, "_dataset_to_tasks"):
            raise RuntimeError(
                "setup_tasks() must be called before accessing dataset_to_tasks"
            )
        return self._dataset_to_tasks


class HydraModelV2(nn.Module):
    def __init__(
        self,
        backbone: BackboneInterface,
        heads: dict[str, HeadInterface],
        freeze_backbone: bool = False,
    ):
        super().__init__()
        self.backbone = backbone
        self.output_heads = torch.nn.ModuleDict(heads)
        self.device = None
        if freeze_backbone:
            for param in self.backbone.parameters():
                param.requires_grad = False

    def forward(self, data):
        # lazily get device from input to use with amp, at least one input must be a tensor to figure out it's device
        if not self.device:
            device_from_tensors = {
                x.device.type for x in data.values() if isinstance(x, torch.Tensor)
            }
            assert (
                len(device_from_tensors) == 1
            ), f"all inputs must be on the same device, found the following devices {device_from_tensors}"
            self.device = device_from_tensors.pop()

        emb = self.backbone(data)
        # Predict all output properties for all structures in the batch for now.
        out = {}
        for k in self.output_heads:
            with torch.autocast(
                device_type=self.device, enabled=self.output_heads[k].use_amp
            ):
                out[k] = self.output_heads[k](data, emb)
        return out

    @property
    def direct_forces(self) -> bool:
        """
        Whether this model uses direct force prediction.
        """
        return getattr(self.backbone, "direct_forces", False)

    def prepare_for_inference(
        self, data: AtomicData, settings: InferenceSettings
    ) -> None:
        """
        Prepare model for inference. Called once on first prediction.
        """
        data_clone = data.clone()
        need_eval = False
        # Backbone may return a new backbone (e.g., after MOLE merge)
        new_backbone = self.backbone.prepare_for_inference(data_clone, settings)
        if new_backbone is not self.backbone:
            self.backbone = new_backbone
            need_eval = True

        new_output_heads = torch.nn.ModuleDict()
        heads_changed = False
        for head_name, head in self.output_heads.items():
            if hasattr(head, "prepare_for_inference"):
                new_head = head.prepare_for_inference(data_clone, settings)
                new_output_heads[head_name] = new_head
                if new_head is not head:
                    heads_changed = True
            else:
                new_output_heads[head_name] = head

        if heads_changed:
            self.output_heads = new_output_heads
            need_eval = True

        if need_eval:
            self.eval()

    def on_predict_check(self, data: AtomicData) -> None:
        """
        Called before each prediction for any per-prediction checks.
        """
        self.backbone.on_predict_check(data)

    def validate_atoms_data(self, atoms: Atoms, task_name: str) -> None:
        """
        Validate and set defaults for calculator input data.
        """
        self.backbone.validate_atoms_data(atoms, task_name)

    def setup_tasks(self, tasks_config: list) -> None:
        """
        Setup tasks from checkpoint config.

        Args:
            tasks_config: List of task configurations from checkpoint
        """
        tasks = [hydra.utils.instantiate(task_config) for task_config in tasks_config]
        self.tasks = {t.name: t for t in tasks}
        self._dataset_to_tasks = _get_dataset_to_tasks_map(tasks)

        # Let backbone validate tasks
        self.backbone.validate_tasks(self._dataset_to_tasks)

    @property
    def dataset_to_tasks(self) -> dict[str, list]:
        """
        Mapping from dataset names to their associated tasks.
        """
        if not hasattr(self, "_dataset_to_tasks"):
            raise RuntimeError(
                "setup_tasks() must be called before accessing dataset_to_tasks"
            )
        return self._dataset_to_tasks
