"""
07_split_sft_rl.py
最终合并和拆分：
1. 合并所有数据（real + filtered speed/chiral + new speed/chiral/reverse）
2. 统一时间戳格式为 HH:MM:SS（精确到1秒）
3. 拆分为 SFT 和 RL 两个数据集
4. 输出最终文件
"""
import json
import os
import re
import collections

OUT_DIR = "/root/final_dataset"

# 输入文件
REAL_DATA = f"{OUT_DIR}/sft/real_data.jsonl"
SPEED_FILTERED = f"{OUT_DIR}/sft/speed_filtered.jsonl"
CHIRAL_FILTERED = f"{OUT_DIR}/sft/chiral_filtered.jsonl"
SPEED_SINGLE_SFT = f"{OUT_DIR}/sft/speed_single_sft.jsonl"
CHIRAL_SINGLE_SFT = f"{OUT_DIR}/sft/chiral_single_sft.jsonl"
REVERSE_SINGLE_SFT = f"{OUT_DIR}/sft/reverse_single_sft.jsonl"

SPEED_SINGLE_RL = f"{OUT_DIR}/rl/speed_single_rl.json"
CHIRAL_SINGLE_RL = f"{OUT_DIR}/rl/chiral_single_rl.json"
REVERSE_SINGLE_RL = f"{OUT_DIR}/rl/reverse_single_rl.json"

# 输出文件
SFT_OUTPUT = f"{OUT_DIR}/sft/video_qa_sft_v5_for_cot.jsonl"
RL_OUTPUT = f"{OUT_DIR}/rl/video_qa_rl_augmented.json"
STATS_OUTPUT = f"{OUT_DIR}/docs/DATA_STATS.json"


def normalize_timestamp_in_response(response):
    """
    将 response 中的时间戳统一为 HH:MM:SS 精确到1秒。
    处理格式：
    - <time>00:00:05.3</time> → <time>00:00:05</time>
    - <time>00:00:05.30</time> → <time>00:00:05</time>
    - <time_range>00:00:01.2 - 00:00:03.8</time_range> → <time_range>00:00:01 - 00:00:04</time_range>
    """
    if not response:
        return response

    def round_timestamp(match):
        ts = match.group(1)
        # 解析 HH:MM:SS.x 格式
        parts = ts.split(':')
        if len(parts) == 3:
            h, m, s_part = parts
            # 去掉小数部分，四舍五入到秒
            if '.' in s_part:
                s_float = float(s_part)
                s = round(s_float)
                # 处理进位
                s_int = int(s)
                m_int = int(m)
                h_int = int(h)
                if s_int >= 60:
                    s_int -= 60
                    m_int += 1
                if m_int >= 60:
                    m_int -= 60
                    h_int += 1
                return f"{h_int:02d}:{m_int:02d}:{s_int:02d}"
            else:
                return ts
        return ts

    # 处理 <time> 标签内的时间戳
    response = re.sub(
        r'<time>(\d+:\d+:\d+(?:\.\d+)?)</time>',
        lambda m: f'<time>{round_timestamp(m)}</time>',
        response
    )

    # 处理 <time_range> 标签内的时间戳
    def fix_time_range(match):
        content = match.group(1)
        # 找到所有时间戳并替换
        fixed = re.sub(
            r'(\d+:\d+:\d+(?:\.\d+)?)',
            lambda m2: round_timestamp(m2),
            content
        )
        return f'<time_range>{fixed}</time_range>'

    response = re.sub(
        r'<time_range>(.*?)</time_range>',
        fix_time_range,
        response
    )

    return response


def load_jsonl(path):
    if not os.path.exists(path):
        print(f"  [WARN] File not found: {path}")
        return []
    with open(path) as f:
        return [json.loads(line) for line in f if line.strip()]


def main():
    print("=== Loading data ===")

    # 加载原有数据
    real_data = load_jsonl(REAL_DATA)
    speed_filtered = load_jsonl(SPEED_FILTERED)
    chiral_filtered = load_jsonl(CHIRAL_FILTERED)

    # 加载新增数据
    speed_single = load_jsonl(SPEED_SINGLE_SFT)
    chiral_single = load_jsonl(CHIRAL_SINGLE_SFT)
    reverse_single = load_jsonl(REVERSE_SINGLE_SFT)

    print(f"  Real data: {len(real_data)}")
    print(f"  Speed filtered (pair): {len(speed_filtered)}")
    print(f"  Chiral filtered (pair): {len(chiral_filtered)}")
    print(f"  Speed single (new): {len(speed_single)}")
    print(f"  Chiral single (new): {len(chiral_single)}")
    print(f"  Reverse single (new): {len(reverse_single)}")

    # === SFT 数据合并 ===
    print("\n=== Building SFT dataset ===")
    sft_all = []

    # 1. 原有真实数据 - 统一时间戳
    for d in real_data:
        if 'response' in d:
            d['response'] = normalize_timestamp_in_response(d['response'])
        sft_all.append(d)

    # 2. 保留的 Speed pair - 统一时间戳
    for d in speed_filtered:
        if 'response' in d:
            d['response'] = normalize_timestamp_in_response(d['response'])
        sft_all.append(d)

    # 3. 保留的 Chiral pair - 统一时间戳
    for d in chiral_filtered:
        if 'response' in d:
            d['response'] = normalize_timestamp_in_response(d['response'])
        sft_all.append(d)

    # 4. 新增单视频数据 (这些还没有 CoT response，标记一下)
    for d in speed_single:
        d['_needs_cot'] = True
        d['_data_source'] = 'Speed_Single'
        sft_all.append(d)

    for d in chiral_single:
        d['_needs_cot'] = True
        d['_data_source'] = 'Chiral_Single'
        sft_all.append(d)

    for d in reverse_single:
        d['_needs_cot'] = True
        d['_data_source'] = 'Reverse_Single'
        sft_all.append(d)

    print(f"  SFT total: {len(sft_all)}")

    # 统计需要 CoT 的
    needs_cot = [d for d in sft_all if d.get('_needs_cot')]
    has_cot = [d for d in sft_all if not d.get('_needs_cot')]
    print(f"  Has CoT: {len(has_cot)}")
    print(f"  Needs CoT generation: {len(needs_cot)}")

    # === RL 数据合并 ===
    print("\n=== Building RL dataset ===")
    rl_all = []

    for rl_file in [SPEED_SINGLE_RL, CHIRAL_SINGLE_RL, REVERSE_SINGLE_RL]:
        if os.path.exists(rl_file):
            with open(rl_file) as f:
                rl_data = json.load(f)
                rl_all.extend(rl_data)
                print(f"  {os.path.basename(rl_file)}: {len(rl_data)}")

    print(f"  RL total: {len(rl_all)}")

    # === 输出 ===
    print("\n=== Writing outputs ===")

    # SFT: 全量（含待CoT的新数据）
    with open(SFT_OUTPUT, "w") as f:
        for d in sft_all:
            f.write(json.dumps(d, ensure_ascii=False) + "\n")
    print(f"  SFT: {SFT_OUTPUT} ({len(sft_all)} entries)")

    # SFT: 单独输出需要 CoT 生成的部分
    cot_needed_output = f"{OUT_DIR}/sft/needs_cot_generation.jsonl"
    with open(cot_needed_output, "w") as f:
        for d in needs_cot:
            f.write(json.dumps(d, ensure_ascii=False) + "\n")
    print(f"  Needs CoT: {cot_needed_output} ({len(needs_cot)} entries)")

    # RL
    with open(RL_OUTPUT, "w") as f:
        json.dump(rl_all, f, ensure_ascii=False, indent=2)
    print(f"  RL: {RL_OUTPUT} ({len(rl_all)} entries)")

    # === 统计 ===
    synth_count = len(speed_filtered) + len(chiral_filtered) + len(speed_single) + len(chiral_single) + len(reverse_single)
    total = len(sft_all)
    synth_ratio = synth_count / total * 100

    stats = {
        "total_sft": total,
        "real_data": len(real_data),
        "synthetic_data": synth_count,
        "synthetic_ratio": f"{synth_ratio:.1f}%",
        "breakdown": {
            "real": len(real_data),
            "speed_pair_kept": len(speed_filtered),
            "chiral_pair_kept": len(chiral_filtered),
            "speed_single_new": len(speed_single),
            "chiral_single_new": len(chiral_single),
            "reverse_single_new": len(reverse_single),
        },
        "rl_total": len(rl_all),
        "rl_breakdown": {
            "speed_single": len([d for d in rl_all if d.get('data_source') == 'Speed_Single']),
            "chiral_single": len([d for d in rl_all if d.get('data_source') == 'Chiral_Single']),
            "reverse_single": len([d for d in rl_all if d.get('data_source') == 'Reverse_Single']),
        },
        "needs_cot_generation": len(needs_cot),
        "timestamp_format": "HH:MM:SS (rounded to 1s precision)",
    }

    with open(STATS_OUTPUT, "w") as f:
        json.dump(stats, f, ensure_ascii=False, indent=2)

    print(f"\n{'='*60}")
    print(f"FINAL STATS")
    print(f"{'='*60}")
    print(f"SFT total:        {total}")
    print(f"  Real data:      {len(real_data)}")
    print(f"  Synthetic:      {synth_count} ({synth_ratio:.1f}%)")
    print(f"    Speed pair:   {len(speed_filtered)}")
    print(f"    Chiral pair:  {len(chiral_filtered)}")
    print(f"    Speed single: {len(speed_single)}")
    print(f"    Chiral single:{len(chiral_single)}")
    print(f"    Reverse:      {len(reverse_single)}")
    print(f"RL total:         {len(rl_all)}")
    print(f"Needs CoT:        {len(needs_cot)}")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()