# Prompt: HumanEval/MBPP eval wrapper for dVLM models
"""
HumanEval and MBPP evaluation for dVLM models (text-only code generation).

Tests whether VLM fine-tuning degrades code generation ability by running
standard HumanEval (164 problems) and MBPP (500 sanitized problems) benchmarks.

The model is prompted in text-only mode (no images).

Usage:
    # HumanEval
    python3 eval_humaneval_mbpp.py \
        --model_path output_models/checkpoint \
        --task humaneval \
        --n_samples 1

    # MBPP
    python3 eval_humaneval_mbpp.py \
        --model_path output_models/checkpoint \
        --task mbpp \
        --n_samples 1

    # Evaluate existing generations
    python3 eval_humaneval_mbpp.py \
        --task humaneval \
        --eval_only \
        --generations_path output_dir/humaneval_generations.jsonl
"""

import argparse
import json
import os
import sys
from pathlib import Path

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
DVLM_DIR = os.path.dirname(os.path.dirname(SCRIPT_DIR))
MODELS_DIR = os.path.join(DVLM_DIR, "models")
PROJECT_ROOT = os.path.dirname(os.path.dirname(DVLM_DIR))

if MODELS_DIR not in sys.path:
    sys.path.insert(0, MODELS_DIR)

# HumanEval data from openai/human-eval package
HUMANEVAL_PATH = None  # Will be resolved at runtime
MBPP_PATH = None


def load_humaneval_problems():
    """Load HumanEval problems from the human_eval package."""
    try:
        from human_eval.data import read_problems
        return read_problems()
    except ImportError:
        # Fallback: load from HuggingFace datasets
        from datasets import load_dataset
        ds = load_dataset("openai_humaneval", split="test")
        problems = {}
        for item in ds:
            problems[item["task_id"]] = {
                "task_id": item["task_id"],
                "prompt": item["prompt"],
                "entry_point": item["entry_point"],
                "canonical_solution": item["canonical_solution"],
                "test": item["test"],
            }
        return problems


def load_mbpp_problems():
    """Load MBPP sanitized problems."""
    try:
        from datasets import load_dataset
        ds = load_dataset("mbpp", "sanitized", split="test")
        problems = {}
        for item in ds:
            task_id = f"Mbpp/{item['task_id']}"
            problems[task_id] = {
                "task_id": task_id,
                "prompt": item["prompt"],
                "code": item["code"],
                "test_list": item["test_list"],
                "test_setup_code": item.get("test_setup_code", ""),
            }
        return problems
    except Exception as e:
        print(f"ERROR loading MBPP: {e}")
        sys.exit(1)


def load_model(args):
    """Load the model for text-only inference."""
    from qwen_vl_mdm_model import QwenVLMDMModel

    model = QwenVLMDMModel(
        model_name=args.model_path,
        processor_name=args.processor_path,
        device="cuda",
        block_size=args.block_size,
        max_tokens=args.max_tokens,
        mask_token="|<MASK>|",
        use_speculative=False,
        force_left_to_right=True,
    )
    return model


def generate_completion(model, prompt: str) -> str:
    """Generate a code completion (text-only, no image)."""
    result = model.generate(prompt=prompt, image_path=None)
    text = result.get("text", "") if isinstance(result, dict) else str(result)
    return text


def run_humaneval(args):
    """Generate completions for HumanEval problems."""
    problems = load_humaneval_problems()
    print(f"Loaded {len(problems)} HumanEval problems")

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    gen_path = output_dir / "humaneval_generations.jsonl"

    model = load_model(args)
    print("Model loaded. Generating completions...")

    samples = []
    for i, (task_id, problem) in enumerate(sorted(problems.items())):
        print(f"[{i+1}/{len(problems)}] {task_id}")
        for _ in range(args.n_samples):
            try:
                completion = generate_completion(model, problem["prompt"])
            except Exception as e:
                print(f"  ERROR: {e}")
                completion = ""
            samples.append({"task_id": task_id, "completion": completion})

    # Save generations
    with open(gen_path, "w") as f:
        for s in samples:
            f.write(json.dumps(s) + "\n")
    print(f"\nSaved {len(samples)} generations to {gen_path}")
    return str(gen_path)


def run_mbpp(args):
    """Generate solutions for MBPP problems."""
    problems = load_mbpp_problems()
    print(f"Loaded {len(problems)} MBPP problems")

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    gen_path = output_dir / "mbpp_generations.jsonl"

    model = load_model(args)
    print("Model loaded. Generating solutions...")

    MBPP_TEMPLATE = (
        "You are an expert Python programmer. Write a Python function to solve "
        "the following problem.\n\n{prompt}\n\nYour code should pass these tests:\n"
        "{tests}\n\nWrite only the Python function, nothing else."
    )

    samples = []
    for i, (task_id, problem) in enumerate(sorted(problems.items())):
        print(f"[{i+1}/{len(problems)}] {task_id}")
        test_str = "\n".join(problem.get("test_list", []))
        formatted_prompt = MBPP_TEMPLATE.format(
            prompt=problem["prompt"], tests=test_str
        )
        for _ in range(args.n_samples):
            try:
                completion = generate_completion(model, formatted_prompt)
            except Exception as e:
                print(f"  ERROR: {e}")
                completion = ""
            samples.append({"task_id": task_id, "completion": completion})

    with open(gen_path, "w") as f:
        for s in samples:
            f.write(json.dumps(s) + "\n")
    print(f"\nSaved {len(samples)} generations to {gen_path}")
    return str(gen_path)


def run_evaluation(args, gen_path: str):
    """Evaluate generated code using execution-based pass@k."""
    output_dir = Path(args.output_dir)

    if args.task == "humaneval":
        try:
            from human_eval.evaluation import evaluate_functional_correctness
            results = evaluate_functional_correctness(
                gen_path, k=[1], n_workers=4
            )
            print("\n" + "=" * 50)
            print("HumanEval Results")
            print("=" * 50)
            for k, v in results.items():
                print(f"  {k}: {v:.4f}")
            print("=" * 50)

            results_path = output_dir / "humaneval_results.json"
            with open(results_path, "w") as f:
                json.dump(results, f, indent=2)
            print(f"Results saved to {results_path}")

        except ImportError:
            print(
                "WARNING: human_eval package not installed. "
                "Install with: pip install human-eval\n"
                "Then run: evaluate_functional_correctness "
                f"{gen_path}\n"
                "NOTE: You must uncomment the execution code in "
                "human_eval/execution.py first."
            )

    elif args.task == "mbpp":
        # MBPP evaluation: execute test cases
        print("\nMBPP evaluation requires code execution.")
        print(f"Generations saved to: {gen_path}")
        print(
            "To evaluate, use bigcode-evaluation-harness:\n"
            f"  cd {PROJECT_ROOT}/third_party/bigcode-eval\n"
            f"  python3 main.py --tasks mbpp --allow_code_execution "
            f"--load_generations_path {gen_path} --model dummy"
        )


def main():
    parser = argparse.ArgumentParser(
        description="HumanEval/MBPP evaluation for dVLM models"
    )
    parser.add_argument("--model_path", default=None, help="Path to model checkpoint")
    parser.add_argument(
        "--processor_path",
        default="Qwen/Qwen2.5-VL-3B-Instruct",
        help="Path to processor/tokenizer",
    )
    parser.add_argument(
        "--task",
        choices=["humaneval", "mbpp"],
        required=True,
        help="Which benchmark to run",
    )
    parser.add_argument(
        "--output_dir", default=None, help="Output directory"
    )
    parser.add_argument(
        "--n_samples", type=int, default=1, help="Number of samples per problem"
    )
    parser.add_argument("--block_size", type=int, default=None, help="MDM block_size")
    parser.add_argument(
        "--max_tokens", type=int, default=512, help="Max generation tokens"
    )
    parser.add_argument(
        "--eval_only",
        action="store_true",
        help="Skip generation, only evaluate existing results",
    )
    parser.add_argument(
        "--generations_path",
        default=None,
        help="Path to existing generations (for --eval_only)",
    )
    args = parser.parse_args()

    # Resolve defaults
    if args.output_dir is None:
        if args.model_path:
            args.output_dir = os.path.join(args.model_path, f"eval_{args.task}")
        else:
            args.output_dir = f"eval_outputs/{args.task}"

    os.makedirs(args.output_dir, exist_ok=True)

    print("=" * 50)
    print(f"{args.task.upper()} Evaluation")
    print("=" * 50)
    print(f"Model: {args.model_path}")
    print(f"Output: {args.output_dir}")
    print(f"Samples per problem: {args.n_samples}")
    print("=" * 50)

    if args.eval_only:
        gen_path = args.generations_path
        if gen_path is None:
            gen_path = os.path.join(
                args.output_dir, f"{args.task}_generations.jsonl"
            )
        if not os.path.exists(gen_path):
            print(f"ERROR: Generations file not found: {gen_path}")
            sys.exit(1)
    else:
        if args.model_path is None:
            print("ERROR: --model_path required for generation")
            sys.exit(1)
        if args.task == "humaneval":
            gen_path = run_humaneval(args)
        else:
            gen_path = run_mbpp(args)

    run_evaluation(args, gen_path)


if __name__ == "__main__":
    main()
