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

class CascadedSIMethod(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_per_seg = int(sampling["n_sde_steps_per_seg"])
        self.n_segments = int(sampling.get("n_segments", 10))
        self.noise_scale = float(sampling.get("noise_scale", 1.0))
        # Segment boundaries — uniform partition of [0, 1] by default, matching
        # the normalized physical forecast horizon.
        if "t_values" in sampling:
            self.t_values = list(sampling["t_values"])
        else:
            self.t_values = [k / self.n_segments for k in range(self.n_segments + 1)]
        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):
        if cfg.get("anchor_ic", False):
            raise ValueError("This release supports the paper's SI models without retained initial-state inputs")
        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_tk, x_tk1, t_k, t_k1 = batch
        return core.local_si_loss(model, x_tk, x_tk1, t_k, t_k1,
                                  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.rollout(model, x_0[i], n_samples=n_samples,
                                 t_values=self.t_values,
                                 n_sde_steps_per_seg=self.n_sde_steps_per_seg,
                                 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")