models/full_model.py

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from config import (
    EMBEDDING_DIM,
    MODEL_DIM,
    PHASE_MODE,
    PHASE_SOURCE,
    MIXING,
    USE_LLM_PRIOR,
    QUANTUM_DEVICE,
)

from models.validator import HypothesisValidator
from models.quantum_layer import QuantumEvolutionLayer
from models.classical_control_layer import ClassicalControlLayer
from models.persistent_reasoner import PersistentReasoner
from models.collapse import CollapseController
from models.answer_selector import EnergyAnswerSelector, bounded_phase
from models.energy_fusion import EnergyFusion
from models.interference import CoherentCollapse

BACKENDS = ("quantum", "classical_control", "none")


def _masked_zscore(x, mask, eps=1e-5):
    """Per-sample standardization over the option axis, ignoring padding."""

    m = mask.to(x.dtype) if mask is not None else torch.ones_like(x)

    n = m.sum(dim=1, keepdim=True).clamp(min=1.0)

    mean = (x * m).sum(dim=1, keepdim=True) / n

    var = (((x - mean) ** 2) * m).sum(dim=1, keepdim=True) / n

    return ((x - mean) / torch.sqrt(var + eps)) * m


class QAIRvNext(nn.Module):
    """
    qAIR-v45.

        Q, H, O  ->  in_proj (384 -> MODEL_DIM)
                 ->  PersistentReasoner        (orthogonal mixing)
                 ->  Quantum / classical layer (state + energy + PHASE)
                 ->  HypothesisValidator       (-> potential, fed back)
                 ->  EnergyAnswerSelector      (E_kn + compatibility phase)
                 ->  EnergyFusion              (-> E_k)
                 ->  CollapseController        (per-sample temperature)
                 ->  CoherentCollapse          (complex amplitudes,
                                                |sum_k c_k v_kn|^2)
                 ->  + LLM log-prior           (optional, ablatable)

    Three things are different from v44 in ways that matter, all of them
    responses to specific measured failures rather than tuning:

    1. **Width.** v44 ran 10.58M parameters over 800 training examples --
       13,223 per example -- and its best validation accuracy was at
       epoch 1, declining monotonically after while train loss fell 54%.
       Everything past `in_proj` now runs at MODEL_DIM (128), roughly 1.1M
       parameters. Below that ratio, "hypotheses don't help" and "the
       model memorizes before it learns to use them" are not
       distinguishable, so no mechanistic claim from v44 was safe.

    2. **Evidence channels.** The cache now carries a supporting AND an
       attacking hypothesis per option, each tagged with its polarity and
       the option it argues about (`align`). v44's support-only
       hypotheses were measured at argmax diag(H.O) = 0.2481 against a
       chance of 0.2500 -- literally zero discriminative signal, because
       asking a competent LLM to justify all four options gets all four
       justified.

    3. **Genuine interference.** The final distribution is
       |sum_k c_k v_kn|^2 over complex amplitudes, so hypotheses can
       cancel. See models/interference.py. `phase_mode="classical"`
       recovers v44's real-amplitude mixture as a control, and
       `interference_ratio` is reported every epoch so an inert phase
       channel shows up as a number rather than hiding behind a plausible
       accuracy.

    The LLM log-prior (`use_llm_prior`) deserves its own warning. It adds
    the generator's own cached per-option log-likelihood to the final
    score. That is legitimate -- it is the only channel by which the LLM's
    judgement reaches the answer without being squeezed through a frozen
    384-d sentence encoder -- but it also means the model can score well
    while ignoring every hypothesis. Always read
    `evaluation/input_ablation.py` alongside accuracy, and always run the
    `use_llm_prior=False` arm before attributing a gain to reasoning.
    """

    def __init__(
        self,
        dim=EMBEDDING_DIM,
        model_dim=MODEL_DIM,
        use_quantum=True,
        use_validator=True,
        persistent_steps=3,
        n_qubits=12,
        backend=None,
        use_question=True,
        validator_feedback=True,
        keep_trajectory=False,
        phase_mode=PHASE_MODE,
        phase_source=PHASE_SOURCE,
        mixing=MIXING,
        use_llm_prior=USE_LLM_PRIOR,
        use_attack=True,
        quantum_device=QUANTUM_DEVICE,
        verbose=False,
    ):

        super().__init__()

        # backend selects what occupies the quantum-layer slot:
        #   "quantum"           -- QuantumEvolutionLayer (PennyLane circuit)
        #   "classical_control" -- ClassicalControlLayer (parameter-matched
        #                          classical MLP in the same slot)
        #   "none"              -- no bottleneck layer at all
        # use_quantum is a backward-compatible alias (True -> "quantum",
        # False -> "none") used only when backend isn't given explicitly.
        # The submodule is stored as self.quantum regardless of backend so
        # training/train.py's quantum./validator.-prefixed lower-LR
        # parameter group applies identically to both -- the comparison
        # between them should differ only in what's inside that slot, not
        # in how it is optimized.
        if backend is None:
            backend = "quantum" if use_quantum else "none"

        if backend not in BACKENDS:
            raise ValueError(f"backend must be one of {BACKENDS}, got {backend!r}")

        d = model_dim or dim

        self.embed_dim = dim
        self.model_dim = d

        self.backend = backend
        self.use_quantum = backend != "none"
        self.use_validator = use_validator
        self.use_question = use_question
        self.use_llm_prior = use_llm_prior
        self.use_attack = use_attack
        self.phase_mode = phase_mode
        self.phase_source = phase_source
        self.mixing = mixing
        self.quantum_device = quantum_device
        self.persistent_steps = persistent_steps

        # The validator computes a `potential` field that was never fed
        # back into the reasoner before v44: forward() called
        # self.reasoner(H) with no potential argument even though
        # PersistentReasoner accepted one, so the validator only ever
        # reached the answer through validator_energy. With feedback on,
        # the reasoner runs twice: once to get hypotheses good enough to
        # validate, then again guided by the resulting potential.
        self.validator_feedback = validator_feedback and use_validator

        # ------------------------------------------------------------
        # Input projection: 384-d frozen embeddings -> model width
        # ------------------------------------------------------------

        if d == dim:
            self.in_proj = nn.Identity()
        else:
            self.in_proj = nn.Sequential(nn.Linear(dim, d), nn.LayerNorm(d))

        # Marks a hypothesis as supporting (+1) or attacking (-1) inside
        # the vector itself, so the reasoner's mixing can tell the two
        # channels apart rather than treating them as interchangeable.
        self.polarity_embed = nn.Parameter(torch.zeros(d))

        self.reasoner = PersistentReasoner(
            d, steps=persistent_steps, keep_trajectory=keep_trajectory,
            mixing=mixing,
        )

        if backend == "quantum":
            self.quantum = QuantumEvolutionLayer(
                d, n_qubits=n_qubits, verbose=verbose,
                quantum_device=quantum_device,
            )
        elif backend == "classical_control":
            self.quantum = ClassicalControlLayer(d, n_qubits=n_qubits, verbose=verbose)

        if use_validator:
            self.validator = HypothesisValidator(d, use_question=use_question)

        self.selector = EnergyAnswerSelector(d, use_question=use_question)

        self.fusion = EnergyFusion()

        self.collapse = CollapseController()

        self.interference = CoherentCollapse(phase_mode=phase_mode)

        # Classical fallback phase head, used when phase_source ==
        # "learned" or when there is no bottleneck layer to read a phase
        # off. One linear map to one bounded scalar per hypothesis.
        self.phase_head = nn.Linear(d, 1)

        # Relative phase between the two evidence directions, initialized
        # to pi.
        #
        # This is the one piece of physics the architecture actually
        # asserts, and it is worth stating plainly: an objection to
        # option n is the NEGATION of the support for option n, so the
        # natural relative phase between them is pi, where their
        # amplitudes cancel. Without it the phase channel starts
        # essentially uniform -- measured at initialization, phases from
        # both the circuit and the learned head are concentrated tightly
        # enough that the coherent sum is purely constructive and
        # `destructive_fraction` is exactly 0. The mechanism the whole
        # v45 claim rests on would begin dead and have to discover
        # cancellation from a gradient signal that only exists once
        # cancellation happens.
        #
        # It is a learnable parameter, not a constant: if support and
        # attack turn out not to be opposites in any useful sense, the
        # model can drive it to 0 and recover the in-phase behaviour.
        # `phase_mode="zero"` and `phase_mode="classical"` bypass it
        # entirely, so it is covered by the existing controls.
        self.polarity_phase = nn.Parameter(torch.tensor(math.pi))

        # Weight on the cached LLM log-prior. softplus keeps it
        # non-negative (a prior the model can trust or ignore, but never
        # invert); softplus(-1.0) = 0.313 at init -- present, not
        # dominant.
        self.llm_prior_weight = nn.Parameter(torch.tensor(-1.0))

    def _run_quantum(self, H):

        if H.is_cuda:

            # Only construct a "cuda" autocast context on an actual CUDA
            # tensor -- on CPU/MPS this context wouldn't do anything
            # anyway (autocast is disabled either way in the training
            # loop for non-CUDA devices), and some PyTorch versions
            # validate device_type at construction even when disabled.
            #
            # Forcing fp32 here applies regardless of which PennyLane
            # device is attached (see config.QUANTUM_DEVICE) -- PennyLane's
            # autograd through lightning.qubit isn't mixed-precision safe,
            # and lightning.gpu's cuStateVec backend is fp32/fp64
            # statevectors, not fp16/bf16, so autocast has nothing valid to
            # do to either one.
            from torch.amp import autocast

            with autocast(device_type="cuda", enabled=False):
                return self.quantum(H.float())

        return self.quantum(H.float())

    def forward(self, H, O, Q=None, y=None, H_mask=None, O_mask=None,
                polarity=None, align=None, llm_logprob=None):
        """
        H           : (B, K, D_embed) hypothesis embeddings, K = 2N
        O           : (B, N, D_embed) option embeddings
        Q           : (B, D_embed)    question embedding
        y           : (B,)  labels, training only -- used solely to build
                      the validator's BCE target, never to compute energy
        H_mask      : (B, K) bool, True for real hypotheses
        O_mask      : (B, N) bool, True for real options
        polarity    : (B, K) +1 support / -1 attack
        align       : (B, K, N) 1 where hypothesis k is about option n
        llm_logprob : (B, N) cached per-option LLM log-likelihood

        Masks matter: collate_fn pads ragged option/hypothesis counts with
        zeros, and before the v44 audit those masks were computed and then
        never passed to the model, so zero-padded options participated in
        every reduction and could be returned as the argmax prediction.
        """

        B, K, _ = H.shape
        N = O.shape[1]

        if H_mask is None:
            H_mask = torch.ones(B, K, dtype=torch.bool, device=H.device)

        if O_mask is None:
            O_mask = torch.ones(B, N, dtype=torch.bool, device=O.device)

        if polarity is None:
            polarity = torch.ones(B, K, device=H.device, dtype=H.dtype)

        # Support-only ablation: drop the attacking channel without
        # touching the cache, so the comparison is over identical data.
        if not self.use_attack:
            H_mask = H_mask & (polarity >= 0)

        # ------------------------------------------------------------
        # Project into model width
        # ------------------------------------------------------------

        H = self.in_proj(H)
        O = self.in_proj(O)
        Q = self.in_proj(Q) if Q is not None else None

        H = H + polarity.unsqueeze(-1) * self.polarity_embed

        H = H * H_mask.unsqueeze(-1).to(H.dtype)

        # ------------------------------------------------------------
        # Reason
        # ------------------------------------------------------------

        H, trajectory, interaction = self.reasoner(H, mask=H_mask)

        quantum_energy = None
        quantum_phase = None

        if self.use_quantum:

            q_state, quantum_energy, quantum_phase = self._run_quantum(H)

            H = H + q_state

        validator_out = None
        validator_energy = None
        potential = None

        if self.use_validator:

            validator_out = self.validator(
                H, O, Q=Q, y=y, H_mask=H_mask, O_mask=O_mask,
                align=align, polarity=polarity,
            )

            potential = validator_out["potential"]

            validator_energy = validator_out["validator_energy"]

            if self.validator_feedback:

                # Second reasoning pass, guided by the validator's
                # potential -- the feedback edge the architecture was
                # described as having but did not until v44.
                H, trajectory2, interaction = self.reasoner(
                    H, potential=potential, mask=H_mask
                )

                trajectory = trajectory + trajectory2

                if self.use_quantum:
                    q_state, quantum_energy, quantum_phase = self._run_quantum(H)
                    H = H + q_state

                validator_out = self.validator(
                    H, O, Q=Q, y=y, H_mask=H_mask, O_mask=O_mask,
                    align=align, polarity=polarity,
                )

                potential = validator_out["potential"]

                validator_energy = validator_out["validator_energy"]

        # ------------------------------------------------------------
        # Pairwise energy + compatibility phase
        # ------------------------------------------------------------

        selector = self.selector(
            H, O, Q=Q, align=align, polarity=polarity,
            H_mask=H_mask, O_mask=O_mask,
        )

        answer_energy = selector["energy"]

        collapse_energy = self.fusion(
            answer_energy,
            quantum_energy=quantum_energy,
            validator_energy=validator_energy,
            option_mask=O_mask,
        )

        collapse_out = self.collapse(collapse_energy, mask=H_mask)

        # ------------------------------------------------------------
        # Hypothesis phase
        # ------------------------------------------------------------

        if self.phase_source == "circuit" and quantum_phase is not None:
            phase_h = quantum_phase
        else:
            phase_h = bounded_phase(self.phase_head(H).squeeze(-1))

        # +0 for support (polarity +1), +polarity_phase for attack (-1).
        phase_h = phase_h + 0.5 * (1.0 - polarity) * self.polarity_phase

        phase_h = phase_h * H_mask.to(phase_h.dtype)

        # ------------------------------------------------------------
        # Coherent collapse -- |sum_k c_k v_kn|^2
        # ------------------------------------------------------------

        coherent = self.interference(
            collapse_energy,
            answer_energy,
            phase_h=phase_h,
            phase_kn=selector["phase"],
            H_mask=H_mask,
            O_mask=O_mask,
            hyp_temperature=collapse_out["temperature"],
        )

        scores = coherent["log_probs"]

        # ------------------------------------------------------------
        # LLM log-prior (explicit, weighted, ablatable)
        # ------------------------------------------------------------

        llm_term = None

        if self.use_llm_prior and llm_logprob is not None:

            # Standardized per sample: raw mean-token log-likelihoods sit
            # around -2 to -8 with a scale that varies with option length,
            # and an unnormalized additive term at that scale would swamp
            # log_probs outright.
            llm_term = _masked_zscore(llm_logprob, O_mask)

            scores = scores + F.softplus(self.llm_prior_weight) * llm_term

        # Padded options must never be predicted or carry probability
        # mass in the cross-entropy.
        neg_inf = torch.finfo(scores.dtype).min / 2
        scores = scores.masked_fill(~O_mask, neg_inf)

        # Renormalized, so `scores` is a genuine log-probability vector
        # and calibration metrics (evaluation ECE) are meaningful.
        #
        # The floor is not cosmetic: without it a padded option lands at
        # exactly -inf, and label smoothing -- which puts a little mass on
        # EVERY class -- multiplies that by epsilon and returns inf. The
        # loss then NaNs on any batch containing a 3-option question,
        # which ARC has. At -30 a padded option carries 9e-14 of the
        # probability mass, i.e. nothing, while staying finite.
        scores = torch.log_softmax(scores, dim=1).clamp(min=-30.0)

        collapse_probs = coherent["collapse_probs"]

        entropy = -(
            collapse_probs * torch.log(collapse_probs + 1e-8)
        ).sum(dim=1).mean()

        return {
            "scores": scores,
            "answer_energy": answer_energy,
            "collapse_energy": collapse_energy,
            "collapse_probs": collapse_probs,
            "quantum_energy": quantum_energy,
            "quantum_phase": quantum_phase,
            "phase_h": phase_h,
            "phase_kn": selector["phase"],
            "interference": coherent["interference"],
            "interference_ratio": coherent["interference_ratio"],
            "phase_effect": coherent["phase_effect"],
            "destructive_fraction": coherent["destructive_fraction"],
            "coherent_contrast": coherent["coherent_contrast"],
            "classical_weights": coherent["classical"],
            "coherent_weights": coherent["coherent"],
            "llm_prior": llm_term,
            "llm_prior_weight": F.softplus(self.llm_prior_weight).detach(),
            "entropy": entropy,
            "diversity": collapse_out["diversity"],
            "spread": collapse_out["spread"],
            "peak": collapse_probs.max(dim=1).values.mean(),
            "collapse_confidence": collapse_out["confidence"],
            "validator_potential": potential,
            "trajectory": trajectory,
            "attention": interaction,
            "validator": validator_out,
            "H_mask": H_mask,
            "O_mask": O_mask,
            # Post-reasoning hypotheses, so training/evaluate.py can
            # distinguish "the hypotheses merged into one vector" from
            # "the selector learned to ignore them". Both make the
            # multi-hypothesis mechanism inert, with different fixes.
            "H_reasoned": H,
        }