#!/usr/bin/env python
"""
E4 实验：对齐即谱工程

目的：验证对齐（SFT/Instruct）会使模型在指令集上的谱更集中（r_eff 下降）

假设:
1. 在指令集上：Instruct 版的 r_eff < Base 版（谱压缩）
2. 在通用集上：两者 r_eff 相近（基础能力不变）
3. 域专一化 (STG) = |r_eff_instruction - r_eff_general|，Instruct 版应更大

用法:
    python experiments/alignment_spectrum.py
"""

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["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


@dataclass
class E4Config:
    """E4 实验配置"""
    # 模型对：Base vs Instruct
    base_model_name: str = "llama3-8b"
    instruct_model_name: str = "llama3-8b-instruct"

    base_model_path: str = "model/weights/llama3-8b"
    instruct_model_path: str = "model/weights/llama3-8b-instruct"

    # 数据集
    instruction_dataset: str = "alpaca"  # 指令集
    general_dataset: str = "flores200"   # 通用集（对照）

    # 实验参数
    batch_size: int = 4
    max_length: int = 512
    pooling: str = "mean"
    sample_size: int = 200  # 每个数据集采样 200 条

    output_dir: str = "experiments/output/e4"
    device: str = "cuda"


@dataclass
class E4Result:
    """E4 实验结果"""
    timestamp: str
    config: dict

    # 指令集结果
    instruction_r_eff: Dict[str, float]  # {model_name: r_eff}
    instruction_rs_cross: float  # Base vs Instruct 的差异

    # 通用集结果（对照）
    general_r_eff: Dict[str, float]
    general_rs_cross: float

    # 域专一化指标
    stg: Dict[str, float]  # {model_name: |r_eff_instr - r_eff_general|}

    # 谱距离
    spd_instruction: float  # Base vs Instruct 在指令集上的 SPD
    spd_general: float      # Base vs Instruct 在通用集上的 SPD

    # 假设验证
    hypothesis_support: Dict[str, Any]


class AlignmentSpectrumAnalyzer:
    """对齐谱分析器"""

    def __init__(self, config: E4Config):
        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]] = {}

    def load_model(self, model_name: str) -> Tuple[AutoTokenizer, AutoModelForCausalLM]:
        """加载模型"""
        if model_name in self.loaded_models:
            return self.loaded_models[model_name]

        model_path = (self.config.base_model_path if model_name == self.config.base_model_name
                     else self.config.instruct_model_path)

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

        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 = 128009  # LLaMA3 default

        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 unload_model(self, model_name: str):
        """卸载模型释放显存"""
        if model_name in self.loaded_models:
            del self.loaded_models[model_name]
            import gc
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            print(f"已卸载：{model_name}")

    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} 不存在，使用备用数据")
            # 备用：生成简单指令
            if dataset_name == "alpaca":
                return [f"Instruction {i}: Please explain the concept of machine learning."
                       for i in range(limit)]
            return [f"Text {i}" for i in range(limit)]

    @torch.no_grad()
    def extract_hidden_states(
        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)

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

            if self.config.pooling == "mean":
                mask_expanded = attention_mask.unsqueeze(-1).float()
                feat = (hidden * mask_expanded).sum(1) / (mask_expanded.sum(1) + 1e-10)
            else:
                feat = hidden[:, 0, :]  # CLS token

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

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

    def compute_r_eff(self, H: torch.Tensor) -> float:
        """计算有效秩"""
        return compute_gram_spectrum(H)

    def compute_spectra(self, H: torch.Tensor) -> np.ndarray:
        """计算谱分布"""
        _, prob = compute_spectrum_array(H)
        return prob

    def compute_spd(self, prob1: np.ndarray, prob2: np.ndarray) -> float:
        """计算谱 Platonic 距离 (Wasserstein-1 距离)"""
        from model.spectrum import wasserstein1_distance

        # 对齐长度
        max_len = max(len(prob1), len(prob2))
        p1_pad = np.pad(prob1, (0, max_len - len(prob1)))
        p2_pad = np.pad(prob2, (0, max_len - len(prob2)))

        return wasserstein1_distance(p1_pad, p2_pad)

    def analyze_pair(
        self,
        dataset_name: str,
        texts: List[str]
    ) -> Dict[str, Any]:
        """
        分析一对模型在特定数据集上的谱

        Returns:
            {
                "r_effs": {model: r_eff},
                "spectra": {model: spectrum},
                "rs_cross": float,  # 方差
                "spd": float
            }
        """
        result = {"r_effs": {}, "spectra": {}}

        # 提取 Base 模型
        print(f"  提取 {self.config.base_model_name}...")
        H_base = self.extract_hidden_states(self.config.base_model_name, texts)
        r_eff_base = self.compute_r_eff(H_base)
        spec_base = self.compute_spectra(H_base)
        result["r_effs"][self.config.base_model_name] = r_eff_base
        result["spectra"][self.config.base_model_name] = spec_base

        # 卸载 Base 模型，加载 Instruct 模型
        self.unload_model(self.config.base_model_name)

        # 提取 Instruct 模型
        print(f"  提取 {self.config.instruct_model_name}...")
        H_instruct = self.extract_hidden_states(self.config.instruct_model_name, texts)
        r_eff_instruct = self.compute_r_eff(H_instruct)
        spec_instruct = self.compute_spectra(H_instruct)
        result["r_effs"][self.config.instruct_model_name] = r_eff_instruct
        result["spectra"][self.config.instruct_model_name] = spec_instruct

        # 计算 RS_cross (方差) 和 SPD
        r_eff_values = list(result["r_effs"].values())
        result["rs_cross"] = float(np.var(r_eff_values))
        result["spd"] = self.compute_spd(spec_base, spec_instruct)

        return result

    def run_full_analysis(self) -> E4Result:
        """运行完整的 E4 分析"""
        print("=" * 60)
        print("E4 实验：对齐即谱工程")
        print("=" * 60)
        print(f"\n模型对：{self.config.base_model_name} vs {self.config.instruct_model_name}")
        print(f"指令集：{self.config.instruction_dataset}")
        print(f"通用集：{self.config.general_dataset}")
        print()

        # 加载数据
        print("[1/4] 加载指令集...")
        instruct_texts = self.load_texts(self.config.instruction_dataset,
                                         limit=self.config.sample_size)

        print("[2/4] 加载通用集...")
        general_texts = self.load_texts(self.config.general_dataset,
                                        limit=self.config.sample_size)

        # 分析指令集
        print("\n[3/4] 分析指令集上的谱...")
        instruct_result = self.analyze_pair(self.config.instruction_dataset, instruct_texts)

        # 卸载所有模型
        self.unload_model(self.config.instruct_model_name)

        # 分析通用集
        print("\n[4/4] 分析通用集上的谱...")
        general_result = self.analyze_pair(self.config.general_dataset, general_texts)

        # 计算域专一化 (STG)
        stg = {}
        for model_name in [self.config.base_model_name, self.config.instruct_model_name]:
            r_instr = instruct_result["r_effs"][model_name]
            r_gen = general_result["r_effs"][model_name]
            stg[model_name] = abs(r_instr - r_gen)

        # 假设验证
        hypothesis = self._evaluate_hypothesis(
            instruct_result, general_result, stg
        )

        # 打包结果
        result = E4Result(
            timestamp=datetime.now().isoformat(),
            config=asdict(self.config),
            instruction_r_eff=instruct_result["r_effs"],
            instruction_rs_cross=instruct_result["rs_cross"],
            general_r_eff=general_result["r_effs"],
            general_rs_cross=general_result["rs_cross"],
            stg=stg,
            spd_instruction=instruct_result["spd"],
            spd_general=general_result["spd"],
            hypothesis_support=hypothesis
        )

        self._save_result(result)
        return result

    def _evaluate_hypothesis(
        self,
        instruct: Dict,
        general: Dict,
        stg: Dict
    ) -> Dict[str, Any]:
        """评估实验假设"""
        base_name = self.config.base_model_name
        instruct_name = self.config.instruct_model_name

        # H1: Instruct 版在指令集上 r_eff 更低
        h1_supported = (instruct["r_effs"][instruct_name] <
                       instruct["r_effs"][base_name])

        # H2: Instruct 版的域专一化 (STG) 更高
        h2_supported = stg[instruct_name] > stg[base_name]

        # H3: 指令集上的谱距离 > 通用集上的谱距离
        # （这意味着对齐主要影响指令理解）
        h3_data_needed = True  # 需要更多模型对才能验证

        return {
            "h1_spectrum_compression": {
                "supported": h1_supported,
                "interpretation": "Instruct 版在指令集上谱更集中" if h1_supported else "假设未获支持"
            },
            "h2_domain_specialization": {
                "supported": h2_supported,
                "interpretation": "Instruct 版域专一化程度更高" if h2_supported else "假设未获支持"
            },
            "overall": "假设获支持" if (h1_supported and h2_supported) else "部分支持"
        }

    def _save_result(self, result: E4Result):
        """保存结果"""
        result_dict = {
            "timestamp": result.timestamp,
            "config": result.config,
            "instruction_r_eff": result.instruction_r_eff,
            "instruction_rs_cross": result.instruction_rs_cross,
            "general_r_eff": result.general_r_eff,
            "general_rs_cross": result.general_rs_cross,
            "stg": result.stg,
            "spd_instruction": result.spd_instruction,
            "spd_general": result.spd_general,
            "hypothesis_support": result.hypothesis_support
        }

        # JSON 结果
        result_file = self.output_dir / "e4_result.json"
        with open(result_file, "w", encoding="utf-8") as f:
            json.dump(result_dict, f, indent=2, ensure_ascii=False)

        # Markdown 报告
        self._generate_report(result)

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

    def _generate_report(self, result: E4Result):
        """生成 Markdown 报告"""
        r = result

        report = f"""# E4 实验报告：对齐即谱工程

**实验时间**: {r.timestamp}
**模型对**: {r.config['base_model_name']} (Base) vs {r.config['instruct_model_name']} (Instruct)

---

## 一、实验目标

验证对齐（SFT/Instruct 微调）对模型谱结构的影响：
- **假设 H1**: Instruct 版在指令集上的 r_eff 更低（谱压缩）
- **假设 H2**: Instruct 版的域专一化程度更高（STG 更大）

---

## 二、实验配置

| 配置项 | 值 |
|--------|-----|
| Base 模型 | {r.config['base_model_name']} |
| Instruct 模型 | {r.config['instruct_model_name']} |
| 指令数据集 | {r.config['instruction_dataset']} |
| 通用数据集 | {r.config['general_dataset']} |
| 样本数量 | {r.config['sample_size']} 条 |

---

## 三、核心结果

### 3.1 指令集上的 r_eff 对比

| 模型 | r_eff |
|------|-------|
| {r.config['base_model_name']} | {r.instruction_r_eff[r.config['base_model_name']]:.4f} |
| {r.config['instruct_model_name']} | {r.instruction_r_eff[r.config['instruct_model_name']]:.4f} |
| **RS_cross** | {r.instruction_rs_cross:.4f} |

### 3.2 通用集上的 r_eff 对比（对照）

| 模型 | r_eff |
|------|-------|
| {r.config['base_model_name']} | {r.general_r_eff[r.config['base_model_name']]:.4f} |
| {r.config['instruct_model_name']} | {r.general_r_eff[r.config['instruct_model_name']]:.4f} |
| **RS_cross** | {r.general_rs_cross:.4f} |

### 3.3 域专一化 (STG)

STG = |r_eff_instruction - r_eff_general|

| 模型 | STG |
|------|-----|
| {r.config['base_model_name']} | {r.stg[r.config['base_model_name']]:.4f} |
| {r.config['instruct_model_name']} | {r.stg[r.config['instruct_model_name']]:.4f} |

### 3.4 谱距离 (SPD)

| 数据集 | SPD (Base vs Instruct) |
|--------|----------------------|
| 指令集 | {r.spd_instruction:.4f} |
| 通用集 | {r.spd_general:.4f} |

---

## 四、假设验证

### H1: 谱压缩假设

**预测**: Instruct 版在指令集上的 r_eff < Base 版

**结果**: {'✅ 支持' if r.hypothesis_support['h1_spectrum_compression']['supported'] else '❌ 不支持'}

{r.hypothesis_support['h1_spectrum_compression']['interpretation']}

### H2: 域专一化假设

**预测**: Instruct 版的 STG > Base 版

**结果**: {'✅ 支持' if r.hypothesis_support['h2_domain_specialization']['supported'] else '❌ 不支持'}

{r.hypothesis_support['h2_domain_specialization']['interpretation']}

---

## 五、总体结论

**实验结论**: {r.hypothesis_support['overall']}

### 解释

"""
        # 添加具体解释
        delta_r = (r.instruction_r_eff[r.config['base_model_name']] -
                  r.instruction_r_eff[r.config['instruct_model_name']])

        if delta_r > 0:
            report += f"""1. **谱压缩效应**: Instruct 微调使 r_eff 降低了 {delta_r:.4f} ({delta_r/r.instruction_r_eff[r.config['base_model_name']]*100:.1f}%)
   - 这表明对齐过程使模型在指令理解上使用更少的表示维度
   - 与"对齐即谱工程"的假设一致

2. **域专一化**:
"""
        else:
            report += f"""1. **谱压缩效应不显著**: Instruct 版 r_eff 反而高 {abs(delta_r):.4f}
   - 可能的原因：指令微调增加了表示的多样性
   - 需要更多模型对验证

2. **域专一化**:
"""

        stg_diff = r.stg[r.config['instruct_model_name']] - r.stg[r.config['base_model_name']]
        if stg_diff > 0:
            report += f"""   - Instruct 版的 STG 比 Base 版高 {stg_diff:.4f}
   - 表明对齐增强了模型对指令域的专一化
"""
        else:
            report += f"""   - Instruct 版的 STG 比 Base 版低 {abs(stg_diff):.4f}
   - 表明对齐没有显著增强域专一化
"""

        report += """
---

## 六、局限性

1. **单一模型对**: 仅使用 LLaMA-3-8B Base/Instruct 一对，结论普适性有限
2. **数据集有限**: 仅使用 alpaca 作为指令集，可能需要更多指令数据验证
3. **缺少 RLHF/DPO**: 仅对比 SFT，没有包含更强的对齐方式

## 七、未来工作

- 增加 Qwen2.5-7B Base/Instruct 对验证
- 添加 DPO/RLHF 版本对比
- 分析逐层谱变化（LCP）

---

*报告生成时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*
"""

        report_file = self.output_dir / "e4_report.md"
        with open(report_file, "w", encoding="utf-8") as f:
            f.write(report)

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


def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(description="E4 实验：对齐即谱工程")
    parser.add_argument("--model-pair", type=str, default="qwen",
                       choices=["qwen", "llama3", "mistral"],
                       help="选择模型对：qwen (Qwen2.5-7B), llama3 (LLaMA-3-8B), 或 mistral (Mistral-7B)")
    parser.add_argument("--sample-size", type=int, default=200,
                       help="采样文本数量")
    parser.add_argument("--output-dir", type=str, default=None,
                       help="输出目录")

    args = parser.parse_args()

    # 根据选择配置模型对
    if args.model_pair == "qwen":
        # Qwen2.5-7B 对（真正的 Base/Instruct）
        config = E4Config(
            base_model_name="Qwen2.5-7B-Base",
            instruct_model_name="Qwen2.5-7B-Instruct",
            base_model_path="model/weights/Qwen2.5-7B-Base",
            instruct_model_path="model/weights/Qwen2.5-7B-Instruct",
            instruction_dataset="alpaca",
            general_dataset="flores200",
            sample_size=args.sample_size,
            batch_size=4,
            output_dir=args.output_dir or "experiments/output/e4_qwen"
        )
    elif args.model_pair == "llama3":
        # LLaMA-3-8B 对（注意：之前发现两个权重相同，仅用于测试）
        config = E4Config(
            base_model_name="llama3-8b",
            instruct_model_name="llama3-8b-instruct",
            base_model_path="model/weights/llama3-8b",
            instruct_model_path="model/weights/llama3-8b-instruct",
            instruction_dataset="alpaca",
            general_dataset="flores200",
            sample_size=args.sample_size,
            batch_size=4,
            output_dir=args.output_dir or "experiments/output/e4_llama3"
        )
    else:  # mistral
        # Mistral-7B 对
        config = E4Config(
            base_model_name="Mistral-7B-v0.3",
            instruct_model_name="Mistral-7B-Instruct-v0.3",
            base_model_path="model/weights/Mistral-7B-v0.3",
            instruct_model_path="model/weights/Mistral-7B-Instruct-v0.3",
            instruction_dataset="alpaca",
            general_dataset="flores200",
            sample_size=args.sample_size,
            batch_size=4,
            output_dir=args.output_dir or "experiments/output/e4_mistral"
        )

    analyzer = AlignmentSpectrumAnalyzer(config)
    result = analyzer.run_full_analysis()

    print("\n" + "=" * 60)
    print("E4 实验完成!")
    print("=" * 60)
    print(f"\n核心发现:")
    print(f"  指令集 r_eff: {result.instruction_r_eff[result.config['base_model_name']]:.4f} (Base) -> "
          f"{result.instruction_r_eff[result.config['instruct_model_name']]:.4f} (Instruct)")
    print(f"  域专一化 STG: {result.stg[result.config['base_model_name']]:.4f} (Base) -> "
          f"{result.stg[result.config['instruct_model_name']]:.4f} (Instruct)")
    print(f"\n结论：{result.hypothesis_support['overall']}")


if __name__ == "__main__":
    main()