#!/usr/bin/env python
"""
方向 9 实验运行脚本 - 整合所有实验流程

实验列表:
- E1: 基础谱收敛验证
- E2: 跨域谱分层验证
- E3: 能力涌现谱突变检测
- E4: 对齐即谱工程
- LCP: 逐层谱收敛轮廓

用法:
    # 运行所有实验
    python experiments/run_direction9_experiments.py --all

    # 运行单个实验
    python experiments/run_direction9_experiments.py --experiment E1

    # 生成可视化
    python experiments/run_direction9_experiments.py --visualize
"""

import os
import sys
import json
import argparse
from pathlib import Path
from datetime import datetime

# 设置 HF 镜像
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

# 项目根目录
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))


def check_prerequisites():
    """检查实验前置条件"""
    print("=" * 60)
    print("检查实验前置条件")
    print("=" * 60)

    issues = []

    # 检查模型
    print("\n[模型检查]")
    required_models = {
        "llama3.2-3b-instruct": "model/weights/llama3.2-3b-instruct",
        "Mistral-7B-v0.3": "model/weights/Mistral-7B-v0.3",
        "llama3-8b": "model/weights/llama3-8b",
        "Qwen2.5-7B-Instruct": "model/weights/Qwen2.5-7B-Instruct",
        "gemma-2-9b-it": "model/weights/gemma-2-9b-it"
    }

    for name, path in required_models.items():
        model_path = PROJECT_ROOT / path
        config_file = model_path / "config.json"
        if config_file.exists():
            print(f"  ✓ {name}")
        else:
            print(f"  ✗ {name} - 缺失")
            issues.append(f"模型缺失：{name}")

    # 检查数据集
    print("\n[数据集检查]")
    from database.corpus import CorpusLoader
    loader = CorpusLoader(PROJECT_ROOT / "database/corpus")

    required_datasets = ["gsm8k", "math", "humaneval", "alpaca", "bigbench", "flores200"]
    available = loader.list_available_domains()

    for ds in required_datasets:
        if ds in available:
            stats = loader.get_corpus_stats(ds)
            print(f"  ✓ {ds}: {stats['num_texts']} 条")
        else:
            print(f"  ✗ {ds} - 缺失")
            issues.append(f"数据集缺失：{ds}")

    # 检查依赖
    print("\n[依赖检查]")
    try:
        import ruptures
        print("  ✓ ruptures")
    except ImportError:
        print("  ✗ ruptures - 未安装")
        issues.append("依赖缺失：ruptures")

    try:
        import seaborn
        print("  ✓ seaborn")
    except ImportError:
        print("  ✗ seaborn - 未安装")
        issues.append("依赖缺失：seaborn")

    # GPU 检查
    print("\n[GPU 检查]")
    import torch
    if torch.cuda.is_available():
        print(f"  ✓ CUDA {torch.version.cuda}")
        print(f"  ✓ GPU: {torch.cuda.get_device_name(0)}")
        print(f"  ✓ 显存：{torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
    else:
        print("  ✗ CUDA 不可用 - 将使用 CPU (慢)")

    print("\n" + "=" * 60)
    if issues:
        print(f"发现 {len(issues)} 个问题:")
        for issue in issues:
            print(f"  - {issue}")
        print("\n建议先解决上述问题再运行实验")
        return False
    else:
        print("✓ 所有前置条件满足")
        return True


def run_experiment_e1(output_dir: str):
    """运行 E1 实验"""
    from experiments.cross_model_convergence import CrossModelConvergence, ExperimentConfig

    config = ExperimentConfig(
        models=[
            "llama3.2-3b-instruct",
            "Mistral-7B-v0.3",
            "llama3-8b",
            "Qwen2.5-7B-Instruct",
            "gemma-2-9b-it"
        ],
        model_paths={
            "llama3.2-3b-instruct": "model/weights/llama3.2-3b-instruct",
            "Mistral-7B-v0.3": "model/weights/Mistral-7B-v0.3",
            "llama3-8b": "model/weights/llama3-8b",
            "Qwen2.5-7B-Instruct": "model/weights/Qwen2.5-7B-Instruct",
            "gemma-2-9b-it": "model/weights/gemma-2-9b-it"
        },
        datasets=["gsm8k", "math", "humaneval", "alpaca", "flores200"],
        output_dir=output_dir,
        batch_size=4,  # 降低 batch size 以防止 OOM
        max_length=512,
        pooling="mean"
    )

    analyzer = CrossModelConvergence(config)
    return analyzer.run_experiment_e1()


def run_experiment_e2(output_dir: str):
    """运行 E2 实验 - 跨域谱分层验证"""
    from experiments.cross_model_convergence import CrossModelConvergence, ExperimentConfig

    # E2 复用 E1 的配置和数据
    config = ExperimentConfig(
        models=[
            "llama3.2-3b-instruct",
            "Mistral-7B-v0.3",
            "llama3-8b",
            "Qwen2.5-7B-Instruct",
            "gemma-2-9b-it"
        ],
        model_paths={
            "llama3.2-3b-instruct": "model/weights/llama3.2-3b-instruct",
            "Mistral-7B-v0.3": "model/weights/Mistral-7B-v0.3",
            "llama3-8b": "model/weights/llama3-8b",
            "Qwen2.5-7B-Instruct": "model/weights/Qwen2.5-7B-Instruct",
            "gemma-2-9b-it": "model/weights/gemma-2-9b-it"
        },
        datasets=["gsm8k", "math", "humaneval", "alpaca", "flores200"],
        output_dir=output_dir,
        batch_size=4,
        max_length=512,
        pooling="mean"
    )

    analyzer = CrossModelConvergence(config)
    return analyzer.run_experiment_e2()


def run_experiment_e3(output_dir: str):
    """运行 E3 实验"""
    from experiments.emergence_detection import EmergenceDetector

    MODEL_SEQUENCE = [
        "llama3.2-3b-instruct",
        "Mistral-7B-v0.3",
        "llama3-8b",
        "Qwen2.5-7B-Instruct",
        "gemma-2-9b-it"
    ]

    MODEL_PATHS = {
        "llama3.2-3b-instruct": "model/weights/llama3.2-3b-instruct",
        "Mistral-7B-v0.3": "model/weights/Mistral-7B-v0.3",
        "llama3-8b": "model/weights/llama3-8b",
        "Qwen2.5-7B-Instruct": "model/weights/Qwen2.5-7B-Instruct",
        "gemma-2-9b-it": "model/weights/gemma-2-9b-it"
    }

    MODEL_MMLU = {
        "llama3.2-3b-instruct": 58.0,
        "Mistral-7B-v0.3": 62.5,
        "llama3-8b": 68.4,
        "Qwen2.5-7B-Instruct": 86.3,
        "gemma-2-9b-it": 82.0
    }

    detector = EmergenceDetector(
        model_sequence=MODEL_SEQUENCE,
        model_paths=MODEL_PATHS,
        model_mmlu=MODEL_MMLU,
        output_dir=output_dir
    )

    return detector.run_full_analysis()


def run_experiment_e5(output_dir: str, dataset: str = "bigbench", sample_size: int = 100):
    """运行 E5 实验 - RS_cross 作为文本难度代理"""
    from experiments.text_difficulty_proxy import TextDifficultyProxy, E5Config

    config = E5Config(
        models=[
            "llama3.2-3b-instruct",
            "Mistral-7B-v0.3",
            "llama3-8b",
            "Qwen2.5-7B-Instruct",
            "gemma-2-9b-it"
        ],
        model_paths={
            "llama3.2-3b-instruct": "model/weights/llama3.2-3b-instruct",
            "Mistral-7B-v0.3": "model/weights/Mistral-7B-v0.3",
            "llama3-8b": "model/weights/llama3-8b",
            "Qwen2.5-7B-Instruct": "model/weights/Qwen2.5-7B-Instruct",
            "gemma-2-9b-it": "model/weights/gemma-2-9b-it"
        },
        model_mmlu={
            "llama3.2-3b-instruct": 58.0,
            "Mistral-7B-v0.3": 62.5,
            "llama3-8b": 68.4,
            "Qwen2.5-7B-Instruct": 86.3,
            "gemma-2-9b-it": 82.0
        },
        dataset=dataset,
        output_dir=output_dir,
        sample_size=sample_size,
        batch_size=4,
        max_length=512
    )

    analyzer = TextDifficultyProxy(config)
    return analyzer


def run_visualization(input_dir: str, output_dir: str):
    """运行可视化"""
    from experiments.visualization_direction9 import Direction9Visualizer

    visualizer = Direction9Visualizer(output_dir)

    # 加载 E1 结果
    e1_result_file = Path(input_dir) / "cross_model/e1_result.json"
    if e1_result_file.exists():
        with open(e1_result_file) as f:
            e1_result = json.load(f)
        print("加载 E1 结果成功")
    else:
        e1_result = {}
        print("未找到 E1 结果文件")

    # 加载 E3 结果
    e3_result_file = Path(input_dir) / "emergence/emergence_results.json"
    if e3_result_file.exists():
        with open(e3_result_file) as f:
            e3_result = json.load(f)
        print("加载 E3 结果成功")
    else:
        e3_result = {}
        print("未找到 E3 结果文件")

    # 生成可视化
    model_mmlu = {
        "llama3.2-3b-instruct": 58.0,
        "Mistral-7B-v0.3": 62.5,
        "llama3-8b": 68.4,
        "Qwen2.5-7B-Instruct": 86.3,
        "gemma-2-9b-it": 82.0
    }

    if e1_result.get('r_effs'):
        visualizer.plot_rs_cross_comparison(e1_result.get('rs_cross', {}))

        # 尝试绘制 MMLU vs r_eff (需要重新组织数据结构)

    if e1_result.get('spd_matrices'):
        for dataset, spd in e1_result['spd_matrices'].items():
            model_names = list(e1_result.get('r_effs', {}).get(dataset, {}).keys())
            if model_names:
                visualizer.plot_spd_heatmap(spd, model_names, dataset)

    if e3_result:
        visualizer.plot_emergence_curve(e3_result)

    # 生成摘要仪表盘
    visualizer.create_summary_dashboard(e1_result, e3_result)


def main():
    parser = argparse.ArgumentParser(description="方向 9 实验运行脚本")
    parser.add_argument("--all", action="store_true", help="运行所有实验")
    parser.add_argument("--experiment", type=str,
                       choices=["E1", "E2", "E3", "E5", "LCP"],
                       help="运行指定实验")
    parser.add_argument("--dataset", type=str, default="bigbench",
                       help="数据集名称 (E5 实验用)")
    parser.add_argument("--sample", type=int, default=100,
                       help="采样文本数量 (E5 实验用)")
    parser.add_argument("--visualize", action="store_true",
                       help="生成可视化")
    parser.add_argument("--output-dir", type=str,
                       default="experiments/output",
                       help="输出目录")
    parser.add_argument("--skip-checks", action="store_true",
                       help="跳过前置检查")

    args = parser.parse_args()

    print("=" * 60)
    print("方向 9 实验运行脚本")
    print(f"时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("=" * 60)

    # 前置检查
    if not args.skip_checks:
        if not check_prerequisites():
            print("\n前置检查未通过，请解决上述问题")
            sys.exit(1)

    # 创建输出目录
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # 运行实验
    if args.all:
        print("\n" + "=" * 60)
        print("运行 E1: 基础谱收敛验证")
        print("=" * 60)
        run_experiment_e1(str(output_dir / "cross_model"))

        print("\n" + "=" * 60)
        print("运行 E3: 能力涌现谱突变检测")
        print("=" * 60)
        run_experiment_e3(str(output_dir / "emergence"))

        print("\n" + "=" * 60)
        print("生成可视化")
        print("=" * 60)
        run_visualization(str(output_dir), str(output_dir / "figures"))

    elif args.experiment:
        if args.experiment == "E1":
            run_experiment_e1(str(output_dir / "cross_model"))
        elif args.experiment == "E2":
            run_experiment_e2(str(output_dir / "cross_model"))
        elif args.experiment == "E3":
            run_experiment_e3(str(output_dir / "emergence"))
        elif args.experiment == "E5":
            analyzer = run_experiment_e5(
                str(output_dir / "e5"),
                dataset=args.dataset,
                sample_size=args.sample
            )
            # 加载文本并运行
            from database.corpus import CorpusLoader
            loader = CorpusLoader(PROJECT_ROOT / "database/corpus")
            texts = loader.load_texts(args.dataset, limit=args.sample)
            analyzer.run_analysis(texts)
            analyzer.save_results()
            analyzer.generate_report()

    elif args.visualize:
        run_visualization(str(output_dir), str(output_dir / "figures"))

    else:
        parser.print_help()


if __name__ == "__main__":
    main()