"""Method interface. EDM, Cascaded SI and Direct SI subclass this."""

from abc import ABC, abstractmethod

import torch


class Method(ABC):
    """Training-method interface.

    A method bundles together (a) the training-time loss, (b) the sampling
    procedure used at eval time, and (c) the optimizer for the model. The
    trainer is method-agnostic — it calls these three hooks and nothing else.
    """

    @abstractmethod
    def loss(self, model: torch.nn.Module, batch: tuple[torch.Tensor, ...]) -> torch.Tensor:
        """Per-batch training loss. batch = (x_0, x_T) tensors."""

    @abstractmethod
    def sample(self, model: torch.nn.Module, x_0: torch.Tensor, n_samples: int,
               **variant_kwargs) -> torch.Tensor:
        """Generate n_samples per x_0 conditioning vector. Returns [B, n_samples, D].

        Methods that support multiple eval variants accept variant-specific
        overrides as keyword arguments (e.g. ``noise_scale=0.25`` for SI). The
        legal kwargs are whatever ``variants_for_eval()`` produces.
        """

    def variants_for_eval(self) -> dict[str, dict]:
        """Eval-time sampling variants this method exposes.

        Returns a dict ``{variant_label: sample_kwargs}``; ``bifurcation.sample`` writes
        one directory per variant.

        Default: one variant called ``default`` with no kwarg overrides — used
        by EDM (the Heun sampler is deterministic; there's no analogous
        noise-scale knob).
        """
        return {"default": {}}

    def configure_optimizer(self, model: torch.nn.Module, cfg: dict) -> torch.optim.Optimizer:
        """Build the optimizer from ``self.optimizer_cfg`` (set by the method's __init__).

        Supported types:
          - ``adam``  (default)  → ``torch.optim.Adam`` ; weight_decay ignored
          - ``adamw``            → ``torch.optim.AdamW``; ``weight_decay`` honoured
                                   (true decoupled decay — the right way to regularise
                                   adaptive-LR optimisers)

        Default (no ``optimizer`` block in the method section, or
        ``optimizer.type == "adam"`` with ``weight_decay`` 0): plain Adam.

        A non-None ``cfg["weight_decay"]`` (the training section, passed by
        the trainer) overrides the method's block and selects AdamW with that decay.
        """
        opt_cfg = getattr(self, "optimizer_cfg", None) or {}
        opt_type = str(opt_cfg.get("type", "adam")).lower()
        wd = float(opt_cfg.get("weight_decay", 0.0))
        params = {"lr": cfg["lr"], "betas": (0.9, 0.999), "eps": 1e-8}
        wd_override = cfg.get("weight_decay")
        if wd_override is not None:
            return torch.optim.AdamW(model.parameters(),
                                     weight_decay=float(wd_override), **params)
        if opt_type == "adam":
            return torch.optim.Adam(model.parameters(), **params)
        if opt_type == "adamw":
            return torch.optim.AdamW(model.parameters(), weight_decay=wd, **params)
        raise ValueError(f"unknown optimizer.type: {opt_type!r}")