"""
Head-LoRA Finetune (Refactored)
==========================================
Target: Align specific attention heads to ignore 'data' tokens using Flash Attention Interception.

Modification: Unselected heads in target layers are preserved via KL loss against
the original model (no LoRA, no blind mask).
"""

import os
import json
import csv
import argparse
import math
import random
import re
from typing import List, Tuple, Dict

import torch
import numpy as np
from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader
from evallib import quick_eval_mmlu, quick_eval_asr_util
from transformers import (
    AutoConfig, AutoTokenizer, AutoModelForCausalLM,
    BitsAndBytesConfig, get_linear_schedule_with_warmup,
)
from transformers.models.llama.modeling_llama import ALL_ATTENTION_FUNCTIONS
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, PeftModel
from lib.tokenize_data_mask import apply_chat_tokenize_with_strip_and_mark
from lib.printcolor import print_tok_color
# -----------------------------------------------------------------------------
# 1. Flash Attention Spy Interceptor (核心機制)
# -----------------------------------------------------------------------------
print("🥷 Injecting Spy Interceptor into Flash Attention 2...")

if "flash_attention_2" not in ALL_ATTENTION_FUNCTIONS:
    raise RuntimeError("Current environment does not support Flash Attention 2!")

orig_flash_attn = ALL_ATTENTION_FUNCTIONS["flash_attention_2"]

def wrapped_flash_attn(module, query, key, value, attention_mask, scaling, **kwargs):
    """
    攔截 Flash Attention，並在 Breakpoint 驗證手算結果與原始 FA2 結果的一致性。
    """
    # -------------------------------------------------------------------------
    # [TASK A] 間諜行動：偷算 Last Token Weight
    # -------------------------------------------------------------------------
    # 1. 切片 (只看最後一個 Token)
    q_last = query[..., -1:, :]

    # 2. 處理 GQA (手動 Expand Key 和 Value)
    num_heads = query.shape[1]
    num_kv_heads = key.shape[1]

    # 用於後面計算 Manual Output
    k_expanded = key
    v_expanded = value

    if num_heads != num_kv_heads:
        n_rep = num_heads // num_kv_heads
        k_expanded = key.repeat_interleave(n_rep, dim=1)
        v_expanded = value.repeat_interleave(n_rep, dim=1) # Value 也要 Expand 才能乘

    # 3. 手動計算 Tiny Attention Weight [Batch, Heads, 1, Seq]
    # Q: [B, H, 1, D], K: [B, H, S, D] -> QK^T: [B, H, 1, S]
    attn_weights_tiny = torch.matmul(q_last, k_expanded.transpose(2, 3)) * scaling

    # 4. 處理 Mask
    if attention_mask is not None:
        if attention_mask.dim() == 4:
            if attention_mask.size(2) > 1:
                mask_slice = attention_mask[..., -1:, :]
                attn_weights_tiny = attn_weights_tiny + mask_slice
            else:
                attn_weights_tiny = attn_weights_tiny + attention_mask
        elif attention_mask.dim() == 2:
            mask_expanded = attention_mask[:, None, None, :]
            if attention_mask.dtype == torch.bool or (attention_mask.max() <= 1.0 and attention_mask.min() >= 0.0):
                min_dtype = torch.finfo(attn_weights_tiny.dtype).min
                attn_weights_tiny = attn_weights_tiny.masked_fill(mask_expanded == 0, min_dtype)
            else:
                attn_weights_tiny = attn_weights_tiny + mask_expanded

    # 5. Softmax (得到機率分佈 P)
    # 注意：這裡轉成 float32 做 softmax 以求精確，實際 FA2 內部可能是 fp16/bf16
    attn_probs = torch.nn.functional.softmax(attn_weights_tiny, dim=-1, dtype=torch.float32).to(query.dtype)

    # 填入 Config (原本的邏輯)
    retrieve_map = getattr(module.config, "retrieve_attn_map", None)
    if retrieve_map is not None and hasattr(module, "layer_idx") and module.layer_idx in retrieve_map:
        target_heads_dict = retrieve_map[module.layer_idx]
        for h_idx in target_heads_dict.keys():
            # 存入 config 用於 Loss 計算
            target_heads_dict[h_idx] = attn_probs[:, h_idx, :, :]

    # -------------------------------------------------------------------------
    # [DEBUG] 驗證環節：手算 Output vs Flash Attention Output
    # -------------------------------------------------------------------------

    # 1. 執行原始 Flash Attention
    orig_result = orig_flash_attn(module, query, key, value, attention_mask, scaling=scaling, **kwargs)
    debug=False
    if debug:
        orig_tensor = orig_result[0]
        # 2. 手動計算 Output ( O = Attention_Probs * V )
        # attn_probs: [B, H, 1, S]
        # v_expanded: [B, H, S, D]
        # manual_out: [B, H, 1, D]
        manual_out = torch.matmul(attn_probs, v_expanded)

        # 3. 對齊形狀以便比較
        # FA2 output 通常是 [Batch, Seq, Heads, Dim]
        # 我們取最後一個 token: [Batch, 1, Heads, Dim]
        orig_out_last = orig_tensor[:, -1:, :, :]

        # 手算結果原本是 [Batch, Heads, 1, Dim]，轉置成 [Batch, 1, Heads, Dim]
        manual_out_check = manual_out.transpose(1, 2)

        # 4. 計算誤差
        diff = (orig_out_last - manual_out_check).abs()
        max_diff = diff.max().item()
        mean_diff = diff.mean().item()

        print(f"\n🕵️ [Layer {getattr(module, 'layer_idx', '?')}] Validation:")
        print(f"   Max Diff: {max_diff:.8f}")
        print(f"   Mean Diff: {mean_diff:.8f}")

        # 如果誤差過大 (例如 > 1e-3)，自動暫停檢查
        # 注意：BF16 下 FlashAttention 和標準 Matmul 會有一定精度差異是正常的
        if max_diff > 1e-2:
            print("⚠️ Warning: Large difference detected!")

        # 呼叫 breakpoint 讓你進去檢查
        # 在 pdb 輸入:
        #   p manual_out_check[0,0,0,:5]
        #   p orig_out_last[0,0,0,:5]
        # 來對比數值
        breakpoint()

    return orig_result

# 替換全局函數
ALL_ATTENTION_FUNCTIONS["flash_attention_2"] = wrapped_flash_attn
print("✅ Spy Interceptor installed successfully!")

# -----------------------------------------------------------------------------
# 2. Utils: Model & Heads
# -----------------------------------------------------------------------------
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)

def get_lora_targets(model, layers: List[int]) -> List[str]:
    mtype = (getattr(model.config, "model_type", "") or "").lower()
    if "llama" in mtype or "mistral" in mtype:
        return [f"model.layers.{i}.self_attn.{p}" for i in layers for p in ("q_proj", "k_proj")]
    # Fallback / Other architectures
    cand = []
    for name, _ in model.named_modules():
        if any(f".{i}." in name for i in layers) and name.split(".")[-1] in {"q_proj", "k_proj"}:
            cand.append(name)
    return cand

def load_model(model_path: str):
    cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
    tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, use_fast=True,padding_side="left")

    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token
        tok.pad_token_id = tok.eos_token_id

    bnb_cfg = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
    )
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        config=cfg,
        quantization_config=bnb_cfg,
        device_map="auto",
        trust_remote_code=True,
        attn_implementation="flash_attention_2",
    )
    return model, tok, tok # use same tokenizer for simplicity

def get_target_structure(heads_file: str, topk_spec: str) -> Dict[int, List[int]]:
    """Load heads from file and parse into {layer: [heads]} structure."""
    if not os.path.exists(heads_file):
        raise FileNotFoundError(f"Heads file not found: {heads_file}")

    with open(heads_file, "r") as f:
        ihead_list = json.load(f)
        if type(ihead_list[0]) == str:
            all_heads = [(str(tag), 0.0) for tag in ihead_list]
        else:
            all_heads = ihead_list

    # Filter Top-K
    count = len(all_heads)
    if topk_spec.endswith("p"):
        count = max(1, math.ceil(float(topk_spec[:-1]) / 100.0 * len(all_heads)))
    else:
        count = int(topk_spec)

    target_heads = all_heads[:min(count, len(all_heads))]

    # Parse into structure
    structure = {}
    for tag, _ in target_heads:
        try:
            # tag format: "L22_H7"
            parts = tag.split('_')
            l = int(parts[0][1:])
            h = int(parts[1][1:])
            structure.setdefault(l, []).append(h)
        except Exception:
            continue

    print(f"🎯 Selected {len(target_heads)} heads from {topk_spec}.")
    return structure

def discover_existing_adapter(out_dir: str):
    if not os.path.isdir(out_dir): return None, 0, 0
    candidates = []
    if os.path.exists(os.path.join(out_dir, "adapter_config.json")):
        candidates.append((0, 0, out_dir))

    for entry in os.listdir(out_dir):
        path = os.path.join(out_dir, entry)
        if os.path.isdir(path) and os.path.exists(os.path.join(path, "adapter_config.json")):
            m = re.match(r"^batch_(\d+)_(\d+)$", entry)
            if m:
                epoch_idx = int(m.group(1))
                batch_idx = int(m.group(2))
                candidates.append((epoch_idx, batch_idx, path))

    if not candidates: return None, 0, 0
    candidates.sort(key=lambda x: (x[0], x[1]))
    return candidates[-1][2], candidates[-1][0], candidates[-1][1]

# -----------------------------------------------------------------------------
# 3. Loss Functions
# -----------------------------------------------------------------------------
def head_attention_loss(tuned_map, base_map, data_mask, lambda_data=1.0, eps=1e-8):
    """Loss for SELECTED heads: KL toward blind teacher + data mass penalty."""
    mask_data  = data_mask.bool()               # [B,S]
    mask_valid = ~mask_data

    if mask_valid.dim() == 2:
        mask_valid = mask_valid.unsqueeze(1)    # [B,1,S]
        mask_data  = mask_data.unsqueeze(1)     # [B,1,S]

    total = 0.0
    n = 0

    for l, heads in tuned_map.items():
        for h, tuned in heads.items():
            base = base_map.get(l, {}).get(h, None)
            if tuned is None or base is None:
                continue

            tuned = tuned.float()   # [B,1,S]
            base  = base.float().to(tuned.device)

            # 1) teacher: base 已經是 blind 分佈（理想上 data=0），但我們也只取 valid 區域
            base_v = base * mask_valid.float()
            base_v = base_v / base_v.sum(dim=-1, keepdim=True).clamp_min(eps)

            # 2) student: tuned 在 valid 區域重新 normalize
            tuned_v = tuned * mask_valid.float()
            tuned_v_sum = tuned_v.sum(dim=-1, keepdim=True).clamp_min(eps)
            tuned_v = tuned_v / tuned_v_sum

            # 3) KL(base || tuned)
            kl = (base_v * (torch.log(base_v + eps) - torch.log(tuned_v + eps))).sum(dim=-1).mean()

            # 4) data mass penalty（不重正規化，直接壓 data 的注意力總量）
            data_mass = (tuned * mask_data.float()).sum(dim=-1).mean()

            total = total + kl + lambda_data * data_mass
            n += 1

    return total / max(n, 1)


def head_preservation_loss(tuned_map, orig_map, eps=1e-8):
    """Loss for UNSELECTED heads: KL toward original model (no LoRA, no mask)."""
    total = 0.0
    n = 0
    for l, heads in orig_map.items():
        for h, orig in heads.items():
            tuned = tuned_map.get(l, {}).get(h, None)
            if tuned is None or orig is None:
                continue
            tuned = tuned.float()
            orig = orig.float().to(tuned.device)
            # KL(orig || tuned) — 讓 tuned 分佈接近 orig
            kl = (orig * (torch.log(orig + eps) - torch.log(tuned + eps))).sum(dim=-1).mean()
            total += kl
            n += 1
    return total / max(n, 1)


# -----------------------------------------------------------------------------
# 4. Data Loading
# -----------------------------------------------------------------------------
class JsonlMessagesDS(Dataset):
    def __init__(self, jsonl_path: str):
        self.samples = []
        with open(jsonl_path, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    obj = json.loads(line)
                    self.samples.append(obj)
        random.shuffle(self.samples)
        print(f"📊 Loaded {len(self.samples)} samples.")

    def __len__(self): return len(self.samples)
    def __getitem__(self, idx): return self.samples[idx]

def collate(batch, tokenizer):
    input_ids, attention_mask, data_mask = apply_chat_tokenize_with_strip_and_mark(
        batch, tokenizer, device="cpu", add_generation_prompt=True,
        encode_kwargs = {"padding_side" :"left" },
        mode="custom_mask == 'inst'",
        custom_mask_identifier={"data": ["<data>", "</data>"], "inst": ["<inst>", "</inst>"]},
        return_tensors="pt"
    )

    return {"input_ids": input_ids, "attention_mask": attention_mask, "data_mask": data_mask}

# -----------------------------------------------------------------------------
# 5. Training Loop
# -----------------------------------------------------------------------------
def save_model(model, tok, out_dir, epoch, batch_idx, data_path):
    """Safely saves model by temporarily cleaning config, then logs eval results."""
    print("Saving checkpoint...")

    # 🚑 CRITICAL: Temporarily remove the tensor map from config to prevent DeepCopy error
    temp_map = getattr(model.config, "retrieve_attn_map", None)
    if temp_map is not None:
        del model.config.retrieve_attn_map

    save_dir = os.path.join(out_dir, f"batch_{epoch}_{batch_idx}")
    model.save_pretrained(save_dir)
    tok.save_pretrained(save_dir)

    print(f"✅ Saved to {save_dir}")

    mllu_result = quick_eval_mmlu(model, tok)
    asr_result = quick_eval_asr_util(
        model,
        tok,
        training_data_path=data_path,
    )
    log_path = os.path.join(out_dir, "training_log.csv")
    write_header = not os.path.exists(log_path)
    with open(log_path, "a", newline="") as f:
        writer = csv.writer(f)
        attack_columns = [
            ("naive", "a_naive", "v_naive"),
            ("ignore", "a_ignore", "v_ignore"),
            ("escape_separation", "a_escape", "v_escape"),
            ("completion_realcmb", "a_cmb", "v_cmb"),
            ("conv_attack", "a_conv", "v_conv"),
            ("none", "a_none", "v_none"),
        ]
        if write_header:
            header = ["epoch", "batch", "mmlu score", "asr_mode", "asr_status"]
            for _, a_col, v_col in attack_columns:
                header.append(a_col)
                header.append(v_col)
            writer.writerow(header)
        row = [epoch, batch_idx, mllu_result.get("accuracy"), asr_result.get("eval_mode"), asr_result.get("status")]
        metrics = asr_result.get("metrics", {}) if asr_result.get("status") == "ok" else {}
        for attack, _, _ in attack_columns:
            values = metrics.get(attack, {})
            row.append(values.get("asr"))
            row.append(values.get("valid_rate"))
        writer.writerow(row)
    return save_dir

def tune_new(model, tok, tok_inf, target_structure, data_path, out_dir, epochs, bs, lr,
             resume_adapter=None, start_epoch=0, start_batch_idx=0, batch_save_interval=-1,
             lambda_preserve=1.0):
    # Setup LoRA
    layers = sorted(target_structure.keys())
    targets = get_lora_targets(model, layers)

    lora_cfg = LoraConfig(r=32, lora_alpha=16, bias="none", target_modules=targets, task_type="CAUSAL_LM")
    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)

    if resume_adapter:
        model = PeftModel.from_pretrained(model, resume_adapter, is_trainable=True)
        print(f"♻️ Resumed LoRA: {resume_adapter}")
    else:
        model = get_peft_model(model, lora_cfg)

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    ratio = trainable / total * 100
    print(f"🔧 Trainable parameters: {trainable:,} / {total:,} ({ratio:.4f}% of total)")

    # -------------------------------------------------------------------------
    # Precompute unselected heads structure for preservation loss
    # -------------------------------------------------------------------------
    num_heads = model.config.num_attention_heads
    unselected_structure = {}
    for l in target_structure.keys():
        selected_set = set(target_structure[l])
        unselected = [h for h in range(num_heads) if h not in selected_set]
        if unselected:
            unselected_structure[l] = unselected

    has_unselected = len(unselected_structure) > 0
    if has_unselected:
        total_unselected = sum(len(v) for v in unselected_structure.values())
        print(f"🛡️ Preservation targets: {total_unselected} unselected heads across {len(unselected_structure)} layers (λ={lambda_preserve})")
    else:
        print("⚠️ No unselected heads found — preservation loss disabled.")

    dl = DataLoader(JsonlMessagesDS(data_path), batch_size=bs, shuffle=True, collate_fn=lambda b: collate(b, tok))
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    sch = get_linear_schedule_with_warmup(opt, int(0.05 * epochs * len(dl)), epochs * len(dl))

    current_idx = start_epoch

    for ep in range(epochs):
        if ep < start_epoch:
            continue
        pbar = tqdm(dl, desc=f"Ep {ep+1}/{epochs}")
        for idx, batch in enumerate(pbar):
            if ep == start_epoch and idx < start_batch_idx:
                continue
            device = next(model.parameters()).device
            MAX_LEN = 50000

            input_ids = batch["input_ids"][:, :MAX_LEN].to(device)
            attn_mask = batch["attention_mask"][:, :MAX_LEN].to(device) # 原始 Mask (看得到全部)
            data_mask = batch["data_mask"][:, :MAX_LEN].to(device)      # Data 標記

            # -------------------------------------------------------
            # [Step 1] 建構 "Blind Mask" 給 Base Model
            # -------------------------------------------------------
            blind_attn_mask = attn_mask * (~data_mask).long()

            # --- A. Base Model Pass (Reference: Blind — no LoRA, with blind mask) ---
            # 收集「選中 head」的 blind attention 作為 teacher
            model.eval()
            base_map_container = {l: {h: None for h in h_list} for l, h_list in target_structure.items()}
            model.config.retrieve_attn_map = base_map_container

            with torch.no_grad(), model.disable_adapter():
                model(input_ids=input_ids, attention_mask=blind_attn_mask, output_attentions=False, use_cache=False)
                base_attns_map = base_map_container

            # --- A2. Original Model Pass (Reference: Full visibility, no LoRA, no blind mask) ---
            # 收集「未選中 head」的原始 attention 作為 preservation teacher
            if has_unselected:
                orig_map_container = {l: {h: None for h in h_list} for l, h_list in unselected_structure.items()}
                model.config.retrieve_attn_map = orig_map_container

                with torch.no_grad(), model.disable_adapter():
                    model(input_ids=input_ids, attention_mask=attn_mask, output_attentions=False, use_cache=False)
                orig_attns_map = orig_map_container
            else:
                orig_attns_map = None

            # --- B. Tuned Model Pass (Training: See Everything) ---
            # 收集「選中 + 未選中」head 的 tuned attention
            model.train()
            tuned_map_container = {}
            for l in target_structure.keys():
                selected_set = set(target_structure[l])
                unselected = unselected_structure.get(l, [])
                tuned_map_container[l] = {h: None for h in list(selected_set) + unselected}
            model.config.retrieve_attn_map = tuned_map_container

            # Tuned Model 仍然傳入原始 attn_mask (它看得到 Data，但我們要訓練它忽略)
            model(input_ids=input_ids, attention_mask=attn_mask, output_attentions=False, use_cache=False)
            tuned_attns_map = model.config.retrieve_attn_map

            # --- C. Loss & Step ---
            # 選中 head: KL toward blind teacher + data mass penalty
            loss_selected = head_attention_loss(tuned_attns_map, base_attns_map, data_mask)

            # 未選中 head: KL toward original model (preservation)
            loss_preserve = torch.tensor(0.0, device=device)
            if has_unselected and orig_attns_map is not None:
                loss_preserve = head_preservation_loss(tuned_attns_map, orig_attns_map)

            loss = loss_selected + lambda_preserve * loss_preserve

            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            sch.step()
            opt.zero_grad()

            del model.config.retrieve_attn_map
            pbar.set_postfix(
                loss=f"{loss.item():.4f}",
                sel=f"{loss_selected.item():.4f}",
                pres=f"{loss_preserve.item():.4f}" if has_unselected else "N/A",
            )
            if batch_save_interval > 0 and (idx + 1) % batch_save_interval == 0:
                save_model(model, tok, out_dir, current_idx, idx, data_path)
        save_model(model, tok, out_dir, current_idx, idx, data_path)
        current_idx += 1

# -----------------------------------------------------------------------------
# 6. Main
# -----------------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_path", required=True)
    parser.add_argument("--data_path", default="../3-1_model_training_preprocess/inj-likechen/crafted_instruction_data_tri_injection_qa.jsonl")
    parser.add_argument("--head_path", default="../2-2_head_identification/head_scoring/llama31-8b_injsq_dev/heads_sorted/user_prop_inst_0.1.json")

    parser.add_argument("--output_dir", default="outputs_lora")
    parser.add_argument("--lora_path", default=None)
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--batch-save-interval", type=int, default=-1)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--topk", type=str, default="18.75p")
    parser.add_argument("--lambda_preserve", type=float, default=1.0,
                        help="Weight for unselected head preservation loss")
    args = parser.parse_args()

    resume_adapter, start_epoch, start_batch_idx = discover_existing_adapter(args.output_dir)
    if args.lora_path:
        resume_adapter, start_epoch, start_batch_idx = args.lora_path, 0, 0

    model, tok, tok_inf = load_model(args.model_path)
    target_structure = get_target_structure(args.head_path, args.topk)
    os.makedirs(args.output_dir, exist_ok=True)
    tune_new(
        model, tok, tok_inf, target_structure,
        args.data_path, args.output_dir,
        args.epochs, args.batch_size, args.lr,
        resume_adapter=resume_adapter, start_epoch=start_epoch, start_batch_idx=start_batch_idx,
        batch_save_interval=args.batch_save_interval,
        lambda_preserve=args.lambda_preserve,
    )

if __name__ == "__main__":
    main()