#!/usr/bin/env python3
"""Minimal attention heatmap toy example.

The default check path uses only the Python standard library. PNG export needs
matplotlib and is intentionally optional so shape/mask semantics can still be
tested in a minimal environment.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
from pathlib import Path


TOKENS = ["[CLS]", "the", "method", "fails", "when", "mask", "flips", "[PAD]"]
PAD_INDEX = len(TOKENS) - 1


def dot(a: list[float], b: list[float]) -> float:
    return sum(x * y for x, y in zip(a, b))


def softmax(scores: list[float]) -> list[float]:
    m = max(scores)
    exps = [math.exp(s - m) for s in scores]
    z = sum(exps)
    return [v / z for v in exps]


def make_qk_vectors(tokens: list[str]) -> tuple[list[list[float]], list[list[float]]]:
    q, k = [], []
    for i, token in enumerate(tokens):
        is_boundary = 1.0 if token in {"[CLS]", "[PAD]"} else 0.0
        is_diagnostic = 1.0 if token in {"method", "mask", "flips"} else 0.0
        position = i / max(1, len(tokens) - 1)
        q.append([1.0, is_boundary, is_diagnostic, position])
        k.append([1.0, is_boundary, is_diagnostic, 1.0 - position])
    return q, k


def attention_matrix(tokens: list[str], pad_index: int) -> list[list[float]]:
    q, k = make_qk_vectors(tokens)
    d_head = len(q[0])
    rows = []
    for query_index, query in enumerate(q):
        scores = []
        for key_index, key in enumerate(k):
            score = dot(query, key) / math.sqrt(d_head)
            if key_index == query_index:
                score += 0.6
            if tokens[key_index] == "[PAD]":
                score = -1e9
            scores.append(score)
        rows.append(softmax(scores))
    assert all(abs(row[pad_index]) < 1e-9 for row in rows)
    return rows


def write_csv(path: Path, tokens: list[str], matrix: list[list[float]]) -> None:
    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["query/key", *tokens])
        for token, row in zip(tokens, matrix):
            writer.writerow([token, *[f"{v:.6f}" for v in row]])


def write_png(path: Path, tokens: list[str], matrix: list[list[float]]) -> None:
    import matplotlib.pyplot as plt  # type: ignore

    fig, ax = plt.subplots(figsize=(10, 6), dpi=160)
    im = ax.imshow(matrix, cmap="viridis", vmin=0.0, vmax=max(max(r) for r in matrix))
    ax.set_xticks(range(len(tokens)), labels=tokens, rotation=45, ha="right")
    ax.set_yticks(range(len(tokens)), labels=tokens)
    ax.set_xlabel("Key token")
    ax.set_ylabel("Query token")
    ax.set_title("Toy attention weights after softmax")
    fig.colorbar(im, ax=ax, fraction=0.035, pad=0.04)
    fig.tight_layout()
    fig.savefig(path)
    plt.close(fig)


def invariant_report(tokens: list[str], matrix: list[list[float]]) -> dict[str, object]:
    row_sums = [sum(row) for row in matrix]
    return {
        "tokens": tokens,
        "shape": [len(matrix), len(matrix[0])],
        "row_sums_close_to_one": all(abs(v - 1.0) < 1e-9 for v in row_sums),
        "masked_pad_column_max": max(row[PAD_INDEX] for row in matrix),
        "max_attention": max(max(row) for row in matrix),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out-dir", default="outputs")
    parser.add_argument("--save-png", action="store_true")
    parser.add_argument("--check-only", action="store_true")
    args = parser.parse_args()

    matrix = attention_matrix(TOKENS, PAD_INDEX)
    report = invariant_report(TOKENS, matrix)
    print(json.dumps(report, ensure_ascii=False, indent=2))

    if args.check_only:
        return

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    write_csv(out_dir / "toy_attention.csv", TOKENS, matrix)
    if args.save_png:
        write_png(out_dir / "toy_attention_heatmap.png", TOKENS, matrix)
        print(f"wrote {out_dir / 'toy_attention_heatmap.png'}")
    print(f"wrote {out_dir / 'toy_attention.csv'}")


if __name__ == "__main__":
    main()