"""Export a cached float32 ONNX model for the static browser demo.

    pip install onnx onnxruntime
    python export_browser.py --model checkpoints/release --out checkpoints/browser
"""

import argparse
import json
import math
from pathlib import Path
import shutil

import numpy as np
import onnxruntime as ort
import torch

from inference import digest, load, prompt
from model import rope


class BrowserModel(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, tokens, cache):
        model = self.model
        past = cache.shape[4]
        length = tokens.shape[1]
        positions = past + torch.arange(length)
        x = model.tok_emb(tokens)
        if model.pos_emb is not None:
            x = x + model.pos_emb(positions)
        cos, sin = model.cos[positions], model.sin[positions]
        mask = torch.arange(past + length)[None, :] <= positions[:, None]
        updated = []
        for i, block in enumerate(model.blocks):
            attention = block.attn
            y = block.attn_norm(x)
            heads = attention.n_head
            q = attention.wq(y).reshape(1, length, heads, -1).transpose(1, 2)
            k = attention.wk(y).reshape(1, length, heads, -1).transpose(1, 2)
            v = attention.wv(y).reshape(1, length, heads, -1).transpose(1, 2)
            q, k = rope(q, cos, sin), rope(k, cos, sin)
            k = torch.cat((cache[i, 0], k), dim=2)
            v = torch.cat((cache[i, 1], v), dim=2)
            scores = (q @ k.transpose(-1, -2)) / math.sqrt(model.tok_emb.embedding_dim // heads)
            scores = scores.masked_fill(~mask, float("-inf"))
            y = (scores.softmax(-1) @ v).transpose(1, 2).reshape(1, length, -1)
            x = x + attention.wo(y)
            x = x + block.mlp(block.mlp_norm(x))
            updated.append(torch.stack((k, v)))
        return model.head(model.norm(x[:, -1])), torch.stack(updated)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default="checkpoints/release")
    parser.add_argument("--out", default="checkpoints/browser")
    args = parser.parse_args()
    out = Path(args.out)
    if out.exists() and any(out.iterdir()):
        raise ValueError("Output directory is not empty. Use a new directory.")
    torch.set_num_threads(4)
    model, tok, config = load(args.model)
    wrapper = BrowserModel(model).eval()
    shape = config["model"]
    empty = torch.empty(shape["n_layer"], 2, 1, shape["n_head"], 0,
                        shape["dim"] // shape["n_head"])
    tokens = torch.tensor([prompt(tok, {"Summary": "A fox finds a seed."})])
    out.mkdir(parents=True, exist_ok=True)
    with torch.no_grad():
        torch.onnx.export(wrapper, (tokens, empty), str(out / "model.onnx"),
                          input_names=["tokens", "cache"], output_names=["logits", "present"],
                          dynamic_axes={"tokens": {1: "length"}, "cache": {4: "past"},
                                        "present": {4: "total"}},
                          opset_version=17, dynamo=False)
    options = ort.SessionOptions()
    options.intra_op_num_threads = 4
    session = ort.InferenceSession(str(out / "model.onnx"), options,
                                   providers=["CPUExecutionProvider"])
    cache = empty.numpy()
    reference_cache = None
    with torch.no_grad():
        for step in range(64):
            logits, cache = session.run(None, {"tokens": tokens.numpy(), "cache": cache})
            reference, reference_cache = model.forward_cached(tokens, reference_cache)
            np.testing.assert_allclose(logits, reference[:, -1].numpy(), atol=1e-4, rtol=1e-4)
            assert logits.argmax() == reference[:, -1].argmax().item(), step
            tokens = torch.tensor([[int(logits.argmax())]])
    shutil.copyfile(Path(args.model) / "tokenizer.model", out / "tokenizer.model")
    config["browser_sha256"] = digest(out / "model.onnx")
    config["browser_bytes"] = (out / "model.onnx").stat().st_size
    (out / "config.json").write_text(json.dumps(config, indent=2) + "\n")
    print(f"64 cached greedy tokens and logits match PyTorch; ONNX: {config['browser_bytes']:,} bytes.")


if __name__ == "__main__":
    main()