"""Layer 1: the embedding - id -> vector.

wte: token table    (vocab_size, d_model)   "what"
wpe: position table (block_size, d_model)   "where"
x = wte(ids) + wpe(pos), then optional dropout.

Independence: imports only torch. Knows nothing about layer 0 (data)
or anything above it. Exposes tie_lm_head for the (later) output layer.
"""
from dataclasses import dataclass

import torch
import torch.nn as nn


@dataclass
class Layer1Config:
    vocab_size: int = 109
    block_size: int = 1024
    d_model: int = 512
    init_std: float = 0.02
    dropout: float = 0.0


class Layer1(nn.Module):
    def __init__(self, cfg: Layer1Config):
        super().__init__()
        self.cfg = cfg
        self.wte = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.wpe = nn.Embedding(cfg.block_size, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)
        self.init_weights()

    def init_weights(self):
        nn.init.normal_(self.wte.weight, mean=0.0, std=self.cfg.init_std)
        nn.init.normal_(self.wpe.weight, mean=0.0, std=self.cfg.init_std)

    def forward(self, ids, pos_offset=0):
        """ids: (B, T) long -> (B, T, d_model). pos_offset: position embedding
        start, used by stateful generation so wpe advances with the stream."""
        B, T = ids.shape
        assert pos_offset + T <= self.cfg.block_size, \
            f"pos_offset+T={pos_offset + T} exceeds block_size {self.cfg.block_size}"
        pos = torch.arange(T, device=ids.device) + pos_offset
        return self.drop(self.wte(ids) + self.wpe(pos))

    def tie_lm_head(self, head):
        """Share wte with the (later) output head: head.weight IS wte.weight."""
        head.weight = self.wte.weight
        return head

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


def _self_test():
    torch.manual_seed(0)
    cfg = Layer1Config()
    l1 = Layer1(cfg)
    B, T = 4, 16
    ids = torch.randint(0, cfg.vocab_size, (B, T))
    x = l1(ids)

    assert x.shape == (B, T, cfg.d_model), "embed shape"

    # same id everywhere -> the wte part is constant across positions
    same = torch.full((1, T), 42, dtype=torch.long)
    xs = l1(same)
    for p in range(T - 1):
        d = (xs[0, p] - l1.wpe.weight[p]) - (xs[0, p + 1] - l1.wpe.weight[p + 1])
        assert d.abs().max().item() < 1e-6, "wte part varies with position"

    # same position, different ids -> the wpe part is constant
    for b in range(B - 1):
        d = (x[b, 0] - l1.wte(ids[b, 0])) - (x[b + 1, 0] - l1.wte(ids[b + 1, 0]))
        assert d.abs().max().item() < 1e-6, "wpe part varies with token"

    # gradients flow through both tables
    x.sum().backward()
    assert l1.wte.weight.grad is not None and l1.wte.weight.grad.abs().sum().item() > 0
    assert l1.wpe.weight.grad is not None and l1.wpe.weight.grad.abs().sum().item() > 0

    # param count: 109*256 + 256*256
    assert l1.n_params == cfg.vocab_size * cfg.d_model + cfg.block_size * cfg.d_model == 93440

    # weight tying is a true share, not a copy
    head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
    l1.tie_lm_head(head)
    assert head.weight is l1.wte.weight, "tie is not a shared tensor"

    # block_size guard
    try:
        l1(torch.zeros((1, cfg.block_size + 1), dtype=torch.long))
        raise SystemExit("expected assertion for T > block_size")
    except AssertionError:
        pass

    print(f"wte {tuple(l1.wte.weight.shape)}  wpe {tuple(l1.wpe.weight.shape)}  "
          f"params {l1.n_params:,}")
    print("layer1 self-test: PASS")


if __name__ == "__main__":
    _self_test()