"""Transparent forward passes for the Latent Space attention course.

Install numpy. Run: python numpy_core.py
Regions are included directly in the lessons, so displayed code stays in sync.
"""
from __future__ import annotations

import numpy as np
from numpy.typing import NDArray

FloatArray = NDArray[np.float64]
IntArray = NDArray[np.int64]
BoolArray = NDArray[np.bool_]
# Type aliases describe dtype; docstrings and assertions specify tensor shapes.

# region embeddings
def embedding_demo() -> tuple[IntArray, FloatArray, FloatArray]:
    vocab = {"<pad>": 0, "<bos>": 1, "<eos>": 2, "the": 3,
             "small": 4, "robot": 5, "moves": 6, "<unk>": 7}
    tokens = "the small robot moves".split()  # Teaching tokenizer only.
    ids = np.array([vocab.get(token, vocab["<unk>"]) for token in tokens])
    table = np.random.default_rng(0).normal(size=(len(vocab), 6))
    x = table[ids]                           # (4, 6), no arithmetic on IDs.
    assert np.array_equal(x[2], table[vocab["robot"]])
    return ids, table, x
# endregion embeddings

# region positions
def positions(length: int, width: int, start: int = 0) -> FloatArray:
    if width <= 0 or width % 2:
        raise ValueError("This example uses a positive, even model width")
    pos = np.arange(start, start + length, dtype=float)[:, None]
    frequency = 10000.0 ** (-np.arange(0, width, 2) / width)
    angles = pos * frequency[None, :]
    pe = np.empty((length, width))
    pe[:, 0::2] = np.sin(angles)
    pe[:, 1::2] = np.cos(angles)
    return pe
# endregion positions

# region attention
def attention(q: FloatArray, k: FloatArray, v: FloatArray, allowed: BoolArray | None = None) -> tuple[FloatArray, FloatArray]:
    """(..., T, dk), (..., S, dk), (..., S, dv). True means visible."""
    if q.shape[-1] != k.shape[-1] or k.shape[-2] != v.shape[-2]:
        raise ValueError("Q/K feature widths and K/V source lengths must match")
    scores = (q @ np.swapaxes(k, -2, -1)) / np.sqrt(q.shape[-1])
    if allowed is not None:
        visible = np.broadcast_to(np.asarray(allowed, dtype=bool), scores.shape)
        if not visible.any(axis=-1).all():
            raise ValueError("Every query needs at least one visible key")
        scores = np.where(visible, scores, -np.inf)
    shifted = scores - scores.max(axis=-1, keepdims=True)
    weights = np.exp(shifted)
    weights /= weights.sum(axis=-1, keepdims=True)
    return weights @ v, weights
# endregion attention

# region masks
def attention_masks(source_ids: IntArray, target_ids: IntArray, pad: int = 0) -> tuple[BoolArray, BoolArray]:
    """IDs: (batch, length). Return source-key and causal-target masks."""
    source_visible = (source_ids != pad)[:, None, None, :]
    target_visible = (target_ids != pad)[:, None, None, :]
    length = target_ids.shape[1]
    causal = np.arange(length)[None, :] <= np.arange(length)[:, None]
    target_visible = target_visible & causal[None, None, :, :]
    return source_visible, target_visible
# endregion masks

# region heads
class MultiHead:
    def __init__(self, width: int, heads: int, seed: int = 0) -> None:
        if heads <= 0 or width % heads:
            raise ValueError("Model width must be divisible by head count")
        self.heads, self.head_width = heads, width // heads
        rng = np.random.default_rng(seed)
        self.wq, self.wk, self.wv, self.wo = [
            rng.normal(size=(width, width)) / np.sqrt(width) for _ in range(4)
        ]

    def split(self, x: FloatArray) -> FloatArray:
        batch, length, _ = x.shape
        return x.reshape(batch, length, self.heads, self.head_width).transpose(0, 2, 1, 3)

    def __call__(self, query: FloatArray, context: FloatArray, allowed: BoolArray | None = None) -> tuple[FloatArray, FloatArray]:
        q = self.split(query @ self.wq)
        k = self.split(context @ self.wk)
        v = self.split(context @ self.wv)
        heads, weights = attention(q, k, v, allowed)
        batch, _, length, _ = heads.shape
        joined = heads.transpose(0, 2, 1, 3).reshape(batch, length, -1)
        return joined @ self.wo, weights
# endregion heads

# region norm_ffn
def layer_norm(x: FloatArray, gamma: FloatArray | None = None, beta: FloatArray | None = None, eps: float = 1e-5) -> FloatArray:
    mean = x.mean(axis=-1, keepdims=True)
    variance = ((x - mean) ** 2).mean(axis=-1, keepdims=True)
    normalized = (x - mean) / np.sqrt(variance + eps)
    gamma = np.ones(x.shape[-1]) if gamma is None else gamma
    beta = np.zeros(x.shape[-1]) if beta is None else beta
    return normalized * gamma + beta


def feed_forward(x: FloatArray, w1: FloatArray, b1: FloatArray, w2: FloatArray, b2: FloatArray) -> FloatArray:
    hidden = np.maximum(x @ w1 + b1, 0.0)  # ReLU, separately per token.
    return hidden @ w2 + b2               # Return to the residual width.
# endregion norm_ffn


def ffn_parameters(width: int, inner: int, seed: int) -> tuple[FloatArray, FloatArray, FloatArray, FloatArray]:
    rng = np.random.default_rng(seed)
    return (rng.normal(size=(width, inner)) / np.sqrt(width), np.zeros(inner),
            rng.normal(size=(inner, width)) / np.sqrt(inner), np.zeros(width))

# region encoder
class EncoderBlock:
    """A pre-norm forward pass. Normalization gain=1 and bias=0 here."""
    def __init__(self, width: int = 32, heads: int = 4, inner: int = 64, seed: int = 10) -> None:
        self.attn = MultiHead(width, heads, seed)
        self.ffn = ffn_parameters(width, inner, seed + 1)

    def __call__(self, x: FloatArray, source_visible: BoolArray) -> FloatArray:
        normalized = layer_norm(x)
        x = x + self.attn(normalized, normalized, source_visible)[0]
        return x + feed_forward(layer_norm(x), *self.ffn)
# endregion encoder

# region decoder
class DecoderBlock:
    def __init__(self, width: int = 32, heads: int = 4, inner: int = 64, seed: int = 20) -> None:
        self.self_attn = MultiHead(width, heads, seed)
        self.cross_attn = MultiHead(width, heads, seed + 1)
        self.ffn = ffn_parameters(width, inner, seed + 2)

    def __call__(self, y: FloatArray, memory: FloatArray, target_visible: BoolArray, source_visible: BoolArray) -> FloatArray:
        normalized = layer_norm(y)
        y = y + self.self_attn(normalized, normalized, target_visible)[0]
        y = y + self.cross_attn(layer_norm(y), memory, source_visible)[0]
        return y + feed_forward(layer_norm(y), *self.ffn)
# endregion decoder

# region architectures
def architecture_demo() -> tuple[FloatArray, FloatArray]:
    source = np.array([[4, 7, 9], [6, 8, 0]])
    prefix = np.array([[1, 4, 7], [1, 6, 8]])
    width, vocabulary = 32, 12
    rng = np.random.default_rng(8)
    embedding = rng.normal(size=(vocabulary, width))
    source_mask, target_mask = attention_masks(source, prefix)
    x = embedding[source] + positions(source.shape[1], width)
    y = embedding[prefix] + positions(prefix.shape[1], width)

    # Encoder-only: pool only real source positions for classification.
    encoded = layer_norm(EncoderBlock()(x, source_mask))
    real = (source != 0)[..., None]
    pooled = (encoded * real).sum(1) / real.sum(1)
    class_logits = pooled @ rng.normal(size=(width, 3))

    # Decoder-only: same two-sublayer structure, with CAUSAL visibility.
    causal_hidden = layer_norm(EncoderBlock(seed=30)(y, target_mask))
    next_token_logits = causal_hidden @ rng.normal(size=(width, vocabulary))
    return class_logits, next_token_logits  # (2, 3), (2, 3, 12)
# endregion architectures

# region forward
def demo_forward() -> FloatArray:
    # PAD=0, BOS=1, EOS=2; source symbols are integers 4..11.
    source = np.array([[4, 7, 9], [6, 8, 0]])
    target = np.array([[1, 4, 7], [1, 6, 8]])
    width, vocabulary = 32, 12
    rng = np.random.default_rng(7)
    embedding = rng.normal(size=(vocabulary, width))
    source_mask, target_mask = attention_masks(source, target)
    x = embedding[source] + positions(source.shape[1], width)
    y = embedding[target] + positions(target.shape[1], width)
    memory = layer_norm(EncoderBlock()(x, source_mask))
    hidden = layer_norm(DecoderBlock()(y, memory, target_mask, source_mask))
    output_projection = rng.normal(size=(width, vocabulary)) / np.sqrt(width)
    logits = hidden @ output_projection
    assert logits.shape == (2, 3, vocabulary)
    return logits  # Random, untrained scores: not meaningful predictions.
# endregion forward

if __name__ == "__main__":
    print("Embedding shape:", embedding_demo()[2].shape)
    print("Vocabulary logits:", demo_forward().shape)