"""Modular-TS training model: a Timestral-style network trained from scratch with DHF.

The network *is* `timestral.model.TimestralModel` (Apache-2.0 architecture code, random
init), so parameter names match the open inference package and a trained checkpoint
exports to `config.json` + `model.safetensors` that `timestral.TimestralForecaster`
loads unchanged. We add `forward_train`, the Decoupled-Horizon-Forecasting training pass
(our reconstruction; Mistral has not published it):

- batch layout (B groups, V rows per group, T = N * patch steps); rows are target,
  past-only covariate, known-future covariate, or padding;
- Enc = all N patch embeddings (causal per patch, past-only scaling + asinh);
  Dec = N copies of the learned forecast query; per group a min-context s, decoder
  blocks of K patches anchored at s, s+K, ...; Dec token in block b reads Enc[< b] plus
  the causal Dec chain of its own block (timestral.masks.build_two_stream_time_mask);
- variate attention per patch slot over the V rows of a group with the covariate
  firewall (covariates never read targets; known-future rows never read causal rows);
- target of Dec token p = true patch p normalized in the causal frame of patches < p
  (the "advancing refine frame" the inference code decodes in), asinh-squashed;
- loss = pinball over the 21 quantile knots on observed target values of active patches.

NEVER load Timestral weights into this model (Modified-MIT licence, see NOTICE).
"""

from __future__ import annotations

import math
import os
import sys

import torch
import torch.nn.functional as F
from torch import Tensor

_HERE = os.path.dirname(os.path.abspath(__file__))
_TS_SRC = os.path.join(os.path.dirname(_HERE), "timestral", "src")
if os.path.isdir(_TS_SRC) and _TS_SRC not in sys.path:
    sys.path.insert(0, _TS_SRC)

from timestral.config import TimestralConfig  # noqa: E402
from timestral.layers import Attention, _apply_rotary  # noqa: E402
from timestral.masks import build_two_stream_time_mask, compute_enc_key_validity  # noqa: E402
from timestral.model import ROLE_KNOWN_COV, ROLE_PAST_COV, ROLE_TARGET, TimestralModel  # noqa: E402
from timestral.scaler import apply_affine_asinh_masked, causal_patch_loc_scale  # noqa: E402

ROLE_PAD = -1
QUANTILES_21 = (0.01, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5,
                0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 0.99)


def make_config(size: str = "2.5M", context_length: int = 4096) -> TimestralConfig:
    presets = {
        # d_model, heads, mlp, layers-groups
        "2.5M": (128, 2, 320, 2),
        "5M": (192, 3, 512, 2),
        "28M": (384, 6, 1024, 3),
        "tiny": (64, 1, 128, 2),  # CPU tests only
    }
    d, h, mlp, k = presets[size]
    return TimestralConfig(
        patch_size=32, d_model=d, num_heads=h, mlp_hidden_dim=mlp,
        block_schedule=("time", "time", "time", "time_variate") * k,
        quantile_knots=QUANTILES_21, block_size=32, context_length=context_length,
        rope_theta=1e4, rope_partial_factor=0.5, norm_eps=1e-6, attn_bias=False, mlp_bias=False,
        time_attn_gate="sdpa_sigmoid", variate_attn_gate="sdpa_sigmoid", minimum_scale=1e-6,
        use_asinh=True, stabilize_with_global=False, covariate_enabled=True,
        split_context_horizon_attn=True, attend_missing_patches=True, quantile_width_factor=3.0,
    )


def count_params(model: torch.nn.Module) -> dict:
    tot = sum(p.numel() for p in model.parameters())
    emb = sum(p.numel() for n, p in model.named_parameters() if n.startswith("representation"))
    head = sum(p.numel() for n, p in model.named_parameters() if n.startswith("head"))
    return {"total": tot, "embed": emb, "trunk": tot - emb - head, "head": head}


def _variate_attention(attn: Attention, x: Tensor, mask: Tensor, B: int, V: int) -> Tensor:
    """Variate attention per (group, slot) over V rows. x: (B*V, S, D); mask: (B*S, 1, V, V)."""
    R, S, D = x.shape
    inp = x.view(B, V, S, D).permute(0, 2, 1, 3).reshape(B * S, V, D)
    q, k, v = attn._qkv(inp)  # (B*S, V, H, hd)
    out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_mask=mask)
    out = out.transpose(1, 2).reshape(B * S, V, D)
    if attn._gate is not None:
        out = out * attn._gate(inp).sigmoid()
    out = attn.wO(out)
    return out.view(B, S, V, D).permute(0, 2, 1, 3).reshape(R, S, D)


def _time_attention(attn: Attention, x: Tensor, mask: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
    q, k, v = attn._qkv(x)
    q, k = _apply_rotary(q, k, cos, sin)
    out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_mask=mask)
    out = out.transpose(1, 2).reshape(x.shape[0], x.shape[1], -1)
    if attn._gate is not None:
        out = out * attn._gate(x).sigmoid()
    return attn.wO(out)


class ModularTS(TimestralModel):
    """TimestralModel + DHF training pass (identical parameter names)."""

    def __init__(self, config: TimestralConfig) -> None:
        super().__init__(config)
        self.register_buffer("q_levels", torch.tensor(config.quantile_knots, dtype=torch.float32), persistent=False)
        self._init_weights()

    def _init_weights(self) -> None:
        for name, p in self.named_parameters():
            if p.dim() == 2:
                torch.nn.init.normal_(p, std=0.02)
                if name.endswith("wO.weight") or name.endswith("w3.weight"):
                    p.data.mul_(1.0 / math.sqrt(2 * len(self.config.block_schedule)))
            elif name.endswith(".bias"):
                torch.nn.init.zeros_(p)
        torch.nn.init.normal_(self.representation.forecast_query, std=0.02)

    # ------------------------------------------------------------------ masks
    def build_masks(self, observed: Tensor, roles: Tensor, s: Tensor):
        """observed (B,V,T) bool, roles (B,V) long, s (B,) long -> masks/positions/key validity."""
        c = self.config
        B, V, T = observed.shape
        P, N, R = c.patch_size, T // c.patch_size, B * V
        dev = observed.device
        obs_pp = observed.view(R, N, P).any(-1)
        not_pad = (roles != ROLE_PAD).view(R)
        key_valid = compute_enc_key_validity(obs_pp, c.attend_missing_patches) & not_pad[:, None]
        known = (roles == ROLE_KNOWN_COV).view(R)
        target = (roles == ROLE_TARGET).view(R)
        s_rows = s.repeat_interleave(V)
        time_mask = build_two_stream_time_mask(s_rows, N, c.block_size, key_valid, bidirectional_rows=known)
        # variate mask (B*2N, 1, V, V)
        slot_valid = torch.cat([key_valid, not_pad[:, None].expand(R, N)], dim=1).view(B, V, 2 * N)
        keyok = slot_valid.permute(0, 2, 1)[:, :, None, :]  # (B, 2N, 1, Vk)
        tg = target.view(B, V)
        kn = known.view(B, V)
        forbid = ((~tg)[:, :, None] & tg[:, None, :]) | (kn[:, :, None] & (~kn)[:, None, :])  # (B, Vq, Vk)
        vmask = keyok & ~forbid[:, None]
        eye = torch.eye(V, dtype=torch.bool, device=dev)
        vmask = (vmask | eye).reshape(B * 2 * N, 1, V, V)
        # positions anchored at each group's first observed patch
        grp_obs = obs_pp.view(B, V, N).any(1)  # (B, N)
        first = torch.where(grp_obs.any(-1), grp_obs.float().argmax(-1), torch.zeros(B, device=dev)).long()
        p = torch.arange(N, device=dev)
        pos = (p[None] - first[:, None]).repeat_interleave(V, 0)
        positions = torch.cat([pos, pos], dim=1)
        return time_mask, vmask, positions, obs_pp, target, known

    # ---------------------------------------------------------------- forward
    def forward_train(self, values: Tensor, observed: Tensor, roles: Tensor, s: Tensor, return_preds: bool = False,
                      teacher: Tensor | None = None, teacher_obs: Tensor | None = None, gt_row: Tensor | None = None,
                      lam: float = 0.5, teacher_q_idx: tuple = (2, 4, 6, 8, 10, 12, 14, 16, 18)):
        """values (B,V,T) float (anything where unobserved), observed (B,V,T) bool,
        roles (B,V) long in {-1,0,1,2}, s (B,) min-context patches. Returns (loss, stats)."""
        c = self.config
        B, V, T = values.shape
        P, N, R = c.patch_size, T // c.patch_size, B * V
        x = torch.where(observed, values, torch.zeros_like(values)).float().view(R, T)
        obs = observed.view(R, T)

        with torch.no_grad():
            time_mask, vmask, positions, obs_pp, target_rows, known_rows = self.build_masks(observed, roles, s)
            loc, scale = causal_patch_loc_scale(x[:, None], obs[:, None], P, c.minimum_scale, False)
            loc, scale = loc[:, 0], scale[:, 0]
            scaled = apply_affine_asinh_masked(x, loc, scale, obs, use_asinh=c.use_asinh)
            scaled, _, _ = self._known_cov_scaling(x, obs, scaled, loc, scale, known_rows)
            # DHF targets in the frame of patches < p
            loc_pp = loc.view(R, N, P)[..., 0]
            sc_pp = scale.view(R, N, P)[..., 0]
            f_loc = torch.cat([loc_pp[:, :1], loc_pp[:, :-1]], dim=1)
            f_sc = torch.cat([sc_pp[:, :1], sc_pp[:, :-1]], dim=1)
            y = torch.asinh((x.view(R, N, P) - f_loc[..., None]) / f_sc[..., None].clamp_min(1e-6))
            cnt_before = torch.cumsum(obs.view(R, N, P).sum(-1), 1) - obs.view(R, N, P).sum(-1)
            p_idx = torch.arange(N, device=x.device)
            active = p_idx[None] >= s.repeat_interleave(V)[:, None]
            w = (obs.view(R, N, P) & target_rows[:, None, None] & active[..., None] & (cnt_before >= 2)[..., None]).float()
            if gt_row is not None:  # rows whose ground truth may not be used (teacher-only inputs)
                w = w * gt_row.view(R, 1, 1).float()
            y = torch.where(w > 0, y, torch.zeros_like(y))
            if teacher is not None:
                # teacher (R, T, Qt) raw-space quantiles (anything where ~teacher_obs) -> same advancing frame + asinh
                Qt = teacher.shape[-1]
                tv = torch.where(teacher_obs.view(R, T, 1), teacher.view(R, T, Qt).float(), torch.zeros(1, device=x.device))
                yt = torch.asinh((tv.view(R, N, P, Qt) - f_loc[..., None, None]) / f_sc[..., None, None].clamp_min(1e-6))
                wt = (teacher_obs.view(R, N, P) & target_rows[:, None, None] & active[..., None] & (cnt_before >= 2)[..., None]).float()
                yt = torch.where(wt[..., None] > 0, yt, torch.zeros_like(yt))

        x_enc = self.representation(scaled, obs)
        x_dec = self.representation.forecast_tokens(R, N)
        h = torch.cat([x_enc, x_dec.to(x_enc.dtype)], dim=1)
        cos, sin = self.transformer.rope(h, positions)
        for layer in self.transformer.layers:
            for norm, attn in zip(layer.attn_norms, layer.attns):
                if attn.is_time:
                    h = h + _time_attention(attn, norm(h), time_mask, cos, sin)
                else:
                    h = h + _variate_attention(attn, norm(h), vmask, B, V)
            h = h + layer.mlp(layer.mlp_norm(h))
        h = self.transformer.out_norm(h)
        preds = self.head(h[:, N:]).view(R, N, P, -1).float()  # (R,N,P,Q) normalized, monotone

        if return_preds:
            return preds, f_loc, f_sc
        err = y[..., None] - preds
        tau = self.q_levels.view(1, 1, 1, -1)
        pin = torch.maximum(tau * err, (tau - 1) * err).mean(-1)
        denom = w.sum().clamp_min(1.0)
        loss = (pin * w).sum() / denom
        with torch.no_grad():
            mae = ((y - preds[..., c.median_index]).abs() * w).sum() / denom
        stats = {"pinball": loss.detach(), "mae_med": mae, "n_tgt": denom}
        if teacher is not None:
            pq = preds[..., list(teacher_q_idx)]
            tq = self.q_levels[list(teacher_q_idx)].view(1, 1, 1, -1)
            et = yt - pq
            kd = torch.maximum(tq * et, (tq - 1) * et).mean(-1)
            kd_loss = (kd * wt).sum() / wt.sum().clamp_min(1.0)
            loss = (1 - lam) * loss + lam * kd_loss
            stats["kd"] = kd_loss.detach()
        return loss, stats