"""Simple neural baselines on the 1-minute branch only: MLP on the origin features and a small TCN over the last 390
one-minute bars, both with the same Student-t head and loss as MRPT (matched information set, matched tuning budget).
PatchTST-style 1m-only = MRPT(branches=['A']) (see configs/ablations.yaml); iTransformer/Chronos/TimesFM: see
docs/MODEL_CARD.md §baselines-not-run for the licensing/compatibility audit."""
from __future__ import annotations

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

from ..data.windows import F_BAR, F_GLOBAL
from ..models.student_t import StudentTHead


class MLPBaseline(nn.Module):
    """Flattened last 60 one-minute bars + global features -> MLP -> per-minute Student-t."""

    def __init__(self, horizon: int = 60, n_last: int = 60, d_hidden: int = 256, dropout: float = 0.1):
        super().__init__()
        self.n_last, self.H = n_last, horizon
        self.net = nn.Sequential(nn.Linear(n_last * F_BAR + F_GLOBAL + 4, d_hidden), nn.GELU(), nn.Dropout(dropout),
                                 nn.Linear(d_hidden, d_hidden), nn.GELU(), nn.Dropout(dropout))
        self.fut = nn.Linear(4, d_hidden)
        self.head = StudentTHead(d_hidden)
        self.cfg = type("C", (), {"head": "student_t", "quantiles": [], "__dict__": {}})()

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

    def forward(self, b):
        x = b["xA"][:, -self.n_last:].reshape(b["xA"].shape[0], -1)
        g = torch.cat([x, b["glob"], b["fut_cal"][:, 0]], -1)
        h = self.net(g).unsqueeze(1) + self.fut(b["fut_cal"].float())
        out = self.head(F.gelu(h), b["sigma"]); out["hidden"] = h
        return out


class TCNBaseline(nn.Module):
    """Dilated causal 1-D convolutions over the 390 one-minute bars; last state + future calendar -> Student-t."""

    def __init__(self, horizon: int = 60, channels: int = 96, n_blocks: int = 6, kernel: int = 3, dropout: float = 0.1):
        super().__init__()
        self.inp = nn.Conv1d(F_BAR + 1, channels, 1)
        self.blocks = nn.ModuleList()
        for i in range(n_blocks):
            d = 2 ** i
            self.blocks.append(nn.Sequential(nn.Conv1d(channels, channels, kernel, dilation=d, padding=(kernel - 1) * d),
                                             nn.GELU(), nn.Dropout(dropout), nn.Conv1d(channels, channels, 1)))
        self.pad = [(kernel - 1) * 2 ** i for i in range(n_blocks)]
        self.glob = nn.Linear(F_GLOBAL, channels)
        self.fut = nn.Linear(4, channels)
        self.pos = nn.Embedding(horizon, channels)
        self.head = StudentTHead(channels)
        self.cfg = type("C", (), {"head": "student_t", "quantiles": [], "__dict__": {}})()

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

    def forward(self, b):
        x = torch.cat([b["xA"], b["mA"].unsqueeze(-1).float()], -1).transpose(1, 2)
        h = self.inp(x)
        for blk, p in zip(self.blocks, self.pad):
            y = blk(h)[:, :, :-p] if p > 0 else blk(h)
            h = h + y
        last = h[:, :, -1] + self.glob(b["glob"].float())
        H = b["fut_cal"].shape[1]
        q = last.unsqueeze(1) + self.fut(b["fut_cal"].float()) + self.pos.weight[:H].unsqueeze(0)
        out = self.head(F.gelu(q), b["sigma"]); out["hidden"] = q
        return out