#!/usr/bin/env python
"""
E5 实验：RS_cross 作为文本难度代理（修正版）

方法说明:
由于单条文本无法计算有意义的 RS_cross（N=1 时协方差矩阵退化），
本实验采用"分组难度分析"方法：

1. 将数据集按文本长度/复杂度分组
2. 在每个组内计算 RS_cross
3. 分析 RS_cross 与组平均难度的相关性

或者采用"模型错误率分组"方法：
1. 先用模型做题，统计每题的准确率
2. 按准确率分组（简单题/中等题/困难题）
3. 计算每组的 RS_cross（用该组所有文本一起计算）

用法:
    python experiments/text_difficulty_proxy.py --dataset bigbench --sample 200
"""

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
from scipy.stats import spearmanr, pearsonr

import sys
sys.path.insert(0, str(Path(__file__).parent.parent))
from model.spectrum import compute_gram_spectrum


@dataclass
class E5Config:
    """E5 实验配置"""
    models: List[str]
    model_paths: Dict[str, str]
    model_mmlu: Dict[str, float]
    dataset: str
    output_dir: str
    sample_size: int = 200
    batch_size: int = 4
    max_length: int = 512


@dataclass
class GroupDifficultyResult:
    """分组难度分析结果"""
    group_name: str  # 如"简单"、"中等"、"困难"
    text_ids: List[int]
    num_texts: int
    r_effs: Dict[str, float]  # {model_name: r_eff}
    rs_cross: float
    avg_text_length: float
    estimated_difficulty: float  # 估计的难度（基于 r_eff）


class TextDifficultyProxy:
    """文本难度代理分析器（分组方法）"""

    def __init__(self, config: E5Config):
        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.group_results: List[GroupDifficultyResult] = []

    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_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}")

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

        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        if tokenizer.eos_token_id is None:
            tokenizer.eos_token_id = 50256

        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

    def extract_hidden_states_batch(
        self,
        model_name: str,
        texts: List[str]
    ) -> torch.Tensor:
        """
        批量提取文本的隐层表示

        Returns:
            H: 表示矩阵 (d, N)
        """
        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)

            with torch.no_grad():
                outputs = model(**inputs)
                hidden = outputs.hidden_states[-1]  # 最后一层
                attention_mask = inputs["attention_mask"]

                # 对有效 token 取均值
                mask_expanded = attention_mask.unsqueeze(-1).float()
                feat = (hidden * mask_expanded).sum(1) / (mask_expanded.sum(1) + 1e-10)

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

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

    def compute_group_rs_cross(
        self,
        texts: List[str],
        model_names: List[str]
    ) -> Tuple[Dict[str, float], float]:
        """
        计算一组文本的 RS_cross

        方法：
        1. 每个模型提取这组文本的表示
        2. 计算每个模型的 r_eff
        3. RS_cross = 跨模型 r_eff 方差

        Returns:
            r_effs: {model_name: r_eff}
            rs_cross: 跨模型 r_eff 方差
        """
        r_effs = {}

        for model_name in tqdm(model_names, desc=f"计算组 r_eff"):
            # 提取表示
            H = self.extract_hidden_states_batch(model_name, texts)

            # 计算有效秩
            r_eff = compute_gram_spectrum(H)
            r_effs[model_name] = float(r_eff)

            # 立即卸载模型释放显存
            self.unload_all_models()

        # RS_cross = 跨模型 r_eff 方差
        values = list(r_effs.values())
        rs_cross = float(np.var(values))

        return r_effs, rs_cross

    def load_texts(self, dataset_name: str, limit: int = 200) -> List[str]:
        """加载数据集文本"""
        from database.corpus import CorpusLoader
        loader = CorpusLoader("database/corpus")

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

    def group_by_text_length(
        self,
        texts: List[str],
        num_groups: int = 5
    ) -> Dict[str, List[int]]:
        """
        按文本长度分组（作为难度的代理）

        假设：长文本 = 更复杂 = 更难
        """
        lengths = [len(t) for t in texts]
        indices = np.argsort(lengths)

        # 等分为 num_groups 组
        group_size = len(indices) // num_groups
        groups = {}

        for i in range(num_groups):
            start = i * group_size
            if i == num_groups - 1:
                end = len(indices)
            else:
                end = (i + 1) * group_size

            group_indices = indices[start:end].tolist()
            group_name = f"group_{i}_({['最短', '较短', '中等', '较长', '最长'][i] if num_groups==5 else i})"
            groups[group_name] = group_indices

        return groups

    def run_analysis(self, texts: List[str]) -> List[GroupDifficultyResult]:
        """
        运行分组难度分析

        方法：按文本长度分组，计算每组的 RS_cross
        """
        results = []
        model_names = self.config.models

        # 按长度分组
        groups = self.group_by_text_length(texts, num_groups=5)

        print(f"\n开始分析 {len(texts)} 条文本...")
        print(f"分组数：{len(groups)}")
        print(f"使用模型：{len(model_names)} 个")

        for group_name, group_indices in groups.items():
            group_texts = [texts[i] for i in group_indices]

            try:
                # 计算该组的 RS_cross
                r_effs, rs_cross = self.compute_group_rs_cross(group_texts, model_names)

                # 计算平均长度
                avg_length = np.mean([len(texts[i]) for i in group_indices])

                # 估计难度（用平均 r_eff 作为代理）
                avg_r_eff = np.mean(list(r_effs.values()))

                result = GroupDifficultyResult(
                    group_name=group_name,
                    text_ids=group_indices,
                    num_texts=len(group_indices),
                    r_effs=r_effs,
                    rs_cross=rs_cross,
                    avg_text_length=avg_length,
                    estimated_difficulty=avg_r_eff
                )
                results.append(result)

                print(f"\n{group_name}:")
                print(f"  文本数：{len(group_indices)}")
                print(f"  平均长度：{avg_length:.1f}")
                print(f"  RS_cross: {rs_cross:.4f}")

            except Exception as e:
                print(f"组 {group_name} 处理失败：{e}")
                continue

        self.group_results = results
        return results

    def compute_correlations(self) -> Dict[str, Any]:
        """
        计算 RS_cross 与文本长度（难度代理）的相关性

        假设：文本越长 = 越复杂 = RS_cross 越高
        """
        if not self.group_results:
            return {"error": "没有结果数据"}

        rs_cross_values = [r.rs_cross for r in self.group_results]
        text_lengths = [r.avg_text_length for r in self.group_results]
        difficulties = [r.estimated_difficulty for r in self.group_results]

        correlations = {}

        # RS_cross vs 文本长度
        if len(rs_cross_values) >= 3:
            spearman_length = spearmanr(rs_cross_values, text_lengths)
            correlations["rs_cross_vs_text_length"] = {
                "spearman_rho": spearman_length.correlation,
                "spearman_pvalue": spearman_length.pvalue,
                "n_samples": len(rs_cross_values),
                "interpretation": "正相关表示长文本（复杂）导致更高 RS_cross"
            }

        # RS_cross vs 估计难度
        if len(rs_cross_values) >= 3:
            spearman_diff = spearmanr(rs_cross_values, difficulties)
            correlations["rs_cross_vs_estimated_difficulty"] = {
                "spearman_rho": spearman_diff.correlation,
                "spearman_pvalue": spearman_diff.pvalue,
                "n_samples": len(rs_cross_values),
                "interpretation": "正相关表示高难度导致更高 RS_cross"
            }

        # RS_cross 随组别的变化趋势
        if len(self.group_results) >= 3:
            group_order = list(range(len(self.group_results)))
            spearman_trend = spearmanr(group_order, rs_cross_values)
            correlations["rs_cross_trend"] = {
                "spearman_rho": spearman_trend.correlation,
                "spearman_pvalue": spearman_trend.pvalue,
                "interpretation": "正相关表示 RS_cross 随文本长度增加而上升"
            }

        return correlations

    def generate_summary(self) -> Dict[str, Any]:
        """生成摘要统计"""
        if not self.group_results:
            return {}

        rs_cross_values = [r.rs_cross for r in self.group_results]

        # 按 RS_cross 排序
        sorted_groups = sorted(
            self.group_results,
            key=lambda x: x.rs_cross,
            reverse=True  # 从高到低
        )

        return {
            "total_groups": len(self.group_results),
            "rs_cross_stats": {
                "mean": float(np.mean(rs_cross_values)),
                "std": float(np.std(rs_cross_values)),
                "min": float(np.min(rs_cross_values)),
                "max": float(np.max(rs_cross_values)),
            },
            "groups_sorted_by_rs_cross": [
                {
                    "name": g.group_name,
                    "rs_cross": g.rs_cross,
                    "num_texts": g.num_texts,
                    "avg_length": g.avg_text_length,
                    "r_effs": g.r_effs
                }
                for g in sorted_groups
            ],
        }

    def save_results(self):
        """保存结果"""
        # 保存详细结果
        results_dict = []
        for r in self.group_results:
            results_dict.append({
                "group_name": r.group_name,
                "text_ids": r.text_ids,
                "num_texts": r.num_texts,
                "r_effs": r.r_effs,
                "rs_cross": r.rs_cross,
                "avg_text_length": r.avg_text_length,
                "estimated_difficulty": r.estimated_difficulty
            })

        results_file = self.output_dir / "e5_group_difficulties.json"
        with open(results_file, "w") as f:
            json.dump(results_dict, f, indent=2, ensure_ascii=False)

        # 保存相关性分析
        correlations = self.compute_correlations()
        corr_file = self.output_dir / "e5_correlations.json"
        with open(corr_file, "w") as f:
            json.dump(correlations, f, indent=2)

        # 保存摘要
        summary = self.generate_summary()
        summary["correlations"] = correlations
        summary_file = self.output_dir / "e5_summary.json"
        with open(summary_file, "w") as f:
            json.dump(summary, f, indent=2, ensure_ascii=False)

        print(f"\n结果已保存到：{self.output_dir}")
        print(f"  - e5_group_difficulties.json (分组结果)")
        print(f"  - e5_correlations.json (相关性分析)")
        print(f"  - e5_summary.json (摘要)")

    def generate_report(self):
        """生成 Markdown 报告"""
        summary = self.generate_summary()
        correlations = self.compute_correlations()

        report = [
            "# E5 实验报告：RS_cross 作为文本难度代理（分组方法）",
            "",
            f"实验时间：{datetime.now().isoformat()}",
            "",
            "## 实验配置",
            f"- 数据集：{self.config.dataset}",
            f"- 文本数量：{sum(g.num_texts for g in self.group_results)}",
            f"- 分组数：{len(self.group_results)}",
            f"- 模型数量：{len(self.config.models)}",
            f"- 模型列表：{', '.join(self.config.models)}",
            "",
            "## 方法说明",
            "",
            "由于单条文本无法计算有意义的 RS_cross（N=1 时协方差矩阵退化），",
            "本实验采用**分组难度分析**方法：",
            "",
            "1. 将数据集按文本长度分组（假设：长文本 = 更复杂 = 更难）",
            "2. 在每个组内计算跨模型 RS_cross",
            "3. 分析 RS_cross 随文本长度的变化趋势",
            "",
            "## RS_cross 统计",
            "",
            "| 统计量 | 值 |",
            "|--------|-----|",
            f"| 均值 | {summary['rs_cross_stats']['mean']:.4f} |",
            f"| 标准差 | {summary['rs_cross_stats']['std']:.4f} |",
            f"| 最小值 | {summary['rs_cross_stats']['min']:.4f} |",
            f"| 最大值 | {summary['rs_cross_stats']['max']:.4f} |",
            "",
            "## 分组结果（按 RS_cross 排序）",
            "",
            "| 组名 | 文本数 | 平均长度 | RS_cross | 平均 r_eff |",
            "|------|--------|----------|----------|------------|",
        ]

        for g in summary.get('groups_sorted_by_rs_cross', []):
            avg_r_eff = np.mean(list(g['r_effs'].values()))
            report.append(
                f"| {g['name']} | {g['num_texts']} | {g['avg_length']:.1f} | "
                f"{g['rs_cross']:.4f} | {avg_r_eff:.2f} |"
            )

        report.extend([
            "",
            "## 相关性分析",
            "",
        ])

        if "rs_cross_vs_text_length" in correlations:
            corr = correlations["rs_cross_vs_text_length"]
            report.extend([
                "### RS_cross vs 文本长度",
                "",
                f"- Spearman ρ: **{corr['spearman_rho']:.4f}** (p={corr['spearman_pvalue']:.4f})",
                f"- 样本数：{corr['n_samples']}",
                f"- 解释：{corr['interpretation']}",
                "",
            ])

            # 评估相关性强度
            rho = abs(corr['spearman_rho'])
            if rho > 0.7:
                strength = "强相关 ✓"
            elif rho > 0.5:
                strength = "中等相关 ✓"
            elif rho > 0.3:
                strength = "弱相关 △"
            else:
                strength = "几乎不相关 ✗"

            report.extend([
                f"**相关性强度**: {strength}",
                "",
            ])

        if "rs_cross_trend" in correlations:
            corr = correlations["rs_cross_trend"]
            report.extend([
                "### RS_cross 随文本长度组的变化趋势",
                "",
                f"- Spearman ρ: **{corr['spearman_rho']:.4f}** (p={corr['spearman_pvalue']:.4f})",
                f"- 解释：{corr['interpretation']}",
                "",
            ])

        report.extend([
            "## 结论",
            "",
        ])

        if "rs_cross_vs_text_length" in correlations:
            rho = correlations["rs_cross_vs_text_length"]["spearman_rho"]
            pval = correlations["rs_cross_vs_text_length"]["spearman_pvalue"]

            if rho > 0.5 and pval < 0.1:
                conclusion = (
                    f"✅ **假设获支持**：RS_cross 与文本长度呈显著正相关 (ρ={rho:.4f}, p={pval:.4f})。\n\n"
                    "这表明**文本复杂度越高，跨模型表示分歧越大**，"
                    "RS_cross 可以作为文本难度的无监督代理指标。"
                )
            elif rho > 0.3:
                conclusion = (
                    f"△ **假设部分支持**：RS_cross 与文本长度呈弱到中等相关 (ρ={rho:.4f}, p={pval:.4f})。\n\n"
                    "趋势存在但不够强，可能需要更多数据点或更细粒度的难度划分。"
                )
            else:
                conclusion = (
                    f"✗ **假设未获支持**：RS_cross 与文本长度几乎不相关 (ρ={rho:.4f}, p={pval:.4f})。\n\n"
                    "可能需要重新审视'文本长度=难度'的假设，或使用其他难度指标（如模型准确率）。"
                )
            report.append(conclusion)

        report.extend([
            "",
            "## 局限性",
            "",
            "1. **文本长度≠难度**: 文本长度只是难度的粗糙代理，有些短文本可能很难（如数学证明），有些长文本可能很简单（如重复叙述）",
            "2. **分组数量少**: 仅 5 组，相关性统计检验力有限",
            "3. **单数据集**: 仅在 bigbench 上验证，需要更多数据集验证普适性",
            "",
            "## 未来改进方向",
            "",
            "1. **使用真实难度标注**: 用模型在题目上的实际准确率作为难度标签",
            "2. **滑动窗口**: 使用滑动窗口代替硬分组，获得更平滑的难度曲线",
            "3. **多维度难度**: 考虑文本的多维度难度（推理、知识、语言复杂度等）",
        ])

        report_file = self.output_dir / "e5_report.md"
        with open(report_file, "w") as f:
            f.write("\n".join(report))

        print(f"报告已保存到：{report_file}")


def main():
    import argparse

    parser = argparse.ArgumentParser(description="E5 实验：RS_cross 作为文本难度代理")
    parser.add_argument("--dataset", type=str, default="bigbench",
                       help="数据集名称")
    parser.add_argument("--sample", type=int, default=200,
                       help="采样文本数量")
    parser.add_argument("--output-dir", type=str,
                       default="experiments/output/e5",
                       help="输出目录")

    args = parser.parse_args()

    # 实验配置
    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=args.dataset,
        output_dir=args.output_dir,
        sample_size=args.sample,
        batch_size=4,
        max_length=512
    )

    analyzer = TextDifficultyProxy(config)

    # 加载文本
    texts = analyzer.load_texts(args.dataset, limit=args.sample)

    if not texts:
        print("未找到文本，退出")
        return

    # 运行分析
    analyzer.run_analysis(texts)

    # 保存结果
    analyzer.save_results()

    # 生成报告
    analyzer.generate_report()


if __name__ == "__main__":
    main()