"""Auto-generated by cascade Toto2Trainer. Loads the trained checkpoint and
decodes the full horizon in one forward pass via contiguous patch masking
(CPM) — no autoregressive sampling. Exposes:

  forecast(history, horizon, num_samples) -> (1, num_samples, horizon)
      the cascade validator contract — sample paths drawn once from the
      decoded quantiles (seeded per window for validator consensus).
  forecast_quantiles(history, horizon) -> (1, horizon, num_q)
  forecast_quantiles_batch(histories, horizon) -> (B, horizon, num_q)
      the quantile head directly — what benchmark CRPS consumes; batched
      across series so eval sweeps amortize the forward passes.
"""

from __future__ import annotations

import hashlib
import importlib.util
import json
import sys
from pathlib import Path

import numpy as np
import torch

# Single-pass CPM decoding is stable to ~768 steps (Toto 2.0 tech report);
# longer horizons block-decode: commit the median per block, then continue.
STABLE_DECODE_STEPS = 768


def _load_model_module(d: Path):
    spec = importlib.util.spec_from_file_location("cascade_ckpt_model", d / "model.py")
    mod = importlib.util.module_from_spec(spec)
    # Register before exec: model.py defines an @dataclass, and the dataclass
    # machinery does sys.modules.get(cls.__module__).__dict__ during class
    # creation — which is None (AttributeError) unless the module is registered.
    sys.modules[spec.name] = mod
    spec.loader.exec_module(mod)
    return mod


class Wrapper:
    def __init__(self, checkpoint_dir, device: str = "cpu"):
        d = Path(checkpoint_dir)
        self.device = device
        cfg_obj = json.loads((d / "config.json").read_text())
        self.m = _load_model_module(d)
        self.cfg = self.m.Toto2Config(**cfg_obj["toto2"])
        self.quantile_levels = [float(v) for v in cfg_obj["quantile_levels"]]
        self.levels = torch.tensor(self.quantile_levels, dtype=torch.float32, device=device)
        self.model = self.m.Toto2Model(self.cfg).to(device).eval()
        from safetensors.torch import load_file
        state = load_file(str(d / "weights.safetensors"))
        self.model.load_state_dict(state)

    # ── CPM decoding ──────────────────────────────────────────────────────────

    def _prep(self, histories):
        """Left-pad (with the first value) or truncate each 1-D history to the
        context window. Returns the real-space context ``(B, window_len)`` in
        float64 — standardization happens per decode block, from full
        precision, so large-level series keep their fluctuations."""
        ps = self.cfg.patch_size
        n_ctx = max(2, self.cfg.context_length // ps)
        window_len = n_ctx * ps
        rows = []
        for h in histories:
            h = np.asarray(h, dtype=np.float64).reshape(-1)
            if h.shape[0] < window_len:
                pad = np.full(window_len - h.shape[0], h[0] if h.size else 0.0)
                h = np.concatenate([pad, h])
            else:
                h = h[-window_len:]
            rows.append(h)
        return torch.as_tensor(np.stack(rows), dtype=torch.float64, device=self.device)

    @torch.no_grad()
    def _decode_block_z(self, z, block: int):
        """One CPM forward pass: append ``block`` masked patches to the
        normalized context ``(B, L)`` and read their z-space quantiles
        ``(B, block*patch_size, num_q)``."""
        ps = self.cfg.patch_size
        # keep as much context as the positional table allows
        ctx_p = min(z.shape[1] // ps, self.cfg.max_patches - block)
        ctx = z[:, -ctx_p * ps :].view(z.shape[0], ctx_p, ps)
        filler = torch.zeros(z.shape[0], block, ps, dtype=ctx.dtype, device=self.device)
        mask = torch.zeros(z.shape[0], ctx_p + block, dtype=ctx.dtype, device=self.device)
        mask[:, ctx_p:] = 1.0
        pred = self.model(torch.cat([ctx, filler], dim=1), mask=mask)
        # position i predicts patch i+1 → the horizon patches come from
        # positions ctx_p-1 .. ctx_p+block-2.
        q = pred[:, ctx_p - 1 : ctx_p + block - 1]          # (B, block, ps, nq)
        q, _ = torch.sort(q, dim=-1)                        # prevent quantile crossing
        return q.reshape(z.shape[0], block * ps, -1)

    @torch.no_grad()
    def _decode_quantiles(self, x, horizon: int):
        """Block-decode real-space quantiles ``(B, horizon, num_q)`` from the
        real-space context ``x`` ``(B, L)``.

        Each block re-runs the causal scaler over history + committed medians
        and unscales with the resulting end-of-context anchor. Committed
        patches are *observed* context for later blocks, and in training the
        causal stats advance through every observed patch — so the anchor must
        advance with them; reusing the pre-horizon anchor would feed blocks ≥ 2
        a scale/location regime the model never sees in training. Clamp bounds
        are fixed from the original context (min/max ± 1e4x anchor scale, per
        the report) so committed medians can't widen them.
        """
        ps = self.cfg.patch_size
        stable = max(1, min(STABLE_DECODE_STEPS // ps, self.cfg.max_patches - 2))
        remaining = -(-int(horizon) // ps)
        lo = hi = None
        out = []
        while remaining > 0:
            block = min(remaining, stable)
            z, loc_t, scale_t = self.m.causal_standardize(x)
            loc = loc_t[:, -1:].double().unsqueeze(-1)      # (B, 1, 1)
            scale = scale_t[:, -1:].double().unsqueeze(-1)
            if lo is None:
                lo = x.min(dim=-1, keepdim=True).values.unsqueeze(-1) - 1e4 * scale
                hi = x.max(dim=-1, keepdim=True).values.unsqueeze(-1) + 1e4 * scale
            qz = self._decode_block_z(z.to(torch.float32), block)
            q = torch.sinh(qz.double()) * scale + loc       # (B, block*ps, nq)
            q = torch.clamp(q, min=lo, max=hi)
            out.append(q)
            remaining -= block
            if remaining > 0:
                x = torch.cat([x, q[..., q.shape[-1] // 2]], dim=1)
        return torch.cat(out, dim=1)[:, : int(horizon)]

    # ── quantile head (benchmark path) ────────────────────────────────────────

    @torch.no_grad()
    def forecast_quantiles_batch(self, histories, horizon: int) -> np.ndarray:
        """Decode ``len(histories)`` series in one batch → real-space quantiles
        ``(B, horizon, num_q)`` at ``self.quantile_levels``. arcsinh + affine
        are monotone increasing, so quantiles map pointwise."""
        q = self._decode_quantiles(self._prep(list(histories)), horizon)
        return q.detach().cpu().numpy().astype(np.float64)

    def forecast_quantiles(self, history, horizon: int) -> np.ndarray:
        return self.forecast_quantiles_batch([history], horizon)

    # ── validator contract (sample paths) ─────────────────────────────────────

    @torch.no_grad()
    def forecast(self, history, horizon: int, num_samples: int) -> np.ndarray:
        hist = np.asarray(history, dtype=np.float64).reshape(-1)
        # Deterministic per-window sampling: seed from the (raw history, horizon,
        # num_samples) so every validator computes identical scores and king vs
        # challenger share the uniform draws (paired Monte-Carlo).
        seed_src = hist.tobytes() + int(horizon).to_bytes(8, "big") + int(num_samples).to_bytes(8, "big")
        seed = int.from_bytes(hashlib.sha256(seed_src).digest()[:8], "big") & ((1 << 63) - 1)
        generator = torch.Generator(device=self.device)
        generator.manual_seed(seed)

        q = self._decode_quantiles(self._prep([hist]), horizon)[0]  # (h, nq) real-space
        # One draw per step per path via the piecewise-linear inverse CDF of the
        # decoded quantiles (already clamped and monotone in the level).
        # Quantiles decode once; samples never feed back.
        nq = q.shape[-1]
        levels = self.levels
        u = torch.rand(int(num_samples), int(horizon), device=self.device, generator=generator)
        idx = torch.searchsorted(levels, u.clamp(levels[0].item(), levels[-1].item()))
        idx = idx.clamp(1, nq - 1)
        i_lo = idx - 1
        i_hi = idx
        qe = q.unsqueeze(0).expand(u.shape[0], -1, -1)      # (ns, h, nq)
        vl = torch.gather(qe, -1, i_lo.unsqueeze(-1)).squeeze(-1)
        vh = torch.gather(qe, -1, i_hi.unsqueeze(-1)).squeeze(-1)
        ql = levels[i_lo].double(); qh = levels[i_hi].double()
        frac = ((u.double() - ql) / (qh - ql).clamp_min(1e-8)).clamp(0, 1)
        out = vl + frac * (vh - vl)                         # (ns, h)
        return out.detach().cpu().numpy().reshape(1, int(num_samples), int(horizon))