import torch
import torch.nn as nn
import torch.nn.functional as F
import copy


def sequence_log_prob(model, x, y):
    """Sum of log P(y_t | x, y_<t>) under `model`, treating x as
    a fixed prompt prefix and y as the continuation to score."""
    full = torch.cat([x, y], dim=1)
    logits = model(full)[:, x.size(1) - 1:-1, :]  # predict each y token
    log_probs = F.log_softmax(logits, dim=-1)
    token_log_probs = log_probs.gather(2, y.unsqueeze(-1)).squeeze(-1)
    return token_log_probs.sum(dim=1)


def dpo_loss(policy, ref_model, x, y_w, y_l, beta=0.1):
    with torch.no_grad():
        ref_logp_w = sequence_log_prob(ref_model, x, y_w)
        ref_logp_l = sequence_log_prob(ref_model, x, y_l)

    policy_logp_w = sequence_log_prob(policy, x, y_w)
    policy_logp_l = sequence_log_prob(policy, x, y_l)

    logits = beta * (
        (policy_logp_w - ref_logp_w) - (policy_logp_l - ref_logp_l)
    )
    loss = -F.logsigmoid(logits).mean()

    with torch.no_grad():
        implicit_reward_w = beta * (policy_logp_w - ref_logp_w)
        implicit_reward_l = beta * (policy_logp_l - ref_logp_l)
        accuracy = (implicit_reward_w > implicit_reward_l).float().mean()

    return loss, accuracy.item()
