"""ARC-V training objective (ARCV_METHOD_DESIGN.md §4.3).

Combines the loss-side ARC-V mechanisms on top of any model that maps images to
logits: the Fourier amplitude-mix (M2) augmentation plus a FACT-style
phase-consistency co-teaching term. MixStyle (M1) lives inside the backbone via
:func:`src.dg.mixstyle.register_mixstyle_hooks` (architectural, no loss term); the
cross-source SupCon (M3) is a LOSO-only add-on for Regime B.

With ``p_f=0`` and ``gamma_pc=0`` this reduces *exactly* to ERM cross-entropy, so a
single code path covers both the ERM baseline and the ARC-V method (toggled by
config) -- the honest apples-to-apples comparison the design demands.

``forward`` is a callable mapping an image batch to class logits, e.g.
``model.forward_image``.
"""
import torch.nn.functional as F

from src.dg.fourier import FourierAmplitudeMix, phase_consistency_loss


class ARCVObjective:
    def __init__(self, p_f: float = 0.5, eta: float = 1.0,
                 eta_weak: float = 0.3, gamma_pc: float = 1.0):
        self.strong = FourierAmplitudeMix(p=p_f, eta=eta)
        self.weak = FourierAmplitudeMix(p=p_f, eta=eta_weak)
        self.gamma_pc = float(gamma_pc)

    def __call__(self, forward, imgs, labels, training: bool = True):
        self.strong.train(training)
        self.weak.train(training)
        logits = forward(self.strong(imgs))
        loss = F.cross_entropy(logits, labels)
        if training and self.gamma_pc > 0.0:
            logits_weak = forward(self.weak(imgs))
            loss = loss + self.gamma_pc * phase_consistency_loss(logits, logits_weak)
        return loss, logits