"""Step 2 / Step 4 benchmark suite: lm_eval on the 3 required tasks + GSM8K
testing split, matching the rubric's "full datasets without setting a limit"
requirement (Check-in 4 Step 2). Same pattern as the class's sample_script.py,
pointed at Qwen2.5-7B-Instruct with 4-bit quantization (needed to fit the T4),
and able to load a LoRA adapter for post-training runs.

Usage:
  python3 benchmark_suite.py --tag smoke --limit 20        # dry run, sanity check
  python3 benchmark_suite.py --tag pre_base                # full, base model (Step 2)
  python3 benchmark_suite.py --tag post_B --adapter ../training/runs/B/best_model  # Step 4

Writes results/<tag>_results.json and results/<tag>_samples.json.
"""
import argparse
import json
import os
import time

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
import lm_eval
import lm_eval.tasks

MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"  # default; override with --model for comparison models
TASKS = ["gsm8k", "logiqa2", "arc_challenge", "mmlu"]
# logiqa2, not logiqa: the installed lm_eval's "logiqa" task uses a legacy HF
# "dataset script" loader, dropped entirely by datasets>=4 (RuntimeError:
# "Dataset scripts are no longer supported"). logiqa2 (baber/logiqa2 on the
# Hub) is the same task family via a proper Hub dataset -- same multiple-choice
# / loglikelihood scoring as arc_challenge and mmlu, so no cost implication.
RESULTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results")


def load_model(adapter_path=None, model_id=MODEL_ID):
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    bnb = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,  # T4 has no bf16
        bnb_4bit_use_double_quant=True,
    )
    model = AutoModelForCausalLM.from_pretrained(
        model_id, quantization_config=bnb, device_map="auto")
    if adapter_path:
        model = PeftModel.from_pretrained(model, adapter_path)
    model.eval()
    return model, tokenizer


def run(tag, adapter_path, limit, tasks=None, model_id=MODEL_ID):
    t0 = time.time()
    model, tokenizer = load_model(adapter_path, model_id)
    print(f"model {model_id} ready in {time.time()-t0:.1f}s")

    task_manager = lm_eval.tasks.TaskManager()
    t0 = time.time()
    results = lm_eval.simple_evaluate(
        model="hf",
        model_args={"pretrained": model, "dtype": "float16", "tokenizer": tokenizer},
        tasks=tasks or TASKS,
        task_manager=task_manager,
        log_samples=True,
        batch_size="auto:4",
        limit=limit,
    )
    print(f"eval done in {time.time()-t0:.1f}s")

    os.makedirs(RESULTS_DIR, exist_ok=True)
    with open(os.path.join(RESULTS_DIR, f"{tag}_results.json"), "w") as f:
        json.dump(results["results"], f, indent=1)
    with open(os.path.join(RESULTS_DIR, f"{tag}_samples.json"), "w") as f:
        json.dump(results["samples"], f, indent=1)
    print(json.dumps(results["results"], indent=1))
    print(f"wrote results/{tag}_results.json and results/{tag}_samples.json")


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--tag", required=True)
    p.add_argument("--adapter", default=None,
                    help="path to a LoRA adapter dir, e.g. ../training/runs/B/best_model")
    p.add_argument("--limit", type=int, default=None,
                    help="cap examples per task; omit for the required full-dataset run")
    p.add_argument("--tasks", default=None,
                    help="comma-separated task override, e.g. gsm8k or mmlu; default is all 4 required tasks")
    p.add_argument("--model", default=MODEL_ID,
                    help="HF model id to evaluate; override for Deliverable 5 comparison models")
    args = p.parse_args()
    task_list = args.tasks.split(",") if args.tasks else None
    run(args.tag, args.adapter, args.limit, task_list, args.model)