"""Test-time ensembling for Modular-TS (v3): context-length ensembles, temporal-hierarchy
(multi-resolution) reconciliation and in-context backtest weighting. Inference only.

Members (spec string, comma separated):
  cN   - the base forecaster with max_context N (e.g. c8192, c2048, c512)
  dK   - downsampled member: context mean-pooled by K (aligned to the forecast origin, so it
         sees up to K x the history at the same token budget), forecast ceil(H/K) coarse
         steps, step-upsampled back. Used as a *level corrector* for the fine forecasts
         (temporal hierarchies / THieF, Athanasopoulos et al. 2017): fine + a*(coarse - pool(fine))
         instead of being averaged in directly, so intra-period shape comes from full res.
Combination of the c-members: quantile averaging (Vincentisation), uniform or, with
backtest=True, per-series softmax weights from each member's scaled MAE on the last H
observed steps of the context (a zero-shot, in-context model selection - it only uses the
series' own history).

    fc = EnsembleForecaster(TimestralForecaster.from_pretrained(exp), "c8192,c2048,d4", alpha=0.5)
"""

from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np
import torch


@dataclass
class _F:
    quantiles: torch.Tensor


def _pool(x: np.ndarray | None, k: int) -> np.ndarray | None:
    """(B,V,T) NaN-aware mean pool by k aligned to the end."""
    if x is None:
        return None
    B, V, T = x.shape
    n = T // k
    if n == 0:
        return None
    y = x[..., T - n * k:].reshape(B, V, n, k)
    with np.errstate(all="ignore"):
        cnt = np.isfinite(y).sum(-1)
        s = np.nansum(y, -1)
        out = np.where(cnt > 0, s / np.maximum(cnt, 1), np.nan)
    return out.astype(np.float32)


class EnsembleForecaster:
    def __init__(self, fc, spec: str = "c8192", alpha: float = 0.5, backtest: bool = False, tau: float = 0.5, bt_windows: int = 1, prior: float = 0.0,
                 min_coarse_ctx: int = 256, min_extra: float = 1.5, min_ratio: float = 8.0):
        self.fc = fc
        self.members = [m.strip() for m in spec.split(",") if m.strip()]
        cl = fc.model.config.context_length  # never exceed the trained context (RoPE does not extrapolate)
        self.ctxs = [min(int(m[1:]), cl) for m in self.members if m[0] == "c"] or [fc.max_context]
        self.ks = [int(m[1:]) for m in self.members if m[0] == "d"]
        self.alpha, self.backtest, self.tau = alpha, backtest, tau
        self.bt_windows, self.prior = bt_windows, prior  # prior: logit bonus for the first (longest) member
        self.min_coarse_ctx, self.min_extra, self.min_ratio = min_coarse_ctx, min_extra, min_ratio  # member ctx >= min_ratio*H
        self.base_ctx = fc.max_context
        self.max_context = max(self.ctxs + [self.base_ctx * k for k in self.ks] + [self.base_ctx])
        self.quantile_levels = fc.quantile_levels
        self.mi = list(fc.quantile_levels).index(0.5) if 0.5 in list(fc.quantile_levels) else len(fc.quantile_levels) // 2

    def _one(self, ctx, h, past, mc):
        self.fc.max_context = mc
        try:
            return self.fc.forecast(ctx, h, past_covariates=past).quantiles.float().cpu()
        finally:
            self.fc.max_context = self.base_ctx

    def _members(self, ctx, h, past):
        T = ctx.shape[-1]
        outs, used = [], set()
        for i, c in enumerate(self.ctxs):
            if i > 0 and c < self.min_ratio * h:  # a short context cannot support a long horizon
                continue
            eff = min(c, T)
            if eff in used:  # a shorter series makes several members identical
                continue
            used.add(eff)
            outs.append(self._one(ctx[..., -c:], h, None if past is None else past[..., -c:], c))
        return outs

    def forecast(self, context, prediction_length, past_covariates=None):
        ctx = np.asarray(context, np.float32)
        if ctx.ndim == 2:
            ctx = ctx[:, None]
        past = None if past_covariates is None else np.asarray(past_covariates, np.float32)
        h = prediction_length
        T = ctx.shape[-1]
        outs = self._members(ctx, h, past)
        if len(outs) == 1:
            q = outs[0]
        elif self.backtest and T - h >= 2 * 32:
            E, ok = 0.0, True
            nw = 0
            for wdw in range(1, self.bt_windows + 1):
                if T - wdw * h < 2 * 32:
                    break
                cut = T - wdw * h
                bt = self._members(ctx[..., :cut], h, None if past is None else past[..., :cut])
                if len(bt) != len(outs):
                    ok = False
                    break
                y = torch.from_numpy(ctx[..., cut:cut + h])
                hist = torch.from_numpy(ctx[..., :cut])
                dd = hist[..., 1:] - hist[..., :-1]
                fin = torch.isfinite(dd)
                scale = (torch.where(fin, dd.abs(), torch.zeros_like(dd)).sum(-1) / fin.sum(-1).clamp_min(1)).clamp_min(1e-8)
                errs = []
                for b in bt:
                    e = (b[..., self.mi] - y).abs()
                    fe = torch.isfinite(e)
                    e = torch.where(fe, e, torch.zeros_like(e)).sum(-1) / fe.sum(-1).clamp_min(1) / scale
                    errs.append(e.mean(-1))
                E = E + torch.stack(errs, 0)
                nw += 1
            if not ok or nw == 0:
                q = torch.stack(outs).mean(0)
            else:
                E = E / nw
                logit = -E / (self.tau * E.min(0, keepdim=True).values.clamp_min(1e-3))
                logit[0] = logit[0] + self.prior
                w = torch.softmax(logit, 0)
                q = (torch.stack(outs) * w[:, :, None, None, None]).sum(0)
        else:
            q = torch.stack(outs).mean(0)
        # temporal-hierarchy level correction from downsampled members
        for k in self.ks:
            hk = math.ceil(h / k)
            pooled = _pool(ctx, k)
            if pooled is None or pooled.shape[-1] < self.min_coarse_ctx:
                continue
            # only worth it when the coarse view sees materially more history than the fine one
            if T < self.min_extra * self.base_ctx and hk < 2:
                continue
            pp = _pool(past, k) if past is not None else None
            cq = self._one(pooled[..., -self.base_ctx:], hk, None if pp is None else pp[..., -self.base_ctx:], self.base_ctx)
            cmed = cq[..., self.mi]  # (B,V,hk)
            fmed = q[..., self.mi]
            pad = hk * k - h
            fm = torch.cat([fmed, fmed[..., -1:].expand(*fmed.shape[:-1], pad)], -1) if pad else fmed
            fpool = fm.view(*fm.shape[:-1], hk, k).mean(-1)
            delta = (cmed - fpool).repeat_interleave(k, -1)[..., :h]
            q = q + self.alpha / len(self.ks) * delta[..., None]
        return _F(q)