"""EDMMethod — wraps the math in core.py."""

import torch

from ..base import Method
from . import core


class EDMMethod(Method):
    """EDM (Karras et al. 2022) for conditional endpoint prediction.

    Hyperparameters
    ---------------
    sigma_data : float
        EDM σ_data; matches the empirical per-feature std of the (normalized)
        training x_T. For data normalized to unit variance per feature this is
        1.0.
    P_mean, P_std : float
        Log-normal training-σ sampler (Table 1 of the EDM paper).
    sampling : dict
        Keys: n_steps, sigma_min, sigma_max, rho. Used by Heun's sampler.
    """

    def __init__(self, sigma_data: float, P_mean: float, P_std: float, sampling: dict,
                 optimizer_cfg: dict | None = None):
        self.sigma_data = float(sigma_data)
        self.P_mean = float(P_mean)
        self.P_std = float(P_std)
        self.n_steps = int(sampling["n_steps"])
        self.sigma_min = float(sampling["sigma_min"])
        self.sigma_max = float(sampling["sigma_max"])
        self.rho = float(sampling["rho"])
        # Picked up by base Method.configure_optimizer; defaults to plain Adam.
        self.optimizer_cfg = optimizer_cfg or {}

    @classmethod
    def from_config(cls, cfg: dict) -> "EDMMethod":
        return cls(
            sigma_data=cfg["sigma_data"],
            P_mean=cfg["P_mean"],
            P_std=cfg["P_std"],
            sampling=cfg["sampling"],
            optimizer_cfg=cfg.get("optimizer"),
        )

    def loss(self, model: torch.nn.Module, batch) -> torch.Tensor:
        x_0, x_T = batch
        return core.edm_loss(model, x_T, x_0, self.sigma_data, self.P_mean, self.P_std)

    def sample(self, model: torch.nn.Module, x_0: torch.Tensor, n_samples: int) -> torch.Tensor:
        """Return Heun-sampled clouds, shape [B, n_samples, D].

        x_0 may be either a single conditioning vector of shape (D,) or a
        batch of shape (B, D). The output is always (B, n_samples, D).
        """
        device = next(model.parameters()).device
        if x_0.ndim == 1:
            x_0 = x_0.unsqueeze(0)
        B = x_0.shape[0]
        clouds = []
        for i in range(B):
            cloud = core.heun_sample(
                model,
                x_0[i],
                n_samples=n_samples,
                N=self.n_steps,
                sigma_min=self.sigma_min,
                sigma_max=self.sigma_max,
                rho=self.rho,
                sigma_data=self.sigma_data,
                device=device,
            )
            clouds.append(cloud)
        return torch.stack(clouds, dim=0)