"""
03_model_finetune.py - 模型微调（LoRA/QLoRA + 公平性约束损失）

使用说明：
    本脚本对基座模型 Qwen1.5-1.8B-Chat 进行参数高效微调，实现四个消融实验变体：
        Model A：原始模型（无任何微调，对照组）
        Model B：仅数据净化（SFT with aligned data）
        Model C：仅KL散度约束损失（原始数据 + 公平性损失）
        Model D：双重纠偏（对齐数据 + 公平性损失）

    KL 约束损失公式：
        L_total = L_task + λ · D_KL(logits_male || logits_female)

运行前提：
    1. 已运行 02_data_debiasing.py，data/debiased_dataset/ 目录已生成
    2. 环境中已安装 CUDA（GPU 训练），或使用 CPU 模式（极慢，仅供测试）
    3. 已安装所有依赖：pip install -r requirements.txt
    4. 建议显存 ≥ 12 GB（QLoRA 4-bit 量化）

运行方式：
    # 训练所有变体（B/C/D）
    python 03_model_finetune.py

    # 仅训练指定变体
    python 03_model_finetune.py --variant B
    python 03_model_finetune.py --variant D

命令行参数：
    --variant   {B,C,D,ALL}  指定训练哪个变体，默认 ALL

输出：
    models/model_b_data_only/       Model B 权重
    models/model_c_loss_only/       Model C 权重
    models/model_d_dual/            Model D 权重

注意：
    - Model A 不需要训练，直接使用原始基座模型权重
    - 若显存不足，可减小 per_device_train_batch_size 或关闭 fp16

作者：研究团队
日期：2025
"""

from __future__ import annotations

import sys
import logging
import argparse
from pathlib import Path
from typing import Optional

import torch
import torch.nn.functional as F
import numpy as np
from datasets import load_from_disk, Dataset, DatasetDict

sys.path.insert(0, str(Path(__file__).parent))
import config

# ============================================================
# 日志配置
# ============================================================
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger(__name__)


# ============================================================
# 工具函数
# ============================================================

def set_seed(seed: int = config.RANDOM_SEED) -> None:
    """固定随机种子，保证实验可重复"""
    import random
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def load_tokenizer_and_model(use_qlora: bool = True):
    """
    加载基座模型和分词器。

    参数：
        use_qlora: 是否使用 QLoRA 4-bit 量化（节省显存）

    返回：
        (tokenizer, model)
    """
    from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

    logger.info(f"加载基座模型: {config.BASE_MODEL_ID}")
    tokenizer = AutoTokenizer.from_pretrained(
        config.BASE_MODEL_ID,
        trust_remote_code=True,
        padding_side="right",
    )
    # 确保 pad_token 存在
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # 量化配置
    bnb_config = None
    if use_qlora and torch.cuda.is_available():
        from transformers import BitsAndBytesConfig
        q_cfg = config.QLORA_QUANTIZATION
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=q_cfg["load_in_4bit"],
            bnb_4bit_quant_type=q_cfg["bnb_4bit_quant_type"],
            bnb_4bit_compute_dtype=getattr(torch, q_cfg["bnb_4bit_compute_dtype"]),
            bnb_4bit_use_double_quant=q_cfg["bnb_4bit_use_double_quant"],
        )

    model = AutoModelForCausalLM.from_pretrained(
        config.BASE_MODEL_ID,
        quantization_config=bnb_config,
        device_map="auto" if torch.cuda.is_available() else "cpu",
        trust_remote_code=True,
    )
    model.config.use_cache = False  # 训练时关闭 KV Cache

    if use_qlora and torch.cuda.is_available():
        from peft import prepare_model_for_kbit_training
        model = prepare_model_for_kbit_training(model)

    return tokenizer, model


def apply_lora(model):
    """为模型添加 LoRA 适配器"""
    from peft import LoraConfig, get_peft_model, TaskType

    lora_cfg = config.LORA_CONFIG
    peft_config = LoraConfig(
        r=lora_cfg["r"],
        lora_alpha=lora_cfg["lora_alpha"],
        target_modules=lora_cfg["target_modules"],
        lora_dropout=lora_cfg["lora_dropout"],
        bias=lora_cfg["bias"],
        task_type=TaskType.CAUSAL_LM,
    )
    model = get_peft_model(model, peft_config)
    model.print_trainable_parameters()
    return model


def load_dataset_for_variant(variant: str) -> DatasetDict:
    """
    根据训练变体加载对应的数据集。

    变体 B / D 使用对齐后的净化数据集；
    变体 C 使用原始有偏数据（从 bias_detection_results.csv 重建）。

    参数：
        variant: "B" | "C" | "D"

    返回：
        HuggingFace DatasetDict
    """
    if variant in ("B", "D"):
        # 对齐数据集
        dataset_path = config.DEBIASED_DATASET_DIR
        if not Path(dataset_path).exists():
            raise FileNotFoundError(
                f"净化数据集不存在: {dataset_path}\n"
                "请先运行 02_data_debiasing.py"
            )
        ds = load_from_disk(dataset_path)
        logger.info(f"[变体 {variant}] 已加载净化数据集: {dataset_path}")
        return ds

    elif variant == "C":
        # 使用原始有偏数据（或模拟数据）重建
        import pandas as pd
        from config import SCENARIOS, PROMPT_TEMPLATE, COMPENSATION_MULTIPLIER

        bias_csv = Path(config.BIAS_RESULT_CSV)
        if bias_csv.exists():
            df = pd.read_csv(bias_csv, encoding="utf-8-sig")
            logger.info(f"[变体 C] 已加载原始偏见数据: {bias_csv}")
        else:
            # 使用模拟数据（importlib动态导入以数字开头的文件名模块）
            logger.warning("[变体 C] 原始偏见数据不存在，使用模拟有偏数据")
            import importlib.util as _ilu
            _spec = _ilu.spec_from_file_location(
                "data_debiasing", Path(__file__).parent / "02_data_debiasing.py"
            )
            _mod = _ilu.module_from_spec(_spec)
            _spec.loader.exec_module(_mod)
            df = _mod.generate_mock_bias_results()

        # 将原始（有偏）数据构建为训练样本
        scenario_map = {s["id"]: s for s in SCENARIOS}
        records = []
        valid = df.dropna(subset=["compensation"])
        grouped = (
            valid
            .groupby(["scenario_id", "gender"])["compensation"]
            .mean()
            .reset_index()
        )
        for _, row in grouped.iterrows():
            scenario = scenario_map.get(row["scenario_id"])
            if scenario is None:
                continue
            instruction = PROMPT_TEMPLATE.format(
                gender=row["gender"],
                position=scenario["position"],
                years=scenario["years"],
                salary=scenario["salary"],
                education=scenario.get("education", "不详"),
                family=scenario.get("family", "不详"),
                reason=scenario["reason"],
                evidence=scenario.get("evidence", ""),
            )
            # 直接使用模型原始输出作为目标（保留偏见，依靠损失函数纠正）
            comp = row["compensation"]
            output_text = f"建议赔偿金：{int(comp)}元"
            text = (
                f"<|im_start|>user\n{instruction}<|im_end|>\n"
                f"<|im_start|>assistant\n{output_text}<|im_end|>"
            )
            records.append({
                "instruction": instruction,
                "input": "",
                "output": output_text,
                "text": text,
                "scenario_id": row["scenario_id"],
                "gender": row["gender"],
            })

        np.random.seed(config.RANDOM_SEED)
        indices = np.random.permutation(len(records))
        split = int(len(records) * 0.85)
        train_ds = Dataset.from_list([records[i] for i in indices[:split]])
        val_ds   = Dataset.from_list([records[i] for i in indices[split:]])
        return DatasetDict({"train": train_ds, "validation": val_ds})

    else:
        raise ValueError(f"未知变体: {variant}，应为 B / C / D")


# ============================================================
# 公平性约束损失（KL 散度正则化）
# ============================================================

class FairnessKLLoss(torch.nn.Module):
    """
    KL 散度约束损失模块。

    对于一个 batch，要求包含成对的（男性输入, 女性输入）样本。
    损失公式：
        L_total = L_task + λ · D_KL(P_male_logits || P_female_logits)

    其中 P 是对 logits 做 softmax 后的概率分布（仅取最后一个 token 位置）。
    """

    def __init__(self, lam: float = config.KL_LAMBDA):
        super().__init__()
        self.lam = lam

    def forward(
        self,
        task_loss: torch.Tensor,
        logits_male: torch.Tensor,
        logits_female: torch.Tensor,
    ) -> torch.Tensor:
        """
        参数：
            task_loss:     标准语言模型任务损失（交叉熵）
            logits_male:   男性输入对应的 logits，形状 [batch, seq_len, vocab_size]
            logits_female: 女性输入对应的 logits，形状 [batch, seq_len, vocab_size]

        返回：
            总损失（task_loss + λ * KL）
        """
        # 取最后一个有效 token 位置的 logits
        p_male   = F.softmax(logits_male[:, -1, :],   dim=-1)
        p_female = F.softmax(logits_female[:, -1, :], dim=-1)

        # D_KL(P_male || P_female)
        kl_div = F.kl_div(
            input=p_female.log(),        # log(Q)
            target=p_male,               # P
            reduction="batchmean",
        )
        return task_loss + self.lam * kl_div


# ============================================================
# 自定义 Trainer（支持 KL 约束损失）
# ============================================================

class FairnessTrainer:
    """
    封装带 KL 约束损失的训练逻辑。
    当 use_fairness_loss=True 时，需要数据集中包含 gender 字段，
    训练时对男女样本分别前向传播后计算 KL 约束。

    为简化实现，此处使用手动训练循环而非 HuggingFace Trainer 子类，
    以确保在各版本 transformers 下的兼容性。
    """

    def __init__(
        self,
        model,
        tokenizer,
        train_dataset: Dataset,
        eval_dataset: Dataset,
        output_dir: str,
        use_fairness_loss: bool = False,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.use_fairness_loss = use_fairness_loss
        self.fairness_loss_fn = FairnessKLLoss(lam=config.KL_LAMBDA)

        # 超参数
        self.args = config.TRAIN_ARGS
        self.device = next(model.parameters()).device

    def _tokenize_batch(self, texts: list[str]) -> dict[str, torch.Tensor]:
        """将文本列表分词并截断填充"""
        encoded = self.tokenizer(
            texts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=1024,
        )
        return {k: v.to(self.device) for k, v in encoded.items()}

    def _forward_pass(
        self,
        batch_texts: list[str],
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        对 batch 执行前向传播。

        返回：
            (task_loss, logits)
        """
        enc = self._tokenize_batch(batch_texts)
        # labels = input_ids（Causal LM 自回归训练）
        outputs = self.model(
            input_ids=enc["input_ids"],
            attention_mask=enc["attention_mask"],
            labels=enc["input_ids"],
        )
        return outputs.loss, outputs.logits

    def _train_epoch(
        self,
        optimizer: torch.optim.Optimizer,
        scheduler,
    ) -> float:
        """执行一个 epoch 的训练，返回平均损失"""
        from torch.utils.data import DataLoader

        self.model.train()
        total_loss = 0.0
        n_steps = 0

        # 构建数据加载器
        loader = DataLoader(
            self.train_dataset,
            batch_size=self.args["per_device_train_batch_size"],
            shuffle=True,
            collate_fn=lambda x: x,  # 保持 list of dict
        )

        grad_acc = self.args["gradient_accumulation_steps"]
        optimizer.zero_grad()

        for step, batch in enumerate(loader):
            texts = [item["text"] for item in batch]

            if self.use_fairness_loss and "gender" in batch[0]:
                # 分离男女样本
                male_texts   = [t for item, t in zip(batch, texts) if item.get("gender") == "男"]
                female_texts = [t for item, t in zip(batch, texts) if item.get("gender") == "女"]

                if male_texts and female_texts:
                    loss_m, logits_m = self._forward_pass(male_texts)
                    loss_f, logits_f = self._forward_pass(female_texts)
                    task_loss = (loss_m + loss_f) / 2
                    loss = self.fairness_loss_fn(task_loss, logits_m, logits_f)
                else:
                    # 若 batch 内只有单一性别，退化为普通任务损失
                    loss, _ = self._forward_pass(texts)
            else:
                loss, _ = self._forward_pass(texts)

            # 梯度累积
            (loss / grad_acc).backward()
            total_loss += loss.item()
            n_steps += 1

            if (step + 1) % grad_acc == 0:
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()

                if n_steps % self.args["logging_steps"] == 0:
                    logger.info(
                        f"  Step {n_steps:4d} | "
                        f"Loss: {total_loss / n_steps:.4f} | "
                        f"LR: {scheduler.get_last_lr()[0]:.2e}"
                    )

        return total_loss / max(n_steps, 1)

    def train(self) -> None:
        """完整训练流程"""
        from transformers import get_cosine_schedule_with_warmup

        # 构建优化器
        optimizer = torch.optim.AdamW(
            filter(lambda p: p.requires_grad, self.model.parameters()),
            lr=self.args["learning_rate"],
        )

        # 计算 warmup steps
        steps_per_epoch = max(
            len(self.train_dataset) // self.args["per_device_train_batch_size"],
            1,
        )
        total_steps = steps_per_epoch * self.args["num_train_epochs"]
        warmup_steps = int(total_steps * self.args["warmup_ratio"])

        scheduler = get_cosine_schedule_with_warmup(
            optimizer,
            num_warmup_steps=warmup_steps,
            num_training_steps=total_steps,
        )

        best_eval_loss = float("inf")

        for epoch in range(1, self.args["num_train_epochs"] + 1):
            logger.info(f"Epoch {epoch}/{self.args['num_train_epochs']}")
            train_loss = self._train_epoch(optimizer, scheduler)
            eval_loss  = self._evaluate()
            logger.info(
                f"  Epoch {epoch} | Train Loss: {train_loss:.4f} | "
                f"Eval Loss: {eval_loss:.4f}"
            )

            # 保存最优模型
            if eval_loss < best_eval_loss:
                best_eval_loss = eval_loss
                self._save_model("best")
                logger.info(f"  新最优模型已保存 (eval_loss={best_eval_loss:.4f})")

        # 保存最终模型
        self._save_model("final")
        logger.info(f"训练完成，最终模型保存至: {self.output_dir}")

    def _evaluate(self) -> float:
        """在验证集上计算平均损失"""
        from torch.utils.data import DataLoader

        self.model.eval()
        total_loss = 0.0
        n_steps = 0

        loader = DataLoader(
            self.eval_dataset,
            batch_size=self.args["per_device_train_batch_size"],
            shuffle=False,
            collate_fn=lambda x: x,
        )

        with torch.no_grad():
            for batch in loader:
                texts = [item["text"] for item in batch]
                loss, _ = self._forward_pass(texts)
                total_loss += loss.item()
                n_steps += 1

        return total_loss / max(n_steps, 1)

    def _save_model(self, suffix: str = "final") -> None:
        """保存 LoRA 适配器权重及分词器"""
        save_path = self.output_dir / suffix
        self.model.save_pretrained(str(save_path))
        self.tokenizer.save_pretrained(str(save_path))


# ============================================================
# 各变体训练入口
# ============================================================

def train_model_b(tokenizer, base_model, dataset: DatasetDict) -> None:
    """
    Model B：仅数据净化（SFT with aligned data）
    使用对齐后数据做标准 SFT，不加公平性损失。
    """
    logger.info("=" * 50)
    logger.info("训练 Model B：仅数据净化（SFT with aligned data）")
    logger.info("=" * 50)

    # 为 B 克隆模型（避免共享参数影响其他变体）
    import copy
    model_b = apply_lora(copy.deepcopy(base_model))

    trainer = FairnessTrainer(
        model=model_b,
        tokenizer=tokenizer,
        train_dataset=dataset["train"],
        eval_dataset=dataset["validation"],
        output_dir=config.ABLATION_MODELS["Model_B_DataOnly"],
        use_fairness_loss=False,
    )
    trainer.train()


def train_model_c(tokenizer, base_model, dataset: DatasetDict) -> None:
    """
    Model C：仅 KL 散度约束损失（原始有偏数据 + 公平性损失）
    """
    logger.info("=" * 50)
    logger.info("训练 Model C：仅 KL 约束损失（原始数据 + 公平性损失）")
    logger.info("=" * 50)

    import copy
    model_c = apply_lora(copy.deepcopy(base_model))

    # 加载原始有偏数据（variant C）
    dataset_c = load_dataset_for_variant("C")

    trainer = FairnessTrainer(
        model=model_c,
        tokenizer=tokenizer,
        train_dataset=dataset_c["train"],
        eval_dataset=dataset_c["validation"],
        output_dir=config.ABLATION_MODELS["Model_C_LossOnly"],
        use_fairness_loss=True,
    )
    trainer.train()


def train_model_d(tokenizer, base_model, dataset: DatasetDict) -> None:
    """
    Model D：双重纠偏（对齐数据 + 公平性损失）
    """
    logger.info("=" * 50)
    logger.info("训练 Model D：双重纠偏（对齐数据 + 公平性损失）")
    logger.info("=" * 50)

    import copy
    model_d = apply_lora(copy.deepcopy(base_model))

    trainer = FairnessTrainer(
        model=model_d,
        tokenizer=tokenizer,
        train_dataset=dataset["train"],
        eval_dataset=dataset["validation"],
        output_dir=config.ABLATION_MODELS["Model_D_Dual"],
        use_fairness_loss=True,
    )
    trainer.train()


# ============================================================
# 主入口
# ============================================================

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="法律LLM性别偏见消融实验微调")
    parser.add_argument(
        "--variant",
        type=str,
        default="ALL",
        choices=["B", "C", "D", "ALL"],
        help="训练哪个变体（B/C/D/ALL），默认 ALL",
    )
    parser.add_argument(
        "--no_qlora",
        action="store_true",
        help="禁用 QLoRA 量化（适用于无 CUDA 环境测试）",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    set_seed(config.RANDOM_SEED)

    logger.info("=" * 60)
    logger.info("模型微调实验（LoRA/QLoRA + 公平性约束损失）启动")
    logger.info(f"基座模型: {config.BASE_MODEL_ID}")
    logger.info(f"训练变体: {args.variant}")
    logger.info(f"CUDA 可用: {torch.cuda.is_available()}")
    logger.info("=" * 60)

    # 加载基座模型（共享，各变体会 deepcopy 后独立修改）
    use_qlora = not args.no_qlora
    tokenizer, base_model = load_tokenizer_and_model(use_qlora=use_qlora)

    # 加载对齐数据集（B / D 使用）
    try:
        aligned_dataset = load_dataset_for_variant("B")
    except FileNotFoundError as e:
        logger.error(str(e))
        logger.info("将使用模拟数据继续（请先运行 02_data_debiasing.py 获取真实数据）")
        # 使用 importlib 动态导入以文件名数字开头的模块
        import importlib.util
        spec = importlib.util.spec_from_file_location(
            "data_debiasing",
            Path(__file__).parent / "02_data_debiasing.py"
        )
        debiasing_mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(debiasing_mod)
        mock_df = debiasing_mod.generate_mock_bias_results()
        pivot = debiasing_mod.build_pivot_table(mock_df)
        samples = debiasing_mod.create_debiased_pairs(pivot)
        aligned_dataset = debiasing_mod.build_and_save_dataset(samples)

    # 根据参数训练对应变体
    variants_to_train = ["B", "C", "D"] if args.variant == "ALL" else [args.variant]

    for variant in variants_to_train:
        if variant == "B":
            train_model_b(tokenizer, base_model, aligned_dataset)
        elif variant == "C":
            train_model_c(tokenizer, base_model, aligned_dataset)
        elif variant == "D":
            train_model_d(tokenizer, base_model, aligned_dataset)

    logger.info("\n所有指定变体训练完成！")
    for name, path in config.ABLATION_MODELS.items():
        p = Path(path)
        if p.exists():
            logger.info(f"  {name}: {path} [存在]")
        else:
            logger.info(f"  {name}: {path} [未训练]")


if __name__ == "__main__":
    main()
