"""Multi-objective preference loss combiner (v0.40.1 Part B runtime).

Closes the v0.40.0 Part D stub-then-live deferral: ``preference_loss_weights``
now actually combines 2-5 preference losses into one backward pass.

Each preference loss reduces to a pure function of the same forward-pass
quantities: ``policy_chosen_logps`` / ``policy_rejected_logps`` and (for
DPO / IPO) ``ref_chosen_logps`` / ``ref_rejected_logps``. Sharing one
forward pass keeps the cost ~equal to single-loss training.

Compatibility matrix (enforced at config-load + at runtime):

* DPO / IPO  — require a frozen reference model (β log-ratio family).
* SimPO / ORPO — reference-free.
* BCO — uses ``prompt + completion + label`` data, *incompatible* with
  paired ``prompt + chosen + rejected`` batches. Rejected at runtime when
  combined with anything else; users wanting to blend BCO with paired
  losses must run them as separate stages.

The helper itself is dependency-light — it only imports torch lazily so it
can be unit-tested on toy tensors without pulling TRL.
"""

from __future__ import annotations

import logging
import math
from typing import TYPE_CHECKING, Dict, Mapping, Optional

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    import torch  # noqa: F401

PAIRED_LOSSES: frozenset = frozenset({"dpo", "simpo", "orpo", "ipo"})
REF_MODEL_LOSSES: frozenset = frozenset({"dpo", "ipo"})
REF_FREE_LOSSES: frozenset = frozenset({"simpo", "orpo"})
UNPAIRED_LOSSES: frozenset = frozenset({"bco"})


def validate_weight_compat(weights: Mapping[str, float]) -> None:
    """Enforce the BCO-incompatible-with-paired rule at runtime.

    Schema-level validation already restricts keys to the allowlist and
    bounds the sum to 1; this guard catches the data-format mismatch that
    only manifests at training time.
    """
    keys = set(weights.keys())
    if "bco" in keys and (keys - {"bco"}):
        raise ValueError(
            "preference_loss_weights cannot mix 'bco' with paired losses "
            "(dpo/simpo/orpo/ipo). BCO consumes prompt+completion+label rows, "
            "while paired losses consume prompt+chosen+rejected. Run BCO as a "
            "separate task=bco stage."
        )


def needs_reference_model(weights: Mapping[str, float]) -> bool:
    """True iff any active loss in the blend uses a frozen reference model."""
    return bool(set(weights) & REF_MODEL_LOSSES)


def _sigmoid(x):
    import torch

    return torch.sigmoid(x)


def _logsigmoid(x):
    import torch

    return torch.nn.functional.logsigmoid(x)


def compute_dpo_term(pol_chosen, pol_rejected, ref_chosen, ref_rejected, beta: float):
    """Standard DPO loss: ``-log σ(β · (Δπ - Δπ_ref))``.

    All log-prob args are summed-token log-likelihoods of the *response*
    only (matching TRL's ``DPOTrainer.compute_reference_log_probs`` shape).
    """
    if ref_chosen is None or ref_rejected is None:
        raise ValueError("DPO requires reference-model log-probs")
    pi_logratio = pol_chosen - pol_rejected
    ref_logratio = ref_chosen - ref_rejected
    logits = beta * (pi_logratio - ref_logratio)
    return -_logsigmoid(logits).mean()


def compute_ipo_term(pol_chosen, pol_rejected, ref_chosen, ref_rejected, beta: float):
    """IPO loss: squared-hinge regularised ``(Δπ - Δπ_ref - 1/(2β))²``."""
    if ref_chosen is None or ref_rejected is None:
        raise ValueError("IPO requires reference-model log-probs")
    if beta <= 0:
        raise ValueError(f"IPO beta must be > 0, got {beta}")
    pi_logratio = pol_chosen - pol_rejected
    ref_logratio = ref_chosen - ref_rejected
    target = 1.0 / (2.0 * beta)
    return ((pi_logratio - ref_logratio - target) ** 2).mean()


def compute_simpo_term(
    pol_chosen,
    pol_rejected,
    beta: float,
    gamma: float,
    chosen_lens=None,
    rejected_lens=None,
):
    """Reference-free length-normalised preference loss (SimPO).

    ``pol_chosen`` / ``pol_rejected`` are summed log-probs; lengths are the
    response token counts used to length-normalise. When ``chosen_lens`` is
    None, falls back to per-sample 1.0 (i.e. acts like length-blind DPO,
    with no reference).
    """
    import torch

    if chosen_lens is None or rejected_lens is None:
        chosen_norm = pol_chosen
        rejected_norm = pol_rejected
    else:
        # Avoid div by zero.
        chosen_lens = torch.clamp(chosen_lens.float(), min=1.0)
        rejected_lens = torch.clamp(rejected_lens.float(), min=1.0)
        chosen_norm = pol_chosen / chosen_lens
        rejected_norm = pol_rejected / rejected_lens
    logits = beta * (chosen_norm - rejected_norm) - gamma
    return -_logsigmoid(logits).mean()


def compute_orpo_term(pol_chosen, pol_rejected, alpha: float):
    """Reference-free odds-ratio preference loss (ORPO).

    Uses the response-log-prob formulation ``-log σ(log(p_w) - log(p_l) +
    log(1-p_l) - log(1-p_w))`` scaled by ``alpha``. Approximates the full
    ORPO loss without the SFT term — caller is expected to mix in SFT via
    its own weight if desired.
    """
    import torch

    log_odds_chosen = pol_chosen - torch.log1p(-torch.exp(pol_chosen).clamp(max=1 - 1e-7))
    log_odds_rejected = pol_rejected - torch.log1p(
        -torch.exp(pol_rejected).clamp(max=1 - 1e-7)
    )
    sigm_term = _logsigmoid(log_odds_chosen - log_odds_rejected)
    return (-alpha * sigm_term).mean()


def combine_losses(
    losses: Dict[str, "torch.Tensor"],
    weights: Mapping[str, float],
) -> "torch.Tensor":
    """Weighted sum of per-loss tensors, validated against ``weights``.

    Raises:
        ValueError: weight dict and loss dict keys differ, or weights don't
            sum to 1 within ±1e-6 (defence-in-depth — schema also enforces).
    """
    if not weights:
        raise ValueError("weights mapping must not be empty")
    if set(losses.keys()) != set(weights.keys()):
        raise ValueError(
            f"loss keys {sorted(losses)} != weight keys {sorted(weights)}"
        )
    # v0.40.1 review fix — defence-in-depth bool rejection (schema also
    # rejects bool, but the runtime path should not silently accept True/False).
    for name, weight in weights.items():
        if isinstance(weight, bool):
            raise TypeError(
                f"preference_loss_weights[{name!r}] must be float, not bool"
            )
    total = sum(weights.values())
    if not math.isclose(total, 1.0, abs_tol=1e-6):
        raise ValueError(f"weights must sum to 1.0 (±1e-6), got {total}")
    out = None
    for name, weight in weights.items():
        contrib = float(weight) * losses[name]
        out = contrib if out is None else out + contrib
    return out


def attach_weighted_preference_combine(trainer: object, weights: Mapping[str, float]) -> bool:
    """v0.53.11 #68 — wrap inner trainer's ``compute_loss`` for TRUE weighted blend.

    Replaces the v0.40.1 primary-loss approximation with a per-batch
    weighted sum across all named losses. After the inner trainer's
    ``compute_loss`` runs (computing the primary forward pass), we:

    1. Read the policy + reference summed log-probs that TRL stashes on
       the trainer / batch (``policy_chosen_logps``, ``policy_rejected_logps``,
       ``ref_chosen_logps``, ``ref_rejected_logps``).
    2. For each weighted loss in ``weights``, compute its term via the
       matching ``compute_*_term`` helper.
    3. Combine via :func:`combine_losses` and return the blended scalar.

    When the TRL trainer does not expose the per-batch log-probs (older
    TRL versions, custom forks), we fall back to the v0.40.1 primary-loss
    scaling — defence-in-depth so a TRL upgrade does not crash training.

    Idempotent: re-attaching detects the ``_soup_weighted_combine`` marker.
    """
    validate_weight_compat(weights)
    if not hasattr(trainer, "compute_loss"):
        return False
    original = trainer.compute_loss
    if getattr(original, "_soup_weighted_combine", False):
        return True
    snapshot = dict(weights)

    def wrapped(model, inputs, return_outputs=False, **kwargs):
        result = original(model, inputs, return_outputs=return_outputs, **kwargs)
        if return_outputs:
            primary_loss, outputs = result
        else:
            primary_loss = result
            outputs = None

        # True weighted-sum path: read per-batch logps from inputs OR the
        # trainer's last-batch state.
        pol_chosen = _read_logps(inputs, "policy_chosen_logps")
        if pol_chosen is None:
            pol_chosen = _read_logps(inputs, "chosen_logps")
        pol_rejected = _read_logps(inputs, "policy_rejected_logps")
        if pol_rejected is None:
            pol_rejected = _read_logps(inputs, "rejected_logps")
        ref_chosen = _read_logps(inputs, "reference_chosen_logps")
        if ref_chosen is None:
            ref_chosen = _read_logps(inputs, "ref_chosen_logps")
        ref_rejected = _read_logps(inputs, "reference_rejected_logps")
        if ref_rejected is None:
            ref_rejected = _read_logps(inputs, "ref_rejected_logps")

        terms: dict = {}
        if pol_chosen is not None and pol_rejected is not None:
            # v0.53.11 review fix (security HIGH) — explicit None check
            # before float(). The previous `or 0.1` idiom triggered tensor
            # truth-value evaluation when TRL stored these as zero-element
            # tensors instead of Python floats.
            beta_attr = getattr(trainer, "beta", None)
            beta = float(beta_attr) if beta_attr is not None else 0.1
            for name in snapshot:
                try:
                    if name == "dpo":
                        if ref_chosen is None or ref_rejected is None:
                            logger.debug(
                                "weighted-combine: dpo term skipped — ref logps missing"
                            )
                            continue
                        terms["dpo"] = compute_dpo_term(
                            pol_chosen, pol_rejected, ref_chosen, ref_rejected, beta
                        )
                    elif name == "ipo":
                        if ref_chosen is None or ref_rejected is None:
                            logger.debug(
                                "weighted-combine: ipo term skipped — ref logps missing"
                            )
                            continue
                        terms["ipo"] = compute_ipo_term(
                            pol_chosen, pol_rejected, ref_chosen, ref_rejected, beta
                        )
                    elif name == "simpo":
                        gamma_attr = getattr(trainer, "simpo_gamma", None)
                        gamma = (
                            float(gamma_attr) if gamma_attr is not None else 1.0
                        )
                        terms["simpo"] = compute_simpo_term(
                            pol_chosen, pol_rejected, beta, gamma
                        )
                    elif name == "orpo":
                        alpha_attr = getattr(trainer, "orpo_alpha", None)
                        alpha = (
                            float(alpha_attr) if alpha_attr is not None else 1.0
                        )
                        terms["orpo"] = compute_orpo_term(
                            pol_chosen, pol_rejected, alpha
                        )
                    elif name == "bco":
                        # BCO is data-format-incompatible — already rejected
                        # by validate_weight_compat. Defensive skip.
                        continue
                except (TypeError, ValueError) as exc:
                    logger.debug(
                        "weighted-combine: %s term skipped — %s", name, exc
                    )
                    continue

        if len(terms) == len(snapshot):
            # All requested terms computed — return true weighted sum.
            blended = combine_losses(terms, snapshot)
        else:
            # Fallback: primary-loss scaling.
            primary_weight = snapshot.get(_pick_primary(snapshot), 1.0)
            blended = primary_loss * float(primary_weight)
        if return_outputs:
            return blended, outputs
        return blended

    wrapped._soup_weighted_combine = True  # type: ignore[attr-defined]
    trainer.compute_loss = wrapped  # type: ignore[assignment]
    return True


def _read_logps(obj, name: str):
    """Read a log-prob tensor from a mapping / object — TRL inputs vary."""
    if obj is None:
        return None
    if hasattr(obj, "get"):
        val = obj.get(name)
        if val is not None:
            return val
    return getattr(obj, name, None)


def _pick_primary(weights: Mapping[str, float]) -> str:
    """Return the loss with the highest weight (deterministic on ties)."""
    return max(sorted(weights), key=lambda k: weights[k])


def describe_blend(weights: Optional[Mapping[str, float]]) -> str:
    """Human-readable summary for advisory output."""
    if not weights:
        return "(none)"
    parts = [f"{w:.2f}·{n}" for n, w in sorted(weights.items())]
    return " + ".join(parts)
