#!/usr/bin/env python3
"""Run MMLU against a loaded baseline, atomic-pruned, or repacked grouped model."""

from __future__ import annotations

import argparse
import json
import sys
from contextlib import nullcontext
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import numpy as np

from heapr.grouped_model import RepackedExpandedGroupedLagunaContext
from heapr.model_utils import (
    build_max_memory,
    load_causal_lm,
    load_tokenizer,
    validate_model_device_placement,
)
from heapr.prune import apply_atomic_mask_to_model, apply_group_mask_to_model, atomic_mask_from_scores
from heapr.utils import require_torch


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-id", required=True)
    parser.add_argument("--output-path", required=True)
    parser.add_argument(
        "--mode",
        choices=["baseline", "atomic", "native-group", "repacked-group"],
        default="baseline",
    )
    parser.add_argument("--scores-path")
    parser.add_argument("--ratio", type=float)
    parser.add_argument("--group-mask")
    parser.add_argument("--group-indices")
    parser.add_argument(
        "--child-budget-mode",
        choices=["backfill", "layer-scaled", "parent-weighted"],
        default="backfill",
    )
    parser.add_argument("--tasks", default="mmlu")
    parser.add_argument("--limit", type=float)
    parser.add_argument("--num-fewshot", type=int)
    parser.add_argument("--max-gen-toks", type=int)
    parser.add_argument("--batch-size", default="1")
    parser.add_argument("--revision")
    parser.add_argument("--dtype", default="bfloat16")
    parser.add_argument("--gpu-memory-per-device")
    parser.add_argument("--max-gpu-memory")
    parser.add_argument("--max-cpu-memory")
    parser.add_argument("--offload-folder")
    parser.add_argument("--allow-cpu-offload", action="store_true")
    return parser.parse_args()


def _batch_size(value: str):
    return int(value) if value.isdigit() else value


def _json_default(value):
    if isinstance(value, np.generic):
        return value.item()
    if isinstance(value, np.ndarray):
        return value.tolist()
    return str(value)


def _generation_kwargs(max_gen_toks: int | None):
    if max_gen_toks is None:
        return None
    return {"max_gen_toks": max_gen_toks}


def main() -> None:
    args = parse_args()
    if args.offload_folder and not args.allow_cpu_offload:
        raise ValueError("--offload-folder requires --allow-cpu-offload")
    if args.mode == "atomic" and (args.scores_path is None or args.ratio is None):
        raise ValueError("--scores-path and --ratio are required for atomic mode")
    if args.mode in {"native-group", "repacked-group"} and (
        args.group_mask is None or args.group_indices is None
    ):
        raise ValueError("--group-mask and --group-indices are required for grouped modes")

    max_memory = build_max_memory(
        gpu_memory_per_device=args.gpu_memory_per_device,
        max_gpu_memory=args.max_gpu_memory,
        max_cpu_memory=args.max_cpu_memory,
        allow_cpu_offload=args.allow_cpu_offload,
    )
    torch = require_torch()
    requested_gpu_count = torch.cuda.device_count() if args.gpu_memory_per_device else None
    model = load_causal_lm(
        args.model_id,
        revision=args.revision,
        dtype=args.dtype,
        max_memory=max_memory,
        offload_folder=args.offload_folder if args.allow_cpu_offload else None,
        use_cache=False,
    )
    if args.max_gen_toks is not None:
        # Laguna ships max_new_tokens=4096 in its generation config. Align it
        # with the lm-eval limit so it does not override the harness max_length.
        model.generation_config.max_new_tokens = args.max_gen_toks
    validate_model_device_placement(
        model,
        allow_cpu_offload=args.allow_cpu_offload,
        requested_gpu_count=requested_gpu_count,
    )
    tokenizer = load_tokenizer(args.model_id, revision=args.revision)

    context = nullcontext(model)
    pruning_summary = {"mode": args.mode}
    if args.mode == "atomic":
        keep_mask = atomic_mask_from_scores(np.load(args.scores_path), args.ratio)
        apply_atomic_mask_to_model(model, keep_mask)
        pruning_summary.update(
            ratio=args.ratio,
            kept=int(keep_mask.sum()),
            pruned=int((~keep_mask).sum()),
        )
    elif args.mode == "native-group":
        keep_mask = np.load(args.group_mask)
        group_indices = np.load(args.group_indices)
        apply_group_mask_to_model(model, keep_mask, group_indices=group_indices)
        pruning_summary.update(
            kept_groups=int(keep_mask.sum()),
            pruned_groups=int((~keep_mask).sum()),
        )
    elif args.mode == "repacked-group":
        keep_mask = np.load(args.group_mask)
        group_indices = np.load(args.group_indices)
        context = RepackedExpandedGroupedLagunaContext(
            model,
            keep_group_masks=keep_mask,
            group_indices=group_indices,
            child_budget_mode=args.child_budget_mode,
        )
        pruning_summary.update(
            kept_groups=int(keep_mask.sum()),
            pruned_groups=int((~keep_mask).sum()),
            child_budget_mode=args.child_budget_mode,
        )

    from lm_eval import evaluator
    from lm_eval.models.huggingface import HFLM

    with context:
        harness_model = HFLM(
            pretrained=model,
            tokenizer=tokenizer,
            backend="causal",
            batch_size=_batch_size(args.batch_size),
        )
        results = evaluator.simple_evaluate(
            model=harness_model,
            tasks=[task.strip() for task in args.tasks.split(",") if task.strip()],
            num_fewshot=args.num_fewshot,
            batch_size=_batch_size(args.batch_size),
            limit=args.limit,
            gen_kwargs=_generation_kwargs(args.max_gen_toks),
        )

    payload = {
        "model_id": args.model_id,
        "pruning": pruning_summary,
        "tasks": args.tasks,
        "limit": args.limit,
        "num_fewshot": args.num_fewshot,
        "max_gen_toks": args.max_gen_toks,
        "batch_size": args.batch_size,
        "results": results,
    }
    output_path = Path(args.output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(
        json.dumps(payload, indent=2, sort_keys=True, default=_json_default) + "\n",
        encoding="utf-8",
    )
    print(output_path)


if __name__ == "__main__":
    main()