"""2Xplat v2: the MVP-Giant appearance expert.

One difference from ``src.model.model``, which the checkpoint depends on: the Gaussian head
emits a single ray-depth channel instead of three xyz channels, and activates scale with
softplus in linear space instead of exp in log space.

Multi-resolution PRoPE lives in the base class, so both models share it.
"""

import torch
import torch.nn.functional as F

from src.model.model import AppearanceExpert, TwoExpertModel


class AppearanceExpertV2(AppearanceExpert):
    """Appearance expert whose head predicts per-pixel depth rather than xyz."""

    def _activate_scale(self, scale: torch.Tensor) -> torch.Tensor:
        """Activate scale with softplus, then return its log.

        v2 was trained with scale in linear space, but ``GaussianRenderer.render``
        exponentiates whatever it is given, so handing back the log makes that
        round-trip exact and keeps the shared renderer untouched.
        """
        return F.softplus(scale - 4.0).clamp(1e-6, 1.0).log()

    def _ray_distance(self, pos: torch.Tensor) -> torch.Tensor:
        """``pos`` is already a single depth channel, so no reduction is needed."""
        return pos.sigmoid() * self.max_dist


class TwoExpertModelV2(TwoExpertModel):
    """2Xplat v2 compositor: depth-parameterised Gaussians."""

    POS_DIM = 1
    APPEARANCE_EXPERT_CLS = AppearanceExpertV2