#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
小狸 AI - 预训练数据生成器 (v5)
==================================
两阶段生成：

Phase 1 - 种子直出:
  5并发 API 直接生成完整多轮对话
  每API调用 = 1条高质量对话 (15-25轮, 2000-4000字)
  速率: ~600条/h, 17h出1万条种子

Phase 2 - 语义扩展:
  基于种子主题，API重新"编剧"出新对话
  每条种子 → 10个变体（换角度/深度/角色）
  速率: ~600条/h, ~7天扩展至10万+条

数据格式: JSONL, 每条: {"text":"...", "domain":"...", "concept":"..."}
"""
import sys, os, json, time, random, hashlib, argparse, re
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from api_correction import call_astron

# ═══════════════════════════════════════════════════════════
# 1. 角色定义
# ═══════════════════════════════════════════════════════════

TEACHER_SYSTEM = "你是一位AI导师，回答专业详细，包含具体数据和原理。每轮回答150-500字。"

# ═══════════════════════════════════════════════════════════
# 2. 深度学习领域图谱
# ═══════════════════════════════════════════════════════════

DOMAINS = {
    "量子物理": ["量子叠加态","量子纠缠","量子隧穿","海森堡测不准原理","波粒二象性","量子计算","量子场论","狄拉克方程"],
    "相对论": ["狭义相对论","质能方程","时空弯曲","引力波","黑洞物理","虫洞理论","暗物质","暗能量"],
    "人工智能": ["深度学习原理","Transformer","强化学习","大语言模型","扩散模型","多模态学习","注意力机制","模型量化"],
    "计算机科学": ["计算理论","操作系统","分布式系统","数据库","编译器","密码学","网络协议","计算机视觉"],
    "分子生物学": ["DNA复制","CRISPR","中心法则","蛋白质折叠","表观遗传学","RNA干扰","免疫系统","合成生物学"],
    "数学": ["微积分","傅里叶分析","微分方程","抽象代数","数论","拓扑学","概率论","图论"],
    "宇宙学": ["大爆炸","恒星演化","星系形成","宇宙微波背景","暗物质粒子","暴胀模型","引力波天文学","哈勃常数危机"],
    "神经科学": ["神经元","突触可塑性","大脑皮层","记忆机制","睡眠与意识","视觉系统","脑机接口","情感神经科学"],
    "化学": ["化学键","反应动力学","热力学","催化原理","光谱学","纳米化学","电化学","高分子化学"],
    "工程科技": ["航天工程","半导体","通信工程","机器人学","电力系统","材料工程","机械设计","核工程"],
    "经济学": ["一般均衡","博弈论","信息不对称","行为经济学","货币理论","计量经济学","拍卖理论","机制设计"],
    "哲学": ["先验哲学","权力意志","存在主义","语言游戏","证伪主义","范式理论","正义论","功利主义"],
    "天文学": ["赫罗图","超新星","中子星","系外行星","星系分类","脉冲星","恒星形成","宇宙射线"],
    "地球科学": ["板块构造","地震学","火山学","古生物学","气候系统","海洋环流","大气动力学","行星地质"],
    "生物技术": ["基因测序","PCR","冷冻电镜","核磁共振","光遗传学","类器官","纳米医学","基因敲除"],
}

ALL_DOMAINS = list(DOMAINS.keys())


# ═══════════════════════════════════════════════════════════
# 3. Phase 1: API直接生成种子
# ═══════════════════════════════════════════════════════════

def build_seed_prompt(domain: str, concept: str) -> str:
    return f"""请围绕【{domain}】领域的【{concept}】概念，生成一段师生深度对话。

要求：
1. 学生：从基础概念问起，逐步深入，包含追问、质疑、要求举例
2. 导师：每轮回答150-500字，专业准确，有数据支撑
3. 对话演进：基础定义→核心机制→实验证据→应用场景→前沿进展
4. 至少15轮对话

每轮格式（严格）：
学生：[问题]
导师：[回答]

不要额外说明，直接开始对话。"""


def generate_seed(domain: str, concept: str) -> dict:
    """一次API调用生成一条种子对话"""
    prompt = build_seed_prompt(domain, concept)
    text = call_astron(prompt, temperature=random.uniform(0.7, 0.9),
                       system_prompt=TEACHER_SYSTEM, max_tokens=8192)
    if not text:
        return None

    # 解析
    turns = []
    for line in text.strip().split('\n'):
        line = line.strip()
        if not line:
            continue
        if re.match(r'^学生[：:]', line):
            content = re.sub(r'^学生[：:]\s*', '', line)
            turns.append(('user', content))
        elif re.match(r'^导师[：:]', line):
            content = re.sub(r'^导师[：:]\s*', '', line)
            turns.append(('assistant', content))

    if len(turns) < 10:
        return None

    lines = []
    for role, content in turns:
        lines.append(f"<|im_start|>{role}\n{content}<|im_end|>")

    return {
        "text": '\n'.join(lines),
        "domain": domain,
        "concept": concept,
    }


# ═══════════════════════════════════════════════════════════
# 4. Phase 2: 语义扩展
# ═══════════════════════════════════════════════════════════

VARIATION_PROMPTS = [
    "用更通俗的方式重新讲一遍，多用日常类比",
    "从更高级的视角讲，假设学生已有基础，重点讲深层机制和数学原理",
    "换个角度，从历史发展脉络讲，突出关键实验和突破性发现",
    "从工程应用角度讲，重点讲实际怎么用、有什么技术瓶颈",
    "从批判性视角讲，有哪些争议、局限性、未解决的问题",
    "假设学生是高中水平，用最通俗的语言和大量类比来解释",
    "从跨学科视角讲，这个概念与其他领域有什么联系",
    "以问答竞赛的方式讲，学生提尖锐问题，导师用证据回击",
]


def generate_variation(seed: dict, seed_text: str, prompt_template: str) -> dict:
    """基于种子主题生成一个语义变体"""
    domain = seed["domain"]
    concept = seed["concept"]
    instruction = prompt_template

    prompt = f"""请围绕【{domain}】领域的【{concept}】概念，生成一段师生深度对话。

{instruction}

要求：
1. 至少15轮对话
2. 导师回答150-500字，专业详细
3. 对话要有递进：定义→机制→证据→应用→前沿

每轮格式（严格）：
学生：[问题]
导师：[回答]

直接开始对话，不要额外说明。"""

    text = call_astron(prompt, temperature=random.uniform(0.75, 0.95),
                       system_prompt=TEACHER_SYSTEM, max_tokens=8192)
    if not text:
        return None

    turns = []
    for line in text.strip().split('\n'):
        line = line.strip()
        if not line:
            continue
        if re.match(r'^学生[：:]', line):
            content = re.sub(r'^学生[：:]\s*', '', line)
            turns.append(('user', content))
        elif re.match(r'^导师[：:]', line):
            content = re.sub(r'^导师[：:]\s*', '', line)
            turns.append(('assistant', content))

    if len(turns) < 10:
        return None

    lines = []
    for role, content in turns:
        lines.append(f"<|im_start|>{role}\n{content}<|im_end|>")

    return {
        "text": '\n'.join(lines),
        "domain": domain,
        "concept": concept,
        "variation": prompt_template[:30],
    }


# ═══════════════════════════════════════════════════════════
# 5. 数据管理
# ═══════════════════════════════════════════════════════════

class DataStore:
    def __init__(self, save_dir: str):
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(parents=True, exist_ok=True)
        self.output_path = self.save_dir / 'dialogues.jsonl'
        self.seed_path = self.save_dir / 'seeds.jsonl'

        self.existing_hashes = set()
        self.count = 0
        self.load()

    def load(self):
        for path in [self.output_path, self.seed_path]:
            if path.exists():
                with open(path, 'r') as f:
                    for line in f:
                        line = line.strip()
                        if line:
                            try:
                                item = json.loads(line)
                                h = hashlib.md5(item['text'][:100].encode()).hexdigest()
                                self.existing_hashes.add(h)
                                self.count += 1
                            except:
                                pass
        print(f"📂 已有: {self.count:,} 条")

    def add_seeds(self, items: list) -> int:
        if not items:
            return 0
        new = []
        for item in items:
            h = hashlib.md5(item['text'][:100].encode()).hexdigest()
            if h not in self.existing_hashes:
                self.existing_hashes.add(h)
                new.append(item)
        if not new:
            return 0
        with open(self.seed_path, 'a') as f:
            for item in new:
                f.write(json.dumps(item, ensure_ascii=False) + '\n')
                self.count += 1
        return len(new)

    def add_variations(self, items: list) -> int:
        if not items:
            return 0
        new = []
        for item in items:
            h = hashlib.md5(item['text'][:100].encode()).hexdigest()
            if h not in self.existing_hashes:
                self.existing_hashes.add(h)
                new.append(item)
        if not new:
            return 0
        with open(self.output_path, 'a') as f:
            for item in new:
                f.write(json.dumps(item, ensure_ascii=False) + '\n')
                self.count += 1
        return len(new)

    def get_seeds(self):
        """读取所有种子"""
        seeds = []
        if self.seed_path.exists():
            with open(self.seed_path, 'r') as f:
                for line in f:
                    line = line.strip()
                    if line:
                        seeds.append(json.loads(line))
        return seeds

    def total_count(self):
        return self.count


# ═══════════════════════════════════════════════════════════
# 6. 质量评分
# ═══════════════════════════════════════════════════════════

def score(text: str) -> float:
    turns = (text.count('<|im_start|>user') + text.count('<|im_start|>assistant')) // 2
    chars = len(text)
    s = 0.0
    if turns >= 20: s += 0.20
    elif turns >= 12: s += 0.15
    elif turns >= 8: s += 0.10
    else: s += 0.05
    if chars >= 4000: s += 0.20
    elif chars >= 2000: s += 0.15
    elif chars >= 1200: s += 0.10
    else: s += 0.05
    avg = chars / max(turns, 1)
    if avg >= 200: s += 0.15
    elif avg >= 100: s += 0.10
    uc = text.count('<|im_start|>user')
    ac = text.count('<|im_start|>assistant')
    ratio = min(uc, ac) / max(uc, ac, 1)
    if ratio >= 0.8: s += 0.10
    unique = len(set(text))
    if unique >= 150: s += 0.10
    return max(0, min(1, s))


# ═══════════════════════════════════════════════════════════
# 7. 卷统计
# ═══════════════════════════════════════════════════════════

def print_stats(store: DataStore, t_start: float, phase: str, target: int):
    elapsed = max(time.time() - t_start, 1)
    rate = store.count / elapsed * 3600
    eta_h = (target - store.count) / max(rate, 0.01)
    
    seed_count = 0
    if store.seed_path.exists():
        with open(store.seed_path) as f:
            seed_count = sum(1 for _ in f if _.strip())
    
    tokens_est = store.count * 6000
    
    print(f"  [{time.strftime('%H:%M:%S')}] "
          f"Phase {phase} | "
          f"累计: {store.count:,}/{target:,} "
          f"({store.count/max(target,1)*100:.1f}%) | "
          f"速率: {rate:.0f}条/h | "
          f"种子: {seed_count} | "
          f"tokens: {tokens_est/1e9:.2f}B | "
          f"预计: {eta_h/24:.1f}d", flush=True)


# ═══════════════════════════════════════════════════════════
# 8. 主函数
# ═══════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--target", type=int, default=100000, help="目标条数")
    parser.add_argument("--phase", type=int, default=1, choices=[1, 2],
                        help="1=生成种子, 2=语义扩展")
    parser.add_argument("--concurrency", type=int, default=5, help="API并发数")
    args = parser.parse_args()

    save_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'pretrain_data')
    store = DataStore(save_dir)

    target = args.target
    if store.count >= target:
        print(f"✅ 已达目标 {target:,} 条")
        return

    print(f"\n{'='*50}")
    print(f"  Phase {args.phase}: {'种子生成' if args.phase == 1 else '语义扩展'}")
    print(f"  目标: {target:,} 条")
    print(f"  并发: {args.concurrency}")
    print(f"  当前: {store.count:,}")
    print(f"{'='*50}\n")

    t_start = time.time()

    if args.phase == 1:
        # Phase 1: 纯API种子生成
        seeds = store.get_seeds()
        # 从种子中的(domain, concept)对去重
        used = set()
        for s in seeds:
            used.add((s.get("domain",""), s.get("concept","")))
        print(f"  已有 {len(used)} 个不同(领域,概念)种子\n")

        while store.count < target:
            batch_items = []
            tasks = []
            for _ in range(args.concurrency):
                # 选一个未用过的(领域,概念)
                for attempt in range(20):
                    d = random.choice(ALL_DOMAINS)
                    c = random.choice(DOMAINS[d])
                    if (d, c) not in used or attempt > 15:
                        break
                used.add((d, c))
                tasks.append((d, c))

            with ThreadPoolExecutor(max_workers=args.concurrency) as ex:
                futures = {ex.submit(generate_seed, d, c): (d, c) for d, c in tasks}
                for f in as_completed(futures):
                    result = f.result()
                    if result and score(result["text"]) >= 0.5:
                        batch_items.append(result)

            if batch_items:
                added = store.add_seeds(batch_items)
                n = len(batch_items)
                print(f"  +{n}条种子 | 领域: {batch_items[0]['domain']}, "
                      f"概念: {batch_items[0]['concept']}...", flush=True)

            print_stats(store, t_start, "1", target)

            if store.count >= target:
                break

    elif args.phase == 2:
        # Phase 2: 语义扩展
        seeds = store.get_seeds()
        if not seeds:
            print("❌ 没有种子数据，先跑 Phase 1")
            return
        print(f"  种子库: {len(seeds)} 条\n")

        used_seeds = set()
        expansion_tasks = []

        while store.count < target:
            if not expansion_tasks:
                # 把所有种子打乱，生成扩展任务
                random.shuffle(seeds)
                for s in seeds:
                    if s["text"] in used_seeds:
                        continue
                    used_seeds.add(s["text"])
                    for vp in VARIATION_PROMPTS:
                        expansion_tasks.append((s, vp))
                print(f"  生成 {len(expansion_tasks)} 个扩展任务\n")

            # 取一批任务
            batch = []
            while len(batch) < args.concurrency and expansion_tasks:
                batch.append(expansion_tasks.pop(0))

            if not batch:
                break

            with ThreadPoolExecutor(max_workers=args.concurrency) as ex:
                futures = {ex.submit(generate_variation, s, s["text"], vp): (s, vp)
                          for s, vp in batch}
                for f in as_completed(futures):
                    result = f.result()
                    if result and score(result["text"]) >= 0.5:
                        store.add_variations([result])

            # 报告
            seed_count = 0
            if store.seed_path.exists():
                with open(store.seed_path) as f:
                    seed_count = sum(1 for _ in f if _.strip())
            elapsed = max(time.time() - t_start, 1)
            rate = store.count / elapsed * 3600
            eta_h = (target - store.count) / max(rate, 0.01) if rate > 0 else 0
            tokens_est = store.count * 6000
            done_tasks = max(1, len(used_seeds) * len(VARIATION_PROMPTS) - len(expansion_tasks))
            print(f"  [{time.strftime('%H:%M:%S')}] "
                  f"累计: {store.count:,}/{target:,} "
                  f"({store.count/max(target,1)*100:.1f}%) | "
                  f"速率: {rate:.0f}条/h | "
                  f"种子: {seed_count} | "
                  f"扩展任务进度: {done_tasks} | "
                  f"tokens: {tokens_est/1e9:.2f}B | "
                  f"预计: {eta_h/24:.1f}d", flush=True)

            if store.count >= target:
                break

    # 最终报告
    total_time = time.time() - t_start
    final_rate = store.count / max(total_time, 1) * 3600
    tokens = store.count * 6000
    path = store.output_path if args.phase == 2 else store.seed_path
    size_mb = os.path.getsize(path) / 1024 / 1024 if path.exists() else 0

    print(f"\n{'='*50}")
    print(f"  ✅ Phase {args.phase} 完成!")
    print(f"  总条数:     {store.count:,}")
    print(f"  文件:       {path}")
    print(f"  大小:       {size_mb:.1f}MB")
    print(f"  估算tokens: {tokens/1e9:.2f}B")
    print(f"  耗时:       {total_time/3600:.1f}h")
    print(f"  平均速率:   {final_rate:.0f}条/h")
    print(f"{'='*50}")

    # 样例
    print("\n📊 样例:")
    with open(path) as f:
        for i, line in enumerate(f):
            if i >= 2:
                break
            if line.strip():
                item = json.loads(line.strip())
                text = item["text"]
                turns = text.count("<|im_start|>") // 2
                print(f"\n  条目{i+1}: {turns}轮 | {len(text):,}字 | {item.get('domain','')}")
                lines = text.split("\n")[:4]
                for l in lines:
                    if l.strip():
                        role = "👤" if "user" in l else "🧑‍🏫"
                        c = l.split("<|im_start|>")[1].split("<|im_end|>")[0].strip()
                        print(f"    {role} {c[:100]}")
                print(f"    ...（{turns}轮）")


if __name__ == "__main__":
    main()