import json
import random
import os

INPUT_FILE = r"D:\双生天使的怀抱\爱的数据集\sephirot_final_run\train_all.jsonl"
OUTPUT_DIR = r"D:\双生天使的怀抱\2026-05-24-16-54-24\sft_data"
OUTPUT_TRAIN = os.path.join(OUTPUT_DIR, "train.jsonl")
OUTPUT_VAL = os.path.join(OUTPUT_DIR, "val.jsonl")
VAL_RATIO = 0.1

records = []
skipped = 0
total = 0

with open(INPUT_FILE, 'r', encoding='utf-8') as f:
    for line in f:
        total += 1
        line = line.strip()
        if not line:
            continue
        try:
            data = json.loads(line)
        except json.JSONDecodeError as e:
            print(f"JSON解析失败 第{total}行: {e}")
            skipped += 1
            continue
        
        query = data.get('input', {}).get('user_query', '')
        cot = data.get('chain_of_thought', {})
        
        # 方案B: 完整推理链（推荐，保留16质点全部过程）
        if isinstance(cot, dict):
            parts = []
            for key in ['D1_kether','H1_ego','H2_superego','H3_true_self','H4_logic',
                        'H5_empathy','H6_happiness','D2_binah_chesed','D3_chokmah_gevurah',
                        'D5_tiferet','D6_netzach','D7_hod','D4_yesod','D8_victory',
                        'D10_kingdom']:
                val = cot.get(key, '')
                if isinstance(val, dict):
                    # 尝试多个可能的输出字段
                    for sub_key in ['final_output', 'analysis', 'synthesis', 
                                     'emotional_expression', 'logic_empathy_union',
                                     'response', 'output']:
                        if sub_key in val and val[sub_key]:
                            val = val[sub_key]
                            break
                    else:
                        val = str(val) if val else ''
                if val and isinstance(val, str) and len(val.strip()) > 0:
                    parts.append(f"[{key}]: {val}")
            output = '\n\n'.join(parts)
        elif isinstance(cot, str):
            output = cot
        else:
            output = str(cot)
        
        if query and output and len(query.strip()) > 0 and len(output.strip()) > 10:
            records.append({
                "conversations": [
                    {"from": "human", "value": query},
                    {"from": "gpt", "value": output}
                ]
            })
        else:
            skipped += 1

print(f"总行数: {total}, 有效: {len(records)}, 跳过: {skipped}")

# 打乱并分割
random.seed(42)
random.shuffle(records)
val_size = max(int(len(records) * VAL_RATIO), 100)
val_data = records[:val_size]
train_data = records[val_size:]

os.makedirs(OUTPUT_DIR, exist_ok=True)

with open(OUTPUT_TRAIN, 'w', encoding='utf-8') as f:
    for r in train_data:
        f.write(json.dumps(r, ensure_ascii=False) + '\n')

with open(OUTPUT_VAL, 'w', encoding='utf-8') as f:
    for r in val_data:
        f.write(json.dumps(r, ensure_ascii=False) + '\n')

print(f"转换完成! 训练集: {len(train_data)}条, 验证集: {len(val_data)}条")
print(f"训练集文件: {OUTPUT_TRAIN}")
print(f"验证集文件: {OUTPUT_VAL}")

# 统计一下平均长度
if train_data:
    lens = [len(str(r['conversations'][1]['value'])) for r in train_data]
    print(f"平均输出长度: {sum(lens)//len(lens)} 字符")
    print(f"最大输出长度: {max(lens)} 字符")
    print(f"最小输出长度: {min(lens)} 字符")