#!/usr/bin/env python
"""
跨模型谱收敛实验流水线 - 方向 9 核心实验

功能:
- 提取多个模型在同一文本集上的隐层表示
- 计算跨模型谱稳定性指标 (RS_cross)
- 计算谱 Platonic 距离矩阵 (SPD)
- 逐层谱收敛轮廓分析 (LCP)

用法:
    # 运行完整 E1 实验
    python experiments/cross_model_convergence.py --experiment E1 --output-dir output/e1

    # 运行逐层分析
    python experiments/cross_model_convergence.py --experiment LCP --layers 10
"""

import os
import json
import torch
import numpy as np
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass, asdict
from datetime import datetime

# 设置环境变量防止显存碎片化
os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from transformers import AutoTokenizer, AutoModelForCausalLM
from tqdm import tqdm

# 导入谱计算模块
import sys
sys.path.insert(0, str(Path(__file__).parent.parent))
from model.spectrum import compute_gram_spectrum, compute_spectrum_array, wasserstein1_distance
from model.extractor import save_H, load_H, RepresentationLoader


@dataclass
class ExperimentConfig:
    """实验配置"""
    models: List[str]
    model_paths: Dict[str, str]
    datasets: List[str]
    output_dir: str
    batch_size: int = 16
    max_length: int = 512
    pooling: str = "mean"
    layer: int = -1  # -1 表示最后一层
    device: str = "cuda"


@dataclass
class ExperimentResult:
    """实验结果"""
    timestamp: str
    config: dict
    r_effs: Dict[str, Dict[str, float]]  # {model: {dataset: r_eff}}
    rs_cross: Dict[str, float]  # {dataset: RS_cross}
    spd_matrices: Dict[str, List[List[float]]]  # {dataset: SPD_matrix}
    spectra: Dict[str, Dict[str, List[float]]]  # {model: {dataset: spectrum}}


class CrossModelConvergence:
    """跨模型谱收敛分析器"""

    def __init__(self, config: ExperimentConfig):
        self.config = config
        self.output_dir = Path(config.output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # 缓存加载的模型
        self.loaded_models: Dict[str, Tuple[AutoTokenizer, AutoModelForCausalLM]] = {}

        # 文本缓存
        self.texts_cache: Dict[str, List[str]] = {}

    def unload_all_models(self):
        """卸载所有模型并清理 GPU 缓存"""
        import gc
        self.loaded_models.clear()
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def load_texts(self, dataset_name: str, limit: int = 500) -> List[str]:
        """加载数据集文本"""
        if dataset_name in self.texts_cache:
            return self.texts_cache[dataset_name][:limit]

        from database.corpus import CorpusLoader
        loader = CorpusLoader("database/corpus")

        try:
            texts = loader.load_texts(dataset_name, limit=limit)
            self.texts_cache[dataset_name] = texts
            print(f"加载 {dataset_name}: {len(texts)} 条文本")
            return texts
        except FileNotFoundError:
            print(f"警告：数据集 {dataset_name} 不存在")
            return []

    def load_model(self, model_name: str):
        """加载模型（带缓存）"""
        if model_name in self.loaded_models:
            return self.loaded_models[model_name]

        model_path = self.config.model_paths.get(
            model_name,
            f"model/weights/{model_name}"
        )

        print(f"加载模型：{model_name} (from {model_path})")

        tokenizer = AutoTokenizer.from_pretrained(
            model_path,
            trust_remote_code=True
        )

        # 设置 pad_token
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        if tokenizer.eos_token_id is None:
            tokenizer.eos_token_id = 50256  # 默认 EOS

        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype=torch.float16,
            device_map="auto",
            output_hidden_states=True,
            trust_remote_code=True
        )
        model.eval()

        self.loaded_models[model_name] = (tokenizer, model)
        return tokenizer, model

    @torch.no_grad()
    def extract_hidden_states(
        self,
        model_name: str,
        texts: List[str],
        layer: int = -1
    ) -> torch.Tensor:
        """
        提取隐层表示

        Returns:
            H: 表示矩阵 (d, N)
        """
        # 先清理缓存再加载新模型
        self.unload_all_models()

        tokenizer, model = self.load_model(model_name)

        all_features = []
        for i in range(0, len(texts), self.config.batch_size):
            batch = texts[i:i + self.config.batch_size]

            inputs = tokenizer(
                batch,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=self.config.max_length
            ).to(model.device)

            outputs = model(**inputs)
            hidden = outputs.hidden_states[layer]  # (B, seq_len, d)
            attention_mask = inputs["attention_mask"]

            if self.config.pooling == "mean":
                # 对有效 token 取均值
                mask_expanded = attention_mask.unsqueeze(-1).float()
                feat = (hidden * mask_expanded).sum(1) / (mask_expanded.sum(1) + 1e-10)
            elif self.config.pooling == "last_token":
                # 取最后一个有效 token
                lengths = attention_mask.sum(1) - 1
                feat = hidden[torch.arange(len(batch)), lengths]
            else:
                # CLS token
                feat = hidden[:, 0, :]

            all_features.append(feat.cpu().float())

        # 转置为 (d, N)
        H = torch.cat(all_features, dim=0).T
        return H

    def compute_rs_cross(
        self,
        dataset_name: str,
        model_names: List[str],
        texts: Optional[List[str]] = None
    ) -> Dict[str, Any]:
        """
        计算跨模型谱稳定性 RS_cross

        Returns:
            {
                "rs_cross": float,
                "r_effs": {model: r_eff},
                "std": float,
                "range": float
            }
        """
        if texts is None:
            texts = self.load_texts(dataset_name)

        if not texts:
            return {"rs_cross": 0, "r_effs": {}, "std": 0, "range": 0}

        r_effs = {}
        spectra = {}

        for model_name in tqdm(model_names, desc=f"RS_cross [{dataset_name}]"):
            # 提取或加载表示
            rep_path = f"representations/{model_name}_{dataset_name}.pt"
            rep_file = self.output_dir.parent / rep_path

            if rep_file.exists():
                H = load_H(model_name, dataset_name,
                          representations_dir=str(rep_file.parent))
            else:
                H = self.extract_hidden_states(model_name, texts)
                save_H(H, model_name, dataset_name,
                      representations_dir=str(rep_file.parent))

            # 计算有效秩
            H_tensor = torch.tensor(H.numpy() if hasattr(H, "cpu") else H) if not isinstance(H, torch.Tensor) else H
            r_eff = compute_gram_spectrum(H_tensor)
            r_effs[model_name] = r_eff

            # 计算谱（用于 SPD）
            _, prob = compute_spectrum_array(H)
            spectra[model_name] = prob

        # 计算 RS_cross = 方差
        values = list(r_effs.values())
        rs_cross = float(np.var(values))
        std = float(np.std(values))
        range_val = float(max(values) - min(values)) if values else 0

        return {
            "rs_cross": rs_cross,
            "r_effs": r_effs,
            "std": std,
            "range": range_val,
            "spectra": spectra
        }

    def compute_spd_matrix(
        self,
        spectra_dict: Dict[str, np.ndarray]
    ) -> np.ndarray:
        """
        计算谱 Platonic 距离矩阵

        Returns:
            SPD 矩阵 (n, n)
        """
        model_names = list(spectra_dict.keys())
        n = len(model_names)
        spd_matrix = np.zeros((n, n))

        for i in range(n):
            for j in range(i + 1, n):
                p_i = spectra_dict[model_names[i]]
                p_j = spectra_dict[model_names[j]]

                # 对齐长度
                max_len = max(len(p_i), len(p_j))
                p_i_pad = np.pad(p_i, (0, max_len - len(p_i)))
                p_j_pad = np.pad(p_j, (0, max_len - len(p_j)))

                d = wasserstein1_distance(p_i_pad, p_j_pad)
                spd_matrix[i, j] = d
                spd_matrix[j, i] = d

        return spd_matrix, model_names

    def compute_layer_convergence_profile(
        self,
        dataset_name: str,
        model_names: List[str],
        n_layers: int,
        sample_layers: int = 10
    ) -> np.ndarray:
        """
        逐层谱收敛轮廓 (LCP)

        Returns:
            RS_cross 值数组 (sample_layers,)
        """
        texts = self.load_texts(dataset_name)
        layer_indices = np.linspace(0, n_layers - 1, sample_layers, dtype=int)

        lcp = []
        for layer_idx in tqdm(layer_indices, desc="LCP"):
            result = self.compute_rs_cross(
                dataset_name, model_names, texts,
                layer=int(layer_idx)
            )
            lcp.append(result["rs_cross"])

        return np.array(lcp), layer_indices

    def run_experiment_e1(self) -> ExperimentResult:
        """运行 E1 基础谱收敛实验"""
        print("=" * 60)
        print("运行 E1: 基础谱收敛验证")
        print("=" * 60)

    def run_experiment_e2(self) -> Dict[str, Any]:
        """
        运行 E2: 跨域谱分层验证

        目的：验证不同领域的谱收敛程度不同
        假设：形式化领域（代码/数学）收敛更好，开放领域（指令/创作）收敛更差
        """
        print("=" * 60)
        print("运行 E2: 跨域谱分层验证")
        print("=" * 60)

        # 领域分类（按形式化程度）
        domain_categories = {
            "highly_formal": ["humaneval", "code"],  # 高度形式化：代码
            "formal": ["gsm8k", "math"],  # 形式化：数学/推理
            "semi_formal": ["flores200", "translation"],  # 半形式化：翻译
            "open": ["alpaca", "instruction"],  # 开放：指令跟随
        }

        # 加载 E1 结果（如果已运行）
        e1_result_file = self.output_dir / "e1_result.json"
        if e1_result_file.exists():
            print(f"从 E1 结果加载数据：{e1_result_file}")
            with open(e1_result_file) as f:
                e1_result = json.load(f)
            rs_cross_all = e1_result.get("rs_cross", {})
            r_effs_all = e1_result.get("r_effs", {})
        else:
            print("未找到 E1 结果，重新计算...")
            # 需要先运行 E1
            self.run_experiment_e1()
            return self.run_experiment_e2()

        # 按类别聚合 RS_cross
        category_rs = {}
        category_domains = {}

        for category, domains in domain_categories.items():
            matching_domains = [d for d in domains if d in rs_cross_all]
            if matching_domains:
                rs_values = [rs_cross_all[d] for d in matching_domains]
                category_rs[category] = {
                    "mean": float(np.mean(rs_values)),
                    "std": float(np.std(rs_values)) if len(rs_values) > 1 else 0,
                    "values": {d: rs_cross_all[d] for d in matching_domains},
                    "domains": matching_domains
                }
                category_domains[category] = matching_domains

        # 排序：按 RS_cross 从小到大（收敛程度从高到低）
        sorted_categories = sorted(
            category_rs.items(),
            key=lambda x: x[1]["mean"]
        )

        # 计算领域间差异
        formal_domains = category_rs.get("highly_formal", {}).get("domains", []) + \
                        category_rs.get("formal", {}).get("domains", [])
        open_domains = category_rs.get("open", {}).get("domains", [])

        formal_rs = [rs_cross_all[d] for d in formal_domains if d in rs_cross_all]
        open_rs = [rs_cross_all[d] for d in open_domains if d in rs_cross_all]

        comparison = {
            "formal_mean": float(np.mean(formal_rs)) if formal_rs else None,
            "open_mean": float(np.mean(open_rs)) if open_rs else None,
            "difference": None,
            "ratio": None
        }

        if comparison["formal_mean"] and comparison["open_mean"]:
            comparison["difference"] = comparison["open_mean"] - comparison["formal_mean"]
            comparison["ratio"] = comparison["open_mean"] / comparison["formal_mean"]

        result = {
            "timestamp": datetime.now().isoformat(),
            "category_rs_cross": category_rs,
            "sorted_by_convergence": [
                {"category": cat, "mean_rs_cross": data["mean"], "domains": data["domains"]}
                for cat, data in sorted_categories
            ],
            "formal_vs_open_comparison": comparison,
            "hypothesis_support": self._evaluate_hypothesis(comparison),
            "raw_rs_cross": rs_cross_all
        }

        # 保存结果
        self._save_e2_result(result)
        return result

    def _evaluate_hypothesis(self, comparison: Dict) -> Dict[str, Any]:
        """评估假设是否得到支持"""
        if not comparison["formal_mean"] or not comparison["open_mean"]:
            return {"supported": False, "reason": "数据不足"}

        # 假设：形式化领域 RS_cross < 开放领域 RS_cross
        supported = comparison["formal_mean"] < comparison["open_mean"]
        effect_size = comparison["ratio"]

        # 效应量判断
        if effect_size is None:
            strength = "unknown"
        elif effect_size > 10:
            strength = "very_strong"
        elif effect_size > 5:
            strength = "strong"
        elif effect_size > 2:
            strength = "moderate"
        else:
            strength = "weak"

        return {
            "supported": supported,
            "formal_mean": comparison["formal_mean"],
            "open_mean": comparison["open_mean"],
            "ratio_open_to_formal": effect_size,
            "effect_strength": strength,
            "interpretation": "形式化领域谱收敛显著优于开放领域" if supported else "假设未获支持"
        }

    def _save_e2_result(self, result: Dict):
        """保存 E2 结果"""
        result_file = self.output_dir / "e2_result.json"
        with open(result_file, "w") as f:
            json.dump(result, f, indent=2)

        print(f"\nE2 结果已保存到：{result_file}")

        # 生成摘要报告
        self._generate_e2_summary(result)

    def _generate_e2_summary(self, result: Dict):
        """生成 E2 摘要报告"""
        summary = [
            "# E2 实验摘要：跨域谱分层验证",
            "",
            f"实验时间：{result['timestamp']}",
            "",
            "## 假设",
            "形式化领域（代码/数学）的谱收敛程度优于开放领域（指令/创作）",
            "",
            "## 按类别 RS_cross（越小越收敛）",
            ""
        ]

        summary.append("| 类别 | 平均 RS_cross | 包含领域 |")
        summary.append("|------|--------------|----------|")
        for cat, data in result["category_rs_cross"].items():
            domains = ", ".join(data["domains"])
            summary.append(f"| {cat} | {data['mean']:.4f} | {domains} |")

        summary.append("")
        summary.append("## 形式化 vs 开放领域对比")
        comp = result["formal_vs_open_comparison"]
        if comp["formal_mean"]:
            summary.append(f"- 形式化领域平均 RS_cross: {comp['formal_mean']:.4f}")
            summary.append(f"- 开放领域平均 RS_cross: {comp['open_mean']:.4f}")
            summary.append(f"- 差异：{comp['difference']:.4f}")
            summary.append(f"- 比率（开放/形式化）：{comp['ratio']:.2f}x")

        summary.append("")
        summary.append("## 假设验证")
        hyp = result["hypothesis_support"]
        status = "✅ 支持" if hyp["supported"] else "❌ 不支持"
        summary.append(f"状态：{status}")
        summary.append(f"效应强度：{hyp.get('effect_strength', 'unknown')}")
        summary.append(f"解释：{hyp['interpretation']}")

        summary.append("")
        summary.append("## 原始数据")
        summary.append("| 领域 | RS_cross |")
        summary.append("|------|----------|")
        for domain, rs in sorted(result["raw_rs_cross"].items(), key=lambda x: x[1]):
            summary.append(f"| {domain} | {rs:.4f} |")

        summary_file = self.output_dir / "e2_summary.md"
        with open(summary_file, "w") as f:
            f.write("\n".join(summary))

        print(f"摘要报告已保存到：{summary_file}")

        config_dict = asdict(self.config)
        result = ExperimentResult(
            timestamp=datetime.now().isoformat(),
            config=config_dict,
            r_effs={},
            rs_cross={},
            spd_matrices={},
            spectra={}
        )

        # 对每个数据集计算 RS_cross
        for dataset_name in self.config.datasets:
            print(f"\n[数据集] {dataset_name}")

            rs_result = self.compute_rs_cross(
                dataset_name,
                self.config.models
            )

            result.r_effs[dataset_name] = rs_result["r_effs"]
            result.rs_cross[dataset_name] = rs_result["rs_cross"]

            # 计算 SPD 矩阵
            spd_matrix, model_names = self.compute_spd_matrix(
                rs_result["spectra"]
            )
            result.spd_matrices[dataset_name] = spd_matrix.tolist()
            result.spectra[dataset_name] = {
                m: s.tolist() for m, s in
                zip(model_names, rs_result["spectra"].values())
            }

            print(f"  RS_cross = {rs_result['rs_cross']:.4f}")
            print(f"  模型 r_eff 范围：[{min(rs_result['r_effs'].values()):.2f}, "
                  f"{max(rs_result['r_effs'].values()):.2f}]")

        # 保存结果
        self._save_result(result)
        return result

    def run_experiment_lcp(
        self,
        dataset_name: str,
        n_layers: int,
        sample_layers: int = 10
    ) -> Dict[str, Any]:
        """运行逐层谱收敛轮廓实验"""
        print("=" * 60)
        print(f"运行 LCP: 逐层谱收敛分析 [{dataset_name}]")
        print("=" * 60)

        lcp, layer_indices = self.compute_layer_convergence_profile(
            dataset_name,
            self.config.models,
            n_layers,
            sample_layers
        )

        result = {
            "timestamp": datetime.now().isoformat(),
            "dataset": dataset_name,
            "models": self.config.models,
            "layer_indices": layer_indices.tolist(),
            "lcp_values": lcp.tolist(),
            "n_layers": n_layers
        }

        # 保存结果
        result_file = self.output_dir / f"lcp_{dataset_name}.json"
        with open(result_file, "w") as f:
            json.dump(result, f, indent=2)

        print(f"\nLCP 结果已保存到：{result_file}")
        print(f"RS_cross 从 {lcp[0]:.4f} (浅层) 到 {lcp[-1]:.4f} (深层)")

        return result

    def _save_result(self, result: ExperimentResult):
        """保存实验结果"""
        result_dict = {
            "timestamp": result.timestamp,
            "config": result.config,
            "r_effs": result.r_effs,
            "rs_cross": result.rs_cross,
            "spd_matrices": result.spd_matrices,
        }

        result_file = self.output_dir / "e1_result.json"
        with open(result_file, "w") as f:
            json.dump(result_dict, f, indent=2)

        print(f"\nE1 结果已保存到：{result_file}")

        # 生成摘要报告
        self._generate_summary(result_dict)

    def _generate_summary(self, result: Dict):
        """生成摘要报告"""
        summary = ["# E1 实验摘要", ""]

        summary.append("## 实验配置")
        summary.append(f"- 时间：{result['timestamp']}")
        summary.append(f"- 模型数：{len(result['config']['models'])}")
        summary.append(f"- 数据集数：{len(result['config']['datasets'])}")
        summary.append("")

        summary.append("## RS_cross 结果")
        summary.append("| 数据集 | RS_cross |")
        summary.append("|--------|----------|")
        for ds, rs in result["rs_cross"].items():
            summary.append(f"| {ds} | {rs:.4f} |")
        summary.append("")

        summary.append("## 各模型 r_eff")
        for ds, r_effs in result["r_effs"].items():
            summary.append(f"### {ds}")
            for model, r_eff in r_effs.items():
                summary.append(f"- {model}: {r_eff:.2f}")
            summary.append("")

        summary_file = self.output_dir / "e1_summary.md"
        with open(summary_file, "w") as f:
            f.write("\n".join(summary))

        print(f"摘要报告已保存到：{summary_file}")


def main():
    import argparse

    parser = argparse.ArgumentParser(description="跨模型谱收敛实验")
    parser.add_argument("--experiment", type=str,
                       choices=["E1", "E2", "LCP"], default="E1",
                       help="实验类型")
    parser.add_argument("--output-dir", type=str,
                       default="experiments/output/cross_model",
                       help="输出目录")
    parser.add_argument("--dataset", type=str,
                       default="gsm8k",
                       help="数据集名称 (LCP 实验用)")
    parser.add_argument("--layers", type=int, default=32,
                       help="模型层数 (LCP 实验用)")
    parser.add_argument("--sample-layers", type=int, default=10,
                       help="采样层数 (LCP 实验用)")

    args = parser.parse_args()

    # 实验配置
    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"],
        output_dir=args.output_dir,
        batch_size=8,
        max_length=512,
        pooling="mean"
    )

    analyzer = CrossModelConvergence(config)

    if args.experiment == "E1":
        analyzer.run_experiment_e1()
    elif args.experiment == "E2":
        analyzer.run_experiment_e2()
    elif args.experiment == "LCP":
        analyzer.run_experiment_lcp(
            args.dataset,
            args.layers,
            args.sample_layers
        )


if __name__ == "__main__":
    main()