# /// script
# requires-python = ">=3.11,<3.14"
# dependencies = [
#   "torch>=2.6,<2.8",
#   "datasets>=3.0",
#   "numpy<3",
#   "onnx>=1.17",
#   "onnxruntime>=1.20",
#   "huggingface_hub>=0.30",
# ]
# ///
"""Train a character-level GRU for in-browser email autocomplete on Enron."""

import argparse
import hashlib
import json
import os
import re
import time
from collections import Counter
from pathlib import Path

import numpy as np
import onnxruntime as ort
import torch
import torch.nn as nn
import torch.nn.functional as F
from datasets import load_dataset
from huggingface_hub import HfApi
from onnxruntime.quantization import QuantType, quantize_dynamic

# Vocabulary: PAD, EOS, newline, then printable ASCII.
PAD, EOS, NEWLINE, FIRST_PRINTABLE = 0, 1, 2, 3
VOCAB_SIZE = FIRST_PRINTABLE + (127 - 32)
IGNORE = -100

ENCODE = np.full(256, 255, dtype=np.uint8)
ENCODE[10] = NEWLINE
ENCODE[32:127] = np.arange(FIRST_PRINTABLE, VOCAB_SIZE, dtype=np.uint8)
DECODE = {NEWLINE: "\n", **{FIRST_PRINTABLE + i: chr(32 + i) for i in range(95)}}

LETTER = np.zeros(VOCAB_SIZE, dtype=bool)
for char in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ":
    LETTER[ENCODE[ord(char)]] = True

KEYBOARD_NEIGHBORS = {
    "q": "wa", "w": "qes", "e": "wrd", "r": "etf", "t": "ryg", "y": "tuh", "u": "yij",
    "i": "uok", "o": "ipl", "p": "ol", "a": "qsz", "s": "awdx", "d": "serf", "f": "drgv",
    "g": "fthb", "h": "gyjn", "j": "hukm", "k": "jil", "l": "kop", "z": "asx", "x": "zsc",
    "c": "xdv", "v": "cfb", "b": "vgn", "n": "bhm", "m": "njk",
}
NEIGHBOR_TABLE = np.tile(np.arange(VOCAB_SIZE, dtype=np.uint8)[:, None], (1, 4))
for key, neighbors in KEYBOARD_NEIGHBORS.items():
    for slot in range(4):
        NEIGHBOR_TABLE[ENCODE[ord(key)], slot] = ENCODE[ord(neighbors[slot % len(neighbors)])]

CUT_REPLY = re.compile(
    r"(?im)^\s*(?:-+\s*original message|-+\s*forwarded by|-+\s*inline attachment|from:\s|to:\s|sent:\s"
    r"|subject:\s|cc:\s|.*\bwrote:\s*$|_{5,}|-{5,}|={5,}|\*{5,}|\d{1,2}/\d{1,2}/\d{2,4}\s+\d{1,2}:\d{2})"
)
SENSITIVE = re.compile(
    r"(?i)https?://|www\.|\S+@\S+|\(?\b\d{3}\)?[-. ]\d{3}[-. ]\d{4}\b|\b\d{3}-\d{2}-\d{4}\b"
)
BULK_SENDER = re.compile(r"(?i)announce|newsletter|mailer|outlook\.team|^all[._]|general\.|chairman|enron\.(?:on|net)")
SMART_CHARS = str.maketrans({"‘": "'", "’": "'", "“": '"', "”": '"', "–": "-", "—": "-", "\t": " "})


def clean_paragraphs(body: str) -> list[str]:
    text = body.replace("\r\n", "\n").replace("=20", " ").replace("=\n", "").translate(SMART_CHARS)
    text = "\n".join(line for line in text.split("\n") if not line.lstrip().startswith(">"))
    cut = CUT_REPLY.search(text)
    if cut:
        text = text[: cut.start()]
    text = re.sub(r"(?m)^-- *\n[\s\S]*", "", text)

    paragraphs = []
    for block in re.split(r"\n\s*\n", text):
        lines = [re.sub(r" +", " ", line).strip() for line in block.split("\n") if line.strip()]
        if not lines:
            continue
        # Enron bodies are hard-wrapped near 70 columns; only keep breaks after short lines.
        joined = lines[0]
        for previous, line in zip(lines, lines[1:]):
            joined += (" " if len(previous) >= 50 else "\n") + line
        letters = sum(char.isalpha() for char in joined)
        if SENSITIVE.search(joined) or letters < 0.6 * len(joined):
            continue
        paragraphs.append(joined)
    return paragraphs


def load_emails(limit_frac: float) -> dict[str, list[np.ndarray]]:
    dataset = load_dataset("corbt/enron-emails", split="train").select_columns(["from", "body"])
    if limit_frac < 1:
        dataset = dataset.shuffle(seed=0).select(range(int(len(dataset) * limit_frac)))

    parsed: list[tuple[str, list[str]]] = []
    for row in dataset:
        sender = (row["from"] or "").lower()
        if "@enron.com" not in sender or BULK_SENDER.search(sender):
            continue
        paragraphs = clean_paragraphs(row["body"] or "")
        if paragraphs:
            parsed.append((sender, paragraphs))

    counts = Counter(p.lower() for _, paragraphs in parsed for p in paragraphs if len(p) >= 60)

    seen: set[bytes] = set()
    splits: dict[str, list[np.ndarray]] = {"train": [], "val": [], "test": []}
    for sender, paragraphs in parsed:
        text = "\n\n".join(p for p in paragraphs if counts[p.lower()] < 5)
        if not 30 <= len(text) <= 6000:
            continue
        digest = hashlib.sha1(re.sub(r"\s+", " ", text.lower()).encode()).digest()
        if digest in seen:
            continue
        seen.add(digest)
        bucket = int(hashlib.sha1(sender.encode()).hexdigest(), 16) % 100
        split = "test" if bucket < 5 else "val" if bucket < 10 else "train"
        ids = ENCODE[np.frombuffer(text.encode("ascii", "ignore"), dtype=np.uint8)]
        splits[split].append(np.append(ids[ids != 255], EOS).astype(np.uint8))
    return splits


def concat_shuffled(emails: list[np.ndarray], rng: np.random.Generator) -> np.ndarray:
    return np.concatenate([emails[i] for i in rng.permutation(len(emails))])


def add_typos(clean: np.ndarray, word_rate: float, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray]:
    """Corrupt the input stream while the targets keep predicting the clean next character."""
    x = clean[:-1].copy()
    y = clean[1:].astype(np.int16)
    is_letter = LETTER[x]
    starts = np.flatnonzero(is_letter[1:] & ~is_letter[:-1]) + 1
    ends = np.flatnonzero(is_letter[:-1] & ~is_letter[1:]) + 1
    if len(ends) and len(starts) and ends[0] <= starts[0]:
        ends = ends[1:]
    count = min(len(starts), len(ends))
    starts, ends = starts[:count], ends[:count]

    lengths = ends - starts
    chosen = (lengths >= 3) & (rng.random(count) < word_rate)
    word_starts, word_lengths = starts[chosen], lengths[chosen]
    position = word_starts + 1 + (rng.random(len(word_starts)) * (word_lengths - 2)).astype(np.int64)
    position = position[position + 1 < len(x)]
    ops = rng.integers(0, 4, size=len(position))
    neighbor = NEIGHBOR_TABLE[x[position], rng.integers(0, 4, size=len(position))]

    substitute, swap, insert, delete = (position[ops == op] for op in range(4))
    x[substitute] = neighbor[ops == 0]
    first, second = x[swap].copy(), x[swap + 1].copy()
    x[swap], x[swap + 1], y[swap] = second, first, IGNORE

    repeats = np.ones(len(x), dtype=np.int64)
    repeats[delete] = 0
    repeats[insert] = 2
    inserted_at = np.cumsum(repeats)[insert] - 1
    x_out, y_out = np.repeat(x, repeats), np.repeat(y, repeats)
    x_out[inserted_at] = neighbor[ops == 2]
    return x_out, y_out


def make_batches(x, y, batch_size, seq_len, rng):
    offset = int(rng.integers(seq_len))
    starts = rng.permutation(np.arange(offset, len(x) - seq_len, seq_len))
    span = np.arange(seq_len)
    for index in range(0, len(starts) - batch_size + 1, batch_size):
        positions = starts[index : index + batch_size][:, None] + span
        yield torch.from_numpy(x[positions].astype(np.int64)), torch.from_numpy(y[positions].astype(np.int64))


class CharGRU(nn.Module):
    def __init__(self, hidden: int, embedding: int = 64):
        super().__init__()
        self.embedding = nn.Embedding(VOCAB_SIZE, embedding, padding_idx=PAD)
        self.gru = nn.GRU(embedding, hidden, batch_first=True)
        self.head = nn.Linear(hidden, VOCAB_SIZE)

    def forward(self, ids, state=None):
        output, state = self.gru(self.embedding(ids), state)
        return self.head(output), state


def sequence_loss(logits, targets, burn_in):
    return F.cross_entropy(logits[:, burn_in:].flatten(0, 1), targets[:, burn_in:].flatten(), ignore_index=IGNORE)


@torch.no_grad()
def evaluate_loss(model, emails, args, device, max_chars=5_000_000):
    rng = np.random.default_rng(1)
    clean = concat_shuffled(emails, rng)[:max_chars]
    x, y = clean[:-1], clean[1:].astype(np.int16)
    total, batches = 0.0, 0
    for inputs, targets in make_batches(x, y, args.batch, args.seq, rng):
        logits, _ = model(inputs.to(device))
        total += sequence_loss(logits, targets.to(device), args.burn_in).item()
        batches += 1
    return total / max(batches, 1)


def decode(ids) -> str:
    return "".join(DECODE.get(int(i), "") for i in ids)


def sample_positions(emails, count, context, horizon, rng):
    contexts, truths = [], []
    attempts = 0
    while len(contexts) < count and attempts < count * 200:
        attempts += 1
        email = emails[int(rng.integers(len(emails)))]
        if len(email) < 64:
            continue
        cursor = int(rng.integers(16, len(email) - 16))
        previous, following = int(email[cursor - 1]), int(email[cursor])
        if not (LETTER[previous] or previous == ENCODE[32]) or not LETTER[following]:
            continue
        window = np.full(context, EOS, dtype=np.int64)
        prefix = email[max(0, cursor - context) : cursor]
        window[-len(prefix):] = prefix
        truth = email[cursor : cursor + horizon]
        ending = np.flatnonzero((truth == EOS) | (truth == NEWLINE))
        truths.append(decode(truth[: ending[0]] if len(ending) else truth))
        contexts.append(window)
    return np.stack(contexts), truths


def build_suggestion(ids, probabilities, threshold, max_words=6):
    text, confidence, best, words = "", 1.0, "", 0
    for token, probability in zip(ids, probabilities):
        token = int(token)
        if token in (EOS, NEWLINE, PAD):
            break
        char = DECODE[token]
        if char == " ":
            if confidence >= threshold:
                best = text
            words += 1
            if words >= max_words:
                break
        confidence *= float(probability)
        if confidence < threshold:
            break
        text += char
        if char in ".,!?":
            best = text
    return best


def score_suggestions(generated, probabilities, truths, thresholds):
    report = {}
    for threshold in thresholds:
        shown, correct, saved = 0, 0, 0
        for ids, probs, truth in zip(generated, probabilities, truths):
            suggestion = build_suggestion(ids, probs, threshold)
            if len(suggestion) < 4:
                continue
            shown += 1
            if truth.startswith(suggestion):
                correct += 1
                saved += len(suggestion)
        report[str(threshold)] = {
            "coverage": round(shown / len(truths), 4),
            "precision": round(correct / shown, 4) if shown else None,
            "avg_chars_when_correct": round(saved / correct, 2) if correct else None,
            "chars_saved_per_position": round(saved / len(truths), 3),
        }
    return report


@torch.no_grad()
def greedy_torch(model, contexts, steps, device, chunk=2048):
    all_ids, all_probs = [], []
    for index in range(0, len(contexts), chunk):
        ids = torch.from_numpy(contexts[index : index + chunk]).to(device)
        logits, state = model(ids)
        step_logits = logits[:, -1]
        batch_ids, batch_probs = [], []
        for _ in range(steps):
            probabilities = step_logits.softmax(-1)
            confidence, token = probabilities.max(-1)
            batch_ids.append(token.cpu().numpy())
            batch_probs.append(confidence.cpu().numpy())
            logits, state = model(token[:, None], state)
            step_logits = logits[:, -1]
        all_ids.append(np.stack(batch_ids, 1))
        all_probs.append(np.stack(batch_probs, 1))
    return np.concatenate(all_ids), np.concatenate(all_probs)


def greedy_onnx(session, contexts, steps, hidden):
    state = np.zeros((len(contexts), hidden), dtype=np.float32)
    logits = None
    for column in range(contexts.shape[1]):
        logits, state = session.run(None, {"input_id": contexts[:, column], "state": state})
    ids, confidences = [], []
    for _ in range(steps):
        probabilities = np.exp(logits - logits.max(-1, keepdims=True))
        probabilities /= probabilities.sum(-1, keepdims=True)
        token = probabilities.argmax(-1)
        ids.append(token)
        confidences.append(probabilities[np.arange(len(token)), token])
        logits, state = session.run(None, {"input_id": token.astype(np.int64), "state": state})
    return np.stack(ids, 1), np.stack(confidences, 1)


class MatMulLinear(nn.Module):
    """Linear layer that exports as MatMul so ONNX Runtime dynamic quantization applies."""

    def __init__(self, weight: torch.Tensor, bias: torch.Tensor):
        super().__init__()
        self.weight = nn.Parameter(weight.t().clone())
        self.bias = nn.Parameter(bias.clone())

    def forward(self, inputs):
        return inputs.matmul(self.weight) + self.bias


class SingleStep(nn.Module):
    """One-character GRU update with weights copied from the trained nn.GRU."""

    def __init__(self, model: CharGRU, hidden: int):
        super().__init__()
        self.embedding = model.embedding
        input_weight, hidden_weight = model.gru.weight_ih_l0.detach(), model.gru.weight_hh_l0.detach()
        input_bias, hidden_bias = model.gru.bias_ih_l0.detach(), model.gru.bias_hh_l0.detach()
        layers = []
        for gate in range(3):
            row = slice(gate * hidden, (gate + 1) * hidden)
            layers.append(MatMulLinear(input_weight[row], input_bias[row]))
            layers.append(MatMulLinear(hidden_weight[row], hidden_bias[row]))
        self.gates = nn.ModuleList(layers)
        self.head = MatMulLinear(model.head.weight.detach(), model.head.bias.detach())

    def forward(self, input_id, state):
        embedded = self.embedding(input_id)
        reset = torch.sigmoid(self.gates[0](embedded) + self.gates[1](state))
        update = torch.sigmoid(self.gates[2](embedded) + self.gates[3](state))
        proposal = torch.tanh(self.gates[4](embedded) + reset * self.gates[5](state))
        next_state = (1 - update) * proposal + update * state
        return self.head(next_state), next_state


def verify_single_step(model, step, hidden):
    ids = torch.randint(3, VOCAB_SIZE, (4, 32))
    logits, _ = model(ids)
    state = torch.zeros(4, hidden)
    outputs = []
    for column in range(ids.shape[1]):
        step_logits, state = step(ids[:, column], state)
        outputs.append(step_logits)
    difference = (torch.stack(outputs, 1) - logits).abs().max().item()
    if difference > 1e-3:
        raise RuntimeError(f"Exported single-step GRU does not match nn.GRU: max difference {difference}")
    return difference


def measure_latency(path, hidden, runs=3000):
    options = ort.SessionOptions()
    options.intra_op_num_threads = 1
    session = ort.InferenceSession(str(path), options, providers=["CPUExecutionProvider"])
    state = np.zeros((1, hidden), dtype=np.float32)
    timings = []
    for run in range(runs):
        token = np.array([FIRST_PRINTABLE + run % 90], dtype=np.int64)
        started = time.perf_counter()
        _, state = session.run(None, {"input_id": token, "state": state})
        timings.append((time.perf_counter() - started) * 1000)
    return {"p50_ms": round(float(np.percentile(timings, 50)), 3), "p95_ms": round(float(np.percentile(timings, 95)), 3)}


def export_and_score(model, args, output, contexts, truths, thresholds):
    model = model.cpu().eval()
    step = SingleStep(model, args.hidden).eval()
    difference = verify_single_step(model, step, args.hidden)
    print(f"single-step GRU matches nn.GRU (max difference {difference:.2e})", flush=True)

    float_path = output / "model.onnx"
    torch.onnx.export(
        step,
        (torch.zeros(1, dtype=torch.long), torch.zeros(1, args.hidden)),
        float_path,
        input_names=["input_id", "state"],
        output_names=["logits", "next_state"],
        dynamic_axes={"input_id": {0: "batch"}, "state": {0: "batch"}, "logits": {0: "batch"}, "next_state": {0: "batch"}},
        opset_version=17,
        dynamo=False,
    )
    quantized_path = output / "model_quint8.onnx"
    quantize_dynamic(float_path, quantized_path, weight_type=QuantType.QUInt8, op_types_to_quantize=["MatMul"])

    subset = min(1000, len(contexts))
    result = {}
    for name, path in (("float", float_path), ("quint8", quantized_path)):
        session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
        ids, probabilities = greedy_onnx(session, contexts[:subset], args.horizon, args.hidden)
        result[name] = {
            "size_mb": round(path.stat().st_size / 1e6, 2),
            "native_cpu_single_thread_step": measure_latency(path, args.hidden),
            "suggestions_1000_positions": score_suggestions(ids, probabilities, truths[:subset], thresholds),
        }
    return result


def show_samples(model, device):
    prompts = [
        "Hi John,\n\nThanks for ", "Please let me know if ", "I will send you ", "Let me kn", "Thnaks for the ",
        "I look forward to ", "Can we schedule a ", "Sorry for the dela", "Plese find attached ", "We need to discuss ",
    ]
    samples = {}
    for prompt in prompts:
        window = np.full(128, EOS, dtype=np.int64)
        encoded = ENCODE[np.frombuffer(prompt.encode("ascii"), dtype=np.uint8)]
        window[-len(encoded):] = encoded
        ids, probabilities = greedy_torch(model, window[None], 48, device)
        samples[prompt] = {
            "continuation": decode(ids[0]),
            "suggestion_at_0.2": build_suggestion(ids[0], probabilities[0], 0.2),
        }
        print(f"{prompt!r} -> {samples[prompt]}", flush=True)
    return samples


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--hidden", type=int, default=512)
    parser.add_argument("--epochs", type=int, default=30)
    parser.add_argument("--batch", type=int, default=256)
    parser.add_argument("--seq", type=int, default=256)
    parser.add_argument("--burn-in", type=int, default=32)
    parser.add_argument("--lr", type=float, default=1.5e-3)
    parser.add_argument("--typo-rate", type=float, default=0.08)
    parser.add_argument("--limit-frac", type=float, default=1.0)
    parser.add_argument("--eval-positions", type=int, default=5000)
    parser.add_argument("--context", type=int, default=128)
    parser.add_argument("--horizon", type=int, default=48)
    parser.add_argument("--repo", default="cfahlgren1/enron-autocomplete")
    parser.add_argument("--no-push", action="store_true")
    return parser.parse_args()


def main():
    args = parse_args()
    torch.manual_seed(0)
    rng = np.random.default_rng(0)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"device={device} {torch.cuda.get_device_name(0) if device.type == 'cuda' else ''} args={vars(args)}", flush=True)
    output = Path("outputs")
    output.mkdir(parents=True, exist_ok=True)

    started = time.time()
    splits = load_emails(args.limit_frac)
    data_stats = {name: {"emails": len(e), "chars": int(sum(map(len, e)))} for name, e in splits.items()}
    print(f"data ready in {time.time() - started:.0f}s: {data_stats}", flush=True)
    if not splits["train"] or not splits["val"] or not splits["test"]:
        raise RuntimeError(f"A split is empty: {data_stats}")

    model = CharGRU(args.hidden).to(device)
    parameters = sum(p.numel() for p in model.parameters())
    print(f"parameters={parameters:,}", flush=True)

    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-5)
    steps_per_epoch = max(data_stats["train"]["chars"] // (args.batch * args.seq), 1)
    total_steps = args.epochs * steps_per_epoch
    scheduler = torch.optim.lr_scheduler.LambdaLR(
        optimizer, lambda step: 0.02 + 0.98 * 0.5 * (1 + np.cos(np.pi * min(step / total_steps, 1)))
    )

    history = []
    best_loss = float("inf")
    best_path = output / "best.pt"
    for epoch in range(args.epochs):
        model.train()
        epoch_started = time.time()
        x, y = add_typos(concat_shuffled(splits["train"], rng), args.typo_rate, rng)
        running, batches, chars = 0.0, 0, 0
        for inputs, targets in make_batches(x, y, args.batch, args.seq, rng):
            logits, _ = model(inputs.to(device, non_blocking=True))
            loss = sequence_loss(logits, targets.to(device, non_blocking=True), args.burn_in)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()
            running += loss.item()
            batches += 1
            chars += inputs.numel()
        if batches == 0:
            raise RuntimeError("No training batches; reduce --batch or --seq for this data size")

        model.eval()
        val_loss = evaluate_loss(model, splits["val"], args, device)
        seconds = time.time() - epoch_started
        record = {
            "epoch": epoch + 1,
            "train_loss": round(running / batches, 4),
            "val_loss": round(val_loss, 4),
            "seconds": round(seconds, 1),
            "chars_per_sec": int(chars / seconds),
        }
        history.append(record)
        print(json.dumps(record), flush=True)
        if val_loss < best_loss:
            best_loss = val_loss
            torch.save({"model": model.state_dict(), "args": vars(args)}, best_path)

    model.load_state_dict(torch.load(best_path, map_location=device)["model"])
    model.eval()

    thresholds = [0.05, 0.1, 0.2, 0.35, 0.5, 0.7]
    eval_rng = np.random.default_rng(42)
    contexts, truths = sample_positions(splits["test"], args.eval_positions, args.context, args.horizon, eval_rng)
    ids, probabilities = greedy_torch(model, contexts, args.horizon, device)
    suggestions = score_suggestions(ids, probabilities, truths, thresholds)
    print(json.dumps({"held_out_sender_suggestions": suggestions}, indent=1), flush=True)

    examples = []
    for index in range(min(25, len(contexts))):
        examples.append({
            "context": decode(contexts[index])[-60:],
            "truth": truths[index],
            "suggestion_at_0.2": build_suggestion(ids[index], probabilities[index], 0.2),
        })
    samples = show_samples(model, device)
    exported = export_and_score(model, args, output, contexts, truths, thresholds)
    print(json.dumps(exported, indent=1), flush=True)

    report = {
        "args": vars(args),
        "device": str(device),
        "parameters": parameters,
        "data": data_stats,
        "history": history,
        "best_val_loss": round(best_loss, 4),
        "test_positions": len(contexts),
        "held_out_sender_suggestions": suggestions,
        "exports": exported,
        "prompt_samples": samples,
        "held_out_examples": examples,
        "total_seconds": round(time.time() - started, 1),
    }
    (output / "report.json").write_text(json.dumps(report, indent=2) + "\n")
    (output / "vocab.json").write_text(json.dumps({
        "pad": PAD, "eos": EOS, "newline": NEWLINE, "first_printable_id": FIRST_PRINTABLE,
        "first_printable_codepoint": 32, "size": VOCAB_SIZE, "hidden": args.hidden,
    }, indent=2) + "\n")

    if args.no_push:
        print(f"Skipping upload; artifacts in {output}", flush=True)
        return
    api = HfApi(token=os.environ["HF_TOKEN"])
    api.create_repo(args.repo, repo_type="model", private=True, exist_ok=True)
    api.upload_folder(folder_path=str(output), repo_id=args.repo, repo_type="model")
    print(f"Uploaded to https://huggingface.co/{args.repo}", flush=True)


if __name__ == "__main__":
    main()