#!/usr/bin/env python3
"""
OpenMath Finetuning 自動評価スクリプト (Llama-3.1-Swallow-8B)

finetuned model でテストデータに対して推論を実行し、
自動採点 (Exact Match, Code Execution, Format Compliance) を全件実行。
結果をJSON + CSV + サマリーレポートで保存。

Usage:
    python3 run_evaluation.py
    python3 run_evaluation.py --model_path /path/to/model --test_data /path/to/test.json
    python3 run_evaluation.py --skip_inference --results_json /path/to/existing_results.json
"""

import os
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"

import argparse
import re
import csv
import gc
import json
import time
from datetime import datetime
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
try:
    from peft import PeftModel
    _PEFT_AVAILABLE = True
except ImportError:
    _PEFT_AVAILABLE = False

from eval_utils import (
    extract_boxed_answer,
    math_equivalent,
    check_code_execution,
    check_format_compliance,
    extract_code_blocks,
    execute_code_safely,
)

# デフォルトパス (Llama-3.1-Swallow-8B)
DEFAULT_FINETUNED_MODEL = "/home/mshohda/Desktop/KambayashiWork/Nextorage-LLM-ChatBot/AiDaptive_Software/output_llama_openmath"
DEFAULT_ORIGINAL_MODEL = "/home/mshohda/Desktop/KambayashiWork/Nextorage-LLM-ChatBot/AiDaptive_Software/models/Llama-3.1-Swallow-8B"
DEFAULT_TEST_DATA = "/home/mshohda/Desktop/KambayashiWork/Nextorage-LLM-ChatBot/Dataset/FinetuiningDataset/OpenMathInstruct/openmath_test_500.json"
DEFAULT_OUTPUT_DIR = "/home/mshohda/Desktop/KambayashiWork/Nextorage-LLM-ChatBot/AiDaptive_Software/vNXUN_2_03_00/commands/finetuning_openmath_llama/evaluation/results"

SYSTEM_PROMPT = (
    "あなたは数学の問題を解く優秀なAIアシスタントです。ステップバイステップで考え、Pythonコードを使って計算し、最終的な答えを明示してください。\n"
    "\n"
    "回答は以下のフォーマットに従ってください：\n"
    "1. Pythonコードは <llm-code> と </llm-code> タグで囲んでください\n"
    "2. コードの実行結果は <llm-code-output> と </llm-code-output> タグで囲んでください\n"
    "3. 最終的な答えは \\boxed{答え} の形式で明示してください\n"
    "\n"
    "回答例：\n"
    "Pythonコードを使用してこの問題を解決しましょう。<llm-code>\n"
    "x = 2 + 3\n"
    "print(x)\n"
    "</llm-code><llm-code-output>\n"
    "5\n"
    "</llm-code-output>\n"
    "したがって、答えは\\boxed{5}です。"
)


def parse_args():
    parser = argparse.ArgumentParser(description="OpenMath Finetuning Evaluation (Llama)")
    parser.add_argument("--model_path", default=DEFAULT_FINETUNED_MODEL,
                        help="Path to finetuned model")
    parser.add_argument("--original_model_path", default=DEFAULT_ORIGINAL_MODEL,
                        help="Path to original model (for tokenizer)")
    parser.add_argument("--test_data", default=DEFAULT_TEST_DATA,
                        help="Path to test dataset JSON")
    parser.add_argument("--output_dir", default=DEFAULT_OUTPUT_DIR,
                        help="Directory to save results")
    parser.add_argument("--max_samples", type=int, default=-1,
                        help="Max samples to evaluate (-1 for all)")
    parser.add_argument("--max_new_tokens", type=int, default=1024,
                        help="Max new tokens for generation")
    parser.add_argument("--gpu_id", type=int, default=0,
                        help="GPU ID to use")
    parser.add_argument("--skip_inference", action="store_true",
                        help="Skip inference, only run scoring on existing results")
    parser.add_argument("--results_json", default=None,
                        help="Path to existing results JSON (for --skip_inference)")
    parser.add_argument("--skip_code_exec", action="store_true",
                        help="Skip code execution checks")
    parser.add_argument("--model_label", default=None,
                        help="Label for this model (e.g. 'finetuned', 'original'). "
                             "Used in output filenames and reports.")
    parser.add_argument("--system_prompt", default=None,
                        help="Custom system prompt (overrides default SYSTEM_PROMPT). "
                             "Use to provide format instructions for baseline models.")
    parser.add_argument("--code_exec_pipeline", action="store_true",
                        help="Enable code execution pipeline: execute <llm-code> blocks "
                             "and inject real outputs into <llm-code-output>")
    parser.add_argument("--max_retries", type=int, default=2,
                        help="Max self-repair retries on code execution error (default: 2)")
    return parser.parse_args()


def build_prompt(question: str) -> str:
    """学習時と同じQAフォーマットでプロンプトを構築 (Llama 3.1形式)"""
    return (
        f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n"
        f"{SYSTEM_PROMPT}<|eot_id|>"
        f"<|start_header_id|>user<|end_header_id|>\n\n"
        f"{question}<|eot_id|>"
        f"<|start_header_id|>assistant<|end_header_id|>\n\n"
    )


def build_repair_prompt(error_msg: str, failed_code: str, attempt: int) -> str:
    """エラー箇所を特定し、具体的な修正指示を生成"""
    # tracebackからエラー行を抽出
    error_line = ""
    line_match = re.search(r'File ".*", line (\d+)', error_msg)
    if line_match:
        line_no = int(line_match.group(1))
        code_lines = failed_code.strip().split('\n')
        adjusted = line_no - 2  # wrapper scriptのオフセット補正
        if 0 <= adjusted < len(code_lines):
            error_line = code_lines[adjusted].strip()

    # 2回目以降はアプローチ転換を促す
    approach_change = ""
    if attempt >= 1:
        approach_change = (
            "前回と同じアプローチでは解決しません。"
            "sympyの代わりにmathモジュールやforループなど、別の方法を試してください。"
        )

    repair = f"<llm-code-output>\nError: {error_msg}\n</llm-code-output>\n"
    repair += "上記のコードでエラーが発生しました。"
    if error_line:
        repair += f"特にこの行が問題です: `{error_line}`\n"
    repair += "エラーメッセージをよく読み、問題の箇所を変更して修正したコードを書いてください。"
    if approach_change:
        repair += f"\n{approach_change}"
    repair += "\n"

    return repair


def postprocess_response(text: str) -> str:
    """回答後の繰り返し生成を除去する。"""
    # 1. ターンマーカー検出
    match = re.search(r'\n[^\n]{0,30}?(?:user|assistant)', text)
    if match:
        text = text[:match.start()].strip()

    # 2. 短フレーズの繰り返し検出 (3回以上同じフレーズが繰り返されたら切り取る)
    match = re.search(r'(.{10,80}?)\1{2,}', text, re.DOTALL)
    if match:
        text = text[:match.start() + len(match.group(1))].strip()

    return text.strip()


def _find_lora_adapter(model_path: str) -> str | None:
    """
    phisonai2 LoRA学習時のアダプタパスを探す。
    model_path: .../output_dir/finetuned_model_YYYY-mm-dd-HH-MM-SS/epoch_N_step_M_#ModelName
    対応するLoRAアダプタ: .../lora_adapters/finetuned_model_.../Lora_epoch_N_step_M_#ModelName
    """
    model_path = Path(model_path)
    # epoch_N_step_M_#ModelName → Lora_epoch_N_step_M_#ModelName
    step_dir_name = model_path.name
    lora_step_name = "Lora_" + step_dir_name

    # output_dir/lora_adapters/finetuned_model_.../Lora_epoch_N_step_M_#...
    run_dir = model_path.parent        # finetuned_model_YYYY-...
    output_dir = run_dir.parent        # output_dir (e.g. /mnt/nvme0/output_llama_lora/lora_01_...)
    candidate = output_dir / "lora_adapters" / run_dir.name / lora_step_name
    if (candidate / "adapter_config.json").exists():
        return str(candidate)
    return None


def load_model(model_path: str, tokenizer_path: str, gpu_id: int):
    """モデルとトークナイザーをロード。LoRAアダプタが存在する場合は自動適用。"""
    # CUDA_VISIBLE_DEVICES でGPUを制限
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    logical_gpu = 0  # CUDA_VISIBLE_DEVICES で絞った後は論理GPU 0

    print(f"Loading tokenizer from: {tokenizer_path}")
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # LoRAアダプタの存在確認
    lora_adapter_path = _find_lora_adapter(model_path)
    if lora_adapter_path:
        if not _PEFT_AVAILABLE:
            raise ImportError("peft is required for LoRA evaluation. Run: pip install peft")
        print(f"LoRA adapter found: {lora_adapter_path}")
        print(f"Loading base model from: {tokenizer_path}")
        model = AutoModelForCausalLM.from_pretrained(
            tokenizer_path,
            torch_dtype=torch.bfloat16,
            device_map={"": logical_gpu},
            trust_remote_code=True,
        )
        print(f"Applying LoRA adapter...")
        model = PeftModel.from_pretrained(model, lora_adapter_path)
        model = model.merge_and_unload()
        print(f"  [OK] LoRA merged into base model")
    else:
        print(f"Loading model from: {model_path}")
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype=torch.bfloat16,
            device_map={"": logical_gpu},
            trust_remote_code=True,
        )

    # vocab_size不一致の場合はリサイズ
    actual_embed_size = model.model.embed_tokens.weight.shape[0]
    tokenizer_vocab_size = len(tokenizer)
    if actual_embed_size != tokenizer_vocab_size:
        print(f"  Resizing embeddings: {actual_embed_size} -> {tokenizer_vocab_size}")
        model.resize_token_embeddings(tokenizer_vocab_size)

    model.config.pad_token_id = tokenizer.pad_token_id
    model.eval()
    print(f"  [OK] Model loaded on GPU {gpu_id} (CUDA_VISIBLE_DEVICES={gpu_id})")
    return model, tokenizer


def compute_perplexity(model, tokenizer, question: str, answer: str) -> float:
    """質問と回答のペアに対するperplexityを計算"""
    prompt = build_prompt(question) + answer
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)

    with torch.no_grad():
        outputs = model(
            input_ids=inputs["input_ids"],
            labels=inputs["input_ids"],
        )

    return torch.exp(outputs.loss).item()


def generate_response(model, tokenizer, question: str, max_new_tokens: int) -> str:
    """プロンプトを生成してモデルから応答を取得"""
    prompt = build_prompt(question)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            repetition_penalty=1.1,
        )

    response_ids = outputs[0][inputs["input_ids"].shape[1]:]
    raw = tokenizer.decode(response_ids, skip_special_tokens=True)
    return postprocess_response(raw)


class StopOnString(StoppingCriteria):
    """指定文字列が生成テキストに出現したら生成を停止する"""

    def __init__(self, stop_string: str, tokenizer, prompt_length: int):
        self.stop_string = stop_string
        self.tokenizer = tokenizer
        self.prompt_length = prompt_length
        # stop_string のトークン数の3倍をチェック窓とする（分割トークン化対策）
        self.check_window = max(
            len(tokenizer.encode(stop_string, add_special_tokens=False)) * 3, 20
        )

    def __call__(self, input_ids, scores, **kwargs):
        if input_ids.shape[1] <= self.prompt_length:
            return False
        start = max(self.prompt_length, input_ids.shape[1] - self.check_window)
        tail = self.tokenizer.decode(input_ids[0][start:], skip_special_tokens=True)
        return self.stop_string in tail


def generate_with_code_exec(
    model, tokenizer, question: str, max_new_tokens: int, max_retries: int = 2
) -> tuple:
    """
    コード実行パイプライン付き推論。

    フロー:
      1. </llm-code> まで生成（コードブロックを含むテキスト）
      2. コードを実際にPythonで実行
      3. 実行結果を <llm-code-output> に注入
      4. 続きを生成して \\boxed{} を含む最終回答を取得
      5. 実行エラー時はエラーを返してリトライ（self-repair）

    Returns:
        (response_text, pipeline_info_dict)
    """
    prompt = build_prompt(question)
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
    prompt_len = input_ids.shape[1]

    pipeline_info = {
        "code_exec_pipeline": True,
        "attempts": [],
        "total_attempts": 0,
        "final_status": "no_code",
    }

    # Phase 1: </llm-code> まで生成
    stop_criteria = StopOnString("</llm-code>", tokenizer, prompt_len)

    with torch.no_grad():
        outputs = model.generate(
            input_ids=input_ids,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            repetition_penalty=1.1,
            stopping_criteria=StoppingCriteriaList([stop_criteria]),
        )

    generated = tokenizer.decode(outputs[0][prompt_len:], skip_special_tokens=True)

    # コードブロックがない場合はそのまま返す
    if "</llm-code>" not in generated:
        pipeline_info["final_status"] = "no_code"
        return postprocess_response(generated), pipeline_info

    # </llm-code> 以降を切り捨て（ハルシネーション防止）
    code_end_idx = generated.index("</llm-code>") + len("</llm-code>")
    response_so_far = generated[:code_end_idx]

    prev_code = None
    for attempt in range(max_retries + 1):
        pipeline_info["total_attempts"] = attempt + 1

        # コードブロック抽出・実行
        code_blocks = extract_code_blocks(response_so_far)
        if not code_blocks:
            pipeline_info["final_status"] = "no_code_block"
            break

        last_code = code_blocks[-1]

        # 同一コード検出（無限ループ防止）
        if prev_code is not None and last_code.strip() == prev_code.strip():
            pipeline_info["final_status"] = "same_code_detected"
            break
        prev_code = last_code

        exec_result = execute_code_safely(last_code)

        attempt_info = {
            "attempt": attempt,
            "code_preview": last_code[:200],
            "success": exec_result["success"],
            "stdout": exec_result["stdout"][:500] if exec_result["stdout"] else "",
            "stderr": exec_result["stderr"][:500] if exec_result["stderr"] else "",
        }
        pipeline_info["attempts"].append(attempt_info)

        if exec_result["success"]:
            # 成功: 実行結果を注入して続きを生成
            output_text = exec_result["stdout"] if exec_result["stdout"] else "None"
            injection = f"<llm-code-output>\n{output_text}\n</llm-code-output>\n"
            response_so_far += injection

            # Phase 2: 最終回答を生成
            full_context = prompt + response_so_far
            context_ids = tokenizer(
                full_context, return_tensors="pt", truncation=True, max_length=4096
            ).input_ids.to(model.device)
            context_len = context_ids.shape[1]

            with torch.no_grad():
                outputs = model.generate(
                    input_ids=context_ids,
                    max_new_tokens=max_new_tokens // 2,
                    do_sample=False,
                    pad_token_id=tokenizer.pad_token_id,
                    eos_token_id=tokenizer.eos_token_id,
                    repetition_penalty=1.1,
                )

            continuation = tokenizer.decode(
                outputs[0][context_len:], skip_special_tokens=True
            )
            response_so_far += continuation
            pipeline_info["final_status"] = "success"
            return postprocess_response(response_so_far), pipeline_info

        else:
            # 実行失敗
            error_msg = exec_result["stderr"] or "Unknown error"
            if len(error_msg) > 300:
                error_msg = error_msg[:300] + "..."

            if attempt < max_retries:
                # Self-repair: エラー箇所を特定した具体的な修正指示を生成
                repair_injection = build_repair_prompt(error_msg, last_code, attempt)
                response_so_far += repair_injection

                full_context = prompt + response_so_far
                context_ids = tokenizer(
                    full_context, return_tensors="pt", truncation=True, max_length=4096
                ).input_ids.to(model.device)
                context_len = context_ids.shape[1]

                stop_criteria = StopOnString("</llm-code>", tokenizer, context_len)

                with torch.no_grad():
                    outputs = model.generate(
                        input_ids=context_ids,
                        max_new_tokens=max_new_tokens // 2,
                        do_sample=False,
                        pad_token_id=tokenizer.pad_token_id,
                        eos_token_id=tokenizer.eos_token_id,
                        repetition_penalty=1.1,
                        stopping_criteria=StoppingCriteriaList([stop_criteria]),
                    )

                repair_text = tokenizer.decode(
                    outputs[0][context_len:], skip_special_tokens=True
                )

                if "</llm-code>" in repair_text:
                    code_end = repair_text.index("</llm-code>") + len("</llm-code>")
                    response_so_far += repair_text[:code_end]
                    # 次のループで再実行（同一コード検出はループ先頭で行う）
                else:
                    # モデルがコードブロックを生成しなかった
                    response_so_far += repair_text
                    pipeline_info["final_status"] = "repair_no_code"
                    return postprocess_response(response_so_far), pipeline_info
            else:
                # リトライ上限到達: エラー出力を注入して続きを生成
                injection = f"<llm-code-output>\nError: {error_msg}\n</llm-code-output>\n"
                response_so_far += injection

                full_context = prompt + response_so_far
                context_ids = tokenizer(
                    full_context, return_tensors="pt", truncation=True, max_length=4096
                ).input_ids.to(model.device)
                context_len = context_ids.shape[1]

                with torch.no_grad():
                    outputs = model.generate(
                        input_ids=context_ids,
                        max_new_tokens=max_new_tokens // 2,
                        do_sample=False,
                        pad_token_id=tokenizer.pad_token_id,
                        eos_token_id=tokenizer.eos_token_id,
                        repetition_penalty=1.1,
                    )

                continuation = tokenizer.decode(
                    outputs[0][context_len:], skip_special_tokens=True
                )
                response_so_far += continuation
                pipeline_info["final_status"] = "error_exhausted"
                return postprocess_response(response_so_far), pipeline_info

    return postprocess_response(response_so_far), pipeline_info


def run_inference(model, tokenizer, test_data: list, max_new_tokens: int,
                  max_samples: int, code_exec_pipeline: bool = False,
                  max_retries: int = 2) -> list:
    """全テストデータに対して推論を実行"""
    samples = test_data if max_samples < 0 else test_data[:max_samples]
    results = []
    total = len(samples)

    mode = "code execution pipeline" if code_exec_pipeline else "standard"
    print(f"\nRunning inference on {total} samples... (mode: {mode})")
    if code_exec_pipeline:
        print(f"  max_retries={max_retries} for self-repair on code execution error")
    start_time = time.time()

    for i, sample in enumerate(samples):
        t0 = time.time()

        if code_exec_pipeline:
            response, pipeline_info = generate_with_code_exec(
                model, tokenizer, sample["question"], max_new_tokens, max_retries
            )
        else:
            response = generate_response(model, tokenizer, sample["question"], max_new_tokens)
            pipeline_info = None

        elapsed = time.time() - t0

        # Perplexity: 参照回答に対するモデルのperplexity
        ppl = compute_perplexity(model, tokenizer, sample["question"], sample["cot_answer"])

        result = {
            "id": i,
            "question": sample["question"],
            "reference_answer": sample["cot_answer"],
            "model_response": response,
            "inference_time_s": round(elapsed, 2),
            "perplexity": round(ppl, 4),
        }
        if pipeline_info:
            result["pipeline_info"] = pipeline_info

        results.append(result)

        if (i + 1) % 10 == 0 or (i + 1) == total:
            elapsed_total = time.time() - start_time
            eta = elapsed_total / (i + 1) * (total - i - 1)
            status = ""
            if pipeline_info:
                status = f" [{pipeline_info['final_status']}]"
            print(f"  [{i+1}/{total}] {elapsed:.1f}s/sample{status}, ETA: {eta/60:.1f}min")

    total_time = time.time() - start_time
    print(f"\nInference completed: {total_time/60:.1f}min total, {total_time/total:.1f}s/sample avg")
    return results


def score_results(results: list, skip_code_exec: bool = False) -> list:
    """推論結果を自動採点"""
    print(f"\nScoring {len(results)} results...")

    for i, r in enumerate(results):
        ref = r["reference_answer"]
        resp = r["model_response"]

        # 1. Boxed Answer Exact Match
        gold_answer = extract_boxed_answer(ref)
        pred_answer = extract_boxed_answer(resp)
        r["gold_boxed"] = gold_answer
        r["pred_boxed"] = pred_answer
        r["exact_match"] = math_equivalent(pred_answer, gold_answer) if (pred_answer and gold_answer) else False

        # 2. Code Execution
        if skip_code_exec:
            r["code_exec"] = {"has_code": False, "skipped": True}
        else:
            r["code_exec"] = check_code_execution(resp)

        # 3. Format Compliance
        r["format"] = check_format_compliance(resp)

        if (i + 1) % 50 == 0:
            print(f"  [{i+1}/{len(results)}] scored")

    return results


def compute_summary(results: list) -> dict:
    """集計サマリーを計算"""
    total = len(results)
    if total == 0:
        return {}

    # Exact Match
    em_correct = sum(1 for r in results if r["exact_match"])
    gold_extracted = sum(1 for r in results if r["gold_boxed"] is not None)
    pred_extracted = sum(1 for r in results if r["pred_boxed"] is not None)

    # Code Execution
    code_results = [r for r in results if r["code_exec"].get("has_code", False)]
    code_success = sum(1 for r in code_results if r["code_exec"]["all_success"])

    # Format Compliance
    format_ok = sum(1 for r in results if r["format"]["format_ok"])
    has_boxed = sum(1 for r in results if r["format"]["has_boxed_answer"])
    no_bad_tokens = sum(1 for r in results if r["format"]["no_bad_tokens"])

    # Perplexity
    ppls = [r["perplexity"] for r in results if "perplexity" in r]
    avg_ppl = sum(ppls) / len(ppls) if ppls else 0

    summary = {
        "total_samples": total,
        "perplexity": {
            "mean": round(avg_ppl, 4),
            "min": round(min(ppls), 4) if ppls else 0,
            "max": round(max(ppls), 4) if ppls else 0,
        },
        "exact_match": {
            "correct": em_correct,
            "accuracy": round(em_correct / total * 100, 1),
            "gold_extracted": gold_extracted,
            "pred_extracted": pred_extracted,
        },
        "code_execution": {
            "samples_with_code": len(code_results),
            "all_success": code_success,
            "success_rate": round(code_success / len(code_results) * 100, 1) if code_results else 0,
        },
        "format_compliance": {
            "format_ok": format_ok,
            "format_ok_rate": round(format_ok / total * 100, 1),
            "has_boxed": has_boxed,
            "has_boxed_rate": round(has_boxed / total * 100, 1),
            "no_bad_tokens": no_bad_tokens,
        },
    }

    # Code Execution Pipeline 統計
    pipeline_results = [r for r in results if "pipeline_info" in r]
    if pipeline_results:
        status_counts = {}
        for r in pipeline_results:
            status = r["pipeline_info"]["final_status"]
            status_counts[status] = status_counts.get(status, 0) + 1
        total_retries = sum(
            r["pipeline_info"]["total_attempts"] - 1
            for r in pipeline_results
            if r["pipeline_info"]["total_attempts"] > 1
        )
        summary["code_exec_pipeline"] = {
            "enabled": True,
            "samples_processed": len(pipeline_results),
            "status_counts": status_counts,
            "total_retries": total_retries,
        }

    return summary


def save_results(results: list, summary: dict, output_dir: str, model_label: str = None):
    """結果を保存"""
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    label = f"_{model_label}" if model_label else ""

    # 1. JSON (全詳細)
    json_path = output_path / f"eval_results{label}_{timestamp}.json"
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump({"summary": summary, "results": results}, f, ensure_ascii=False, indent=2)
    print(f"  JSON: {json_path}")

    # 2. CSV (スコアシート)
    csv_path = output_path / f"eval_scores{label}_{timestamp}.csv"
    with open(csv_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow([
            "id", "question_preview", "exact_match",
            "gold_boxed", "pred_boxed",
            "perplexity",
            "has_code", "code_success",
            "format_ok", "has_boxed", "no_bad_tokens",
            "inference_time_s",
        ])
        for r in results:
            writer.writerow([
                r["id"],
                r["question"][:80],
                int(r["exact_match"]),
                r["gold_boxed"] or "",
                r["pred_boxed"] or "",
                r.get("perplexity", ""),
                int(r["code_exec"].get("has_code", False)),
                int(r["code_exec"].get("all_success", False)) if r["code_exec"].get("has_code") else "N/A",
                int(r["format"]["format_ok"]),
                int(r["format"]["has_boxed_answer"]),
                int(r["format"]["no_bad_tokens"]),
                r.get("inference_time_s", ""),
            ])
    print(f"  CSV:  {csv_path}")

    # 3. サマリーレポート (テキスト)
    report_path = output_path / f"eval_summary{label}_{timestamp}.txt"
    with open(report_path, "w", encoding="utf-8") as f:
        f.write("=" * 60 + "\n")
        f.write("OpenMath Finetuning Evaluation Summary (Llama-3.1-Swallow-8B)\n")
        if model_label:
            f.write(f"Model: {model_label}\n")
        f.write(f"Date: {timestamp}\n")
        f.write("=" * 60 + "\n\n")

        s = summary
        f.write(f"Total Samples: {s['total_samples']}\n\n")

        f.write("--- Perplexity ---\n")
        f.write(f"  Mean: {s['perplexity']['mean']}\n")
        f.write(f"  Min:  {s['perplexity']['min']},  Max: {s['perplexity']['max']}\n\n")

        f.write("--- Exact Match (Primary Metric) ---\n")
        f.write(f"  Accuracy: {s['exact_match']['accuracy']}% ({s['exact_match']['correct']}/{s['total_samples']})\n")
        f.write(f"  Gold answer extracted: {s['exact_match']['gold_extracted']}/{s['total_samples']}\n")
        f.write(f"  Pred answer extracted: {s['exact_match']['pred_extracted']}/{s['total_samples']}\n\n")

        f.write("--- Code Execution ---\n")
        f.write(f"  Samples with code: {s['code_execution']['samples_with_code']}\n")
        f.write(f"  Success rate: {s['code_execution']['success_rate']}% ({s['code_execution']['all_success']}/{s['code_execution']['samples_with_code']})\n\n")

        f.write("--- Format Compliance ---\n")
        f.write(f"  Format OK: {s['format_compliance']['format_ok_rate']}% ({s['format_compliance']['format_ok']}/{s['total_samples']})\n")
        f.write(f"  Has boxed answer: {s['format_compliance']['has_boxed_rate']}%\n")
        f.write(f"  No bad tokens: {s['format_compliance']['no_bad_tokens']}/{s['total_samples']}\n\n")

        if "code_exec_pipeline" in s:
            p = s["code_exec_pipeline"]
            f.write("--- Code Execution Pipeline ---\n")
            f.write(f"  Samples processed: {p['samples_processed']}\n")
            for status, count in sorted(p["status_counts"].items()):
                f.write(f"    {status}: {count}\n")
            f.write(f"  Total self-repair retries: {p['total_retries']}\n")
    print(f"  Report: {report_path}")

    return json_path, csv_path, report_path


def print_summary(summary: dict, model_label: str = None):
    """サマリーをコンソールに表示"""
    s = summary
    print("\n" + "=" * 60)
    title = f"EVALUATION SUMMARY [{model_label}]" if model_label else "EVALUATION SUMMARY"
    print(title)
    print("=" * 60)
    print(f"Total Samples: {s['total_samples']}")
    print()
    print(f"  Perplexity (mean):    {s['perplexity']['mean']}")
    print(f"  Exact Match Accuracy: {s['exact_match']['accuracy']}%"
          f" ({s['exact_match']['correct']}/{s['total_samples']})")
    print(f"  Code Execution Rate:  {s['code_execution']['success_rate']}%"
          f" ({s['code_execution']['all_success']}/{s['code_execution']['samples_with_code']} samples with code)")
    print(f"  Format Compliance:    {s['format_compliance']['format_ok_rate']}%")
    if "code_exec_pipeline" in s:
        p = s["code_exec_pipeline"]
        print()
        print("  --- Code Execution Pipeline ---")
        print(f"  Samples processed: {p['samples_processed']}")
        for status, count in sorted(p["status_counts"].items()):
            print(f"    {status}: {count}")
        print(f"  Total self-repair retries: {p['total_retries']}")
    print("=" * 60)


def main():
    global SYSTEM_PROMPT
    args = parse_args()

    # カスタムシステムプロンプトが指定された場合、グローバル変数を上書き
    if args.system_prompt is not None:
        SYSTEM_PROMPT = args.system_prompt
        print(f"[INFO] Using custom system prompt ({len(SYSTEM_PROMPT)} chars)")

    print("=" * 60)
    print("OpenMath Finetuning Evaluation (Llama-3.1-Swallow-8B)")
    print(f"Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("=" * 60)

    if args.skip_inference:
        # 既存の推論結果を読み込んでスコアリングのみ
        results_path = args.results_json
        if not results_path:
            print("[ERROR] --results_json is required when using --skip_inference")
            return
        print(f"\nLoading existing results: {results_path}")
        with open(results_path, encoding="utf-8") as f:
            data = json.load(f)
        results = data["results"] if "results" in data else data
    else:
        # テストデータ読み込み
        print(f"\nTest data: {args.test_data}")
        with open(args.test_data, encoding="utf-8") as f:
            test_data = json.load(f)
        print(f"  {len(test_data)} samples loaded")

        # モデルロード
        print(f"\nModel: {args.model_path}")
        model, tokenizer = load_model(args.model_path, args.original_model_path, args.gpu_id)

        # 推論実行
        results = run_inference(
            model, tokenizer, test_data, args.max_new_tokens, args.max_samples,
            code_exec_pipeline=args.code_exec_pipeline,
            max_retries=args.max_retries,
        )

        # GPUメモリ解放
        del model
        gc.collect()
        torch.cuda.empty_cache()
        print("  [OK] Model unloaded")

    # スコアリング
    results = score_results(results, skip_code_exec=args.skip_code_exec)

    # サマリー計算
    summary = compute_summary(results)

    # model_label の自動推定
    model_label = args.model_label
    if model_label is None and not args.skip_inference:
        if args.model_path == DEFAULT_ORIGINAL_MODEL:
            model_label = "llama_original"
        elif args.model_path == DEFAULT_FINETUNED_MODEL:
            model_label = "llama_finetuned"

    # サマリーに metadata を追加
    summary["model_label"] = model_label
    summary["model_path"] = args.model_path if not args.skip_inference else "(skip_inference)"

    # 保存
    print("\nSaving results...")
    save_results(results, summary, args.output_dir, model_label)

    # サマリー表示
    print_summary(summary, model_label)


if __name__ == "__main__":
    main()