"""5.11 在一张 NVIDIA 显卡上用 PEFT 给小模型训 LoRA。和 train_mlx.py 用同一份数据、同一种提示格式。

只在“改过的版本”那部分算损失（提示部分的 label 设成 -100），训练集过 2 遍。
LoRA 权重存在 adapters/peft/，几十 MB。

用法：
    python train_peft.py                             # 默认 Qwen/Qwen3-1.7B，约 8 GB 显存
    python train_peft.py --model Qwen/Qwen3-0.6B
    python train_peft.py --data my_data --epochs 3
"""
import argparse
import math
import random
import time

import torch
from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer

from common import HERE, read_pairs, user_message


def encode(tok, pair, max_len):
    """返回 (input_ids, labels)：提示部分的 label 是 -100，不参与损失。"""
    prompt = tok.apply_chat_template([{'role': 'user', 'content': user_message(pair['ai'])}], tokenize=False,
                                     add_generation_prompt=True, enable_thinking=False)
    p_ids = tok(prompt, add_special_tokens=False)['input_ids']
    a_ids = tok(pair['original'] + '<|im_end|>\n', add_special_tokens=False)['input_ids']
    ids = (p_ids + a_ids)[:max_len]
    labels = ([-100] * len(p_ids) + a_ids)[:max_len]
    return ids, labels


def batches(examples, size, pad_id, shuffle):
    order = list(range(len(examples)))
    if shuffle:
        random.shuffle(order)
    for i in range(0, len(order), size):
        chunk = [examples[j] for j in order[i:i + size]]
        n = max(len(ids) for ids, _ in chunk)
        ids = torch.tensor([x + [pad_id] * (n - len(x)) for x, _ in chunk])
        labels = torch.tensor([y + [-100] * (n - len(y)) for _, y in chunk])
        mask = torch.tensor([[1] * len(x) + [0] * (n - len(x)) for x, _ in chunk])
        yield ids, labels, mask


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', default='Qwen/Qwen3-1.7B')
    ap.add_argument('--data')
    ap.add_argument('--epochs', type=int, default=2)
    ap.add_argument('--batch-size', type=int, default=2)
    ap.add_argument('--learning-rate', type=float, default=1e-4)
    ap.add_argument('--rank', type=int, default=8)
    ap.add_argument('--max-len', type=int, default=1024)
    ap.add_argument('--adapter', default=str(HERE / 'adapters' / 'peft'))
    args = ap.parse_args()
    random.seed(0)
    torch.manual_seed(0)

    tok = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16).to('cuda')
    model.gradient_checkpointing_enable()
    model.enable_input_require_grads()
    model = get_peft_model(model, LoraConfig(r=args.rank, lora_alpha=2 * args.rank, lora_dropout=0.0,
                                             target_modules='all-linear', task_type='CAUSAL_LM'))
    model.print_trainable_parameters()

    train = [encode(tok, p, args.max_len) for p in read_pairs('train', args.data)]
    valid = [encode(tok, p, args.max_len) for p in read_pairs('valid', args.data)]
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=args.learning_rate)
    steps = args.epochs * math.ceil(len(train) / args.batch_size)
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / 20) * max(0.0, 1 - s / steps))

    def val_loss():
        model.eval()
        losses = []
        with torch.no_grad():
            for ids, labels, mask in batches(valid, args.batch_size, tok.pad_token_id, False):
                out = model(input_ids=ids.cuda(), labels=labels.cuda(), attention_mask=mask.cuda())
                losses.append(out.loss.item())
        model.train()
        return sum(losses) / len(losses)

    print(f'训练 {len(train)} 对，验证 {len(valid)} 对，共 {steps} 步；开始前验证集损失 {val_loss():.3f}')
    t0, step = time.time(), 0
    model.train()
    for epoch in range(args.epochs):
        for ids, labels, mask in batches(train, args.batch_size, tok.pad_token_id, True):
            loss = model(input_ids=ids.cuda(), labels=labels.cuda(), attention_mask=mask.cuda()).loss
            loss.backward()
            opt.step()
            sched.step()
            opt.zero_grad()
            step += 1
            if step % 50 == 0 or step == steps:
                print(f'  步 {step}/{steps}，训练损失 {loss.item():.3f}，已用 {time.time() - t0:.0f} 秒')
        print(f'第 {epoch + 1} 遍结束，验证集损失 {val_loss():.3f}')
    model.save_pretrained(args.adapter)
    print(f'\nLoRA 已存到 adapters/peft/，用时 {time.time() - t0:.0f} 秒，显存峰值 {torch.cuda.max_memory_allocated() / 2**30:.1f} GB。'
          '\n下一步：python evaluate.py --backend torch')


if __name__ == '__main__':
    main()