#!/usr/bin/env python3
"""
DM-GRPO Final Evaluation Script
================================
Unified evaluation with LOCKED system prompt (matches grpo.py training prompt).
DO NOT modify SYSTEM_MESSAGE without updating the training prompt in grpo.py.

Usage:
    python eval.py --model_path /path/to/model --file_name exp1 --benchmarks tempcompass
    python eval.py --model_path /path/to/model --file_name exp1 --benchmarks all --nframes 32
"""

import os
import json
import re
import math
import argparse
from tqdm import tqdm
import torch
from transformers import AutoProcessor, AutoTokenizer
from vllm import LLM, SamplingParams
from qwen_vl_utils import process_vision_info


# ========================== LOCKED System Prompt ==========================
# This MUST match grpo.py SYSTEM_PROMPT exactly. Do NOT modify independently.
# Source: grpo.py L489-L508 (SFT-v9 aligned, 0.1s precision)
# Last verified: 2026-04-22
SYSTEM_MESSAGE = (
    "A conversation between User and Assistant. The user asks a question about "
    "a video, and the Assistant solves it. The assistant first analyzes the video "
    "by establishing reference frames with <time> tags and tracking temporal "
    "changes with <time_range> and <delta> tags, then summarizes the reasoning "
    "in <reasoning_summary> tags. The analysis and answer are enclosed within "
    "<think> </think> and <answer> </answer> tags, respectively, i.e., "
    "<think> analysis here </think><answer> answer here </answer>"
)

QUESTION_TEMPLATE = "{Question}"

TYPE_TEMPLATE = {
    "multiple choice": " Please provide only the single option letter (e.g., A, B, C, D, etc.) within the <answer> </answer> tags.",
    "numerical":       " Please provide the numerical value (e.g., 42 or 3.14) within the <answer> </answer> tags.",
    "OCR":             " Please transcribe text from the image/video clearly and provide your text answer within the <answer> </answer> tags.",
    "free-form":       " Please provide your text answer within the <answer> </answer> tags.",
    "regression":      " Please provide the numerical value (e.g., 42 or 3.14) within the <answer> </answer> tags.",
}


# ========================== Frame Timestamp Hint ==========================
_FRAME_FACTOR = 2
_DEFAULT_FPS = 2.0
_FPS_MIN_FRAMES = 4


def _smart_nframes(total_frames: int, video_fps: float,
                   sample_fps: float = _DEFAULT_FPS,
                   max_frames: int = 16) -> int:
    min_frames = math.ceil(_FPS_MIN_FRAMES / _FRAME_FACTOR) * _FRAME_FACTOR
    max_f = math.floor(min(max_frames, total_frames) / _FRAME_FACTOR) * _FRAME_FACTOR
    nframes = total_frames / video_fps * sample_fps
    nframes = min(min(max(nframes, min_frames), max_f), total_frames)
    nframes = math.floor(nframes / _FRAME_FACTOR) * _FRAME_FACTOR
    nframes = max(nframes, _FRAME_FACTOR)
    nframes = min(nframes, total_frames)
    return int(nframes)


def _get_frame_timestamps(video_path: str, max_frames: int = 16) -> str:
    try:
        import decord
        vr = decord.VideoReader(video_path)
        total_frames = len(vr)
        video_fps = vr.get_avg_fps()
        if total_frames <= 0 or video_fps <= 0:
            return ""
        duration = total_frames / video_fps
        nframes = _smart_nframes(total_frames, video_fps, max_frames=max_frames)
        indices = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
        timestamps = [idx / video_fps for idx in indices]
        sample_fps = nframes / duration if duration > 0 else _DEFAULT_FPS

        def fmt(t):
            h = int(t // 3600)
            m = int((t % 3600) // 60)
            s = t % 60
            return f"{h:02d}:{m:02d}:{s:04.1f}"

        ts_list = ", ".join(fmt(t) for t in timestamps)
        interval = 1.0 / sample_fps if sample_fps > 0 else 0
        return (
            f"\nThe video frames you see have been sampled at approximately "
            f"{sample_fps:.1f} FPS (roughly one frame every {interval:.1f}s). "
            f"The {nframes} frames correspond to the following timestamps: {ts_list}. "
            f"Use only these timestamps when citing specific moments in the video.\n"
        )
    except Exception:
        return ""


# ========================== Benchmark Registry ==========================
BENCHMARK_REGISTRY = {
    "mvbench":              "eval_mvbench.json",
    "tempcompass":          "eval_tempcompass.json",
    "chirality_in_action":  "eval_chirality_in_action.json",
    "ssv2_validation":      "eval_ssv2_validation.json",
    "tempcompass_ts":       "eval_tempcompass_ts.json",
    "vsibench":             "eval_vsibench.json",
    "mmvu":                 "eval_mmvu.json",
    "videomme":             "eval_videomme.json",
    "videommmu":            "eval_videommmu.json",
}

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
EVAL_DATA_DIR = os.path.join(SCRIPT_DIR, "Evaluation")
EVAL_RESULTS_DIR = os.path.join(SCRIPT_DIR, "eval_results")


# ========================== Argument Parsing ==========================
def parse_args():
    parser = argparse.ArgumentParser(
        description="DM-GRPO Final Evaluation",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=f"Available benchmarks: {', '.join(BENCHMARK_REGISTRY.keys())}",
    )
    parser.add_argument("--model_path", type=str, required=True)
    parser.add_argument("--file_name", type=str, required=True)
    parser.add_argument("--benchmarks", type=str, nargs="+", required=True)
    parser.add_argument("--bsz", type=int, default=64)
    parser.add_argument("--max_tokens", type=int, default=4096)
    parser.add_argument("--max_model_len", type=int, default=16384)
    parser.add_argument("--gpu_memory_utilization", type=float, default=0.8)
    parser.add_argument("--temperature", type=float, default=0.1)
    parser.add_argument("--repetition_penalty", type=float, default=1.0)
    parser.add_argument("--top_p", type=float, default=0.001)
    parser.add_argument("--nframes", type=int, default=None)
    return parser.parse_args()


# ========================== Utility Functions ==========================
def extract_think(output_str):
    match = re.search(r'<think>\s*(.*?)\s*</think>', output_str, re.DOTALL)
    return match.group(1).strip() if match else ""


def extract_answer(text):
    match = re.search(r'<answer>\s*(.*?)\s*</answer>', text, re.DOTALL)
    return match.group(1).strip() if match else ""


def normalize_number(num_str):
    try:
        return float(num_str.replace(',', ''))
    except Exception:
        return None


def mean_relative_accuracy(pred, target, start=0.5, end=0.95, interval=0.05):
    if not torch.is_tensor(pred):
        pred = torch.tensor(pred, dtype=torch.float32)
    if not torch.is_tensor(target):
        target = torch.tensor(target, dtype=torch.float32)
    rel_error = torch.abs(pred - target) / (torch.abs(target) + 1e-8)
    thresholds = torch.arange(start, end + interval / 2, interval, dtype=torch.float32)
    return (rel_error < (1 - thresholds)).float().mean().item()


def reward_fn(sample, model_output, question_type):
    try:
        output_ans = extract_answer(model_output)
        if output_ans == '':
            output_ans = model_output
        gt_ans = extract_answer(sample.get("solution", ""))

        if question_type == "multiple choice":
            out_letter = re.match(r'^([A-Za-z])', output_ans.strip())
            gt_letter = re.match(r'^([A-Za-z])', gt_ans.strip())
            if out_letter and gt_letter:
                return 1.0 if out_letter.group(1).upper() == gt_letter.group(1).upper() else 0.0
            return 1.0 if output_ans.strip() == gt_ans.strip() else 0.0
        elif question_type == "numerical":
            if ("." in gt_ans) != ("." in output_ans):
                return 0.0
            gt_number = normalize_number(gt_ans)
            out_number = normalize_number(output_ans)
            if gt_number is None or out_number is None:
                return 0.0
            return 1.0 if round(gt_number, 2) == round(out_number, 2) else 0.0
        elif question_type == "regression":
            gt_number = normalize_number(gt_ans)
            out_number = normalize_number(output_ans)
            if gt_number is None or out_number is None:
                return 0.0
            return mean_relative_accuracy(out_number, gt_number)
        else:
            return 0.0
    except Exception:
        return 0.0


# ========================== Data Loading & Message Building ==========================
def load_data(prompt_path):
    if prompt_path.endswith('.jsonl'):
        with open(prompt_path, "r", encoding="utf-8") as f:
            return [json.loads(line) for line in f]
    elif prompt_path.endswith('.json'):
        with open(prompt_path, "r", encoding="utf-8") as f:
            return json.load(f)
    raise ValueError("Input file must be .json or .jsonl")


def build_messages(data, args=None):
    messages = []
    for x in data:
        if x["problem_type"] == "multiple choice":
            question = x["problem"] + "\nOptions:\n" + "".join(op + "\n" for op in x["options"])
        else:
            question = x["problem"]

        media_path = x["path"]
        if media_path.startswith("./Evaluation"):
            media_path = os.path.join(SCRIPT_DIR, media_path[2:])
        elif not os.path.isabs(media_path):
            media_path = os.path.join(EVAL_DATA_DIR, media_path)

        if not os.path.exists(media_path):
            base, _ = os.path.splitext(media_path)
            for alt_ext in (".mp4", ".avi", ".webm", ".mkv", ".mov"):
                alt_path = base + alt_ext
                if os.path.exists(alt_path):
                    media_path = alt_path
                    break
            if not os.path.exists(media_path) and os.path.isdir(base):
                media_path = base

        media_type = x["data_type"]
        if media_type == "frame" and os.path.isdir(media_path):
            MAX_FRAMES = 16
            frames = sorted([
                os.path.join(media_path, f) for f in os.listdir(media_path)
                if f.lower().endswith((".jpg", ".jpeg", ".png", ".bmp"))
            ])
            if len(frames) > MAX_FRAMES:
                indices = [int(i * (len(frames) - 1) / (MAX_FRAMES - 1)) for i in range(MAX_FRAMES)]
                frames = [frames[i] for i in indices]
            media_content = [{"type": "image", "image": f} for f in frames]
        else:
            media_content = [{"type": media_type, media_type: media_path}]
            if media_type == "video" and args is not None and args.nframes is not None:
                media_content[0]["nframes"] = args.nframes

        sys_msg = SYSTEM_MESSAGE
        if media_type == "video" and args is not None and args.nframes is not None:
            ts_hint = _get_frame_timestamps(media_path, max_frames=args.nframes)
            sys_msg = SYSTEM_MESSAGE + ts_hint

        msg = [
            {"role": "system", "content": [{"type": "text", "text": sys_msg}]},
            {"role": "user", "content": media_content + [
                {"type": "text", "text": QUESTION_TEMPLATE.format(Question=question) + TYPE_TEMPLATE[x["problem_type"]]}
            ]},
        ]
        messages.append(msg)
    return messages


# ========================== Evaluation Loop ==========================
def evaluate_benchmark(llm, processor, sampling_params, dataset_name, data, messages, output_path, bsz):
    final_output = []
    start_idx = 0
    if os.path.exists(output_path):
        try:
            with open(output_path, "r", encoding="utf-8") as f:
                existing = json.load(f)
                final_output = existing.get("results", [])
                start_idx = len(final_output)
                print(f"[{dataset_name}] Resuming from sample index {start_idx}")
        except Exception as e:
            print(f"[{dataset_name}] Error reading existing output: {e}")

    mean_acc = []
    mean_mra = []
    for item in final_output:
        reward = item.get("reward", 0.0)
        if item.get("problem_type", "") != "regression":
            mean_acc.append(reward)
        else:
            mean_mra.append(reward)

    for i in tqdm(range(start_idx, len(messages), bsz), desc=f"[{dataset_name}] Processing"):
        batch_messages = messages[i:i + bsz]
        prompts = [
            processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
            for msg in batch_messages
        ]
        try:
            image_inputs, video_inputs, video_kwargs = process_vision_info(
                batch_messages, return_video_kwargs=True
            )
            image_idx = 0
            video_idx = 0
            llm_inputs = []
            for idx, prompt in enumerate(prompts):
                content_items = batch_messages[idx][1]["content"]
                img_count = sum(1 for c in content_items if c.get("type") == "image")
                vid_count = sum(1 for c in content_items if c.get("type") == "video")
                sample_mm_data = {}
                sample_video_kw = {}
                if img_count > 0:
                    sample_mm_data["image"] = image_inputs[image_idx] if img_count == 1 else image_inputs[image_idx:image_idx + img_count]
                    image_idx += img_count
                if vid_count > 0:
                    sample_mm_data["video"] = video_inputs[video_idx]
                    for key, value in video_kwargs.items():
                        sample_video_kw[key] = value[video_idx]
                    video_idx += vid_count
                llm_inputs.append({
                    "prompt": prompt,
                    "multi_modal_data": sample_mm_data,
                    "mm_processor_kwargs": sample_video_kw,
                })
            outputs = llm.generate(llm_inputs, sampling_params=sampling_params)
            batch_output_text = [out.outputs[0].text for out in outputs]
        except Exception as e:
            print(f"[{dataset_name}] Error at batch starting index {i}: {e}")
            batch_output_text = ["<answer>error</answer>"] * len(batch_messages)

        for j, (sample, model_output) in enumerate(zip(data[i:i + bsz], batch_output_text), start=i):
            think_chain = extract_think(model_output)
            final_ans = extract_answer(model_output)
            if final_ans == "":
                final_ans = model_output
            sample["output"] = model_output
            sample["prediction"] = final_ans
            q_type = sample.get("problem_type", "")
            sample["reward"] = reward_fn(sample, model_output, q_type)
            sample["correct"] = sample["reward"] == 1.0
            if q_type != "regression":
                mean_acc.append(sample["reward"])
            else:
                mean_mra.append(sample["reward"])
            if think_chain:
                sample["process"] = f"<think>{think_chain}</think>"
            final_output.append(sample)

        try:
            with open(output_path, "w", encoding="utf-8") as f:
                json.dump({"results": final_output}, f, indent=2, ensure_ascii=False)
        except Exception as e:
            print(f"[{dataset_name}] Error writing checkpoint: {e}")

    final_acc = {"mean_acc": 0.0, "mean_mra": 0.0}
    if mean_acc:
        final_acc["mean_acc"] = torch.tensor(mean_acc).mean().item()
    if mean_mra:
        final_acc["mean_mra"] = torch.tensor(mean_mra).mean().item()

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump({"results": final_output, "final_acc": [final_acc]}, f, indent=2, ensure_ascii=False)

    return final_acc


# ========================== Main ==========================
def main():
    args = parse_args()

    if args.benchmarks == ["all"]:
        benchmark_names = list(BENCHMARK_REGISTRY.keys())
    else:
        benchmark_names = [n for n in args.benchmarks if n in BENCHMARK_REGISTRY]
        unknown = [n for n in args.benchmarks if n not in BENCHMARK_REGISTRY]
        if unknown:
            print(f"WARNING: Unknown benchmarks: {unknown}")
        if not benchmark_names:
            print("ERROR: No valid benchmarks specified.")
            return

    print(f"Model:      {args.model_path}")
    print(f"Benchmarks: {benchmark_names}")
    print(f"Frames:     {args.nframes or 'default'}")
    print(f"rep_penalty:{args.repetition_penalty}")
    print()

    llm = LLM(
        model=args.model_path,
        tensor_parallel_size=torch.cuda.device_count(),
        max_model_len=args.max_model_len,
        gpu_memory_utilization=args.gpu_memory_utilization,
        limit_mm_per_prompt={"image": 16, "video": 1},
    )
    sampling_params = SamplingParams(
        temperature=args.temperature,
        top_p=args.top_p,
        max_tokens=args.max_tokens,
        repetition_penalty=args.repetition_penalty,
        stop_token_ids=[],
    )
    processor = AutoProcessor.from_pretrained(args.model_path)
    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    tokenizer.padding_side = "left"
    processor.tokenizer = tokenizer

    os.makedirs(EVAL_RESULTS_DIR, exist_ok=True)

    summary = {}
    for dataset_name in benchmark_names:
        print(f"\n{'='*60}")
        print(f"  Evaluating: {dataset_name}")
        print(f"{'='*60}")

        prompt_path = os.path.join(EVAL_DATA_DIR, BENCHMARK_REGISTRY[dataset_name])
        output_path = os.path.join(EVAL_RESULTS_DIR, f"eval_{dataset_name}_{args.file_name}_greedy_output.json")

        data = load_data(prompt_path)
        messages = build_messages(data, args)
        acc = evaluate_benchmark(llm, processor, sampling_params, dataset_name, data, messages, output_path, args.bsz)
        summary[dataset_name] = acc
        print(f"[{dataset_name}] Accuracy: {acc['mean_acc']:.4f}  MRA: {acc['mean_mra']:.4f}")

    print(f"\n{'='*60}")
    print("  Summary")
    print(f"{'='*60}")
    for name, acc in summary.items():
        mra_str = f"  MRA: {acc['mean_mra']:.4f}" if acc["mean_mra"] > 0 else ""
        print(f"  {name:25s}  Acc: {acc['mean_acc']:.4f}{mra_str}")

    summary_path = os.path.join(EVAL_RESULTS_DIR, f"summary_{args.file_name}.json")
    with open(summary_path, "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)
    print(f"\nSummary saved to {summary_path}")


if __name__ == "__main__":
    main()