#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
高温多样化生成 — 把训练集全部问题用 Astron 高温度重答一次
温度=1.2，追求多样化回答，丰富训练数据
"""
import sys, os, json, time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1)

from api_correction import batch_generate

MODEL_DIR = os.path.dirname(os.path.abspath(__file__))
TEMPERATURE = 1.0       # 高温度(上限1.0)，产出多样化回答
BATCH_SIZE = 500        # 每批 500 条，分批保存防丢
MAX_WORKERS = 5         # Astron 5并发

def extract_questions(data):
    """从训练数据中提取所有问题"""
    questions = []
    for item in data:
        if isinstance(item, dict):
            q = item.get('question', '') or item.get('input', '')
            if q:
                questions.append(q)
    return questions

def main():
    # 1. 加载训练数据
    train_path = os.path.join(MODEL_DIR, 'xiaoli_train_v2.json')
    print(f"加载训练数据: {train_path}")
    with open(train_path, encoding='utf-8') as f:
        train_data = json.load(f)
    
    questions = extract_questions(train_data)
    total = len(questions)
    print(f"提取到 {total} 条问题，温度={TEMPERATURE}，{MAX_WORKERS}并发")
    
    # 2. 创建输出目录
    out_dir = os.path.join(MODEL_DIR, 'high_temp_outputs')
    os.makedirs(out_dir, exist_ok=True)
    
    # 3. 分批处理
    all_results = []
    existing = set()  # 去重用
    
    # 检查是否有已完成的中间结果
    interim_files = sorted([f for f in os.listdir(out_dir) if f.startswith('batch_') and f.endswith('.json')])
    if interim_files:
        for fname in interim_files:
            with open(os.path.join(out_dir, fname), encoding='utf-8') as f:
                batch_data = json.load(f)
            for r in batch_data:
                key = (r['question'][:100], r['answer'][:100])
                if key not in existing:
                    existing.add(key)
                    all_results.append(r)
        print(f"发现已有中间结果: {len(all_results)} 条（来自 {len(interim_files)} 个文件）")
        # 跳过已处理的 question
        processed_questions = set(r['question'] for r in all_results)
        remaining = [q for q in questions if q not in processed_questions]
        print(f"已处理 {len(processed_questions)} 条，剩余 {len(remaining)} 条")
        questions = remaining
    
    total_batches = (len(questions) + BATCH_SIZE - 1) // BATCH_SIZE
    t_start = time.time()
    
    for batch_idx in range(total_batches):
        batch_qs = questions[batch_idx * BATCH_SIZE : (batch_idx + 1) * BATCH_SIZE]
        print(f"\n{'='*60}")
        print(f"批次 {batch_idx + 1}/{total_batches} ({len(batch_qs)}条, 温度{TEMPERATURE})")
        print(f"{'='*60}")
        
        t0 = time.time()
        results = batch_generate(batch_qs, temperature=TEMPERATURE, max_workers=MAX_WORKERS)
        elapsed = time.time() - t0
        
        # 标注入参温度
        for r in results:
            r['temperature'] = TEMPERATURE
        
        # 去重后加入总集
        new_in_batch = 0
        for r in results:
            key = (r['question'][:100], r['answer'][:100])
            if key not in existing:
                existing.add(key)
                all_results.append(r)
                new_in_batch += 1
        
        # 立即保存本批结果
        batch_file = os.path.join(out_dir, f'batch_{batch_idx + 1:04d}.json')
        with open(batch_file, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        
        # 追加保存总进度
        progress_file = os.path.join(out_dir, 'all_progress.json')
        with open(progress_file, 'w', encoding='utf-8') as f:
            json.dump(all_results, f, ensure_ascii=False, indent=2)
        
        total_elapsed = time.time() - t_start
        rate = len(all_results) / max(total_elapsed, 1) * 60
        eta_remaining = (total - len(all_results)) / max(rate, 0.1) / 60 if rate > 0 else 0
        print(f"  本批新增 {new_in_batch} 条 (去重后) | 累计 {len(all_results)} 条 | "
              f"{rate:.0f}条/分 | 已用 {total_elapsed/60:.1f}分 | 预计剩余 {eta_remaining:.1f}分")
    
    # 4. 去重统计
    print(f"\n{'='*60}")
    print(f"全部生成完成!")
    print(f"生成结果: {len(all_results)} 条 (去重后)")
    print(f"用时: {(time.time() - t_start)/60:.1f} 分钟")
    
    # 5. 保存最终去重结果（独立存放，不追加到训练集）
    final_file = os.path.join('/data/user/work', 'high_temp_final.json')
    with open(final_file, 'w', encoding='utf-8') as f:
        json.dump(all_results, f, ensure_ascii=False, indent=2)
    print(f"\n最终结果已保存: {final_file}")
    print(f"如需使用，可手动合并到训练集")
    
    print(f"\n{'='*60}")
    print(f"全部完成! 新增 {len(all_results)} 条多样化回答")
    print(f"{'='*60}")

if __name__ == '__main__':
    main()