"""Reproduce the constructed numerical examples in Article 04.

Standalone standard-library implementation, not a pretrained embedding model.
Run with: python assets/code/first-principles-04-examples.py

The published SVG figures are static educational illustrations. This script
checks their declared values, not a measured training experiment.
"""

from __future__ import annotations

import json
import math
from collections import Counter

VOCAB: tuple[str, ...] = ("I", "like", "cats", "dogs", "pizza")
SENTENCES: tuple[tuple[str, ...], ...] = (
    ("I", "like", "cats"),
    ("I", "like", "dogs"),
    ("I", "like", "cats"),
    ("I", "like", "pizza"),
    ("I", "like", "cats"),
)
INPUT_EMBEDDINGS: tuple[tuple[float, float], ...] = (
    (0.1, 0.4),
    (0.3, -0.2),
    (0.0, 0.0),
    (-0.4, 0.2),
    (0.2, 0.5),
)
OUTPUT_EMBEDDINGS: tuple[tuple[float, float], ...] = (
    (-1.0, 0.0),
    (1.0, 0.0),
    (0.0, 1.0),
    (0.0, -1.0),
    (0.0, 0.0),
)


def toy_bpe() -> list[dict[str, object]]:
    """Four merges using first-encountered tie-breaking and word boundaries."""
    words: list[tuple[list[str], int]] = [
        (list("cat"), 5),
        (list("cats"), 3),
        (list("dog"), 4),
        (list("dogs"), 2),
    ]
    merges: list[dict[str, object]] = []
    for _ in range(4):
        counts: dict[tuple[str, str], int] = {}
        for symbols, frequency in words:
            for pair in zip(symbols, symbols[1:]):
                counts[pair] = counts.get(pair, 0) + frequency
        pair = max(counts, key=counts.__getitem__)
        merges.append({"pair": list(pair), "count": counts[pair]})
        updated: list[tuple[list[str], int]] = []
        for symbols, frequency in words:
            out: list[str] = []
            i = 0
            while i < len(symbols):
                if i + 1 < len(symbols) and (symbols[i], symbols[i + 1]) == pair:
                    out.append(symbols[i] + symbols[i + 1])
                    i += 2
                else:
                    out.append(symbols[i])
                    i += 1
            updated.append((out, frequency))
        words = updated
    assert [row["count"] for row in merges] == [8, 8, 6, 6]
    assert words[1][0] == ["cat", "s"]
    assert words[3][0] == ["dog", "s"]
    return merges


def cooccurrences() -> list[list[int]]:
    lookup = {word: i for i, word in enumerate(VOCAB)}
    matrix = [[0 for _ in VOCAB] for _ in VOCAB]
    for sentence in SENTENCES:
        for pos, word in enumerate(sentence):
            for neighbor in (pos - 1, pos + 1):
                if 0 <= neighbor < len(sentence):
                    matrix[lookup[word]][lookup[sentence[neighbor]]] += 1
    assert matrix == [
        [0, 5, 0, 0, 0],
        [5, 0, 3, 1, 1],
        [0, 3, 0, 0, 0],
        [0, 1, 0, 0, 0],
        [0, 1, 0, 0, 0],
    ]
    return matrix


def prediction(embedding: tuple[float, float]) -> tuple[list[float], float]:
    logits = [embedding[0] * u[0] + embedding[1] * u[1] for u in OUTPUT_EMBEDDINGS]
    weights = [math.exp(x - max(logits)) for x in logits]
    total = sum(weights)
    probabilities = [weight / total for weight in weights]
    return probabilities, -math.log(probabilities[VOCAB.index("like")])


def single_embedding_step() -> dict[str, object]:
    initial = INPUT_EMBEDDINGS[VOCAB.index("cats")]
    q_before, loss_before = prediction(initial)
    actual = VOCAB.index("like")
    gradient = tuple(
        sum((p - float(j == actual)) * OUTPUT_EMBEDDINGS[j][k] for j, p in enumerate(q_before))
        for k in range(2)
    )
    learning_rate = 0.5
    updated = tuple(initial[k] - learning_rate * gradient[k] for k in range(2))
    q_after, loss_after = prediction(updated)
    assert gradient == (-1.0, 0.0)
    assert updated == (0.5, 0.0)
    assert math.isclose(q_after[actual], 0.313728303139719, abs_tol=1e-12)
    assert math.isclose(loss_after, 1.159227944266887, abs_tol=1e-12)
    assert loss_after < loss_before

    epsilon = 1e-6
    finite_difference = [
        (
            prediction(tuple(initial[k] + (epsilon if j == k else 0.0) for k in range(2)))[1]
            - prediction(tuple(initial[k] - (epsilon if j == k else 0.0) for k in range(2)))[1]
        )
        / (2.0 * epsilon)
        for j in range(2)
    ]
    assert all(math.isclose(a, b, abs_tol=1e-7) for a, b in zip(finite_difference, gradient))
    return {
        "input_before": list(initial),
        "input_after": list(updated),
        "gradient": list(gradient),
        "before_probabilities": q_before,
        "after_probabilities": q_after,
        "before_loss": loss_before,
        "after_loss": loss_after,
        "finite_difference": finite_difference,
    }


def positional_code(position: int, coordinate: int, dimension: int = 16) -> float:
    if dimension % 2:
        raise ValueError("Sinusoidal positional encoding requires an even dimension")
    frequency = 10000 ** (2 * (coordinate // 2) / dimension)
    angle = position / frequency
    return math.sin(angle) if coordinate % 2 == 0 else math.cos(angle)


def main() -> None:
    assert [INPUT_EMBEDDINGS[i] for i in [0, 1, 2]] == [
        (0.1, 0.4), (0.3, -0.2), (0.0, 0.0)
    ]
    assert positional_code(0, 0) == 0.0
    assert positional_code(0, 1) == 1.0

    results = {
        "description": "Constructed teaching examples, not pretrained vectors or experimental data",
        "bpe_merges": toy_bpe(),
        "vocabulary": list(VOCAB),
        "word_context_counts": cooccurrences(),
        "single_update": single_embedding_step(),
        "character_trigrams": {
            "cat": ["<ca", "cat", "at>"],
            "cats": ["<ca", "cat", "ats", "ts>"],
        },
        "positional": {
            "positions": 64,
            "dimensions": 16,
            "samples": {
                "0": [positional_code(0, k) for k in range(16)],
                "1": [positional_code(1, k) for k in range(16)],
            },
        },
    }
    print(json.dumps(results, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()