"""The assembler: layer 0 + layer 1 + 9x layer 2 + layer 11 = 12 layers.

The only file that knows the whole stack. forward(ids) -> logits (B, T, vocab).
"""
from dataclasses import dataclass, field

import torch.nn as nn

from .layer_0 import Layer0, Layer0Config
from .layer_1 import Layer1, Layer1Config
from .layer_2 import Layer2, Layer2Config
from .layer_11 import Layer11, Layer11Config


@dataclass
class NotioConfig:
    layer0: Layer0Config = field(default_factory=Layer0Config)
    layer1: Layer1Config = field(default_factory=Layer1Config)
    block: Layer2Config = field(default_factory=Layer2Config)
    head: Layer11Config = field(default_factory=Layer11Config)
    n_blocks: int = 2   # layers 2..10; + head = layer 11 -> 12 layers total


class Notio(nn.Module):
    def __init__(self, cfg: NotioConfig):
        super().__init__()
        self.cfg = cfg
        # cross-layer contracts (the one place they can be checked)
        assert cfg.layer1.block_size >= cfg.layer0.block_size,             "layer 1 capacity must be >= layer 0 sequence length"
        assert cfg.layer1.d_model == cfg.block.d_model == cfg.head.d_model,             "d_model must match everywhere"
        assert cfg.layer1.vocab_size == cfg.head.vocab_size,             "vocab_size must match layer 1 and head"

        self.layer0 = Layer0(cfg.layer0)      # data (not an nn.Module)
        self.layer1 = Layer1(cfg.layer1)
        self.blocks = nn.ModuleList([Layer2(cfg.block) for _ in range(cfg.n_blocks)])
        self.head = Layer11(cfg.head)
        self.tie_head()
        self.init_weights()

    def init_weights(self):
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=self.cfg.layer1.init_std)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
        # residual-path projections get the GPT-2 depth scaling
        for name, p in self.named_parameters():
            if name.endswith("c_proj.weight"):
                nn.init.normal_(p, mean=0.0,
                                std=self.cfg.layer1.init_std / (2 * self.cfg.n_blocks) ** 0.5)

    def tie_head(self):
        self.head.tie_to(self.layer1)

    def forward(self, ids, states=None, pos_offset=0):
        """ids: (B, T) long -> (logits (B, T, vocab), states list).

        states: per-block recurrent state from the previous window, or None
        for zeros. The caller decides when to detach/reset (truncated BPTT).
        pos_offset: position embedding offset for stateful generation.
        """
        x = self.layer1(ids, pos_offset)
        new_states = []
        for i, block in enumerate(self.blocks):
            st = None if states is None else states[i]
            x, st = block(x, st)
            new_states.append(st)
        return self.head(x), new_states

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

    def sample_batch(self):
        return self.layer0.sample_batch()