"""TimeTron — a 33.5M-parameter time-series foundation model.

Distilled from Chronos-2 by latent-space knowledge distillation: the student matches the
teacher's internal representations rather than its forecasts.

    from modeling_timetron import TimeTron
    model = TimeTron.from_pretrained("timetron-v2-33m")
    q = model.predict(context, prediction_length=128)   # (B, H, 21) in the input's own scale

Self-contained: torch + numpy only.
"""
import json
import math
import os
from dataclasses import dataclass, field

import torch
import torch.nn as nn
import torch.nn.functional as F

from dataclasses import dataclass, field

import torch
import torch.nn as nn
import torch.nn.functional as F

QUANTILE_LEVELS_21 = [0.01, 0.05] + [round(0.1 + 0.05 * i, 2) for i in range(17)] + [0.95, 0.99]


# ---------------------------------------------------------------- input path
def patchify(x, mask, patch_size=32):
    """x,(B,L) mask,(B,L) 1=padding -> x_patched,(B,N,P) m_patched,(B,N,P) patch_mask,(B,N)."""
    B, L = x.shape
    P = patch_size
    mask = mask.to(torch.long)
    N = (L + P - 1) // P
    pad = N * P - L
    x_patched = F.pad(x, (pad, 0), value=0.0).reshape(B, N, P)
    m_patched = F.pad(mask, (pad, 0), value=1).reshape(B, N, P)
    patch_mask = m_patched.amin(dim=-1)
    return x_patched, m_patched, patch_mask


class CausalPatchNormV2(nn.Module):
    """Cumulative per-patch stats + σ-relative floor + asinh squash (+ training dither)."""

    def __init__(self, sigma_min=1e-3, dither=0.01):
        super().__init__()
        self.sigma_min = sigma_min
        self.dither = dither

    def forward(self, x, mask):
        valid = 1.0 - mask.float()
        x_valid = x * valid
        count = valid.sum(-1).cumsum(-1)
        safe = count.clamp(min=1.0)
        S1 = x_valid.sum(-1).cumsum(-1)
        S2 = (x_valid ** 2).sum(-1).cumsum(-1)
        mu = S1 / safe
        var = (S2 / safe - mu ** 2).clamp(min=0.0)
        cum_abs = x_valid.abs().sum(-1).cumsum(-1) / safe
        floor = torch.maximum(torch.full_like(cum_abs, self.sigma_min), 0.05 * cum_abs)
        sigma = torch.maximum((var + 1e-8).sqrt(), floor)
        xn = torch.asinh((x - mu.unsqueeze(-1)) / sigma.unsqueeze(-1))   # smooth, unbounded-safe
        if self.training and self.dither > 0:
            xn = xn + self.dither * torch.randn_like(xn)
        return xn, mu, sigma


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-4):                 # v1 backward-amplifier fix baked in
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        rms = (x.pow(2).mean(-1, keepdim=True) + self.eps).sqrt()
        return x / rms * self.weight


class ResidualBlock(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim, bias=False)
        self.fc2 = nn.Linear(hidden_dim, output_dim, bias=False)
        self.skip = nn.Linear(input_dim, output_dim, bias=False)
        self.act = nn.SiLU()

    def forward(self, x):
        return self.fc2(self.act(self.fc1(x))) + self.skip(x)


# ---------------------------------------------------------------- RoPE (v1, unchanged)
class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim, max_position_embeddings=16384, rope_theta=10000.0):
        super().__init__()
        i = torch.arange(0, head_dim, 2).float()
        inv_freq = 1.0 / (rope_theta ** (i / head_dim))
        positions = torch.arange(max_position_embeddings).float()
        angles = torch.outer(positions, inv_freq)
        self.register_buffer("cos_cached", torch.cat([angles.cos(), angles.cos()], dim=-1))
        self.register_buffer("sin_cached", torch.cat([angles.sin(), angles.sin()], dim=-1))

    def forward(self):
        return self.cos_cached, self.sin_cached


def rotate_half(x):
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
    return torch.cat([-x2, x1], dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
    cos = cos[position_ids].unsqueeze(1)
    sin = sin[position_ids].unsqueeze(1)
    return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin


# ---------------------------------------------------------------- attention
class Attention(nn.Module):
    """v1 attention (QK-norm, per-dim scale, RoPE) with a causal switch — v2 runs bidirectional."""

    def __init__(self, hidden_size, num_heads, head_dim, causal=False, dropout=0.0):
        super().__init__()
        assert num_heads * head_dim == hidden_size
        self.h, self.hd, self.hidden = num_heads, head_dim, hidden_size
        self.causal = causal
        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.q_norm = RMSNorm(head_dim)
        self.k_norm = RMSNorm(head_dim)
        self.scale = nn.Parameter(torch.ones(head_dim))
        self.rotary = RotaryEmbedding(head_dim=head_dim)
        self.dropout = dropout

    def forward(self, x, token_mask, position_ids):
        B, N, _ = x.shape
        q = self.q_proj(x).view(B, N, self.h, self.hd).transpose(1, 2)
        k = self.k_proj(x).view(B, N, self.h, self.hd).transpose(1, 2)
        v = self.v_proj(x).view(B, N, self.h, self.hd).transpose(1, 2)
        q, k = self.q_norm(q), self.k_norm(k)
        cos, sin = self.rotary()
        q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
        q = q * self.scale
        blocked = token_mask.bool()[:, None, None, :]
        if self.causal:
            blocked = blocked | torch.ones(N, N, dtype=torch.bool, device=x.device).triu(1)
        blocked = blocked & ~blocked.all(dim=-1, keepdim=True)
        attn_mask = torch.zeros(blocked.shape, dtype=q.dtype, device=x.device).masked_fill(blocked, float("-inf"))
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask,
                                             dropout_p=self.dropout if self.training else 0.0, scale=1.0)
        return self.o_proj(out.transpose(1, 2).contiguous().view(B, N, self.hidden))


class GroupAttention(nn.Module):
    """Chronos-2-style group attention: attends ACROSS series of the same group at each
    token position. o_proj zero-initialized -> exact identity until the multivariate phase
    turns it on. Skipped entirely (no compute) when group_ids is None."""

    def __init__(self, hidden_size, num_heads, head_dim):
        super().__init__()
        self.h, self.hd, self.hidden = num_heads, head_dim, hidden_size
        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        nn.init.zeros_(self.o_proj.weight)                    # neutral start (drift-gate lesson)
        self.q_norm = RMSNorm(head_dim)
        self.k_norm = RMSNorm(head_dim)

    def forward(self, x, group_ids):
        if group_ids is None:
            return torch.zeros_like(x)
        B, N, d = x.shape
        y = x.transpose(0, 1)                                  # (N, B, d): attend across batch per position
        q = self.q_proj(y).view(N, B, self.h, self.hd).transpose(1, 2)
        k = self.k_proj(y).view(N, B, self.h, self.hd).transpose(1, 2)
        v = self.v_proj(y).view(N, B, self.h, self.hd).transpose(1, 2)
        q, k = self.q_norm(q), self.k_norm(k)
        same = group_ids[:, None] == group_ids[None, :]        # (B, B)
        attn_mask = torch.zeros(B, B, dtype=q.dtype, device=x.device).masked_fill(~same, float("-inf"))
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
        return self.o_proj(out.transpose(1, 2).contiguous().view(N, B, d)).transpose(0, 1)


class MLP(nn.Module):
    def __init__(self, hidden_size, intermediate_size):
        super().__init__()
        self.fc1 = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.fc2 = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, x):
        return self.fc2(F.silu(self.fc1(x)))


class V2Layer(nn.Module):
    def __init__(self, cfg, with_group=False):
        super().__init__()
        self.input_layernorm = RMSNorm(cfg.hidden_size)
        self.self_attn = Attention(cfg.hidden_size, cfg.num_attention_heads, cfg.head_dim,
                                   causal=False, dropout=cfg.attention_dropout)
        self.group_attn = GroupAttention(cfg.hidden_size, cfg.num_attention_heads, cfg.head_dim) if with_group else None
        self.group_norm = RMSNorm(cfg.hidden_size) if with_group else None
        self.post_attention_layernorm = RMSNorm(cfg.hidden_size)
        self.mixer = MLP(cfg.hidden_size, cfg.intermediate_size)

    def forward(self, x, token_mask, position_ids, group_ids=None):
        x = x + self.self_attn(self.input_layernorm(x), token_mask, position_ids)
        if self.group_attn is not None:
            x = x + self.group_attn(self.group_norm(x), group_ids)
        x = x + self.mixer(self.post_attention_layernorm(x))
        return x


# ---------------------------------------------------------------- the model
@dataclass
class StudentV2Config:
    patch_length: int = 32
    max_context: int = 2048                 # time-index normalization constant C
    quantile_levels: list = field(default_factory=lambda: list(QUANTILE_LEVELS_21))
    teacher_dim: int = 768                  # Chronos-2 latents
    hidden_size: int = 512
    num_hidden_layers: int = 12
    num_attention_heads: int = 16
    head_dim: int = 32
    intermediate_size: int = 1280
    attention_dropout: float = 0.0
    group_layers: tuple = (5, 11)           # 0-indexed layer positions carrying group attention
    kd_taps: tuple = (3, 7, 11)             # 0-indexed layers tapped for latent distillation
    max_future_patches: int = 12            # up to 384-step native horizon


class StudentV2(nn.Module):
    def __init__(self, cfg: StudentV2Config = StudentV2Config()):
        super().__init__()
        self.cfg = cfg
        P = cfg.patch_length
        self.norm_layer = CausalPatchNormV2()
        self.input_embedding = ResidualBlock(3 * P, cfg.hidden_size, cfg.hidden_size)   # [vals, tidx, mask]
        self.reg_token = nn.Parameter(torch.zeros(1, 1, cfg.hidden_size))
        nn.init.trunc_normal_(self.reg_token, std=0.02)
        self.layers = nn.ModuleList(
            V2Layer(cfg, with_group=(i in cfg.group_layers)) for i in range(cfg.num_hidden_layers))
        self.final_norm = RMSNorm(cfg.hidden_size)
        nq = len(cfg.quantile_levels)
        self.quantile_head = ResidualBlock(cfg.hidden_size, 2 * cfg.hidden_size, P * nq)
        self.latent_projs = nn.ModuleList(
            nn.Linear(cfg.hidden_size, cfg.teacher_dim) for _ in cfg.kd_taps)
        self.mask_token = nn.Parameter(torch.zeros(1, 1, cfg.hidden_size))
        nn.init.trunc_normal_(self.mask_token, std=0.02)

    def _embed_context(self, x, mask):
        P = self.cfg.patch_length
        x_patched, m_patched, patch_mask = patchify(x, mask, P)
        B, Nc, _ = x_patched.shape
        xn, mu, sigma = self.norm_layer(x_patched, m_patched)
        xn = xn * (1.0 - m_patched.float())
        t = torch.arange(-Nc * P + 1, 1, device=x.device, dtype=xn.dtype) / self.cfg.max_context
        tidx = t.view(1, Nc, P).expand(B, Nc, P)
        obs = 1.0 - m_patched.float()
        h = self.input_embedding(torch.cat([xn, tidx, obs], dim=-1))
        return h, patch_mask, mu, sigma, Nc

    def forward(self, x, mask, k_future=4, aug_mask=None, group_ids=None,
                future_values=None, future_observed=None):
        """x (B, L) L%32==0; k_future future patches decoded natively.

        future_values (B, k_future*P): KNOWN future values, for covariate-informed tasks. They
        reuse the context's own [values, tidx, observed] embedding path and the context anchor
        (mu, sigma), so no new parameters are involved. future_observed (B, k_future*P) marks
        where the value is known; patches with nothing known keep the mask token exactly as
        before, so future_values=None reproduces the univariate/multivariate path bit-for-bit.

        Returns quantiles (B, k_future*P, n_q) in asinh-normalized anchor space."""
        cfg = self.cfg
        P = cfg.patch_length
        B = x.shape[0]
        h_ctx, patch_mask, mu, sigma, Nc = self._embed_context(x, mask)
        if aug_mask is not None:
            h_ctx = torch.where(aug_mask.unsqueeze(-1), self.mask_token.to(h_ctx.dtype), h_ctx)
        # future tokens: value channel (known covariates, else 0), future time index, observed flag
        tf = torch.arange(1, k_future * P + 1, device=x.device, dtype=h_ctx.dtype) / cfg.max_context
        if future_values is None:
            fv = torch.zeros(B, k_future, P, device=x.device, dtype=h_ctx.dtype)
            fo = torch.zeros(B, k_future, P, device=x.device, dtype=h_ctx.dtype)
        else:
            fo = (torch.ones_like(future_values) if future_observed is None
                  else future_observed).to(h_ctx.dtype).view(B, k_future, P)
            fv = torch.asinh((future_values - mu[:, -1:]) / sigma[:, -1:]).clamp(-4.0, 4.0)
            fv = fv.to(h_ctx.dtype).view(B, k_future, P) * fo          # unknown slots stay 0
        fut_feats = torch.cat([fv, tf.view(1, k_future, P).expand(B, k_future, P), fo], dim=-1)
        # mask token only where the future is genuinely unknown
        h_fut = self.input_embedding(fut_feats) + \
            self.mask_token.to(h_ctx.dtype) * (1.0 - fo.amax(dim=-1, keepdim=True))
        seq = torch.cat([h_ctx, self.reg_token.expand(B, 1, -1).to(h_ctx.dtype), h_fut], dim=1)
        N = seq.shape[1]
        token_mask = torch.cat([patch_mask,
                                torch.zeros(B, 1 + k_future, dtype=patch_mask.dtype, device=x.device)], dim=1)
        pos = torch.arange(N, device=x.device).unsqueeze(0).expand(B, N)
        taps = {}
        for i, layer in enumerate(self.layers):
            seq = layer(seq, token_mask, pos, group_ids)
            if i in cfg.kd_taps:
                taps[i] = seq[:, :Nc]
        seq = self.final_norm(seq)
        taps[cfg.kd_taps[-1]] = seq[:, :Nc]                      # final tap post-norm
        latents = [proj(taps[t]) for t, proj in zip(cfg.kd_taps, self.latent_projs)]
        h_future = seq[:, Nc + 1:]
        nq = len(cfg.quantile_levels)
        q = self.quantile_head(h_future).view(B, k_future * P, nq)
        return {"quantiles": q, "hidden": seq[:, :Nc], "reg": seq[:, Nc],
                "latents": latents, "mu": mu, "sigma": sigma, "patch_mask": patch_mask}

    def num_params(self):
        return sum(p.numel() for p in self.parameters())

# --------------------------------------------------------------------- public API
QUANTILE_LEVELS = QUANTILE_LEVELS_21


class TimeTron(nn.Module):
    """Thin wrapper over StudentV2 adding load/save and a batched `predict`."""

    def __init__(self, config: StudentV2Config = None):
        super().__init__()
        self.config = config or StudentV2Config()
        self.model = StudentV2(self.config)
        self.median_index = self.config.quantile_levels.index(0.5)
        self.native_horizon = self.config.max_future_patches * self.config.patch_length
        self.eval()   # inference must be deterministic: CausalPatchNormV2 dithers in train mode

    # ---- persistence -------------------------------------------------
    @classmethod
    def from_pretrained(cls, path_or_repo: str, device: str = "cpu"):
        """`path_or_repo` may be a local directory or a Hugging Face repo id."""
        d = path_or_repo
        if not os.path.isdir(d):
            from huggingface_hub import snapshot_download
            d = snapshot_download(path_or_repo)
        with open(os.path.join(d, "config.json")) as f:
            raw = json.load(f)
        cfg = StudentV2Config(**{k: v for k, v in raw.items()
                                 if k in StudentV2Config.__dataclass_fields__})
        cfg.group_layers = tuple(cfg.group_layers)
        cfg.kd_taps = tuple(cfg.kd_taps)
        self = cls(cfg)
        w = os.path.join(d, "model.safetensors")
        if os.path.exists(w):
            from safetensors.torch import load_file
            state = load_file(w)
        else:
            state = torch.load(os.path.join(d, "pytorch_model.bin"),
                               map_location="cpu", weights_only=False)
            state = state.get("model", state)
        missing, unexpected = self.model.load_state_dict(state, strict=False)
        assert not missing and not unexpected, (missing, unexpected)
        return self.to(device).eval()

    def save_pretrained(self, directory: str):
        os.makedirs(directory, exist_ok=True)
        cfg = {k: (list(v) if isinstance(v, tuple) else v)
               for k, v in self.config.__dict__.items()}
        cfg["architectures"] = ["TimeTron"]
        cfg["model_type"] = "timetron"
        with open(os.path.join(directory, "config.json"), "w") as f:
            json.dump(cfg, f, indent=2)
        from safetensors.torch import save_file
        save_file({k: v.contiguous() for k, v in self.model.state_dict().items()},
                  os.path.join(directory, "model.safetensors"))

    # ---- inference ---------------------------------------------------
    @torch.no_grad()
    def predict(self, context, prediction_length: int = 128, group_ids=None):
        """context: (B, L) tensor/array of raw values, L a multiple of 32 (it is cropped
        and left-padded for you). Returns (B, prediction_length, 21) in the input's scale.

        Horizons up to `native_horizon` (384) decode in one pass; beyond that the median
        path is fed back autoregressively in whole chunks.

        `group_ids`: optional (B,) integer tensor. Rows sharing an id attend to each other
        (in-context learning across related series). Leaving it None is bit-identical to a
        model without the group-attention branch. NOTE: on the released checkpoint this
        branch is untrained -- see the model card.
        """
        was_training = self.training
        self.eval()   # dither is a *training* augmentation; never let it reach a forecast
        try:
            return self._predict(context, prediction_length, group_ids)
        finally:
            if was_training:
                self.train()

    @torch.no_grad()
    def _predict(self, context, prediction_length, group_ids):
        dev = next(self.parameters()).device
        x = torch.as_tensor(context, dtype=torch.float32, device=dev)
        if x.dim() == 1:
            x = x[None, :]
        L = max(32, min(self.config.max_context, (x.shape[1] // 32) * 32))
        x = x[:, -L:] if x.shape[1] >= L else F.pad(x, (L - x.shape[1], 0), mode="replicate")
        g = None if group_ids is None else torch.as_tensor(group_ids, dtype=torch.long, device=dev)

        outs, done = [], 0
        cur = x
        while done < prediction_length:
            need = min(prediction_length - done, self.native_horizon)
            kf = (need + 31) // 32
            o = self.model(cur, torch.zeros_like(cur, dtype=torch.long),
                           k_future=kf, group_ids=g)
            y = (o["mu"][:, -1:].unsqueeze(-1)
                 + o["sigma"][:, -1:].unsqueeze(-1)
                 * torch.sinh(o["quantiles"].clamp(-4.0, 4.0)))[:, :need]
            outs.append(y)
            done += need
            if done < prediction_length:
                cur = torch.cat([cur, y[:, :, self.median_index]], dim=1)
                keep = max(32, min(self.config.max_context, (cur.shape[1] // 32) * 32))
                cur = cur[:, -keep:]
        return torch.cat(outs, dim=1)

    @torch.no_grad()
    def predict_median(self, context, prediction_length: int = 128, group_ids=None):
        return self.predict(context, prediction_length, group_ids)[..., self.median_index]

    def num_params(self):
        return sum(p.numel() for p in self.model.parameters())