"""Helpers shared by all three parts: tokenizer training, seeding, logging."""

from __future__ import annotations

import json
import random
import time
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
import torch

if TYPE_CHECKING:            # only for the return annotation below
    from tokenizers import Tokenizer

# Repository root (this file lives at <root>/src/utils.py).
ROOT = Path(__file__).resolve().parent.parent
ASSETS = ROOT / "assets"
RUNS = ROOT / "runs"          # checkpoints + metrics; git-ignored


def set_seed(seed: int = 1337):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def device_of(prefer: str = "cuda") -> torch.device:
    return torch.device(prefer if torch.cuda.is_available() else "cpu")


def human(n: float) -> str:
    for unit in ["", "K", "M", "B"]:
        if abs(n) < 1000:
            return f"{n:.2f}{unit}"
        n /= 1000.0
    return f"{n:.2f}T"


# ---------------------------------------------------------------------------
# Tokenizer
# ---------------------------------------------------------------------------
SPECIALS = ["<pad>", "<bos>", "<eos>", "<unk>", "<vi>", "<ja>", "<en>"]


def train_bpe(text_iter, vocab_size: int, save_path: Path,
              specials=SPECIALS) -> "Tokenizer":
    """Train a byte-level BPE tokenizer and save it to ``save_path``.

    Byte-level is the point here: Vietnamese diacritics and Japanese kana and
    kanji are never romanised in the dataset, so a character- or word-level
    vocabulary built from an ASCII assumption would fall apart.  Working over
    raw UTF-8 bytes means every string is representable with no <unk>, and BPE
    then learns whatever multi-byte sequences are frequent enough to earn a
    slot -- which for Japanese ends up being common kanji and kana bigrams.
    """
    from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders

    save_path = Path(save_path)
    save_path.parent.mkdir(parents=True, exist_ok=True)

    tok = Tokenizer(models.BPE(unk_token=None))
    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tok.decoder = decoders.ByteLevel()
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=list(specials),
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
        show_progress=True,
    )
    tok.train_from_iterator(text_iter, trainer=trainer)
    tok.save(str(save_path))
    return tok


def load_tokenizer(path: Path):
    from tokenizers import Tokenizer
    return Tokenizer.from_file(str(path))


# ---------------------------------------------------------------------------
# Metrics I/O
# ---------------------------------------------------------------------------
def save_json(obj, path: Path):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w") as f:
        json.dump(obj, f, indent=2, default=float)


def load_json(path: Path):
    with open(path) as f:
        return json.load(f)


class Timer:
    def __init__(self):
        self.t0 = time.perf_counter()

    def lap(self) -> float:
        t = time.perf_counter()
        dt = t - self.t0
        self.t0 = t
        return dt

    @property
    def elapsed(self) -> float:
        return time.perf_counter() - self.t0


# ---------------------------------------------------------------------------
# Plot style -- one place so every figure in the report matches
# ---------------------------------------------------------------------------
def use_plot_style():
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    plt.rcParams.update({
        "figure.dpi": 140,
        "savefig.dpi": 140,
        "font.size": 9,
        "axes.grid": True,
        "grid.alpha": 0.25,
        "grid.linewidth": 0.6,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "legend.frameon": False,
        "figure.autolayout": True,
    })
    return plt


# One colour per entity, assigned in a fixed order and never cycled, so a
# configuration keeps its colour in every figure of the report.
#
# These five are a subset of the documented categorical palette chosen to clear
# the *all-pairs* floors, not just the adjacent ones: the five series in these
# charts are lines that cross each other, so any pair can end up adjacent on
# screen.  Verified with ``python -m src.palette_check`` --
# worst CVD dE 13.0 (floor 8.0), worst normal-vision dE 16.3 (hard floor 15.0).
# Slots 2 and 3 sit under 3:1 contrast on the light surface, so the relief rule
# applies and every figure using them carries direct labels or a value table.
SERIES_COLORS = ["#2a78d6", "#eda100", "#e87ba4", "#008300", "#4a3aa7"]

PALETTE = {
    "c1_dense":          SERIES_COLORS[0],
    "c2_moe_4e1a":       SERIES_COLORS[1],
    "c3_moe_4e2a":       SERIES_COLORS[2],
    "c4_moe_shared":     SERIES_COLORS[3],
    "c5_moe_4e2a_wide":  SERIES_COLORS[4],
    "adamw":             SERIES_COLORS[0],
    "mars":              SERIES_COLORS[1],
    "lion":              SERIES_COLORS[2],
    "muon":              SERIES_COLORS[3],
    "sophia":            SERIES_COLORS[4],
}

# Short, readable names for axes and legends.
LABELS = {
    "c1_dense":          "C1 dense MLP",
    "c2_moe_4e1a":       "C2 MoE 4e/top-1",
    "c3_moe_4e2a":       "C3 MoE 4e/top-2",
    "c4_moe_shared":     "C4 1 shared + 3e/top-1",
    "c5_moe_4e2a_wide":  "C5 MoE 4e/top-2 wide",
    "adamw":             "AdamW",
    "mars":              "MARS",
    "lion":              "Lion",
    "muon":              "Muon",
    "sophia":            "Sophia-G",
}

# Sequential ramp (one hue, light -> dark) for plain magnitude heatmaps.
SEQUENTIAL = ["#eef4fc", "#cfe0f7", "#a8c8f0", "#7fade6", "#5590d9",
              "#2a78d6", "#1f5da8", "#154079"]

# Diverging ramp for the routing heatmaps.  Expert share has a meaningful
# midpoint -- 1/E, the perfectly balanced share -- and the question being asked
# is which side of it a cell falls on, so this is a polarity encoding, not a
# magnitude one.  Blue <-> red with a neutral grey centre, equal steps per arm.
DIVERGING = ["#154079", "#1f5da8", "#2a78d6", "#7fade6", "#cfe0f7",
             "#f0efec",
             "#f7d5cf", "#eda28f", "#e0664f", "#c33a2b", "#8f2418"]