"""composer_adv.py — VeRL custom advantage estimator with SDPO + replay shaping.

Architecture spec: docs/INTEGRATION_ARCHITECTURE.md § "Recipe B".
Verified extension point: @register_adv_est decorator + DataProto.batch /
non_tensor_batch fields (DeepWiki audit of volcengine/verl, 2026-05-25).

Pattern:
  - Register a new advantage estimator alongside VeRL's built-in `grpo`.
  - At rollout time, the rollout worker stashes hint-conditioned teacher logprobs
    (channel 2) and N-teacher action distributions (channel 3) into the DataProto.
  - At advantage compute time, we read those fields and shape the GRPO advantage.

This pattern is identical to how VeRL already handles distillation rollouts
(per the DeepWiki audit: "teacher log-probabilities are stashed on the rollout
output and later concatenated into the per-batch DataProto for the student
training step").
"""

from __future__ import annotations

import torch

# These imports work when VeRL is installed — they're not skeleton imports.
# Verified via DeepWiki: the path is verl.trainer.ppo.core_algos.
try:
    from verl.trainer.ppo import core_algos  # type: ignore
    from verl.trainer.ppo.core_algos import register_adv_est  # type: ignore
except ImportError:  # pragma: no cover — fallback so module imports without VeRL
    core_algos = None  # type: ignore
    def register_adv_est(name: str):  # type: ignore
        def deco(fn):
            return fn
        return deco


@register_adv_est("grpo_composer")
def compute_grpo_composer_advantage(
    token_level_rewards: torch.Tensor,
    eos_mask: torch.Tensor,
    index: torch.Tensor,
    *,
    # Channel 2 (SDPO) extras — None when alpha_sdpo == 0
    sdpo_teacher_logprobs: torch.Tensor | None = None,
    sdpo_error_mask: torch.Tensor | None = None,
    old_log_prob: torch.Tensor | None = None,
    alpha_sdpo: float = 0.0,
    # Channel 3 (trace-replay) extras — None when beta_replay == 0
    teacher_consensus_prm: torch.Tensor | None = None,
    beta_replay: float = 0.0,
    **_kwargs,
) -> torch.Tensor:
    """GRPO advantage with SDPO + N-teacher trace-replay shaping.

    The base GRPO outcome advantage is computed as in VeRL's built-in `grpo`
    estimator. Then two additive shaping terms are layered on top:

      base_adv = compute_grpo_outcome_advantage(token_level_rewards, eos_mask, index)
      sdpo_term  = α_sdpo  · (teacher_lp - student_lp) · error_mask
      replay_term = β_replay · teacher_consensus_prm
      adv = base_adv + sdpo_term + replay_term

    Args:
        token_level_rewards: per-token reward signal (RLVR or shaped) [B, T].
        eos_mask: per-token EOS mask [B, T].
        index: group/prompt index for GRPO grouping [B].
        sdpo_teacher_logprobs: per-token logprob from hint-conditioned forward.
            None when alpha_sdpo == 0. Required when alpha_sdpo != 0.
        sdpo_error_mask: per-token mask, 1 at error-turn tokens, 0 elsewhere.
        old_log_prob: per-token logprob of the student under the current policy
            (already in DataProto.batch by VeRL convention).
        alpha_sdpo: weight on the SDPO advantage shaping. 0 to disable.
        teacher_consensus_prm: per-token Process-Reward-Model signal derived from
            N-teacher consensus disagreement. None when beta_replay == 0.
        beta_replay: weight on the trace-replay PRM shaping. 0 to disable.

    Returns:
        Shaped advantage tensor [B, T].
    """
    if core_algos is None:
        raise RuntimeError(
            "VeRL not installed. Install via `pip install verl` and ensure "
            "`from verl.trainer.ppo import core_algos` works before using this estimator."
        )

    # Base GRPO advantage (call VeRL's built-in)
    base_adv = core_algos.compute_grpo_outcome_advantage(
        token_level_rewards=token_level_rewards,
        eos_mask=eos_mask,
        index=index,
    )

    # Channel 2 shaping (SDPO)
    if alpha_sdpo != 0.0 and sdpo_teacher_logprobs is not None:
        if old_log_prob is None or sdpo_error_mask is None:
            raise ValueError(
                "alpha_sdpo != 0 requires sdpo_teacher_logprobs, sdpo_error_mask, "
                "and old_log_prob. Check the rollout worker is attaching them."
            )
        sdpo_term = alpha_sdpo * (sdpo_teacher_logprobs - old_log_prob)
        sdpo_term = sdpo_term * sdpo_error_mask
        base_adv = base_adv + sdpo_term

    # Channel 3 shaping (trace-replay PRM)
    if beta_replay != 0.0 and teacher_consensus_prm is not None:
        base_adv = base_adv + beta_replay * teacher_consensus_prm

    return base_adv


__all__ = ["compute_grpo_composer_advantage"]