#!/usr/bin/env python3
"""Portable NeuronAI-2B / Alloma-style Uzbek benchmark runner.

Examples:
  python benchmark.py --limit 200 --output quick-results.json
  python benchmark.py --limit 0 --comet --output full-results.json

`--limit 0` evaluates every example. The default 200-example run is a quick,
deterministically sampled sanity check and must not be compared with the full
scores in the model card.
"""

from __future__ import annotations

import argparse
import json
import re
from datetime import datetime, timezone
from pathlib import Path

import torch
from datasets import concatenate_datasets, load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer


MODEL_ID = "NeuronUz/NeuronAI-2B"
LETTERS = "ABCD"

TRANSLATION_PROMPTS = {
    "uz-en": (
        "Translate each Uzbek sentence into English.\n\n"
        "1991-yilning 1-sentabrida O'zbekiston mustaqilligini e'lon qildi.\n"
        "-> On 1 September 1991, Uzbekistan declared its independence.\n\n"
        "Tadqiqotchilar yangi usul samaradorligi 47 foizga oshganini aniqladilar.\n"
        "-> Researchers found that the new method improved efficiency by 47 percent.\n\n"
        "{text}\n->"
    ),
    "en-uz": "Translate into Uzbek:\n\n{text}",
}

SENTIMENT_PROMPT = (
    "Given the following Uzbek text, determine the sentiment as either "
    "'Positive' or 'Negative'. Respond with only one label.\n\nText: {text}\n\nLabel:"
)

NEWS_PROMPT = """Classify the given Uzbek news article into one category. Respond with only the category number.

0 - Siyosat
1 - Iqtisodiyot
2 - Texnologiya
3 - Sport
4 - Madaniyat
5 - Salomatlik
6 - Oila va Jamiyat
7 - Ta'lim
8 - Ekologiya
9 - Xorijiy Yangiliklar

Article: {text}

Answer:"""

MCQ_SUFFIX = {
    "uz": "Variantlarni diqqat bilan solishtiring. Javobni A, B, C yoki D harfi bilan boshlang.",
    "en": "Compare the options carefully. Start with the answer letter A, B, C, or D.",
}

MCQ_TASKS = {
    "mmlu-en": ("cais/mmlu", "all", "test", "en"),
    "mmlu-uz": ("murodbek/MMLU-uz", "default", "test", "uz"),
    "tumlu": ("jafarisbarov/TUMLU-mini", "uzbek", "test", "uz"),
}


def choose_rows(dataset, limit: int, seed: int):
    if limit and len(dataset) > limit:
        return dataset.shuffle(seed=seed).select(range(limit))
    return dataset


def strip_thinking(text: str) -> str:
    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()


class Generator:
    def __init__(self, args: argparse.Namespace):
        self.backend = args.backend
        self.tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
        self.tokenizer.padding_side = "left"
        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id

        if args.backend == "vllm":
            from vllm import LLM, SamplingParams

            self.sampling_cls = SamplingParams
            self.model = LLM(
                model=args.model,
                dtype=args.dtype,
                trust_remote_code=True,
                gpu_memory_utilization=args.gpu_memory_utilization,
                max_model_len=args.max_model_len,
                language_model_only=True,
                mamba_block_size=16,
                mamba_cache_mode="align",
            )
        else:
            dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
            self.model = AutoModelForCausalLM.from_pretrained(
                args.model,
                dtype=dtype,
                device_map="auto",
                trust_remote_code=True,
            ).eval()

    def render(self, prompt: str) -> str:
        messages = [{"role": "user", "content": prompt}]
        try:
            return self.tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
                enable_thinking=False,
            )
        except TypeError:
            return self.tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )

    def generate(self, prompts: list[str], max_new_tokens: int, batch_size: int) -> list[str]:
        rendered = [self.render(prompt) for prompt in prompts]
        if self.backend == "vllm":
            params = self.sampling_cls(temperature=0.0, max_tokens=max_new_tokens)
            outputs = self.model.generate(rendered, params)
            return [strip_thinking(item.outputs[0].text) for item in outputs]

        results: list[str] = []
        for start in range(0, len(rendered), batch_size):
            batch = rendered[start : start + batch_size]
            encoded = self.tokenizer(
                batch,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=4096,
            ).to(self.model.device)
            prompt_width = encoded["input_ids"].shape[1]
            with torch.inference_mode():
                output = self.model.generate(
                    **encoded,
                    max_new_tokens=max_new_tokens,
                    do_sample=False,
                )
            results.extend(
                strip_thinking(text)
                for text in self.tokenizer.batch_decode(
                    output[:, prompt_width:], skip_special_tokens=True
                )
            )
        return results


def load_flores(direction: str, limit: int, seed: int) -> list[dict[str, str]]:
    langs = {"uz-en": ("uzn_Latn", "eng_Latn"), "en-uz": ("eng_Latn", "uzn_Latn")}
    src_lang, ref_lang = langs[direction]
    src = concatenate_datasets([
        load_dataset("openlanguagedata/flores_plus", src_lang, split="dev"),
        load_dataset("openlanguagedata/flores_plus", src_lang, split="devtest"),
    ])
    ref = concatenate_datasets([
        load_dataset("openlanguagedata/flores_plus", ref_lang, split="dev"),
        load_dataset("openlanguagedata/flores_plus", ref_lang, split="devtest"),
    ])
    pairs = [
        {"src": src[index]["text"].strip(), "ref": ref[index]["text"].strip()}
        for index in range(min(len(src), len(ref)))
    ]
    if limit and len(pairs) > limit:
        import random

        random.Random(seed).shuffle(pairs)
        pairs = pairs[:limit]
    return pairs


def score_comet(sources: list[str], hypotheses: list[str], references: list[str]) -> float:
    from comet import download_model, load_from_checkpoint

    checkpoint = download_model("Unbabel/wmt22-comet-da")
    model = load_from_checkpoint(checkpoint)
    rows = [
        {"src": src, "mt": hypothesis, "ref": reference}
        for src, hypothesis, reference in zip(sources, hypotheses, references, strict=True)
    ]
    return float(model.predict(rows, batch_size=8, gpus=1 if torch.cuda.is_available() else 0).system_score)


def evaluate_flores(generator: Generator, args: argparse.Namespace) -> dict:
    import sacrebleu

    results = {}
    for direction in ("uz-en", "en-uz"):
        pairs = load_flores(direction, args.limit, args.seed)
        prompts = [TRANSLATION_PROMPTS[direction].format(text=row["src"]) for row in pairs]
        hypotheses = generator.generate(prompts, max_new_tokens=160, batch_size=args.batch_size)
        references = [row["ref"] for row in pairs]
        sources = [row["src"] for row in pairs]
        row = {
            "total": len(pairs),
            "bleu": float(sacrebleu.corpus_bleu(hypotheses, [references]).score),
            "samples": [
                {"source": src, "prediction": hyp, "reference": ref}
                for src, hyp, ref in zip(sources[:3], hypotheses[:3], references[:3])
            ],
        }
        if args.comet:
            row["comet"] = score_comet(sources, hypotheses, references)
        results[direction] = row
        print(f"FLORES+ {direction}: BLEU={row['bleu']:.2f}" + (f", COMET={row['comet']:.4f}" if args.comet else ""))
    return results


def label_to_int(raw, names: list[str]) -> int | None:
    if isinstance(raw, int) and 0 <= raw < len(names):
        return raw
    cleaned = str(raw).strip().casefold().replace("’", "'")
    for index, name in enumerate(names):
        if cleaned == name.casefold():
            return index
    return None


def evaluate_classification(generator: Generator, args: argparse.Namespace, task: str) -> dict:
    if task == "sentiment":
        dataset = load_dataset("behbudiy/uzbek-sentiment-analysis", split="train")
        names = ["Negative", "Positive"]
        rows = [
            {"text": row["text"], "gold": label_to_int(row["label"], names)}
            for row in choose_rows(dataset, args.limit, args.seed)
        ]
        prompt_template = SENTIMENT_PROMPT
        parser = lambda text: 1 if text.casefold().startswith("positive") else (0 if text.casefold().startswith("negative") else None)
    else:
        dataset = load_dataset("risqaliyevds/uzbek-zero-shot-classification", split="train")
        names = ["Siyosat", "Iqtisodiyot", "Texnologiya", "Sport", "Madaniyat",
                 "Salomatlik", "Oila va Jamiyat", "Ta'lim", "Ekologiya", "Xorijiy Yangiliklar"]
        rows = [
            {"text": row["text"], "gold": label_to_int(row["class"], names)}
            for row in choose_rows(dataset, args.limit, args.seed)
        ]
        prompt_template = NEWS_PROMPT
        parser = lambda text: int(match.group()) if (match := re.search(r"\d", text)) else None

    rows = [row for row in rows if row["gold"] is not None]
    prompts = [prompt_template.format(text=row["text"][: args.max_text_chars]) for row in rows]
    outputs = generator.generate(prompts, max_new_tokens=8, batch_size=args.batch_size)
    predictions = [parser(output.strip()) for output in outputs]
    correct = sum(prediction == row["gold"] for prediction, row in zip(predictions, rows, strict=True))
    invalid = sum(prediction is None for prediction in predictions)
    result = {
        "accuracy": correct / len(rows),
        "correct": correct,
        "total": len(rows),
        "invalid_rate": invalid / len(rows),
    }
    print(f"{task}: accuracy={result['accuracy']:.2%} ({correct}/{len(rows)}), invalid={invalid}")
    return result


def answer_letter(raw) -> str | None:
    if isinstance(raw, int) and 0 <= raw < 4:
        return LETTERS[raw]
    cleaned = str(raw).strip().upper()
    return cleaned[0] if cleaned and cleaned[0] in LETTERS else None


def evaluate_mcq(generator: Generator, args: argparse.Namespace, task: str) -> dict:
    dataset_name, config, split, language = MCQ_TASKS[task]
    dataset = choose_rows(load_dataset(dataset_name, config, split=split), args.limit, args.seed)
    rows = []
    for row in dataset:
        choices = row.get("choices") or [row.get(f"option_{letter.lower()}") for letter in LETTERS]
        choices = [str(choice) for choice in choices if choice is not None]
        gold = answer_letter(row.get("answer"))
        if row.get("question") and len(choices) >= 4 and gold:
            rows.append({"question": row["question"], "choices": choices[:4], "gold": gold})
    prompts = []
    for row in rows:
        choices = "\n".join(f"{letter}) {choice}" for letter, choice in zip(LETTERS, row["choices"]))
        prompts.append(f"{row['question']}\n\n{choices}\n\n{MCQ_SUFFIX[language]}")
    outputs = generator.generate(prompts, max_new_tokens=12, batch_size=args.batch_size)
    predictions = []
    for output in outputs:
        match = re.search(r"[ABCD]", output.upper())
        predictions.append(match.group() if match else None)
    correct = sum(prediction == row["gold"] for prediction, row in zip(predictions, rows, strict=True))
    invalid = sum(prediction is None for prediction in predictions)
    result = {
        "accuracy": correct / len(rows),
        "correct": correct,
        "total": len(rows),
        "invalid_rate": invalid / len(rows),
    }
    print(f"{task}: accuracy={result['accuracy']:.2%} ({correct}/{len(rows)}), invalid={invalid}")
    return result


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default=MODEL_ID)
    parser.add_argument("--tasks", default="flores,sentiment,news,mmlu-en,mmlu-uz,tumlu")
    parser.add_argument("--backend", choices=["vllm", "transformers"], default="vllm")
    parser.add_argument("--limit", type=int, default=200, help="Examples per dataset; 0 means full dataset.")
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--max-text-chars", type=int, default=4000)
    parser.add_argument("--max-model-len", type=int, default=4096)
    parser.add_argument("--gpu-memory-utilization", type=float, default=0.85)
    parser.add_argument("--dtype", choices=["bfloat16", "float16"], default="bfloat16")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--comet", action="store_true", help="Download WMT22-COMET-DA and score FLORES+.")
    parser.add_argument("--output", type=Path, default=Path("neuronai-2b-benchmark.json"))
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.limit < 0:
        raise ValueError("--limit must be 0 or greater")
    tasks = {task.strip() for task in args.tasks.split(",") if task.strip()}
    unknown = tasks - {"flores", "sentiment", "news", *MCQ_TASKS}
    if unknown:
        raise ValueError(f"Unknown tasks: {sorted(unknown)}")
    generator = Generator(args)
    results = {
        "model": args.model,
        "timestamp": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "limit_per_dataset": args.limit,
        "seed": args.seed,
        "backend": args.backend,
        "results": {},
    }
    if "flores" in tasks:
        results["results"]["flores"] = evaluate_flores(generator, args)
    for task in ("sentiment", "news"):
        if task in tasks:
            results["results"][task] = evaluate_classification(generator, args, task)
    for task in MCQ_TASKS:
        if task in tasks:
            results["results"][task] = evaluate_mcq(generator, args, task)
    args.output.write_text(json.dumps(results, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()