"""MRPT: Multi-Resolution Patch Transformer with a direct 60-step decoder and Student-t / Gaussian / quantile heads."""
from __future__ import annotations

from dataclasses import dataclass, field, asdict

import torch
import torch.nn as nn

from ..data.windows import F_BAR, F_GLOBAL
from .patches import PatchEmbed, sinusoidal
from .fusion import SharedFusion, CrossAttentionFusion
from .regime import RegimeEmbedding
from .student_t import StudentTHead
from .joint_path import LowRankCopulaHead

BRANCH_SPECS = {"A": dict(patch=5, res=0), "B": dict(patch=3, res=1), "C": dict(patch=2, res=2)}


@dataclass
class MRPTConfig:
    d_model: int = 256
    n_layers: int = 6
    n_heads: int = 8
    d_ff: int = 1024
    dropout: float = 0.1
    dec_layers: int = 2
    horizon: int = 60
    branches: list[str] = field(default_factory=lambda: ["A", "B", "C"])
    fusion: str = "shared"            # shared | cross
    head: str = "student_t"           # student_t | gaussian | quantile
    regime: str = "stats"             # stats | learned | none
    n_symbols: int = 0                # > 0 enables the ticker-identity embedding (ablation only)
    copula_rank: int = 0              # > 0 enables the Stage-B joint path head
    quantiles: list[float] = field(default_factory=lambda: [0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975])

    @classmethod
    def preset(cls, name: str, **kw) -> "MRPTConfig":
        base = {"small": dict(d_model=192, n_layers=4, n_heads=6, d_ff=768, dec_layers=2),
                "base": dict(d_model=384, n_layers=8, n_heads=8, d_ff=1536, dec_layers=2),
                "spec": dict(d_model=256, n_layers=6, n_heads=8, d_ff=1024, dec_layers=2),
                "tiny": dict(d_model=64, n_layers=2, n_heads=4, d_ff=128, dec_layers=1)}[name]
        base.update(kw)
        return cls(**base)


class MRPT(nn.Module):
    def __init__(self, cfg: MRPTConfig):
        super().__init__()
        self.cfg = cfg
        d = cfg.d_model
        self.embeds = nn.ModuleDict({b: PatchEmbed(F_BAR, BRANCH_SPECS[b]["patch"], d, 3, cfg.dropout)
                                     for b in cfg.branches})
        if cfg.fusion == "shared":
            self.fusion = SharedFusion(d, cfg.n_layers, cfg.n_heads, cfg.d_ff, cfg.dropout)
        else:
            self.fusion = CrossAttentionFusion(d, cfg.n_layers, cfg.n_heads, cfg.d_ff, cfg.dropout, len(cfg.branches))
        self.glob_proj = nn.Sequential(nn.Linear(F_GLOBAL, d), nn.GELU(), nn.Linear(d, d))
        self.regime = RegimeEmbedding(F_GLOBAL, d, cfg.regime)
        self.fut_pos = nn.Embedding(cfg.horizon, d)
        self.fut_cal = nn.Linear(4, d)
        self.sym_emb = nn.Embedding(cfg.n_symbols, d) if cfg.n_symbols > 0 else None
        dec_layer = nn.TransformerDecoderLayer(d, cfg.n_heads, cfg.d_ff, cfg.dropout, activation="gelu",
                                               batch_first=True, norm_first=True)
        self.decoder = nn.TransformerDecoder(dec_layer, cfg.dec_layers)
        self.dec_norm = nn.LayerNorm(d)
        if cfg.head == "student_t":
            self.head = StudentTHead(d, learn_nu=True)
        elif cfg.head == "gaussian":
            self.head = StudentTHead(d, learn_nu=False, nu_fixed=float("inf"))
        elif cfg.head == "quantile":
            self.head = nn.Linear(d, len(cfg.quantiles))
        else:
            raise ValueError(cfg.head)
        self.copula = LowRankCopulaHead(d, cfg.horizon, cfg.copula_rank) if cfg.copula_rank > 0 else None

    def n_params(self, trainable_only: bool = True) -> int:
        return sum(p.numel() for p in self.parameters() if p.requires_grad or not trainable_only)

    def encode(self, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        streams = []
        for b in self.cfg.branches:
            tok, valid = self.embeds[b](batch["x" + b], batch["m" + b].bool(), BRANCH_SPECS[b]["res"])
            streams.append((tok, valid))
        # global context token (history-only statistics) prepended to the first stream
        g = self.glob_proj(batch["glob"].float()).unsqueeze(1)
        t0, v0 = streams[0]
        streams[0] = (torch.cat([g, t0], 1), torch.cat([torch.ones_like(v0[:, :1]), v0], 1))
        return self.fusion(streams)

    def decode(self, memory: torch.Tensor, mem_valid: torch.Tensor, batch: dict[str, torch.Tensor]) -> torch.Tensor:
        B = memory.shape[0]
        H = self.cfg.horizon
        q = self.fut_pos.weight[:H].unsqueeze(0).expand(B, H, -1) + self.fut_cal(batch["fut_cal"].float())
        r = self.regime(batch["glob"].float())
        if r is not None:
            q = q + r.unsqueeze(1)
        if self.sym_emb is not None:
            q = q + self.sym_emb(batch["sym"]).unsqueeze(1)
        kpm = ~mem_valid
        if kpm.all(dim=1).any():
            kpm = kpm.clone(); kpm[kpm.all(dim=1), 0] = False
        h = self.decoder(q, memory, memory_key_padding_mask=kpm)
        return self.dec_norm(h)

    def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        memory, valid = self.encode(batch)
        h = self.decode(memory, valid, batch)
        if self.cfg.head == "quantile":
            s = batch["sigma"].float().view(-1, 1, 1)
            q = self.head(h).float() * s
            return {"quantiles": torch.cummax(q, dim=-1).values if False else q, "hidden": h}
        out = self.head(h, batch["sigma"])
        out["hidden"] = h
        return out