#!/usr/bin/env python
"""
下载所有领域的数据集（每个领域 200 条样本）

领域列表:
- news_en:    英文新闻 (CNN/DailyMail)
- news_zh:    中文新闻 (维基百科中文)
- academic:   学术语料 (维基百科科学类)
- code:       代码语料 (Cosmopedia)
- dialogue:   对话语料 (OpenAssistant/oasst1)
- literature: 文学语料 (维基百科人文类)

用法:
    python database/download_200_samples.py
"""

import os
import re
import json
from pathlib import Path
from typing import List, Optional

# 设置 HuggingFace 镜像（国内加速）
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from datasets import load_dataset


def preprocess_text(text: str, lang: str = "en", min_len: int = 30, max_len: int = 300) -> Optional[str]:
    """预处理文本：去特殊符号、控制长度"""
    text = text.strip()
    if not text:
        return None

    if lang == "en":
        text = re.sub(r'[^\w\s.,!?;:()"\']', '', text)
        words = text.split()
        if len(words) < min_len or len(words) > max_len:
            return None
        return ' '.join(words[:max_len]).strip()
    else:
        text = re.sub(r'[^\u4e00-\u9fff0-9.,!?;:()"\',，。！？、；：""（）【】《》]', '', text)
        chars = list(text)
        if len(chars) < min_len or len(chars) > max_len:
            return None
        return ''.join(chars[:max_len]).strip()


def save_corpus(texts: List[str], domain: str, output_dir: str = "database/corpus") -> bool:
    """保存语料到 JSONL 文件"""
    if not texts:
        print(f"  ⚠️ 警告：没有语料可保存")
        return False

    output_path = Path(output_dir) / domain / "texts.jsonl"
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with open(output_path, "w", encoding="utf-8") as f:
        for i, text in enumerate(texts):
            record = {
                "text": text,
                "metadata": {
                    "index": i,
                    "domain": domain,
                }
            }
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

    print(f"  ✅ 已保存至：{output_path} ({len(texts)} 条)")
    return True


# ==================== 各领域加载函数 ====================

def load_news_en(sample_num: int = 200) -> List[str]:
    """加载英文新闻语料 (CNN/DailyMail)"""
    corpus = []
    print("加载 CNN/DailyMail 数据集...")

    try:
        ds = load_dataset("cnn_dailymail", "3.0.0", split="train")

        count = 0
        for item in ds:
            if count >= 2000:
                break
            count += 1

            text = item["article"]
            processed = preprocess_text(text, lang="en", min_len=30, max_len=300)
            if processed and processed not in corpus:
                corpus.append(processed)
            if len(corpus) >= sample_num:
                break

        print(f"  成功提取 {len(corpus)} 条英文新闻")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


def load_news_zh(sample_num: int = 200) -> List[str]:
    """加载中文新闻语料 (维基百科中文)"""
    corpus = []
    print("加载维基百科中文数据集...")

    try:
        ds = load_dataset("wikimedia/wikipedia", "20231101.zh", split="train", streaming=True)

        count = 0
        for item in ds:
            if count >= 5000:
                break
            count += 1

            text = item.get("text", "")
            if text:
                processed = preprocess_text(text, lang="zh", min_len=50, max_len=200)
                if processed and processed not in corpus:
                    corpus.append(processed)
            if len(corpus) >= sample_num:
                break

        print(f"  成功提取 {len(corpus)} 条中文维基百科文章")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


def load_academic(sample_num: int = 200) -> List[str]:
    """加载学术语料 (维基百科科学类)"""
    corpus = []
    print("加载维基百科科学文章...")

    science_keywords = [
        "algorithm", "neural network", "machine learning", "computer science", "mathematics",
        "statistics", "artificial intelligence", "physics", "chemistry", "biology", "quantum",
        "deep learning", "optimization", "data structure", "complexity", "relativity",
        "thermodynamics", "electromagnetism", "genetics", "evolution", "ecosystem",
        "calculus", "topology", "algebra", "geometry", "probability", "enzyme", "protein",
        "dna", "rna", "cell", "molecule", "atom", "electron", "photon", "galaxy", "star"
    ]

    try:
        ds = load_dataset("wikimedia/wikipedia", "20231101.en", split="train", streaming=True)

        count = 0
        for item in ds:
            if count >= 50000:
                break
            count += 1

            text = item.get("text", "")
            title = item.get("title", "")

            # 从标题或内容前 500 字符中筛选
            if any(kw in title.lower() or kw in text[:500].lower() for kw in science_keywords):
                processed = preprocess_text(text, lang="en")
                if processed and processed not in corpus:
                    corpus.append(processed)

            if len(corpus) >= sample_num:
                break

        print(f"  成功提取 {len(corpus)} 篇维基百科科学文章")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


def load_code(sample_num: int = 200) -> List[str]:
    """加载代码语料 (Cosmopedia Openstax)"""
    corpus = []
    print("加载 Cosmopedia Openstax 代码数据集...")

    try:
        ds = load_dataset("HuggingFaceTB/cosmopedia", "openstax", split="train", streaming=True)

        count = 0
        for item in ds:
            if count >= 2000:
                break
            count += 1

            text = item.get("text", "")
            if text and ('import ' in text or 'def ' in text or 'class ' in text or 'print(' in text or 'for ' in text):
                lines = text.strip().split('\n')
                if 5 <= len(lines) <= 80:
                    code_clean = '\n'.join(line[:120] for line in lines[:40])
                    if code_clean and code_clean not in corpus:
                        corpus.append(code_clean)

            if len(corpus) >= sample_num:
                break

        print(f"  成功提取 {len(corpus)} 段代码")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


def load_dialogue(sample_num: int = 200) -> List[str]:
    """加载对话语料 (OpenAssistant/oasst1)"""
    corpus = []
    print("加载 OpenAssistant/oasst1 数据集...")

    try:
        ds = load_dataset("OpenAssistant/oasst1", split="train")

        messages = {}
        for item in ds:
            msg_id = item.get("message_id", "")
            parent_id = item.get("parent_id", "")
            text = item.get("text", "")
            role = item.get("role", "")

            if parent_id not in messages:
                messages[parent_id] = []
            messages[parent_id].append({
                "id": msg_id,
                "role": role,
                "text": text
            })

        for root_msgs in messages.values():
            if len(root_msgs) >= 2:
                dialogue_parts = []
                for msg in root_msgs[:4]:
                    dialogue_parts.append(f"{msg['role']}: {msg['text']}")

                dialogue = " ".join(dialogue_parts)
                processed = preprocess_text(dialogue, lang="en", min_len=30, max_len=300)
                if processed and processed not in corpus:
                    corpus.append(processed)

            if len(corpus) >= sample_num:
                break

        print(f"  成功提取 {len(corpus)} 段对话")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


def load_literature(sample_num: int = 200) -> List[str]:
    """加载文学语料 (维基百科人文类)"""
    corpus = []
    print("加载维基百科人文文章...")

    literature_keywords = [
        "novel", "poetry", "fiction", "literature", "shakespeare",
        "dickens", "austen", "tolkien", "hemingway", "orwell",
        "pride and prejudice", "great gatsby", "1984", "lord of the rings",
        "romeo and juliet", "hamlet", "macbeth", "jane eyre", "wuthering heights",
        "moby dick", "war and peace", "crime and punishment", "ulysses",
        "writer", "author", "book", "story", "character", "plot", "narrative",
        "prose", "verse", "sonnet", "drama", "tragedy", "comedy", "english literature"
    ]

    try:
        ds = load_dataset("wikimedia/wikipedia", "20231101.en", split="train", streaming=True)

        count = 0
        for item in ds:
            if count >= 50000:
                break
            count += 1

            text = item.get("text", "")
            title = item.get("title", "")

            # 从标题或内容前 500 字符中筛选
            if any(kw in title.lower() or kw in text[:500].lower() for kw in literature_keywords):
                processed = preprocess_text(text, lang="en")
                if processed and processed not in corpus:
                    corpus.append(processed)

            if len(corpus) >= sample_num:
                break

        print(f"  成功提取 {len(corpus)} 段文学文本")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


# ==================== 主流程 ====================

DOMAIN_LOADERS = {
    "news_en": load_news_en,
    "news_zh": load_news_zh,
    "academic": load_academic,
    "code": load_code,
    "dialogue": load_dialogue,
    "literature": load_literature,
}


def main():
    print("=" * 60)
    print("数据集下载 - 完整语料库（每个领域 200 条样本）")
    print("=" * 60)
    print()

    domains = list(DOMAIN_LOADERS.keys())
    sample_num = 200

    print(f"目标领域：{domains}")
    print(f"每个领域采样数：{sample_num}")
    print(f"输出目录：database/corpus/")
    print()

    results = {}

    for domain in domains:
        print(f"\n{'='*40}")
        print(f"[{domain}]")
        print(f"{'='*40}")

        loader = DOMAIN_LOADERS[domain]
        texts = loader(sample_num)

        if texts:
            success = save_corpus(texts, domain)
            results[domain] = len(texts) if success else 0
        else:
            print(f"  ⚠️ 警告：未能获取有效语料")
            results[domain] = 0

    print("\n" + "=" * 60)
    print("下载完成! 汇总:")
    print("=" * 60)

    for domain, count in results.items():
        status = "✅" if count > 0 else "❌"
        print(f"  {status} {domain}: {count} 条")

    total = sum(results.values())
    print(f"\n总计：{total} 条")

    print()
    print("下一步：提取模型表示")
    print("  python model/extract_representations.py")


if __name__ == "__main__":
    main()