"""
数据清洗脚本：从 07_oracle_results_nothink + 05_delta_frame_mapping 生成 RL 训练数据

过滤策略：
1. 非多选题（answer 不符合 A/A.xxx 格式）→ 丢弃
2. oracle_results 为空（推理失败）→ 丢弃
3. 单 pattern 条目（1个delta）→ 直接保留
4. 多 pattern 条目（2+个delta）→ 过滤掉含异常 oracle_answer 的 pattern
   异常定义：oracle_answer 不是单个字母 A-E
"""

import json
from pathlib import Path

DATA_DIR = Path(__file__).resolve().parent

ORACLE_FILE = DATA_DIR / "07_oracle_results_nothink.json"
DELTA_FILE  = DATA_DIR / "05_delta_frame_mapping.json"
OUTPUT_FILE = DATA_DIR / "08_rl_training_data.json"

VALID_ANSWERS = set("ABCDE")


def is_multiple_choice(answer: str) -> bool:
    """判断是否为标准多选题格式：A / A. xxx / A xxx"""
    ans = answer.strip()
    return (len(ans) == 1 and ans in VALID_ANSWERS) or \
           (len(ans) >= 2 and ans[0] in VALID_ANSWERS and ans[1] in ". ")


def is_valid_oracle_answer(answer: str) -> bool:
    """oracle_answer 是否为合法的单字母选项"""
    return answer.strip() in VALID_ANSWERS


def main():
    # 加载 delta frame mapping
    with open(DELTA_FILE) as f:
        delta_data = json.load(f)
    delta_map = {m["problem_id"]: m for m in delta_data["mappings"]}
    print(f"[Delta] 加载 {len(delta_map)} 条")

    # 加载 oracle results，同时构建 oracle 查找表
    with open(ORACLE_FILE) as f:
        oracle_data = json.load(f)
    oracle_results = oracle_data["results"]
    print(f"[Oracle] 加载 {len(oracle_results)} 条")

    # 构建查找表：(problem_id, mask_delta_id) -> oracle entry
    oracle_map = {}
    for r in oracle_results:
        pid = r["problem_id"]
        for o in r.get("oracle_results", []):
            oracle_map[(pid, o["mask_delta_id"])] = o

    stats = {
        "total_problems": len(oracle_results),
        "skipped_not_mc": 0,
        "skipped_empty_oracle": 0,
        "skipped_no_delta": 0,
        "output_problems": 0,
        "single_pattern_kept": 0,
        "multi_pattern_kept": 0,
        "multi_pattern_abnormal_removed": 0,
        "total_mask_patterns": 0,
        "oracle_correct_count": 0,
    }

    training_data = []

    for problem in oracle_results:
        pid = problem["problem_id"]

        # 过滤1：非多选题
        if not is_multiple_choice(problem["gt_answer"]):
            stats["skipped_not_mc"] += 1
            continue

        # 过滤2：oracle_results 为空
        if not problem["oracle_results"]:
            stats["skipped_empty_oracle"] += 1
            continue

        # 过滤3：没有 delta mapping
        delta_item = delta_map.get(pid)
        if delta_item is None:
            stats["skipped_no_delta"] += 1
            continue

        delta_mask_patterns = delta_item.get("mask_patterns", [])
        n_patterns = len(delta_mask_patterns)

        # 构建 mask_patterns
        mask_patterns = []
        for p in delta_mask_patterns:
            mid = p["mask_delta_id"]
            oracle_entry = oracle_map.get((pid, mid))
            oracle_correct = oracle_entry["oracle_correct"] if oracle_entry else None

            # 多 pattern 时，过滤掉异常 oracle_answer 的 pattern
            if n_patterns > 1 and oracle_entry is not None:
                if not is_valid_oracle_answer(oracle_entry["oracle_answer"]):
                    stats["multi_pattern_abnormal_removed"] += 1
                    continue

            pattern = {
                "mask_delta_id": mid,
                "zero_frame_indices": p["zero_frame_indices"],
                "remaining_frames": p["remaining_frames"],
                "mask_ratio": p["mask_ratio"],
                "oracle_correct": oracle_correct,
            }
            mask_patterns.append(pattern)
            stats["total_mask_patterns"] += 1
            if oracle_correct:
                stats["oracle_correct_count"] += 1

        # 过滤掉清洗后 mask_patterns 为空的条目
        if not mask_patterns:
            continue

        entry = {
            "problem_id": pid,
            "video": problem["video_path"],
            "question": problem["question"],
            "answer": problem["gt_answer"],
            "n_sampled_frames": problem["n_sampled_frames"],
            "mask_patterns": mask_patterns,
        }
        training_data.append(entry)
        stats["output_problems"] += 1

        if n_patterns == 1:
            stats["single_pattern_kept"] += 1
        else:
            stats["multi_pattern_kept"] += 1

    # 写出
    n_total_p = stats["total_mask_patterns"]
    output = {
        "metadata": {
            "total_entries": len(training_data),
            "total_mask_patterns": n_total_p,
            "oracle_correct_count": stats["oracle_correct_count"],
            "oracle_correct_rate": round(stats["oracle_correct_count"] / n_total_p, 3) if n_total_p else 0,
            "oracle_source": ORACLE_FILE.name,
            "delta_source": DELTA_FILE.name,
        },
        "data": training_data,
    }

    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
        json.dump(output, f, ensure_ascii=False, indent=2)

    print(f"\n{'='*50}")
    print(f"清洗完成，输出到 {OUTPUT_FILE.name}")
    print(f"{'='*50}")
    for k, v in stats.items():
        print(f"  {k}: {v}")


if __name__ == "__main__":
    main()