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

Key features:
- LoRA head masking on Q proj: only selected heads' dimensions are updated.
- LoRA head masking on K proj: only KV groups corresponding to selected heads
  are updated (GQA-aware).
- Preservation loss: heads that share a KV group with selected heads but are
  NOT themselves selected get a KL loss to stay close to the original model,
  since their K changed but they weren't intended training targets.
- Heads whose KV group is completely untouched need no loss at all.
"""

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

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 peft.tuners.lora.layer import Linear as LoraLinear
from lib.tokenize_data_mask import apply_chat_tokenize_with_strip_and_mark
from lib.printcolor import print_tok_color, print_data_color_in_batch

# -----------------------------------------------------------------------------
# 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):
    q_last = query[..., -1:, :]

    num_heads = query.shape[1]
    num_kv_heads = key.shape[1]

    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)

    attn_weights_tiny = torch.matmul(q_last, k_expanded.transpose(2, 3)) * scaling

    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

    attn_probs = torch.nn.functional.softmax(attn_weights_tiny, dim=-1, dtype=torch.float32).to(query.dtype)

    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():
            target_heads_dict[h_idx] = attn_probs[:, h_idx, :, :]

    orig_result = orig_flash_attn(module, query, key, value, attention_mask, scaling=scaling, **kwargs)
    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")]
    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

def get_target_structure(heads_file: str, topk_spec: str) -> Dict[int, List[int]]:
    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

    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))]

    structure = {}
    for tag, _ in target_heads:
        try:
            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. GQA-Aware Head Analysis
# -----------------------------------------------------------------------------
def compute_gqa_affected_heads(target_structure: Dict[int, List[int]],
                               num_q_heads: int,
                               num_kv_heads: int) -> Dict[int, List[int]]:
    """
    For each layer, find Q heads that are NOT selected but share a KV group
    with at least one selected head. These heads will have their K changed
    by LoRA but were not intended as training targets.

    Returns: {layer_idx: [affected_but_unselected_head_indices]}
    """
    group_size = num_q_heads // num_kv_heads

    affected_structure = {}
    for l, selected_heads in target_structure.items():
        selected_set = set(selected_heads)

        # Find which KV groups are touched
        touched_kv_groups = set()
        for h in selected_heads:
            touched_kv_groups.add(h // group_size)

        # Find unselected heads in those touched groups
        affected = []
        for kv_g in touched_kv_groups:
            for offset in range(group_size):
                q_head = kv_g * group_size + offset
                if q_head not in selected_set:
                    affected.append(q_head)

        if affected:
            affected_structure[l] = affected

    return affected_structure

# -----------------------------------------------------------------------------
# 4. LoRA Head Masking (Monkey-Patch)
# -----------------------------------------------------------------------------
def apply_head_mask_to_lora(model, target_structure: Dict[int, List[int]],
                            num_q_heads: int, num_kv_heads: int, head_dim: int):
    """
    Monkey-patch LoRA layers so that:
    - q_proj: LoRA output is zeroed for unselected Q heads
    - k_proj: LoRA output is zeroed for KV groups not associated with any selected head
    """
    group_size = num_q_heads // num_kv_heads

    for name, module in model.named_modules():
        if not isinstance(module, LoraLinear):
            continue

        # Parse layer index from name
        layer_idx = None
        parts = name.split(".")
        for i, part in enumerate(parts):
            if part == "layers" and i + 1 < len(parts) and parts[i + 1].isdigit():
                layer_idx = int(parts[i + 1])
                break

        if layer_idx is None or layer_idx not in target_structure:
            continue

        selected_heads = set(target_structure[layer_idx])
        is_q = name.endswith("q_proj")
        is_k = name.endswith("k_proj")

        if not (is_q or is_k):
            continue

        if is_q:
            mask = torch.zeros(num_q_heads * head_dim)
            for h in selected_heads:
                mask[h * head_dim : (h + 1) * head_dim] = 1.0
            desc = f"Q mask: {len(selected_heads)}/{num_q_heads} heads"
        else:
            touched_kv_groups = set()
            for h in selected_heads:
                touched_kv_groups.add(h // group_size)

            mask = torch.zeros(num_kv_heads * head_dim)
            for g in touched_kv_groups:
                mask[g * head_dim : (g + 1) * head_dim] = 1.0
            desc = f"K mask: {len(touched_kv_groups)}/{num_kv_heads} KV groups"

        module.register_buffer("head_mask", mask)

        orig_forward = module.forward

        def make_masked_forward(orig_fn, mod):
            def masked_forward(x, *args, **kwargs):
                result = orig_fn(x, *args, **kwargs)
                base_out = torch.nn.functional.linear(x, mod.base_layer.weight, mod.base_layer.bias)
                lora_delta = result - base_out
                masked_delta = lora_delta * mod.head_mask.to(lora_delta.device)
                return base_out + masked_delta
            return masked_forward

        module.forward = make_masked_forward(orig_forward, module)
        print(f"  🎭 Patched {name}: {desc}")

# -----------------------------------------------------------------------------
# 5. 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()
    mask_valid = ~mask_data

    if mask_valid.dim() == 2:
        mask_valid = mask_valid.unsqueeze(1)
        mask_data = mask_data.unsqueeze(1)

    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()
            base = base.float().to(tuned.device)

            base_v = base * mask_valid.float()
            base_v = base_v / base_v.sum(dim=-1, keepdim=True).clamp_min(eps)

            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

            kl = (base_v * (torch.log(base_v + eps) - torch.log(tuned_v + eps))).sum(dim=-1).mean()
            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 KV-group-affected but unselected heads: KL toward original model."""
    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 * (torch.log(orig + eps) - torch.log(tuned + eps))).sum(dim=-1).mean()
            total += kl
            n += 1
    return total / max(n, 1)


# -----------------------------------------------------------------------------
# 6. 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}

# -----------------------------------------------------------------------------
# 7. Training Loop
# -----------------------------------------------------------------------------
def save_model(model, tok, out_dir, epoch, batch_idx, data_path):
    print("Saving checkpoint...")

    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):
    # -------------------------------------------------------------------------
    # GQA analysis
    # -------------------------------------------------------------------------
    num_q_heads = model.config.num_attention_heads
    num_kv_heads = getattr(model.config, "num_key_value_heads", num_q_heads)
    head_dim = model.config.hidden_size // num_q_heads
    group_size = num_q_heads // num_kv_heads

    print(f"📐 GQA config: {num_q_heads} Q heads, {num_kv_heads} KV heads, group_size={group_size}")

    affected_structure = compute_gqa_affected_heads(target_structure, num_q_heads, num_kv_heads)
    has_affected = len(affected_structure) > 0
    if has_affected:
        total_affected = sum(len(v) for v in affected_structure.values())
        print(f"🛡️ Preservation targets: {total_affected} KV-group-affected unselected heads across {len(affected_structure)} layers (λ={lambda_preserve})")
        for l, heads in sorted(affected_structure.items()):
            print(f"   Layer {l}: affected heads {heads} (selected: {target_structure[l]})")
    else:
        print("✅ No KV-group collateral — all touched KV groups are fully selected. No preservation loss needed.")

    # -------------------------------------------------------------------------
    # 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)

    # -------------------------------------------------------------------------
    # Apply head mask to LoRA layers (physical isolation)
    # -------------------------------------------------------------------------
    print("🎭 Applying GQA-aware head masks to LoRA layers...")
    apply_head_mask_to_lora(model, target_structure, num_q_heads, num_kv_heads, head_dim)

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

    # -------------------------------------------------------------------------
    # Dataloader & Optimizer
    # -------------------------------------------------------------------------
    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)
            data_mask = batch["data_mask"][:, :MAX_LEN].to(device)

            blind_attn_mask = attn_mask * (~data_mask).long()

            print_data_color_in_batch(5,tok,input_ids,attn_mask)
            print_data_color_in_batch(5,tok,input_ids,data_mask)
            os.exit(0)

            # --- A. Base Model Pass (no LoRA, blind mask) ---
            # Teacher for selected heads
            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 (no LoRA, full mask) ---
            # Teacher for KV-group-affected unselected heads
            if has_affected:
                orig_map_container = {l: {h: None for h in h_list} for l, h_list in affected_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 (with LoRA, full mask) ---
            # Collect attention for both selected and affected heads
            model.train()
            tuned_map_container = {}
            for l in target_structure.keys():
                all_heads_needed = set(target_structure[l])
                if l in affected_structure:
                    all_heads_needed.update(affected_structure[l])
                tuned_map_container[l] = {h: None for h in all_heads_needed}
            model.config.retrieve_attn_map = tuned_map_container

            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 ---
            loss_selected = head_attention_loss(tuned_attns_map, base_attns_map, data_mask)

            loss_preserve = torch.tensor(0.0, device=device)
            if has_affected 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_affected 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

# -----------------------------------------------------------------------------
# 8. 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 KV-group-affected 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()