#!/usr/bin/env python
"""
下载 GLUE Benchmark 完整数据集（不采样）

GLUE 子集说明:
- MNLI: 多段落自然语言推理 (392,702 训练样本)
- QNLI: 问答自然语言推理 (108,431 训练样本)
- SST-2: 情感分析 (67,349 训练样本)
- STS-B: 语义相似度 (8,628 训练样本)
- CoLA: 语法判断 (8,551 训练样本)
- MRPC: 释义识别 (3,668 训练样本)
- RTE: 文本蕴含 (2,490 训练样本)

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

import os
import json
from pathlib import Path
from typing import List, Optional
from datasets import load_dataset

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


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

    # 清理特殊字符
    text = re.sub(r'[^\w\s.,!?;:()"\']', '', text)
    words = text.split()

    if len(words) < min_len:
        return None

    # 截断到最大长度
    text = ' '.join(words[:max_len]).strip()
    return text if text else None


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,
                    "source": "GLUE Benchmark"
                }
            }
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

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


def load_mnli() -> List[str]:
    """
    MNLI (Multi-Genre Natural Language Inference)
    多段落自然语言推理 - 最大子集

    格式：premise, hypothesis, label
    使用 premise 字段作为语料
    """
    corpus = []
    print("加载 MNLI (Multi-Genre Natural Language Inference)...")
    print("  预期规模：~392K 训练样本")

    try:
        ds = load_dataset("glue", "mnli", split="train")

        for item in ds:
            premise = item.get("premise", "")
            hypothesis = item.get("hypothesis", "")

            # 合并 premise 和 hypothesis 作为完整语料
            full_text = f"{premise} {hypothesis}"
            processed = preprocess_text(full_text, lang="en", min_len=20, max_len=300)

            if processed and processed not in corpus:
                corpus.append(processed)

        print(f"  ✅ 成功提取 {len(corpus):,} 条 MNLI 语料")
    except Exception as e:
        print(f"  ❌ 加载失败：{e}")

    return corpus


def load_qnli() -> List[str]:
    """
    QNLI (Question Natural Language Inference)
    问答自然语言推理

    格式：question, sentence, label
    使用 sentence 字段作为语料（维基百科句子）
    """
    corpus = []
    print("加载 QNLI (Question NLI)...")
    print("  预期规模：~108K 训练样本")

    try:
        ds = load_dataset("glue", "qnli", split="train")

        for item in ds:
            sentence = item.get("sentence", "")
            question = item.get("question", "")

            # 合并句子和问题
            full_text = f"{sentence} {question}"
            processed = preprocess_text(full_text, lang="en", min_len=20, max_len=300)

            if processed and processed not in corpus:
                corpus.append(processed)

        print(f"  ✅ 成功提取 {len(corpus):,} 条 QNLI 语料")
    except Exception as e:
        print(f"  ❌ 加载失败：{e}")

    return corpus


def load_sst2() -> List[str]:
    """
    SST-2 (Stanford Sentiment Treebank)
    情感分析

    格式：sentence, label
    使用完整句子
    """
    corpus = []
    print("加载 SST-2 (Stanford Sentiment Treebank)...")
    print("  预期规模：~67K 训练样本")

    try:
        ds = load_dataset("glue", "sst2", split="train")

        for item in ds:
            sentence = item.get("sentence", "")
            processed = preprocess_text(sentence, lang="en", min_len=10, max_len=200)

            if processed and processed not in corpus:
                corpus.append(processed)

        print(f"  ✅ 成功提取 {len(corpus):,} 条 SST-2 语料")
    except Exception as e:
        print(f"  ❌ 加载失败：{e}")

    return corpus


def load_stsb() -> List[str]:
    """
    STS-B (Semantic Textual Similarity Benchmark)
    语义相似度

    格式：sentence1, sentence2, score
    合并两个句子作为语料
    """
    corpus = []
    print("加载 STS-B (Semantic Textual Similarity)...")
    print("  预期规模：~8.6K 训练样本")

    try:
        ds = load_dataset("glue", "stsb", split="train")

        for item in ds:
            sentence1 = item.get("sentence1", "")
            sentence2 = item.get("sentence2", "")

            full_text = f"{sentence1} {sentence2}"
            processed = preprocess_text(full_text, lang="en", min_len=15, max_len=250)

            if processed and processed not in corpus:
                corpus.append(processed)

        print(f"  ✅ 成功提取 {len(corpus):,} 条 STS-B 语料")
    except Exception as e:
        print(f"  ❌ 加载失败：{e}")

    return corpus


def load_cola() -> List[str]:
    """
    CoLA (Corpus of Linguistic Acceptability)
    语法判断

    格式：sentence, label
    """
    corpus = []
    print("加载 CoLA (Corpus of Linguistic Acceptability)...")
    print("  预期规模：~8.5K 训练样本")

    try:
        ds = load_dataset("glue", "cola", split="train")

        for item in ds:
            sentence = item.get("sentence", "")
            processed = preprocess_text(sentence, lang="en", min_len=10, max_len=200)

            if processed and processed not in corpus:
                corpus.append(processed)

        print(f"  ✅ 成功提取 {len(corpus):,} 条 CoLA 语料")
    except Exception as e:
        print(f"  ❌ 加载失败：{e}")

    return corpus


def load_mrpc() -> List[str]:
    """
    MRPC (Microsoft Research Paraphrase Corpus)
    释义识别

    格式：sentence1, sentence2, label
    """
    corpus = []
    print("加载 MRPC (Paraphrase Corpus)...")
    print("  预期规模：~3.6K 训练样本")

    try:
        ds = load_dataset("glue", "mrpc", split="train")

        for item in ds:
            sentence1 = item.get("sentence1", "")
            sentence2 = item.get("sentence2", "")

            full_text = f"{sentence1} {sentence2}"
            processed = preprocess_text(full_text, lang="en", min_len=15, max_len=250)

            if processed and processed not in corpus:
                corpus.append(processed)

        print(f"  ✅ 成功提取 {len(corpus):,} 条 MRPC 语料")
    except Exception as e:
        print(f"  ❌ 加载失败：{e}")

    return corpus


def load_rte() -> List[str]:
    """
    RTE (Recognizing Textual Entailment)
    文本蕴含

    格式：sentence1, sentence2, label
    """
    corpus = []
    print("加载 RTE (Textual Entailment)...")
    print("  预期规模：~2.5K 训练样本")

    try:
        ds = load_dataset("glue", "rte", split="train")

        for item in ds:
            sentence1 = item.get("sentence1", "")
            sentence2 = item.get("sentence2", "")

            full_text = f"{sentence1} {sentence2}"
            processed = preprocess_text(full_text, lang="en", min_len=15, max_len=250)

            if processed and processed not in corpus:
                corpus.append(processed)

        print(f"  ✅ 成功提取 {len(corpus):,} 条 RTE 语料")
    except Exception as e:
        print(f"  ❌ 加载失败：{e}")

    return corpus


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

DOMAIN_LOADERS = {
    "glue_mnli": load_mnli,      # ~392K - 多段落推理
    "glue_qnli": load_qnli,      # ~108K - 问答推理
    "glue_sst2": load_sst2,      # ~67K  - 情感分析
    "glue_stsb": load_stsb,      # ~8.6K - 语义相似度
    "glue_cola": load_cola,      # ~8.5K - 语法判断
    "glue_mrpc": load_mrpc,      # ~3.6K - 释义识别
    "glue_rte": load_rte,        # ~2.5K - 文本蕴含
}


def main():
    print("=" * 70)
    print("GLUE Benchmark - 完整数据集下载（不采样）")
    print("=" * 70)
    print()
    print("GLUE 是 NLP 领域公认的标准评测基准，包含 7 个子任务")
    print("数据将用于模型表示提取和谱分析实验")
    print()

    domains = list(DOMAIN_LOADERS.keys())

    print(f"可用子集：{len(domains)} 个")
    for domain in domains:
        print(f"  - {domain}")
    print()
    print(f"输出目录：database/corpus/")
    print()

    results = {}

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

        loader = DOMAIN_LOADERS[domain]
        texts = loader()

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

    # 汇总
    print("\n" + "=" * 70)
    print("下载完成！汇总:")
    print("=" * 70)

    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("  1. 查看数据：head database/corpus/glue_mnli/texts.jsonl")
    print("  2. 提取表示：python model/extract_representations.py")
    print("  3. 运行实验：python experiments/verify.py")


if __name__ == "__main__":
    main()