import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
import json
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
from modelscope import snapshot_download

# ==================== 1. 动态批次填充 + 严格 Prompt Loss Masking 数据集 ====================
class RAN2SFTDataset(Dataset):
    def __init__(self, jsonl_path, tokenizer, max_length=2048):
        self.data = []
        self.tokenizer = tokenizer
        self.max_length = max_length
        with open(jsonl_path, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip():
                    self.data.append(json.loads(line.strip()))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        messages = self.data[idx]["messages"]

        # 1. 构造 Prompt 前缀（包含系统预设 system + 用户提问 user）
        prompt_text = self.tokenizer.apply_chat_template(
            messages[:-1],
            tokenize=False,
            add_generation_prompt=True
        )
        prompt_ids = self.tokenizer.encode(prompt_text, add_special_tokens=False)
        prompt_len = len(prompt_ids)

        # 2. 构造完整对话文本（包含完整输出 assistant）
        full_text = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=False
        )
        full_ids = self.tokenizer.encode(full_text, add_special_tokens=False)
        if not full_ids or full_ids[-1] != self.tokenizer.eos_token_id:
            full_ids.append(self.tokenizer.eos_token_id)

        # 最大长度保护截断
        if len(full_ids) > self.max_length:
            full_ids = full_ids[:self.max_length]

        # 3. 严格实现 SFT Prompt Loss Masking
        # Prompt 部分（系统提示 + 用户提问）对应标签置为 -100，仅对 Assistant 答案计算交叉熵损失
        labels = list(full_ids)
        mask_len = min(prompt_len, len(labels))
        for i in range(mask_len):
            labels[i] = -100

        return {
            "input_ids": full_ids,
            "attention_mask": [1] * len(full_ids),
            "labels": labels
        }

def collate_fn(batch, tokenizer):
    max_len = max(len(x["input_ids"]) for x in batch)
    # 向上对齐到 8 的倍数，加速 Tensor Core 显存访问
    max_len = (max_len + 7) // 8 * 8

    pad_id = tokenizer.pad_token_id or 0
    input_ids_b, attention_mask_b, labels_b = [], [], []

    for x in batch:
        cur_len = len(x["input_ids"])
        pad_len = max_len - cur_len

        input_ids_b.append(x["input_ids"] + [pad_id] * pad_len)
        attention_mask_b.append(x["attention_mask"] + [0] * pad_len)
        labels_b.append(x["labels"] + [-100] * pad_len)

    return (
        torch.tensor(input_ids_b, dtype=torch.long),
        torch.tensor(attention_mask_b, dtype=torch.long),
        torch.tensor(labels_b, dtype=torch.long)
    )

def main():
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    print(f"[Device] 训练硬件: {device} ({torch.cuda.get_device_name(0) if 'cuda' in device else 'CPU'})", flush=True)

    local_model_dir = r"C:\Users\ZoOoM\.cache\modelscope\hub\models\Qwen\Qwen2___5-3B-Instruct"
    if os.path.exists(local_model_dir):
        print(f"[Model] 本地权重直连秒启: {local_model_dir}", flush=True)
        model_dir = local_model_dir
    else:
        print(f"[Model] 从 ModelScope 下载权重...", flush=True)
        model_dir = snapshot_download("Qwen/Qwen2.5-3B-Instruct")

    tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # ==================== 2. 配置 4-bit QLoRA (NF4 + SDPA + 双重量化) ====================
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=torch.bfloat16
    )

    base_model = AutoModelForCausalLM.from_pretrained(
        model_dir,
        quantization_config=bnb_config,
        device_map="auto",
        attn_implementation="sdpa",
        trust_remote_code=True
    )
    base_model.config.use_cache = False
    base_model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    base_model = prepare_model_for_kbit_training(base_model, use_gradient_checkpointing=False)

    # 配置 LoRA 适配层 (r=16, alpha=32, 针对核心注意力投影矩阵 q, k, v, o)
    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]
    )
    model = get_peft_model(base_model, peft_config)

    trainable_p, all_p = model.get_nb_trainable_parameters()
    print(f"[Params] 总参数: {all_p / 1e6:.2f} M | LoRA 可训参数: {trainable_p / 1e6:.2f} M | 占比: {100 * trainable_p / all_p:.2f}%", flush=True)

    # ==================== 3. 加载 382 条专业训练集 ====================
    train_path = "dataset/ran2_train_382.jsonl"
    max_length = 1536
    dataset = RAN2SFTDataset(train_path, tokenizer, max_length=max_length)
    
    batch_size = 1
    accumulation_steps = 8  # 等效批大小 = 8
    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        collate_fn=lambda b: collate_fn(b, tokenizer)
    )

    epochs = 3
    total_steps = (len(dataloader) * epochs) // accumulation_steps
    import bitsandbytes as bnb
    optimizer = bnb.optim.PagedAdamW8bit(model.parameters(), lr=2e-4, weight_decay=0.01)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps, eta_min=2e-5)

    print(f"[Training] 开始 3GPP-Agent 4-bit QLoRA 严谨微调 (极速低显存版):", flush=True)
    print(f"  - 训练集样本数: {len(dataset)} 条 (独立盲测集 50 条已严格封存)")
    print(f"  - 轮数 (Epochs): {epochs}")
    print(f"  - 上下文长度 (Max Length): {max_length}")
    print(f"  - Batch Size: {batch_size} (梯度累积: {accumulation_steps}, 等效 Batch: {batch_size * accumulation_steps})")
    print(f"  - 总优化步数: {total_steps} 步", flush=True)
    
    model.train()

    step = 0
    opt_step = 0
    running_loss = 0.0
    sample_count = 0

    optimizer.zero_grad()
    for epoch in range(epochs):
        print(f"\n--- Epoch {epoch + 1}/{epochs} ---", flush=True)
        for input_ids, attention_mask, labels in dataloader:
            step += 1
            input_ids = input_ids.to(device)
            attention_mask = attention_mask.to(device)
            labels = labels.to(device)

            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
            loss = outputs.loss / accumulation_steps
            loss.backward()

            running_loss += outputs.loss.item()
            sample_count += 1

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

                if opt_step % 2 == 0 or opt_step == total_steps:
                    avg_loss = running_loss / sample_count
                    vram_gb = torch.cuda.memory_allocated(0) / (1024 ** 3)
                    print(f"Step [{opt_step:3d}/{total_steps:3d}] | Loss: {avg_loss:.4f} | LR: {scheduler.get_last_lr()[0]:.6f} | 显存: {vram_gb:.2f} GB", flush=True)
                    running_loss = 0.0
                    sample_count = 0
                    torch.cuda.empty_cache()

    # ==================== 4. 保存 3B QLoRA 适配器权重 ====================
    save_dir = "output/qwen3b_ran2_qlora"
    os.makedirs(save_dir, exist_ok=True)
    model.save_pretrained(save_dir)
    tokenizer.save_pretrained(save_dir)
    print(f"\n[Saved] Qwen2.5-3B 3GPP-Agent QLoRA 权重已成功保存至: {save_dir}", flush=True)

if __name__ == "__main__":
    main()