"""
Direct Preference Optimization (and ORPO / KTO variants) on preference pairs.

The policy is initialized from the SFT checkpoint; a frozen deep copy of it serves as the
DPO/KTO reference (ORPO is reference-free). Reports implicit-reward accuracy on held-out
preferences and GSM8K dev accuracy.

    PYTHONPATH=. python scripts/train_dpo.py --loss_type dpo --beta 0.1
    PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_dpo.py
"""

from __future__ import annotations

import time

import torch

from config.post_training_config import DPOConfig
from data_loader.preference_dataset import get_preference_iterator
from src.post_training.cli import parse_config_with_json
from src.post_training.distributed import ddp_setup, ddp_wrap, cleanup, reduce_scalar
from src.post_training.dpo import dpo_loss, orpo_loss, kto_loss, implicit_accuracy
from src.post_training.logging_utils import MetricsLogger
from src.post_training.optim import configure_optimizer, cosine_lr
from src.post_training.rollout import sequence_logprobs
from src.post_training.utils import (
    amp_autocast, load_backbone_from_ckpt, make_frozen_copy, save_stage_ckpt, set_seed, unwrap,
)

TEST_PATH = "/ephemeral/data/preferences_test.jsonl"


def _logps(model, ids, mask, requires_grad):
    return sequence_logprobs(model, ids, mask, requires_grad=requires_grad)


def _compute_losses(policy, ref, batch, cfg, ctx):
    B = batch["chosen_ids"].size(0)
    ids = torch.cat([batch["chosen_ids"], batch["rejected_ids"]], dim=0)
    mask = torch.cat([batch["chosen_mask"], batch["rejected_mask"]], dim=0)

    with amp_autocast(cfg.amp_dtype, ctx.device):
        psum, pn = _logps(policy, ids, mask, requires_grad=True)
    pc, pr, ncn, nrn = psum[:B], psum[B:], pn[:B], pn[B:]

    if cfg.loss_type == "orpo":
        return orpo_loss(pc, pr, ncn, nrn, orpo_lambda=cfg.orpo_lambda)

    with torch.no_grad(), amp_autocast(cfg.amp_dtype, ctx.device):
        rsum, _ = _logps(ref, ids, mask, requires_grad=False)
    rc, rr = rsum[:B], rsum[B:]
    if cfg.loss_type == "kto":
        return kto_loss(pc, pr, rc, rr, beta=cfg.beta)
    return dpo_loss(pc, pr, rc, rr, beta=cfg.beta)


@torch.no_grad()
def eval_implicit_acc(policy, ref, cfg, ctx, max_batches: int = 100) -> tuple[float, float]:
    policy.eval()
    it = get_preference_iterator(TEST_PATH, cfg.batch_size, cfg.max_len, device=ctx.device,
                                 rank=ctx.rank, world_size=ctx.world_size, shuffle=False, infinite=False)
    acc, marg, n = 0.0, 0.0, 0
    for batch in it:
        loss, cr, rr = _compute_losses(policy, ref, batch, cfg, ctx)
        acc += implicit_accuracy(cr, rr).item()
        marg += (cr - rr).mean().item()
        n += 1
        if n >= max_batches:
            break
    policy.train()
    return acc / max(1, n), marg / max(1, n)


def main():
    cfg, _ = parse_config_with_json(DPOConfig, "configs/dpo.json")
    ctx = ddp_setup(cfg.device)
    set_seed(cfg.seed + ctx.rank)

    policy = load_backbone_from_ckpt(cfg, cfg.sft_ckpt, ctx.device)
    ref = make_frozen_copy(policy, device=ctx.device) if cfg.loss_type != "orpo" else None
    policy = ddp_wrap(policy, ctx)
    optimizer = configure_optimizer(unwrap(policy), cfg.lr, cfg.weight_decay)

    with open(cfg.pref_path) as f:
        n_rows = sum(1 for line in f if line.strip())
    total_steps = max(1, (n_rows // (cfg.batch_size * ctx.world_size)) * cfg.epochs)

    logger = None
    if ctx.is_main:
        print(f"DPO[{cfg.loss_type}] from {cfg.sft_ckpt} | {n_rows} pairs | total_steps={total_steps} | beta={cfg.beta}")
        logger = MetricsLogger(f"dpo_{cfg.loss_type}", cfg.log_dir, use_wandb=cfg.use_wandb, wandb_project=cfg.wandb_project)

    train_it = get_preference_iterator(cfg.pref_path, cfg.batch_size, cfg.max_len, device=ctx.device,
                                       rank=ctx.rank, world_size=ctx.world_size, shuffle=True, infinite=True)

    policy.train()
    t0 = time.perf_counter()
    for step in range(total_steps):
        lr = cosine_lr(step, warmup_steps=cfg.warmup_steps, max_steps=total_steps, lr=cfg.lr, min_lr=cfg.lr * 0.1)
        for g in optimizer.param_groups:
            g["lr"] = lr

        batch = next(train_it)
        loss, cr, rr = _compute_losses(policy, ref, batch, cfg, ctx)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(policy.parameters(), cfg.grad_clip)
        optimizer.step()

        if ctx.is_main and step % 20 == 0:
            acc = implicit_accuracy(cr, rr).item()
            dt = time.perf_counter() - t0; t0 = time.perf_counter()
            print(f"step {step}/{total_steps} | loss {loss.item():.4f} | acc {acc:.3f} | "
                  f"r_chosen {cr.mean().item():.3f} r_rejected {rr.mean().item():.3f} | {dt:.1f}s/20")
            if logger:
                logger.log(step, {"train_loss": loss.item(), "train_acc": acc,
                                  "r_chosen": cr.mean().item(), "r_rejected": rr.mean().item(), "lr": lr})

        if step > 0 and step % cfg.eval_steps == 0:
            acc, marg = eval_implicit_acc(policy, ref, cfg, ctx)
            acc, marg = reduce_scalar(acc, ctx), reduce_scalar(marg, ctx)
            if ctx.is_main:
                print(f"  [eval] step {step} | test_acc {acc:.3f} | margin {marg:.3f}")
                if logger:
                    logger.log(step, {"test_acc": acc, "test_margin": marg})

        if ctx.is_main and step > 0 and step % cfg.save_every == 0:
            save_stage_ckpt(cfg.out_ckpt, policy, optimizer, stage=f"dpo_{cfg.loss_type}", cfg=cfg, step=step,
                            metrics={"train_loss": loss.item()})

    if ctx.is_main:
        # Unwrap for the final eval: other ranks are already at cleanup(), so a collective on
        # the DDP-wrapped policy here would hang (NCCL timeout). The periodic eval runs on all ranks.
        acc, marg = eval_implicit_acc(unwrap(policy), ref, cfg, ctx)
        save_stage_ckpt(cfg.out_ckpt, policy, optimizer, stage=f"dpo_{cfg.loss_type}", cfg=cfg, step=total_steps,
                        metrics={"test_acc": acc, "test_margin": marg})
        print(f"Done DPO[{cfg.loss_type}]. test_acc {acc:.3f} margin {marg:.3f} -> {cfg.out_ckpt}")
        if logger:
            logger.close()
    cleanup(ctx)


if __name__ == "__main__":
    main()