"""GiftEvalPretrain stream v2 (Modular-TS v3): resolution + mixup augmentation on real data.

Adds to FastGiftEvalPretrain (run1):
  * TTM-style "diverse resolution sampling" (Ekambaram et al., arXiv 2401.03955): with prob
    `ds_p` a univariate batch is drawn at a coarser resolution, mean-pooling k in {2,4,8}
    consecutive raw steps (NaN-aware), creating new frequencies from the same corpus;
  * TSMixup (Chronos, arXiv 2403.07815): with prob `mix_p` rows are replaced by a convex
    combination of 2-3 standardised real series (Dirichlet weights), rescaled;
  * flatter dataset temperature (ds_temp 0.3 instead of 0.5) so small datasets are seen more.
Contamination filter and GPU-resident layout are inherited unchanged.
"""

from __future__ import annotations

import torch

from modular_ts.realdata import FastGiftEvalPretrain


class FastGiftEvalPretrain2(FastGiftEvalPretrain):
    def __init__(self, *a, ds_p=0.3, ds_factors=(2, 4, 8), mix_p=0.2, **k):
        k.setdefault("ds_temp", 0.3)
        super().__init__(*a, **k)
        self.ds_p, self.ds_factors, self.mix_p = ds_p, ds_factors, mix_p

    def _gather_k(self, off, L, kf):
        """Gather T*kf raw steps ending at a random point and mean-pool by kf."""
        T, R = self.T, off.shape[0]
        TK = T * kf
        st = torch.where(L >= TK, (self._rand(R) * (L - TK + 1)).long(), L - TK)
        ar = torch.arange(TK, device=self.dev)
        pos = st[:, None] + ar[None]
        valid = (pos >= 0) & (pos < L[:, None])
        idx = (off[:, None] + pos.clamp(min=0)).clamp_max(self.flat.numel() - 1)
        v = self.flat[idx]
        ok = valid & torch.isfinite(v)
        v = torch.where(ok, v, torch.zeros_like(v)).view(R, T, kf)
        c = ok.view(R, T, kf).sum(-1)
        out = v.sum(-1) / c.clamp_min(1)
        obs = c > 0
        return torch.where(obs, out, torch.zeros_like(out)), obs

    def _mixup(self, v, o):
        R, T = v.shape
        both = o & o.roll(1, 0) & o.roll(2, 0)
        def zs(x, m):
            n = m.sum(-1, keepdim=True).clamp_min(1)
            mu = (x * m).sum(-1, keepdim=True) / n
            sd = (((x - mu) ** 2) * m).sum(-1, keepdim=True).div(n).sqrt().clamp_min(1e-6)
            return (x - mu) / sd, sd, mu
        z0, sd0, mu0 = zs(v, o)
        z1, _, _ = zs(v.roll(1, 0), o.roll(1, 0))
        z2, _, _ = zs(v.roll(2, 0), o.roll(2, 0))
        w = torch.distributions.Dirichlet(torch.full((3,), 1.5, device=self.dev)).sample((R,))
        w[:, 2] *= (self._rand(R) < 0.5).float()
        w = w / w.sum(-1, keepdim=True)
        m = w[:, :1] * z0 + w[:, 1:2] * z1 + w[:, 2:3] * z2
        m = m * sd0 + mu0
        sel = (self._rand(R) < self.mix_p)[:, None]
        return torch.where(sel, torch.where(both, m, 0.0), v), torch.where(sel, both, o)

    @torch.no_grad()
    def next_batch(self, V=None):
        if (self.mv is not None and self.rng.random() < self.mv_frac) or self.rng.random() >= self.ds_p:
            b = super().next_batch(V)
            if b["values"].shape[1] == 1 and self.mix_p > 0:
                v, o = self._mixup(b["values"][:, 0], b["observed"][:, 0])
                b["values"], b["observed"] = v[:, None], o[:, None]
                b["s"] = self._s(b["observed"])
            return b
        kf = self.rng.choice(self.ds_factors)
        Bg = max(1, self.tok // self.N)
        c = self._pick(self.u_w, self.u_start, self.u_cnt, Bg)
        off, L = self.u_off[c], self.u_len[c]
        v, o = self._gather_k(off, L, kf)
        o = o & (torch.arange(self.T, device=self.dev)[None] >= (self.T - self._keep_len(Bg))[:, None])
        v = torch.where(o, v, torch.zeros_like(v))
        if self.mix_p > 0:
            v, o = self._mixup(v, o)
        v, o = v[:, None], o[:, None]
        return {"values": v, "observed": o, "roles": torch.zeros(Bg, 1, dtype=torch.long, device=self.dev), "s": self._s(o)}