# Automated EvalPlus runner for HumanEval and MBPP benchmarks.
# Using the vLLM backend in greedy mode.

import os
import sys
import subprocess
import time
import json
import re
from pathlib import Path
from datetime import datetime
from huggingface_hub import snapshot_download

MODELS = [
    {
        "name": "Quintus-1.7B",
        "id": "iamrahulreddy/Quintus",
        "is_local": False
    },
    {
        "name": "Qwen3-1.7B-Instruct",
        "id": "Qwen/Qwen3-1.7B",
        "is_local": False
    },
    {
        "name": "Qwen3-1.7B-Base",
        "id": "Qwen/Qwen3-1.7B-Base",
        "is_local": False
    }
]

DATASETS = [
    "humaneval", "mbpp",          # EvalPlus benchmarks
    "gsm8k", "winogrande",        # lm-eval fast benchmarks
    "arc_challenge", "boolq", "piqa"
]
EVALPLUS_DATASETS = {"humaneval", "mbpp"}

LM_EVAL_SHOTS = {
    "gsm8k": "10",
    "winogrande": "5",
    "arc_challenge": "25",
    "boolq": "0",
    "piqa": "0"
}

HF_TOKEN = os.environ.get("HF_TOKEN")
TRUST_REMOTE_CODE = os.environ.get("QUINTUS_TRUST_REMOTE_CODE", "").strip().lower() in {"1", "true", "yes", "on"}

def extract_lm_eval_score(results_dir: Path, task: str) -> str:
    """Finds and extracts the primary score from JSON files outputted by lm-evaluation-harness."""
    for json_path in sorted(results_dir.rglob("*.json"), reverse=True):
        try:
            with open(json_path, encoding="utf-8") as fh:
                data = json.load(fh)
            task_results = data.get("results", {})
            for candidate in (task, f"leaderboard_{task}"):
                if candidate in task_results:
                    task_data = task_results[candidate]
                    # Try common metric names
                    for metric in ["acc,none", "acc_norm,none", "exact_match,strict-match", "exact_match,none"]:
                        if metric in task_data:
                            return f"{task_data[metric]*100:.1f}"
        except Exception:
            continue
    return "N/A"

def is_noise(line: str) -> bool:
    l = line.strip()
    if not l:
        return False
    # Progress bar indicators & block characters
    if any(c in l for c in ["█", "━", "╸", "•", "━━━━━━━━"]):
        return True
    # vLLM, ray, flash_attn, huggingface setup/warnings logs
    noise_keywords = [
        "INFO ", "WARNING ", "DEBUG ", "ERROR ", "(EngineCore",
        "Loading safetensors", "Capturing CUDA graphs",
        "Codegen:", "Downloading dataset", "downloading dataset",
        "Initializing a decoder", "Unknown vLLM environment",
        "world_size=", "Using V2 Model Runner", "Model loading took",
        "Using FLASH_ATTN", "Using FlashAttention", "Kernel JIT monitor",
        "autotuner.py", "autotuning", "Autotuning", "loading weights",
        "Loading weights", "Failed to get device capability", "Sanitized code outputs",
        "Raw outputs will be saved", "init engine", "Dynamo bytecode",
        "Directly load the compiled graph", "Directly load AOT compilation", "torch.compile took"
    ]
    if any(k.lower() in l.lower() for k in noise_keywords):
        return True
    # TQDM lines (e.g. 100%|... [00:17<00:00, 9.45it/s])
    if "%|" in l and ("it/s" in l or "s/it" in l):
        return True
    return False

def main():
    print("=" * 80)
    print("              EVALPLUS BENCHMARK RUNNER (HUMANEVAL & MBPP)")
    print("=" * 80)
    print(f"Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"Models to evaluate: {[m['name'] for m in MODELS]}")
    print(f"Datasets: {DATASETS}")
    print("=" * 80)

    # Set optional HF token and runtime configuration.
    if HF_TOKEN:
        os.environ["HF_TOKEN"] = HF_TOKEN
    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
    os.environ["VLLM_MAX_MODEL_LEN"] = "4096"

    # Step 1: Pre-download and prepare model caches
    print("\n--- STAGE 1: WARMING UP MODEL WEIGHTS CACHE ---")

    # Cache all models
    for model in MODELS:
        if model["is_local"]:
            continue
        print(f"\n[DOWNLOADING] Fetching cache for {model['name']} ({model['id']})...")
        try:
            snapshot_download(
                repo_id=model["id"],
                token=HF_TOKEN or None
            )
            print(f"[DOWNLOAD SUCCESS] {model['name']} is cached and ready.")
        except Exception as e:
            print(f"[DOWNLOAD WARNING] Could not pre-download model {model['name']} via snapshot_download: {e}")
            print("The evaluation run will attempt to download it directly during execution.")

    print("\n--- STAGE 2: SEQUENTIAL EVALPLUS EVALUATION ---")
    results = []

    # Run evaluations sequentially
    for model in MODELS:
        # Resolve path
        model_path = str(Path(model["id"]).resolve()) if model["is_local"] else model["id"]

        for dataset in DATASETS:
            print(f"\n[STARTING] Evaluating {model['name']} on {dataset}...")
            print("-" * 60)

            if dataset in EVALPLUS_DATASETS:
                cmd = [
                    sys.executable, "-m", "evalplus.evaluate",
                    "--model", model_path,
                    "--dataset", dataset,
                    "--backend", "vllm",
                    "--greedy"
                ]
            else:
                shots = LM_EVAL_SHOTS.get(dataset, "0")
                out_dir = Path("eval_results") / model["name"] / dataset
                out_dir.mkdir(parents=True, exist_ok=True)

                model_args = (
                    f"pretrained={model_path},dtype=bfloat16,"
                    f"trust_remote_code={str(TRUST_REMOTE_CODE).lower()},"
                    "gpu_memory_utilization=0.9,max_model_len=4096"
                )
                cmd = [
                    sys.executable, "-m", "lm_eval",
                    "--model", "vllm",
                    "--model_args", model_args,
                    "--tasks", dataset,
                    "--num_fewshot", shots,
                    "--batch_size", "auto",
                    "--output_path", str(out_dir),
                    "--log_samples"
                ]
                if dataset == "gsm8k":
                    cmd.extend(["--gen_kwargs", "max_gen_toks=512"])

            print(f"Running command: {' '.join(cmd)}")
            start_time = time.time()

            try:
                # Run the command and stream output
                process = subprocess.Popen(
                    cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1
                )

                # Stream and capture output (filtering out vLLM and progress bar noise)
                stdout_text = ""
                for line in process.stdout:
                    stdout_text += line
                    if not is_noise(line):
                        print(line, end="")

                process.wait()
                duration = time.time() - start_time
                time.sleep(5)  # Let OS/driver fully reclaim GPU VRAM before starting next subprocess

                score_str = "N/A"
                if process.returncode == 0:
                    print(f"[SUCCESS] Completed {model['name']} on {dataset} in {duration:.1f} seconds.")

                    # Parse scores
                    if dataset in EVALPLUS_DATASETS:
                        # Find all pass@1 scores
                        matches = re.findall(r"pass@1:\s+([0-9.]+)", stdout_text)
                        if len(matches) >= 2:
                            val0 = float(matches[0])
                            val1 = float(matches[1])
                            if val0 <= 1.0: val0 *= 100
                            if val1 <= 1.0: val1 *= 100
                            score_str = f"Base: {val0:.1f} | Plus: {val1:.1f}"
                        elif len(matches) == 1:
                            val0 = float(matches[0])
                            if val0 <= 1.0: val0 *= 100
                            score_str = f"Base: {val0:.1f}"
                    else:
                        score_str = extract_lm_eval_score(out_dir, dataset)

                    results.append({
                        "model": model["name"],
                        "dataset": dataset,
                        "status": "Success",
                        "score": score_str,
                        "duration": f"{duration/60:.1f} min"
                    })
                else:
                    print(f"[ERROR] command failed with exit code {process.returncode}")
                    results.append({
                        "model": model["name"],
                        "dataset": dataset,
                        "status": f"Failed ({process.returncode})",
                        "score": "ERROR",
                        "duration": f"{duration/60:.1f} min"
                    })

            except Exception as e:
                duration = time.time() - start_time
                print(f"[ERROR] Failed to run benchmark: {e}")
                results.append({
                    "model": model["name"],
                    "dataset": dataset,
                    "status": f"Error",
                    "score": "ERROR",
                    "duration": f"{duration/60:.1f} min"
                })
            print("-" * 60)

    # Print and save summary report
    report_lines = []
    report_lines.append("\n" + "=" * 100)
    report_lines.append("                       BENCHMARK RUN SUMMARY")
    report_lines.append("=" * 100)
    report_lines.append(f"| {'Model':<30} | {'Dataset':<15} | {'Score':<25} | {'Status':<10} | {'Time':<8} |")
    report_lines.append(f"|{'-'*32}|{'-'*17}|{'-'*27}|{'-'*12}|{'-'*10}|")
    for r in results:
        report_lines.append(f"| {r['model']:<30} | {r['dataset']:<15} | {r['score']:<25} | {r['status']:<10} | {r['duration']:<8} |")
    report_lines.append("=" * 100)

    report_text = "\n".join(report_lines)
    print(report_text)
    print("\nNote: Results are saved in the default EvalPlus directory and eval_results/.")

    # Save to file
    with open("qwen_quintus_scores.txt", "w", encoding="utf-8") as f:
        f.write(report_text + "\n")
    print("\n[SUCCESS] Final score report saved to 'qwen_quintus_scores.txt'")

if __name__ == "__main__":
    main()