from __future__ import annotations
import torch
from ..base import Method
from . import core

class DirectSIMethod(Method):
    def __init__(self, sigma_train, t_eps, sampling, beta_power=1,
                 eval_noise_scales=None, optimizer_cfg=None):
        self.sigma_train = float(sigma_train)
        self.t_eps = float(t_eps)
        self.beta_power = int(beta_power)
        self.n_sde_steps = int(sampling["n_sde_steps"])
        self.noise_scale = float(sampling.get("noise_scale", 1.0))
        self.eval_noise_scales = ([float(s) for s in eval_noise_scales]
                                  if eval_noise_scales is not None else [self.noise_scale])
        self.optimizer_cfg = optimizer_cfg or {}

    @classmethod
    def from_config(cls, cfg):
        return cls(
            sigma_train=cfg["sigma_train"],
            t_eps=cfg.get("t_eps", 0.0),
            sampling=cfg["sampling"],
            beta_power=cfg.get("beta_power", 1),
            eval_noise_scales=cfg.get("eval_noise_scales"),
            optimizer_cfg=cfg.get("optimizer"),
        )

    def variants_for_eval(self):
        return {f"noise_scale_{_fmt(s)}": {"noise_scale": s} for s in self.eval_noise_scales}

    def loss(self, model, batch):
        x_0, x_T = batch
        return core.global_si_loss(model, x_0, x_T, sigma=self.sigma_train,
                                   t_eps=self.t_eps, beta_power=self.beta_power)

    def sample(self, model, x_0, n_samples, noise_scale=None):
        if noise_scale is None:
            noise_scale = self.noise_scale
        if x_0.ndim == 1:
            x_0 = x_0.unsqueeze(0)
        B = x_0.shape[0]
        out = []
        for i in range(B):
            cloud = core.sde_sample(model, x_0[i], n_samples=n_samples,
                                    n_steps=self.n_sde_steps, sigma=self.sigma_train,
                                    noise_scale=noise_scale)
            out.append(cloud)
        return torch.stack(out, dim=0)

def _fmt(s):
    if s == int(s):
        return f"{int(s)}p0"
    return f"{s}".replace(".", "p")