"""
Loss functions for SFT and DPO.

SFT:
    Standard causal LM loss on response tokens.

DPO:
    Direct Preference Optimization:
        -log sigmoid(beta * [(log pi chosen - log ref chosen)
                            - (log pi rejected - log ref rejected)])
"""

import torch
import torch.nn.functional as F


def causal_lm_loss(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
    """
    Calculates causal LM loss.

    Input:
        logits: [batch, seq, vocab]
        labels: [batch, seq]
            Labels use -100 for ignored positions.

    Output:
        scalar loss.

    Important:
        logits[:, t] predicts labels[:, t+1]
    """
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = labels[:, 1:].contiguous()

    loss = F.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1),
        ignore_index=-100,
    )

    return loss


def sequence_logprobs(
    model,
    input_ids: torch.Tensor,
    labels: torch.Tensor,
) -> torch.Tensor:
    """
    Computes sequence log probability of only supervised answer tokens.

    Input:
        model: policy or reference model.
        input_ids: [batch, seq]
        labels: [batch, seq], -100 on ignored tokens.

    Output:
        log probability per sequence: [batch]
    """
    logits, _ = model(input_ids, targets=None)

    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = labels[:, 1:].contiguous()

    log_probs = F.log_softmax(shift_logits, dim=-1)

    mask = shift_labels != -100

    safe_labels = shift_labels.masked_fill(~mask, 0)

    token_log_probs = torch.gather(
        log_probs,
        dim=-1,
        index=safe_labels.unsqueeze(-1),
    ).squeeze(-1)

    token_log_probs = token_log_probs * mask

    return token_log_probs.sum(dim=-1)


def dpo_loss(
    policy_chosen_logps: torch.Tensor,
    policy_rejected_logps: torch.Tensor,
    ref_chosen_logps: torch.Tensor,
    ref_rejected_logps: torch.Tensor,
    beta: float,
) -> tuple[torch.Tensor, dict]:
    """
    Direct Preference Optimization loss.

    Input:
        policy_chosen_logps: log pi_theta(chosen)
        policy_rejected_logps: log pi_theta(rejected)
        ref_chosen_logps: log pi_ref(chosen)
        ref_rejected_logps: log pi_ref(rejected)
        beta: KL trade-off parameter.

    Output:
        loss
        metrics dictionary
    """
    policy_logratios = policy_chosen_logps - policy_rejected_logps
    ref_logratios = ref_chosen_logps - ref_rejected_logps

    logits = beta * (policy_logratios - ref_logratios)

    losses = -F.logsigmoid(logits)

    chosen_rewards = beta * (policy_chosen_logps - ref_chosen_logps)
    rejected_rewards = beta * (policy_rejected_logps - ref_rejected_logps)

    reward_accuracy = (chosen_rewards > rejected_rewards).float().mean()

    metrics = {
        "reward_accuracy": reward_accuracy.detach().item(),
        "chosen_reward": chosen_rewards.mean().detach().item(),
        "rejected_reward": rejected_rewards.mean().detach().item(),
        "margin": (chosen_rewards - rejected_rewards).mean().detach().item(),
    }

    return losses.mean(), metrics