#!/usr/bin/env python
"""
下载 50 条样本用于 MVP 实验测试

下载以下 2 个核心领域的语料：
- news_en:  英文新闻 (CNN/DailyMail)
- academic: 学术论文摘要 (arXiv)

每个领域 50 条样本，保存到 database/corpus/
"""

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 = 50, max_len: int = 200) -> Optional[str]:
    """
    预处理文本：去特殊符号、控制长度

    Args:
        text: 原始文本
        lang: 语言（en/zh）
        min_len: 最小长度（英文=词数，中文=字符数）
        max_len: 最大长度

    Returns:
        预处理后的文本，None 表示不符合长度被过滤
    """
    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
        text = ' '.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
        text = ''.join(chars[:max_len]).strip()

    return text if text else None


def load_news_en(sample_num: int = 50) -> 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 >= 500:  # 最多遍历 500 条
                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_academic(sample_num: int = 50) -> List[str]:
    """
    加载学术论文语料

    备用方案：如果 arXiv 无法加载，使用维基百科科学类文章代替
    """
    corpus = []
    print("加载学术语料...")

    # 方案 1：尝试 arXiv
    try:
        print("  尝试加载 arXiv 数据集...")
        ds = load_dataset("CShorten/arxiv-minimal", split="train", streaming=True)

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

            abstract = item.get("abstract", "")
            categories = item.get("categories", [])

            if any(cat in categories for cat in ["cs.CL", "cs.LG", "stat.ML", "cs.AI"]):
                processed = preprocess_text(abstract, lang="en", min_len=30, max_len=300)
                if processed and processed not in corpus:
                    corpus.append(processed)

            if len(corpus) >= sample_num:
                break

        if corpus:
            print(f"  成功提取 {len(corpus)} 篇 arXiv 论文摘要")
            return corpus
    except Exception as e:
        print(f"  arXiv 加载失败：{e}")

    # 方案 2：使用维基百科（更可靠）
    print("  使用备用方案：维基百科科学文章...")
    try:
        ds = load_dataset("wikimedia/wikipedia", "20231101.en", split="train", streaming=True)

        # 科学相关的关键词
        science_keywords = ["algorithm", "neural network", "machine learning",
                           "computer science", "mathematics", "statistics",
                           "artificial intelligence", "data science"]

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

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

            # 筛选科学相关文章
            if any(kw in title.lower() or kw in text[:500].lower() for kw in science_keywords):
                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 save_corpus(texts: List[str], domain: str, output_dir: str = "database/corpus") -> None:
    """保存语料到 JSONL 文件"""
    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}")


def main():
    print("=" * 50)
    print("MVP 实验数据集下载（50 条样本）")
    print("=" * 50)
    print()

    # 下载英文新闻
    print("\n[news_en]")
    news_texts = load_news_en(50)
    if news_texts:
        save_corpus(news_texts, "news_en")
    else:
        print("  警告：未能获取有效语料")

    # 下载学术论文
    print("\n[academic]")
    academic_texts = load_academic(50)
    if academic_texts:
        save_corpus(academic_texts, "academic")
    else:
        print("  警告：未能获取有效语料")

    print("\n" + "=" * 50)
    print("下载完成!")
    print("=" * 50)
    print()
    print("输出文件:")
    print("  - database/corpus/news_en/texts.jsonl")
    print("  - database/corpus/academic/texts.jsonl")
    print()
    print("下一步：提取模型表示")
    print("  python model/extract_representations.py")


if __name__ == "__main__":
    main()