"""Transcript-window baseline: "the context window is the mind".

A causal decoder-only transformer over the last W transcript tokens (the same
per-event token stream PNSR sees, concatenated), parameter-matched to PNSR,
with the IDENTICAL factorized head system and cache record bank.  TXE is the
same class with W = one event (current-event-only trained null).
"""
from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn

from .modules import Block, HeadSystem, RecordEncoder


@dataclass
class TXConfig:
    vocab: int = 8192
    d: int = 384
    heads: int = 8
    layers: int = 10   # param-matched to PNSR (see PNSRConfig.enc_layers note)
    window: int = 768


class TX(nn.Module):
    def __init__(self, cfg: TXConfig):
        super().__init__()
        self.cfg = cfg
        d = cfg.d
        self.tok_emb = nn.Embedding(cfg.vocab, d)
        self.pos_emb = nn.Parameter(torch.zeros(cfg.window, d))
        self.blocks = nn.ModuleList(Block(d, cfg.heads) for _ in range(cfg.layers))
        self.ln = nn.LayerNorm(d)
        self.recenc = RecordEncoder(d, self.tok_emb)
        self.heads = HeadSystem(d)
        mask = torch.triu(torch.ones(cfg.window, cfg.window, dtype=torch.bool), 1)
        self.register_buffer("causal", mask, persistent=False)

    def forward(self, tok, rec_bank, live_mask):
        """tok: [B,W] left-padded with 0; final position is the query point."""
        B, W = tok.shape
        pad = tok == 0
        x = self.tok_emb(tok.long()) + self.pos_emb[:W]
        am = self.causal[:W, :W]
        for blk in self.blocks:
            x = blk(x, attn_mask=am, self_padding_mask=pad)
        h = self.ln(x[:, -1])                      # final token = query position
        return self.heads(None, h, None, rec_bank, live_mask)

    def param_report(self):
        return {"total": sum(p.numel() for p in self.parameters())}