#!/usr/bin/env python3
"""Text continuation for KerzgrAI/Rondo-Exp and KerzgrAI/Monody-Exp.

Both are proof-of-concept base models: they continue text and follow no
instructions. Rondo-Exp is masked diffusion: the continuation starts as masks
and is revealed in random order, with up to --drafts tokens drafted and checked
per forward (exact self-speculative decoding). Monody-Exp is its
autoregressive twin and samples one token per forward with a cache.

Run from a downloaded repo, which this script reads its model from:
  python inference.py --prompt "The water cycle is"
  python inference.py --prompt "Photosynthesis" --max-new-tokens 128 --temperature 0.8
Or anywhere, downloading the model first:
  python inference.py --repo KerzgrAI/Monody-Exp --prompt "The water cycle is"
"""

from __future__ import annotations

import argparse
import importlib
import subprocess
import sys
import time
import warnings
from pathlib import Path
from types import ModuleType

import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from tokenizers import Tokenizer

REPOS = ("KerzgrAI/Rondo-Exp", "KerzgrAI/Monody-Exp")
# The models trained on this flash-linear-attention commit (GDN-2 and Parallax attention).
FLA_COMMIT = "cbb0a72efb55c18ca0ef4f298298317573ad2cb3"
FLA_REPO = "https://github.com/fla-org/flash-linear-attention.git"
FLA_MODULES = (
    "fla.ops.gdn2",
    "fla.layers.gdn2",
    "fla.modules.rotary",
    "fla.ops.parallax.parallel",
)
TINY_GDN_FILES = (
    "tiny_gdn/__init__.py",
    "tiny_gdn/config.py",
    "tiny_gdn/detokenize.py",
    "tiny_gdn/fugue.py",
    "tiny_gdn/fugue_sampling.py",
    "tiny_gdn/fused_ops.py",
    "tiny_gdn/gated_residual_ops.py",
    "tiny_gdn/haiku_layers.py",
    "tiny_gdn/kdn/__init__.py",
    "tiny_gdn/kdn/kernels.py",
    "tiny_gdn/kdn/layer.py",
    "tiny_gdn/kdn/recurrence.py",
    "tiny_gdn/model.py",
    "tiny_gdn/ngram.py",
    "tiny_gdn/nn_common.py",
    "tiny_gdn/parallax_ops.py",
    "tiny_gdn/senryu2.py",
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Rondo-Exp / Monody-Exp text continuation"
    )
    parser.add_argument("--prompt", required=True, help="Text to continue")
    parser.add_argument(
        "--repo",
        default=None,
        help=f"Hub repo or local folder to load ({' or '.join(REPOS)}); default: this script's folder",
    )
    parser.add_argument("--max-new-tokens", type=int, default=256)
    parser.add_argument("--temperature", type=float, default=1.0, help="0 is greedy")
    parser.add_argument("--top-p", type=float, default=0.9)
    parser.add_argument("--top-k", type=int, default=0, help="0 keeps every token")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument(
        "--drafts",
        type=int,
        default=8,
        help="Rondo-Exp: tokens drafted per forward; 0 reveals one token per forward",
    )
    args = parser.parse_args()
    if args.max_new_tokens <= 0:
        parser.error("--max-new-tokens must be positive")
    if (
        args.temperature < 0
        or not 0 < args.top_p <= 1
        or args.top_k < 0
        or args.drafts < 0
    ):
        parser.error(
            "--temperature >= 0, 0 < --top-p <= 1, --top-k >= 0 and --drafts >= 0"
        )
    return args


def model_directory(repo: str | None) -> Path:
    here = Path(__file__).resolve().parent
    if repo is None:
        if (here / "config.json").is_file():
            return here
        raise SystemExit(
            f"No model next to {Path(__file__).name}; pass --repo {REPOS[0]} or --repo {REPOS[1]}"
        )
    if Path(repo).is_dir():
        return Path(repo).resolve()
    return Path(snapshot_download(repo))


def ensure_fla() -> None:
    try:
        for name in FLA_MODULES:
            importlib.import_module(name)
        return
    except Exception as error:  # noqa: BLE001
        print(
            f"Installing flash-linear-attention {FLA_COMMIT[:7]} ({error}) …",
            file=sys.stderr,
            flush=True,
        )
    pip = [sys.executable, "-m", "pip", "install"]
    subprocess.check_call([*pip, "einops", "transformers>=4.45.0"])
    subprocess.check_call(
        [*pip, "--no-deps", "--force-reinstall", f"git+{FLA_REPO}@{FLA_COMMIT}"]
    )
    for name in [
        name for name in sys.modules if name == "fla" or name.startswith("fla.")
    ]:
        del sys.modules[name]
    importlib.invalidate_caches()
    for name in FLA_MODULES:
        importlib.import_module(name)


def load_model(
    directory: Path, device: torch.device
) -> tuple[torch.nn.Module, ModuleType | None]:
    """The model in ``directory`` on ``device`` in bfloat16, and fugue_sampling when it is Rondo."""

    missing = [name for name in TINY_GDN_FILES if not (directory / name).is_file()]
    if missing:
        raise SystemExit(f"{directory} lacks the bundled code: {', '.join(missing)}")
    sys.path.insert(0, str(directory))
    config = importlib.import_module("tiny_gdn.config").TinyGDNConfig.from_json(
        directory / "config.json"
    )
    if config.model_type == "fugue":
        model_class = importlib.import_module("tiny_gdn.fugue").FugueForMaskedDiffusion
        sampling = importlib.import_module("tiny_gdn.fugue_sampling")
    else:
        model_class = importlib.import_module("tiny_gdn.senryu2").Senryu2ForCausalLM
        sampling = None
    with torch.device(device):
        model = model_class(config)
    model = model.to(dtype=torch.bfloat16)
    model.load_state_dict(
        load_file(directory / "model.safetensors", device="cpu"), strict=True
    )
    return model.eval(), sampling


def sample(
    logits: torch.Tensor, args: argparse.Namespace, generator: torch.Generator
) -> int:
    logits = logits.float()
    if args.temperature == 0:
        return int(logits.argmax())
    logits = logits / args.temperature
    if 0 < args.top_k < logits.numel():
        logits = logits.masked_fill(
            logits < torch.topk(logits, args.top_k).values[-1], -torch.inf
        )
    if args.top_p < 1:
        ordered, order = torch.sort(logits, descending=True)
        drop = torch.softmax(ordered, dim=-1).cumsum(-1) > args.top_p
        drop[1:] = drop[:-1].clone()
        drop[0] = False
        logits = logits.scatter(0, order, ordered.masked_fill(drop, -torch.inf))
    return int(torch.multinomial(torch.softmax(logits, dim=-1), 1, generator=generator))


@torch.inference_mode()
def continue_autoregressive(model, prompt: list[int], args, generator, emit) -> int:
    eos = int(model.config.eos_token_id)
    room = model.config.max_position_embeddings - len(prompt)
    output = model(
        torch.tensor([prompt], device=generator.device),
        return_logits=True,
        logits_to_keep=1,
        use_cache=True,
    )
    produced = 0
    for _ in range(min(args.max_new_tokens, room)):
        token = sample(output.logits[0, -1], args, generator)
        if token == eos:
            break
        emit(token)
        produced += 1
        output = model(
            torch.tensor([[token]], device=generator.device),
            return_logits=True,
            logits_to_keep=1,
            past_key_values=output.past_key_values,
            use_cache=True,
        )
    return produced


@torch.inference_mode()
def continue_diffusion(
    model, sampling: ModuleType, prompt: list[int], args, generator, emit
) -> int:
    eos = int(model.config.eos_token_id)
    mask = model.mask_token_id
    settings = sampling.DiffusionSampling(
        new_tokens=args.max_new_tokens,
        block_length=None,
        threshold=None,
        order="random",
        steps_per_block=None if args.drafts else args.max_new_tokens,
        draft_tokens=args.drafts,
        temperature=args.temperature,
        top_k=args.top_k,
        top_p=args.top_p,
    )
    shown = 0
    forwards = 0

    def show(canvas: torch.Tensor) -> bool:
        nonlocal shown
        committed = sampling.committed_prefix(canvas, mask).tolist()
        for token in committed[shown:]:
            if token == eos:
                return True
            emit(token)
            shown += 1
        return False

    def on_step(state) -> bool:
        nonlocal forwards
        forwards = state.step
        return show(state.tokens[0])

    canvas = sampling.diffusion_generate(
        model,
        torch.tensor([prompt], device=generator.device),
        settings,
        generator=generator,
        eos_token_id=eos,
        on_step=on_step,
    )
    show(canvas[0])
    print(f"\n[{forwards} forwards]", file=sys.stderr, end="", flush=True)
    return shown


def main() -> None:
    args = parse_args()
    warnings.filterwarnings("ignore", message=r"tl\.make_block_ptr is deprecated")
    if not torch.cuda.is_available():
        raise SystemExit(
            "An NVIDIA GPU is required: the GDN-2 kernels run in Triton on CUDA."
        )
    device = torch.device("cuda")
    directory = model_directory(args.repo)
    ensure_fla()
    model, sampling = load_model(directory, device)
    tokenizer = Tokenizer.from_file(str(directory / "tokenizer.json"))
    decoder = importlib.import_module("tiny_gdn.detokenize").IncrementalUtf8Decoder(
        tokenizer
    )
    prompt = tokenizer.encode(args.prompt, add_special_tokens=False).ids
    if not prompt:
        raise SystemExit("The prompt produced no tokens")
    room = model.config.max_position_embeddings - args.max_new_tokens
    if room < 1:
        raise SystemExit(
            f"--max-new-tokens leaves no room in {model.config.max_position_embeddings} positions"
        )
    prompt = prompt[-room:]
    generator = torch.Generator(device=device).manual_seed(args.seed)

    def emit(token: int) -> None:
        print(decoder.push(token), end="", flush=True)

    print(args.prompt, end="", flush=True)
    started = time.perf_counter()
    if sampling is None:
        produced = continue_autoregressive(model, prompt, args, generator, emit)
    else:
        produced = continue_diffusion(model, sampling, prompt, args, generator, emit)
    shown = len(decoder.text)
    print(decoder.finalize()[shown:], end="")
    elapsed = time.perf_counter() - started
    print(f"\n[{produced} tokens in {elapsed:.1f}s]", file=sys.stderr, flush=True)


if __name__ == "__main__":
    main()