from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import torch
from torch import Tensor
from torch.utils.data import DataLoader

import qqtools as qt

__all__ = ["qTaskBase", "PotentialTaskBase"]

OPTIONAL_METHODS = [
    # === functional hook ===
    "get_loss_fn",
    "batch_cache",
    "epoch_metric",
    "pipe_middle_ware",
    "to",
    # === lifetime hook ===
    "onEpochStart",
    "onEpochEnd",
    "onBatchEnd",
    "on_better_model",
    "state_dict",
    "load_state_dict",
]

REQUIRED_METHODS = []


class qTaskBase(ABC):
    """
    Optional methods (can be implemented by users):

    - def get_loss_fn(self, args) -> Callable:
    - def batch_cache(out, batch_data) -> Dict[str, Tensor]: For cross-batch caching
    - def epoch_metric(gathered_cache: Dict[str, Tensor]) -> Dict: For cross-batch metrics calculation
    """

    train_loader: DataLoader
    val_loader: DataLoader
    test_loader: DataLoader
    meta: Dict[str, Any]  # flow with epoch runner and saved in checkpoint

    _opt_impl = set()  # implemented optional methods
    _opt_todo = set()  # un-implemented optional methods
    _required = set(REQUIRED_METHODS)

    def __init__(self):
        """
        set 3 dataloaders:
            - self.train_loader
            - self.val_loader
            - self.test_loader
        """
        for name in OPTIONAL_METHODS:
            if qt.is_override(self, name, qTaskBase):
                self._opt_impl.add(name)
            else:
                self._opt_todo.add(name)

        for name in REQUIRED_METHODS:
            assert qt.is_override(self, name, qTaskBase)

    def has_implemented(self, method_name: str) -> bool:
        """Check if an optional method has been overridden by the subclass."""
        return method_name in self._opt_impl

    @abstractmethod
    def batch_forward(self, model, batch_data) -> Dict[str, Tensor]:
        """data format & model interface

        Returns:
            out (dict): dict-like
        """
        pass

    @abstractmethod
    def batch_loss(self, out, batch_data) -> Dict[str, Tuple[Tensor, int]]:
        """
        Returns:
            dict:  `{ 'loss': (loss_value, sample_count) }`
        """
        pass

    @abstractmethod
    def batch_metric(self, out, batch_data) -> Dict[str, Tuple[Tensor, int]]:
        """
        Returns:
            dict:  `{ 'metric_name': (metric_value, sample_count) }`
        """
        pass

    @abstractmethod
    def post_metrics_to_value(self, result) -> float:
        """
        In cases where multiple metrics are available, the error metric must be prioritized
        to identify the optimal validation performance.

        Args:
            result: Dict[metric_name, metric_avg]

        Returns:
            float: The performance error value.
        """
        pass

    @staticmethod
    def pre_batch_forward(batch_data):
        """
        Args:
            batch_data (dict):

        Returns:
            dict: > dict of preprocessed data
        """
        return batch_data

    @staticmethod
    def post_batch_forward(output, batch_data):
        """
        Args:
           output :  return value of `batch_forward`

        Returns:
            dict: revised output
        """
        return output

    # other conventions:
    def get_loss_fn(self, args) -> Callable:
        raise NotImplementedError

    def batch_cache(self, out, batch_data) -> Dict[str, Union[Tensor, float]]:
        """Return tensors that need to be cached across the epoch.

        Use this to collect metrics like true/predicted labels that need
        to be aggregated at the end of an epoch (e.g., for AUC, Global MAE).
        The framework will safely detach and move them to CPU automatically.

        Warning:
            Do NOT cache high-dimensional tensors (e.g., full feature maps)
            as this will cause system RAM Out-Of-Memory errors!

        Args:
            out (dict): Model outputs
            batch_data (dict-like qData): Batch input data

        Returns:
            dict: `{'key_name': tensor_or_scalar}`
        """
        raise NotImplementedError

    def epoch_metric(self, gathered_cache: Dict[str, Tensor]) -> Dict:
        """
        Receive combined tensors across all batches and specific DDP nodes
        at the end of an epoch, and calculate global metrics.

        `gathered_cache` contains the concatenated tensors provided by `batch_cache`.
        Tensors will be aggregated across all GPUs if in DDP mode.

        Args:
            gathered_cache (dict): Combined dictionary of cached tensors

        Returns:
            dict: dict of epoch level metrics `{'metric_name': value}`
        """
        raise NotImplementedError

    def on_better_model(self, epoch_agent, current_state) -> None:
        raise NotImplementedError

    def to(self, device) -> None:
        raise NotImplementedError

    def state_dict(self) -> Dict:
        raise NotImplementedError

    def load_state_dict(self, d: Dict) -> None:
        raise NotImplementedError

    # def pipe_middle_ware(self, pipe: qPipeline):
    # pass


class PotentialTaskBase(qTaskBase):
    """

    Inherited from qTaskBase, designed for potential energy surface prediction tasks.
    Conventions:
        - implement `batch_forward`, `batch_loss`, `batch_metric`, `post_metrics_to_value` following qTaskBase;
        - implement `init_loader()` to set `train_loader`, `val_loader`, `test_loader`;
        - implement `norm_factor` to set `self.meta["norm_factor"]`;
        - call `super().__init__(args)` in `__init__()`;
        - `args` should follow `qargs-convention`;
    """

    def __init__(self, args, training=True):
        super().__init__()
        self.extract_args(args)

        if training and hasattr(self, "init_loader"):
            self.init_loader()

    @abstractmethod
    def norm_factor(self) -> Tuple[float, float]:
        raise NotImplementedError

    def extract_args(self, args):
        self.args = args.copy()

        batch_size = args.task.dataloader.batch_size
        eval_batch_size = args.task.dataloader.eval_batch_size or batch_size
        num_workers = args.task.dataloader.num_workers
        pin_memory = args.task.dataloader.pin_memory
        distributed = args.distributed
        standarize = args.task.get("standarize", False)
        with_force = args.task.get("with_force", False)

        loader_meta = {
            "batch_size": batch_size,
            "eval_batch_size": eval_batch_size,
            "num_workers": num_workers,
            "pin_memory": pin_memory,
            "distributed": distributed,
        }

        meta = {
            "standarize": standarize,
            "with_force": with_force,
        }

        self.loader_meta = loader_meta
        self.meta = meta

        self.standarize = standarize
        self.with_force = with_force

        prefix = self.__class__.__name__
        print(f"[{prefix}] loader_meta: {self.loader_meta}")
        if self.standarize:
            self.meta["norm_factor"] = self.norm_factor
        else:
            self.meta["norm_factor"] = (0, 1)
        print(f"[{prefix}] standarize: {self.standarize} norm_factor: {self.meta["norm_factor"]}")

    def batch_forward(self, model, batch_data, **meta) -> Dict[str, Tensor]:
        with_force = self.with_force
        z, pos, batch = batch_data["z"], batch_data["pos"], batch_data["batch"]
        out_force: Optional[Tensor] = None
        if with_force:
            pos.requires_grad_()
            result: Union[Tensor, Tuple[Tensor, Tensor]] = model(z, pos, batch)
            if isinstance(result, tuple):
                out_energy, out_force = result  # adaption for models that output force
            else:
                out_energy = result
                out_force = -torch.autograd.grad(
                    out_energy,
                    pos,
                    grad_outputs=torch.ones_like(out_energy),
                    create_graph=True,
                    allow_unused=True,
                )[0]

        else:
            out_energy = model(z, pos, batch)

        if out_energy.shape[-1] == 1:
            out_energy = out_energy.view(-1)

        out = {"energy": out_energy}
        if with_force and out_force is not None:
            if out_force.shape[-1] != 3:
                out_force = out_force.view(-1, 3)
            out["force"] = out_force

        return out

    def batch_metric(self, out, batch_data, **meta) -> Dict[str, Tuple[Tensor, int]]:
        mean, std = meta["norm_factor"]
        ret = dict()
        with torch.no_grad():
            e_out = out["energy"].detach().flatten()
            e_target = batch_data.energy.flatten()
            e_pred = e_out * std + mean
            e_mae = torch.nn.functional.l1_loss(e_pred, e_target)
            ret["e_mae"] = (e_mae, e_out.shape[0])

            if self.with_force:
                f_out = out["force"].flatten().detach()
                f_target = batch_data.force.flatten()
                f_pred = f_out * std + mean
                f_err = torch.nn.functional.l1_loss(f_pred, f_target)
                ret["f_mae"] = (f_err, f_out.numel)

        return ret

    def batch_loss(self, out, batch_data, loss_fn, **meta) -> Dict[str, Tuple[Tensor, int]]:
        mean, std = meta["norm_factor"]
        ret = dict()

        e_target = batch_data.energy.flatten()
        e_out = out["energy"] * std + mean
        e_loss: Tensor = loss_fn["energy"](e_out, (e_target - mean) / std)
        ret["e_loss"] = (e_loss, e_out.shape[0])
        if self.with_force:
            f_out = out["force"].view(-1, 3)
            f_target = batch_data.force.view(-1, 3)
            f_loss = loss_fn["force"](f_out, f_target / std)
            ret["f_loss"] = (f_loss, f_out.numel)
            ew, fw = self.ef_weight
            ret["loss"] = (ew * e_loss + fw * f_loss, e_out.shape[0])
        else:
            ret["loss"] = ret["e_loss"]
        return ret

    def post_metrics_to_value(self, result) -> float:
        if self.with_force:
            if hasattr(self, "ef_weight") and self.ef_weight is not None:
                ew, fw = self.ef_weight
            else:
                ew, fw = 0.8, 0.2
            ret = ew * result["e_mae"] + fw * result["f_mae"]
        else:
            ret = result["e_mae"]
        if torch.is_tensor(ret):
            ret = ret.detach().cpu().numpy()
        return ret


'''
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import torch
from torch import Tensor
from qqtools.pipeline import qTaskBase

class TemplateTask(qTaskBase):

    def __init__(self, args):
        pass

    def batch_forward(self, model, batch_data) -> Dict[str, Tensor]:
        """
        Returns:
            out (dict): dict-like
        """
        pass

    def batch_metric(self, out, batch_data) -> Dict[str, Tuple[Tensor, int]]:
        """
        Returns:
            dict:  `{ 'metric_name': (metric_value, sample_count) }`
        """
        pass

    def batch_loss(self, out, batch_data, loss_fn) -> Dict[str, Tuple[Tensor, int]]:
        """
        Returns:
            dict:  `{ 'loss': (loss_value, sample_count) }`
        """
        pass

    def post_metrics_to_value(self, result) -> float:
        pass
'''
