#!/usr/bin/env python3
"""Reproducible GSM8K and GPQA accuracy runner for Ascend offload comparisons."""

import argparse
import hashlib
import json
import random
import re
import sys
import time
from collections.abc import Mapping
from decimal import Decimal, InvalidOperation
from pathlib import Path


DEFAULT_DATASET_REVISION = "master"
DATASET_SOURCES = {
    "gsm8k": {
        "name": "modelscope/gsm8k",
        "subset": "main",
        "revision": DEFAULT_DATASET_REVISION,
        "split": "test",
    },
    "gpqa": {
        "name": "modelscope/gpqa",
        "subset": "gpqa_diamond",
        "revision": DEFAULT_DATASET_REVISION,
        "split": None,
    },
}
NUMBER_RE = re.compile(r"[+-]?\d[\d,]*(?:\.\d+)?")
EXPLICIT_GPQA_ANSWER_RE = re.compile(
    r"^\s*(?:(?:the\s+)?answer\s*(?:is\s*|:\s*))?[\[(]?([ABCD])[\])]?[.!]?\s*$",
    re.IGNORECASE,
)


def extract_gsm8k_answer(text):
    """Return the final signed numeric answer, preferring GSM8K's #### marker."""
    answer_text = str(text).rsplit("####", 1)[-1]
    matches = NUMBER_RE.findall(answer_text)
    return matches[-1].replace(",", "") if matches else None


def normalize_numeric_answer(value):
    answer = extract_gsm8k_answer(value)
    if answer is None:
        return None
    try:
        decimal = Decimal(answer)
    except InvalidOperation:
        return None
    normalized = format(decimal.normalize(), "f")
    return normalized.rstrip("0").rstrip(".") if "." in normalized else normalized


def score_gsm8k(gold, prediction):
    gold_answer = normalize_numeric_answer(gold)
    prediction_answer = normalize_numeric_answer(prediction)
    return gold_answer is not None and gold_answer == prediction_answer


def normalize_gpqa_row(row):
    """Normalize GPQA's title-case and snake-case data schemas."""
    normalized_keys = {
        key.lower().replace(" ", "_"): value for key, value in row.items()
    }
    question = normalized_keys.get("question")
    correct = normalized_keys.get("correct_answer")
    incorrect = [normalized_keys.get(f"incorrect_answer_{index}") for index in range(1, 4)]
    if not isinstance(question, str) or not isinstance(correct, str) or not all(
        isinstance(answer, str) for answer in incorrect
    ):
        raise ValueError("GPQA row is missing a question or all four labeled answers")
    return {
        "question": question,
        "correct_answer": correct,
        "incorrect_answers": incorrect,
    }


def build_gpqa_choices(row, seed):
    """Build a question-stable A-D mapping without mutating global random state."""
    normalized = normalize_gpqa_row(row)
    choices = [(normalized["correct_answer"], True)] + [
        (answer, False) for answer in normalized["incorrect_answers"]
    ]
    digest = hashlib.sha256(f"{seed}{normalized['question']}".encode("utf-8")).digest()
    random.Random(int.from_bytes(digest, "big")).shuffle(choices)
    correct_index = next(index for index, (_, is_correct) in enumerate(choices) if is_correct)
    return {
        "question": normalized["question"],
        "choices": [answer for answer, _ in choices],
        "correct_index": correct_index,
    }


def parse_gpqa_answer(text):
    match = EXPLICIT_GPQA_ANSWER_RE.fullmatch(str(text))
    return match.group(1).upper() if match else None


def dataset_sha256(path):
    digest = hashlib.sha256()
    with Path(path).open("rb") as file:
        for chunk in iter(lambda: file.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def jsonl_path(data_dir, dataset):
    return Path(data_dir) / f"{dataset}.jsonl"


def manifest_path(data_dir, dataset):
    return Path(data_dir) / f"{dataset}.manifest.json"


def _rows_from_modelscope(dataset):
    return [dict(row) for row in dataset]


def _validated_gpqa_rows(dataset):
    rows = _rows_from_modelscope(dataset)
    if not rows:
        raise ValueError("GPQA split is empty")
    for row in rows:
        normalize_gpqa_row(row)
    return rows


def _load_modelscope_rows(dataset_name, revision=DEFAULT_DATASET_REVISION):
    """Load ModelScope only when local materialized data is unavailable."""
    from modelscope.msdatasets import MsDataset

    source = dict(DATASET_SOURCES[dataset_name])
    source["revision"] = revision
    if dataset_name == "gsm8k":
        dataset = MsDataset.load(
            source["name"],
            subset_name=source["subset"],
            version=source["revision"],
            split=source["split"],
            trust_remote_code=True,
        )
        return _rows_from_modelscope(dataset), source

    # GPQA publishes labeled configurations under gpqa_diamond. Probe split metadata
    # through the dataset object, then accept only the first split with labeled rows.
    metadata = MsDataset.load(
        source["name"],
        subset_name=source["subset"],
        version=source["revision"],
        trust_remote_code=True,
    )
    errors = []
    if isinstance(metadata, Mapping):
        for split, split_dataset in metadata.items():
            try:
                rows = _validated_gpqa_rows(split_dataset)
                source["split"] = str(split)
                return rows, source
            except (TypeError, ValueError) as error:
                errors.append(f"{split}: {error}")
        raise RuntimeError(
            "Unable to discover a labeled GPQA split in mapping: " + "; ".join(errors)
        )

    split_names = getattr(getattr(metadata, "info", None), "splits", {}) or {}
    candidates = list(split_names) or ["test", "train", "validation"]
    for split in candidates:
        try:
            rows = _validated_gpqa_rows(
                MsDataset.load(
                    source["name"],
                    subset_name=source["subset"],
                    version=source["revision"],
                    split=split,
                    trust_remote_code=True,
                )
            )
            source["split"] = split
            return rows, source
        except Exception as error:  # Keep a complete diagnostic if no labeled split exists.
            errors.append(f"{split}: {error}")
    raise RuntimeError("Unable to discover a labeled GPQA split: " + "; ".join(errors))


def _validated_jsonl_row_count(path, dataset):
    row_count = 0
    with Path(path).open(encoding="utf-8") as file:
        for line_number, line in enumerate(file, 1):
            if not line.strip():
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError as error:
                raise RuntimeError(f"Invalid JSONL at {path}:{line_number}") from error
            if not isinstance(row, dict):
                raise RuntimeError(f"Invalid {dataset} row fields at {path}:{line_number}")
            if dataset == "gsm8k":
                if not isinstance(row.get("question"), str) or not isinstance(
                    row.get("answer"), str
                ):
                    raise RuntimeError(
                        f"Invalid GSM8K row fields at {path}:{line_number}"
                    )
            else:
                try:
                    normalize_gpqa_row(row)
                except ValueError as error:
                    raise RuntimeError(
                        f"Invalid GPQA row fields at {path}:{line_number}"
                    ) from error
            row_count += 1
    if row_count == 0:
        raise RuntimeError(f"No examples found in {path}")
    return row_count


def _write_dataset_manifest(metadata_path, dataset, path, source, data_hash, row_count):
    manifest = {
        "dataset": dataset,
        "dataset_path": str(path.resolve()),
        "dataset_sha256": data_hash,
        "dataset_source": source,
        "row_count": row_count,
    }
    metadata_path.write_text(
        json.dumps(manifest, ensure_ascii=False, sort_keys=True, indent=2) + "\n",
        encoding="utf-8",
    )


def materialize_dataset(data_dir, dataset, revision=DEFAULT_DATASET_REVISION):
    path = jsonl_path(data_dir, dataset)
    metadata_path = manifest_path(data_dir, dataset)
    if path.exists():
        if not path.is_file() or path.stat().st_size == 0:
            raise RuntimeError(f"Local dataset is invalid: {path}")
        if not metadata_path.is_file():
            row_count = _validated_jsonl_row_count(path, dataset)
            source = {
                "name": "local",
                "subset": "local",
                "revision": "local",
                "split": "local",
            }
            data_hash = dataset_sha256(path)
            _write_dataset_manifest(
                metadata_path, dataset, path, source, data_hash, row_count
            )
            return path, source, data_hash
        try:
            manifest = json.loads(metadata_path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError) as error:
            raise RuntimeError(f"Invalid dataset manifest: {metadata_path}") from error
        required_keys = {
            "dataset",
            "dataset_path",
            "dataset_source",
            "dataset_sha256",
            "row_count",
        }
        if not required_keys.issubset(manifest) or manifest["dataset"] != dataset:
            raise RuntimeError(f"Dataset manifest does not describe {dataset}: {metadata_path}")
        if manifest["dataset_path"] != str(path.resolve()):
            raise RuntimeError(f"Dataset path does not match manifest: {path}")
        claimed_row_count = manifest["row_count"]
        if type(claimed_row_count) is not int or claimed_row_count <= 0:
            raise RuntimeError(f"Dataset manifest row_count must be a positive integer: {metadata_path}")
        source = manifest["dataset_source"]
        if not isinstance(source, dict) or not {
            "name",
            "subset",
            "revision",
            "split",
        }.issubset(source):
            raise RuntimeError(f"Dataset manifest source metadata is incomplete: {metadata_path}")
        actual_row_count = _validated_jsonl_row_count(path, dataset)
        if claimed_row_count != actual_row_count:
            raise RuntimeError(
                f"Dataset row_count does not match manifest: {claimed_row_count} != {actual_row_count}"
            )
        actual_hash = dataset_sha256(path)
        if manifest["dataset_sha256"] != actual_hash:
            raise RuntimeError(f"Dataset hash does not match manifest: {path}")
        return path, source, actual_hash

    rows, source = _load_modelscope_rows(dataset, revision=revision)
    if not rows:
        raise RuntimeError(f"ModelScope returned no rows for {dataset}")
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="\n") as file:
        for row in rows:
            file.write(json.dumps(row, ensure_ascii=False, sort_keys=True, separators=(",", ":")))
            file.write("\n")
    row_count = _validated_jsonl_row_count(path, dataset)
    data_hash = dataset_sha256(path)
    _write_dataset_manifest(metadata_path, dataset, path, source, data_hash, row_count)
    return path, source, data_hash


def load_local_rows(path, limit):
    rows = []
    with Path(path).open(encoding="utf-8") as file:
        for line_number, line in enumerate(file, 1):
            if not line.strip():
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError as error:
                raise RuntimeError(f"Invalid JSONL at {path}:{line_number}") from error
            if limit and len(rows) >= limit:
                break
    if not rows:
        raise RuntimeError(f"No examples found in {path}")
    return rows


def prompt_for_example(dataset, row, seed):
    if dataset == "gsm8k":
        return str(row["question"]), str(row["answer"])
    choices = build_gpqa_choices(row, seed)
    options = "\n".join(
        f"{chr(65 + index)}. {choice}" for index, choice in enumerate(choices["choices"])
    )
    return f"{choices['question']}\n\n{options}\n\nAnswer:", chr(65 + choices["correct_index"])


def generate_prompts(args, prompts):
    """Generate eagerly, importing the vLLM runtime only for an actual run."""
    from vllm import LLM, SamplingParams

    llm = LLM(
        model=args.model,
        enforce_eager=True,
        max_model_len=args.max_model_len,
        gpu_memory_utilization=args.gpu_memory_utilization,
        offload_group_size=args.offload_group_size,
        offload_num_in_group=args.offload_num_in_group,
    )
    sampling = SamplingParams(temperature=0, max_tokens=args.max_tokens)
    outputs = llm.generate(prompts, sampling)
    return [output.outputs[0].text if output.outputs else "" for output in outputs]


def build_result_metadata(
    args, dataset_source, dataset_sha256, correct, total, elapsed_seconds, output_path
):
    return {
        "model": args.model,
        "dataset": args.dataset,
        "dataset_source": dataset_source,
        "dataset_sha256": dataset_sha256,
        "seed": args.seed,
        "sampling_settings": {
            "temperature": args.temperature,
            "max_tokens": args.max_tokens,
            "enforce_eager": True,
        },
        "engine_settings": {
            "gpu_memory_utilization": args.gpu_memory_utilization,
        },
        "offload_settings": {
            "group_size": args.offload_group_size,
            "num_in_group": args.offload_num_in_group,
        },
        "correct": correct,
        "total": total,
        "accuracy": correct / total if total else 0.0,
        "elapsed_seconds": elapsed_seconds,
        "output_path": str(output_path),
    }


def run_dataset(args, dataset):
    args.dataset = dataset
    data_path, source, data_hash = materialize_dataset(
        args.data_dir, dataset, revision=args.dataset_revision
    )
    rows = load_local_rows(data_path, args.limit)
    prompts_and_gold = [prompt_for_example(dataset, row, args.seed) for row in rows]
    prompts = [prompt for prompt, _ in prompts_and_gold]
    started = time.monotonic()
    predictions = generate_prompts(args, prompts)
    elapsed_seconds = time.monotonic() - started
    if len(predictions) != len(rows):
        raise RuntimeError("Generation returned a different number of predictions")

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    examples_path = output_dir / f"{dataset}-examples.jsonl"
    correct = 0
    with examples_path.open("w", encoding="utf-8", newline="\n") as file:
        for row, (_, gold), prediction in zip(rows, prompts_and_gold, predictions):
            is_correct = score_gsm8k(gold, prediction) if dataset == "gsm8k" else parse_gpqa_answer(prediction) == gold
            correct += is_correct
            file.write(json.dumps({"row": row, "gold": gold, "prediction": prediction, "correct": is_correct}, ensure_ascii=False, sort_keys=True))
            file.write("\n")

    result = build_result_metadata(args, source, data_hash, correct, len(rows), elapsed_seconds, examples_path)
    result_path = output_dir / f"{dataset}-result.json"
    result_path.write_text(json.dumps(result, ensure_ascii=False, sort_keys=True, indent=2) + "\n", encoding="utf-8")
    return result


def parse_args(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True)
    parser.add_argument("--dataset", choices=("gsm8k", "gpqa", "both"), default="both")
    parser.add_argument("--data-dir", type=Path, default=Path("/data/accuracy-data"))
    parser.add_argument("--output-dir", type=Path, default=Path("/data/accuracy-results"))
    parser.add_argument("--limit", type=int, default=0, help="0 means the full dataset")
    parser.add_argument("--seed", default="0")
    parser.add_argument("--dataset-revision", default=DEFAULT_DATASET_REVISION)
    parser.add_argument("--max-tokens", type=int, default=128)
    parser.add_argument("--max-model-len", type=int, default=4096)
    parser.add_argument("--gpu-memory-utilization", type=float, default=0.5)
    parser.add_argument("--offload-group-size", type=int, default=0)
    parser.add_argument("--offload-num-in-group", type=int, default=1)
    args = parser.parse_args(argv)
    args.temperature = 0
    if args.limit < 0:
        parser.error("--limit must be non-negative; 0 means the full dataset")
    if args.max_tokens <= 0 or args.max_model_len <= 0:
        parser.error("--max-tokens and --max-model-len must be positive")
    if not 0 < args.gpu_memory_utilization <= 1:
        parser.error("--gpu-memory-utilization must be in (0, 1]")
    if args.offload_group_size < 0:
        parser.error("--offload-group-size must be non-negative")
    if args.offload_num_in_group <= 0:
        parser.error("--offload-num-in-group must be positive")
    if args.offload_group_size == 0 and args.offload_num_in_group != 1:
        parser.error("group size 0 requires --offload-num-in-group 1")
    if args.offload_group_size and args.offload_num_in_group > args.offload_group_size:
        parser.error("--offload-num-in-group cannot exceed --offload-group-size")
    return args


def main(argv=None):
    args = parse_args(argv)
    datasets = ("gsm8k", "gpqa") if args.dataset == "both" else (args.dataset,)
    results = [run_dataset(args, dataset) for dataset in datasets]
    print(json.dumps(results, ensure_ascii=False, sort_keys=True))


if __name__ == "__main__":
    try:
        main()
    except Exception as error:
        print(f"accuracy reproduction failed: {error}", file=sys.stderr)
        raise
