#!/usr/bin/env python3
"""Generate a planet name with the single-step ONNX model."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import onnxruntime as ort


ROOT = Path(__file__).resolve().parent
DEFAULT_MODEL = ROOT / "planet_namer_fp16.onnx"
DEFAULT_VOCAB = ROOT / "vocab.json"


def generate_name(
    stats: list[float],
    *,
    model_path: Path = DEFAULT_MODEL,
    vocab_path: Path = DEFAULT_VOCAB,
    temperature: float = 0.8,
    seed: int | None = None,
) -> str:
    """Generate one name from seven normalized stats in vocab stat order."""
    with vocab_path.open(encoding="utf-8") as handle:
        vocab = json.load(handle)

    stat_order = vocab["stat_order"]
    if len(stats) != len(stat_order):
        raise ValueError(
            f"expected {len(stat_order)} stats ({', '.join(stat_order)}), "
            f"received {len(stats)}"
        )
    if any(not 0.0 <= value <= 1.0 for value in stats):
        raise ValueError("every stat must be between 0 and 1 inclusive")
    if temperature < 0:
        raise ValueError("temperature must be non-negative")

    idx_to_char = {int(key): value for key, value in vocab["idx_to_char"].items()}
    hidden_size = int(vocab.get("hidden_size", 192))
    max_len = int(vocab.get("max_len", 20))
    pad_idx = int(vocab.get("pad_idx", 0))
    sos_idx = int(vocab.get("sos_idx", 1))
    eos_idx = int(vocab.get("eos_idx", 2))

    session = ort.InferenceSession(
        str(model_path), providers=["CPUExecutionProvider"]
    )
    stats_array = np.asarray([stats], dtype=np.float32)
    zero_stats = np.zeros_like(stats_array)
    char_in = np.asarray([[sos_idx]], dtype=np.int64)
    h_in = np.zeros((1, 1, hidden_size), dtype=np.float32)
    c_in = np.zeros((1, 1, hidden_size), dtype=np.float32)
    rng = np.random.default_rng(seed)
    characters: list[str] = []

    for step in range(max_len):
        logits, h_in, c_in = session.run(
            None,
            {
                "stats_init": stats_array if step == 0 else zero_stats,
                "stats": stats_array,
                "char_in": char_in,
                "h_in": h_in,
                "c_in": c_in,
            },
        )

        if temperature < 1e-6:
            char_idx = int(np.argmax(logits[0]))
        else:
            scaled = logits[0].astype(np.float64) / temperature
            scaled -= scaled.max()
            probabilities = np.exp(scaled)
            probabilities /= probabilities.sum()
            char_idx = int(rng.choice(len(probabilities), p=probabilities))

        if char_idx in (pad_idx, eos_idx):
            break
        if char_idx != sos_idx:
            character = idx_to_char.get(char_idx)
            if character and not character.startswith("<"):
                characters.append(character)
        char_in = np.asarray([[char_idx]], dtype=np.int64)

    return "".join(characters)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--stats",
        nargs=7,
        required=True,
        type=float,
        metavar=("ATM", "GRAV", "RES", "LIFE", "TEMP", "WATER", "RAD"),
        help="seven normalized values in the documented order",
    )
    parser.add_argument("--temperature", type=float, default=0.8)
    parser.add_argument("--seed", type=int)
    parser.add_argument("--model", type=Path, default=DEFAULT_MODEL)
    parser.add_argument("--vocab", type=Path, default=DEFAULT_VOCAB)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    print(
        generate_name(
            args.stats,
            model_path=args.model,
            vocab_path=args.vocab,
            temperature=args.temperature,
            seed=args.seed,
        )
    )


if __name__ == "__main__":
    main()