"""
DPO (Direct Preference Optimization) trainer for DriveMind-V.
Trains the model to prefer safe driving actions over risky ones.

Dataset format (JSONL):
  {"prompt": "...", "chosen": "safe action", "rejected": "risky action"}

DPO loss:
  L = -E[log σ(β * (log π(yw|x) - log πref(yw|x)) - β * (log π(yl|x) - log πref(yl|x)))]

Usage:
  uv run python trainer/train_dpo_driving.py \
    --data_path dataset/driving_preferences.jsonl \
    --use_lora --epochs 3
"""
import os
import sys
import copy
import json

__package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import argparse
import time
import warnings
import torch
import torch.nn.functional as F
from contextlib import nullcontext
from torch import optim
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer
from model.model_vlm import MiniMindVLM, VLMConfig
from trainer.trainer_utils import Logger, setup_seed, get_lr
from trainer.lora_utils import get_lora_config, apply_lora

warnings.filterwarnings('ignore')


# ── Dataset ──────────────────────────────────────────────────────────────────

class DrivingPreferenceDataset(Dataset):
    """
    Loads (prompt, chosen, rejected) pairs and tokenizes them.
    Returns:
        chosen_ids:    [prompt + chosen] token ids
        rejected_ids:  [prompt + rejected] token ids
        prompt_len:    length of prompt tokens (so we can mask the loss)
    """
    def __init__(self, jsonl_path: str, tokenizer, max_length: int = 512):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.samples = []

        with open(jsonl_path) as f:
            for line in f:
                line = line.strip()
                if line:
                    self.samples.append(json.loads(line))

        Logger(f"Loaded {len(self.samples)} preference pairs from {jsonl_path}")

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

    def _tokenize(self, text: str):
        return self.tokenizer(
            text,
            truncation=True,
            max_length=self.max_length,
            return_tensors='pt',
            add_special_tokens=False,
        )['input_ids'].squeeze(0)

    def __getitem__(self, idx):
        sample = self.samples[idx]
        prompt = sample['prompt']
        chosen = sample['chosen']
        rejected = sample['rejected']

        prompt_ids = self._tokenize(prompt)
        chosen_ids = self._tokenize(prompt + chosen)
        rejected_ids = self._tokenize(prompt + rejected)

        prompt_len = len(prompt_ids)

        return {
            'chosen_ids': chosen_ids,
            'rejected_ids': rejected_ids,
            'prompt_len': prompt_len,
        }


def collate_fn(batch, pad_id=0):
    """Pad sequences to same length within a batch."""
    chosen_ids = [x['chosen_ids'] for x in batch]
    rejected_ids = [x['rejected_ids'] for x in batch]
    prompt_lens = [x['prompt_len'] for x in batch]

    def pad(seqs):
        max_len = max(len(s) for s in seqs)
        padded = torch.zeros(len(seqs), max_len, dtype=torch.long)
        mask = torch.zeros(len(seqs), max_len, dtype=torch.bool)
        for i, s in enumerate(seqs):
            padded[i, :len(s)] = s
            mask[i, :len(s)] = True
        return padded, mask

    chosen_ids, chosen_mask = pad(chosen_ids)
    rejected_ids, rejected_mask = pad(rejected_ids)
    prompt_lens = torch.tensor(prompt_lens, dtype=torch.long)

    return chosen_ids, chosen_mask, rejected_ids, rejected_mask, prompt_lens


# ── DPO loss ─────────────────────────────────────────────────────────────────

def compute_log_probs(model, input_ids, attention_mask, prompt_len):
    """
    Compute sum of log-probs over response tokens only (not prompt).

    Args:
        model: language model
        input_ids: (B, L)
        attention_mask: (B, L)
        prompt_len: (B,) — number of prompt tokens per sample

    Returns:
        log_probs: (B,) — sum of log-probs over response tokens
    """
    with torch.no_grad() if not model.training else nullcontext():
        # Model expects float attention_mask (0.0/1.0), not bool
        attn_mask_float = attention_mask.float()
        outputs = model(input_ids=input_ids, attention_mask=attn_mask_float)
        logits = outputs.logits  # (B, L, V)

    # Shift: logits[i] predicts input_ids[i+1]
    shift_logits = logits[:, :-1, :].contiguous()          # (B, L-1, V)
    shift_labels = input_ids[:, 1:].contiguous()            # (B, L-1)
    shift_mask = attention_mask[:, 1:].contiguous()         # (B, L-1)

    log_probs_all = F.log_softmax(shift_logits, dim=-1)     # (B, L-1, V)
    # Gather log-prob of the actual next token
    token_log_probs = log_probs_all.gather(
        2, shift_labels.unsqueeze(-1)
    ).squeeze(-1)                                            # (B, L-1)

    # Mask: only response tokens (after prompt)
    response_mask = shift_mask.clone()
    for i in range(len(prompt_len)):
        # prompt_len[i]-1 because we shifted by 1
        response_start = max(0, prompt_len[i].item() - 1)
        response_mask[i, :response_start] = 0

    # Sum log-probs over response tokens
    token_log_probs = token_log_probs * response_mask.float()
    log_probs = token_log_probs.sum(dim=-1)                  # (B,)

    return log_probs


def dpo_loss(policy_chosen_logp, policy_rejected_logp,
             ref_chosen_logp, ref_rejected_logp, beta: float = 0.1):
    """
    DPO loss from Rafailov et al. (2023).

    L = -E[log σ(β * ((log π(yw|x) - log πref(yw|x)) - (log π(yl|x) - log πref(yl|x))))]
    """
    chosen_rewards = beta * (policy_chosen_logp - ref_chosen_logp)
    rejected_rewards = beta * (policy_rejected_logp - ref_rejected_logp)
    loss = -F.logsigmoid(chosen_rewards - rejected_rewards).mean()

    # Preference accuracy: fraction of samples where chosen > rejected
    with torch.no_grad():
        pref_acc = (chosen_rewards > rejected_rewards).float().mean()

    return loss, chosen_rewards.mean().detach(), rejected_rewards.mean().detach(), pref_acc


# ── Training ──────────────────────────────────────────────────────────────────

def train_epoch(epoch, loader, iters, policy_model, ref_model, optimizer,
                scaler, autocast_ctx, args, wandb=None):
    policy_model.train()
    start_time = time.time()
    total_loss = 0.0
    total_acc = 0.0

    for step, (chosen_ids, chosen_mask, rejected_ids, rejected_mask, prompt_lens) in \
            enumerate(loader, start=1):

        chosen_ids = chosen_ids.to(args.device)
        chosen_mask = chosen_mask.to(args.device)
        rejected_ids = rejected_ids.to(args.device)
        rejected_mask = rejected_mask.to(args.device)
        prompt_lens = prompt_lens.to(args.device)

        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
        for pg in optimizer.param_groups:
            pg['lr'] = lr

        with autocast_ctx:
            # Policy log-probs (with gradient)
            policy_chosen_logp = compute_log_probs(
                policy_model, chosen_ids, chosen_mask, prompt_lens)
            policy_rejected_logp = compute_log_probs(
                policy_model, rejected_ids, rejected_mask, prompt_lens)

            # Reference log-probs (no gradient)
            with torch.no_grad():
                ref_chosen_logp = compute_log_probs(
                    ref_model, chosen_ids, chosen_mask, prompt_lens)
                ref_rejected_logp = compute_log_probs(
                    ref_model, rejected_ids, rejected_mask, prompt_lens)

            loss, chosen_r, rejected_r, pref_acc = dpo_loss(
                policy_chosen_logp, policy_rejected_logp,
                ref_chosen_logp, ref_rejected_logp,
                beta=args.beta,
            )
            loss = loss / args.accumulation_steps

        scaler.scale(loss).backward()

        if step % args.accumulation_steps == 0:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(policy_model.parameters(), args.grad_clip)
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad(set_to_none=True)

        total_loss += loss.item() * args.accumulation_steps
        total_acc += pref_acc.item()

        if step % args.log_interval == 0 or step == iters:
            elapsed = time.time() - start_time
            avg_loss = total_loss / step
            avg_acc = total_acc / step
            Logger(
                f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) '
                f'loss:{avg_loss:.4f} pref_acc:{avg_acc:.3f} '
                f'chosen_r:{chosen_r:.3f} rejected_r:{rejected_r:.3f} '
                f'lr:{lr:.2e} time:{elapsed/60:.1f}min'
            )
            if wandb:
                wandb.log({
                    'loss': avg_loss,
                    'preference_accuracy': avg_acc,
                    'chosen_reward': chosen_r.item(),
                    'rejected_reward': rejected_r.item(),
                    'lr': lr,
                    'epoch': epoch + 1,
                })

    return total_loss / iters, total_acc / iters


# ── Main ──────────────────────────────────────────────────────────────────────

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='DriveMind-V DPO Training')
    parser.add_argument('--data_path', type=str, default='dataset/driving_preferences.jsonl')
    parser.add_argument('--out_dir', type=str, default='out')
    parser.add_argument('--epochs', type=int, default=3)
    parser.add_argument('--batch_size', type=int, default=4)
    parser.add_argument('--learning_rate', type=float, default=5e-5)
    parser.add_argument('--beta', type=float, default=0.1,
                        help='DPO temperature (higher = closer to reference)')
    parser.add_argument('--max_length', type=int, default=256)
    parser.add_argument('--device', type=str,
                        default='cuda:0' if torch.cuda.is_available() else 'cpu')
    parser.add_argument('--dtype', type=str,
                        default='bfloat16' if torch.cuda.is_bf16_supported() else 'float16')
    parser.add_argument('--grad_clip', type=float, default=1.0)
    parser.add_argument('--accumulation_steps', type=int, default=1)
    parser.add_argument('--log_interval', type=int, default=5)
    parser.add_argument('--num_workers', type=int, default=0)
    parser.add_argument('--hidden_size', type=int, default=512)
    parser.add_argument('--num_hidden_layers', type=int, default=8)
    parser.add_argument('--from_weight', type=str, default='sft_vlm',
                        help='Base checkpoint prefix (loads out/<from_weight>_<hidden_size>.pth)')
    # LoRA
    parser.add_argument('--use_lora', action='store_true')
    parser.add_argument('--lora_r', type=int, default=16)
    parser.add_argument('--lora_alpha', type=int, default=32)
    parser.add_argument('--lora_dropout', type=float, default=0.05)
    # WandB
    parser.add_argument('--use_wandb', action='store_true')

    args = parser.parse_args()
    setup_seed(42)

    # ── Tokenizer ──
    tokenizer = AutoTokenizer.from_pretrained('./model')

    # ── Model config ──
    vlm_config = VLMConfig(
        hidden_size=args.hidden_size,
        num_hidden_layers=args.num_hidden_layers,
    )

    # ── Load policy model ──
    weight_path = f'{args.out_dir}/{args.from_weight}_{vlm_config.hidden_size}.pth'
    Logger(f'Loading base model from {weight_path}')

    policy_model = MiniMindVLM(vlm_config,
                               vision_model_path='./model/vision_model/clip-vit-base-patch16')
    if os.path.exists(weight_path):
        weights = torch.load(weight_path, map_location=args.device)
        policy_model.load_state_dict(weights, strict=False)
        Logger(f'Loaded weights from {weight_path}')
    else:
        Logger(f'WARNING: {weight_path} not found — using random init')

    policy_model = policy_model.to(args.device)

    # ── Frozen reference model (deep copy before LoRA) ──
    Logger('Creating frozen reference model...')
    ref_model = copy.deepcopy(policy_model)
    ref_model.eval()
    for p in ref_model.parameters():
        p.requires_grad = False
    Logger('Reference model frozen.')

    # ── Apply LoRA to policy model ──
    if args.use_lora:
        Logger('\nApplying LoRA to policy model...')
        lora_config = get_lora_config(
            r=args.lora_r,
            lora_alpha=args.lora_alpha,
            lora_dropout=args.lora_dropout,
        )
        policy_model = apply_lora(policy_model, lora_config, verbose=True)

    # ── Dataset & loader ──
    dataset = DrivingPreferenceDataset(args.data_path, tokenizer, args.max_length)
    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        collate_fn=collate_fn,
        pin_memory=(args.device != 'cpu'),
    )

    # ── Optimizer & AMP ──
    optimizer = optim.AdamW(
        filter(lambda p: p.requires_grad, policy_model.parameters()),
        lr=args.learning_rate,
        weight_decay=0.01,
    )

    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
    if args.dtype == 'bfloat16':
        autocast_ctx = torch.cuda.amp.autocast(dtype=torch.bfloat16)
    elif args.dtype == 'float16':
        autocast_ctx = torch.cuda.amp.autocast(dtype=torch.float16)
    else:
        autocast_ctx = nullcontext()

    # ── WandB ──
    wandb = None
    if args.use_wandb:
        import wandb as wb
        wb.init(project='drivemind-v', name=f'dpo_lora_r{args.lora_r}_beta{args.beta}',
                config=vars(args))
        wandb = wb

    # ── Train ──
    iters = len(loader)
    Logger(f'\nStarting DPO training: {args.epochs} epochs x {iters} steps')
    Logger(f'Dataset: {len(dataset)} pairs | Batch: {args.batch_size} | Beta: {args.beta}')
    Logger(f'LoRA: {args.use_lora} | Device: {args.device} | dtype: {args.dtype}\n')

    for epoch in range(args.epochs):
        avg_loss, avg_acc = train_epoch(
            epoch, loader, iters,
            policy_model, ref_model,
            optimizer, scaler, autocast_ctx,
            args, wandb,
        )
        Logger(f'Epoch {epoch+1} complete — avg_loss: {avg_loss:.4f}  avg_pref_acc: {avg_acc:.3f}')

    # ── Save ──
    os.makedirs(args.out_dir, exist_ok=True)
    lora_tag = '_lora' if args.use_lora else ''
    save_path = f'{args.out_dir}/dpo_driving{lora_tag}_{vlm_config.hidden_size}'

    if args.use_lora:
        policy_model.save_pretrained(save_path + '_adapters')
        Logger(f'LoRA adapters saved -> {save_path}_adapters')
    else:
        raw = getattr(policy_model, '_orig_mod', policy_model)
        state_dict = {k: v.half().cpu() for k, v in raw.state_dict().items()
                      if not k.startswith('vision_encoder.')}
        ckp = save_path + '.pth'
        torch.save(state_dict, ckp)
        Logger(f'Model saved -> {ckp}')

    if wandb:
        wandb.finish()

    Logger('\nDPO training complete!')