"""
wheeler.py -- the Wheeler-DeWitt channel-mixing block for Wheeler-LM, the family's
quantum-gravity mixer (structural analog of Quazimoto's Kuramoto block and Mycel's
growth block). It is built from the Wheeler-DeWitt equation of canonical quantum
gravity -- the "wavefunction of the universe":

    hbar^2 G_ijmn (d/dg_ij)(dPsi/dg_mn) + R sqrt(g) Psi = 0         (Hamiltonian constraint)
    (dPsi/dg_mn)_|m = 0                                             (momentum / diffeo constraint)

What makes this different from every other mixer in the family:
  * NO TIME. It is not an evolution dPsi/dt = ...; it is a CONSTRAINT H Psi = 0. The
    state does not flow forward in an external time -- it must merely SATISFY the
    constraint. We realise this as an equilibrium/constraint block whose "time" is only
    the iteration index, plus an auxiliary loss <H^2> trained toward zero.
  * LORENTZIAN supermetric. The DeWitt supermetric G_ijmn is INDEFINITE -- one timelike
    direction (the volume / scale mode) and the rest spacelike. So the kinetic operator
    G d d is a WAVE operator, not a diffusion. We give the K minisuperspace modes a
    learnable indefinite metric G^{-1}, initialised diag(-1, +1, ..., +1).
  * CURVATURE POTENTIAL. The potential term is R sqrt(g) (scalar curvature), realised as
    a per-mode curvature weighting r that the token produces.
  * DIFFEO INVARIANCE. The momentum constraint says Psi depends only on the geometry, not
    the coordinates -> the readout uses only INVARIANTS of the mode vector (the timelike
    component and the norm of the spacelike part), never the raw coordinate values.

Mechanics (per token): hidden -> (Psi, Pi) minisuperspace amplitude + conjugate momentum
-> a few leapfrog steps of the Wheeler-DeWitt wave dynamics

    Pi  <- Pi - dt * (r ⊙ Psi)          # Pi_dot = -dU/dPsi   with U = 1/2 r Psi^2
    Psi <- Psi + dt * (G^{-1} @ Pi)      # Psi_dot = G^{-1} Pi   (indefinite -> wave)

-> readout of diffeo-invariants -> hidden, behind a family gate. The Hamiltonian
   H = 1/2 Pi^T G^{-1} Pi + 1/2 sum_k r_k Psi_k^2 is exposed as last_constraint so the
   trunk can add <H^2> to the loss (the block is pressured onto the physical H=0 surface).
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

from family import soft_clamp, RMSNorm
import instrument as _viz


class WheelerDeWittBlock(nn.Module):
    """Wheeler-DeWitt wave mixer: hidden -> minisuperspace (Psi, Pi) -> leapfrog wave
    dynamics under a Lorentzian DeWitt supermetric with a curvature potential -> readout
    of diffeo-invariants. Drop-in for QuazimotoBlock.forward(x, ring_ctl, phase_seed)."""

    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.K = cfg.wdw_modes
        self.norm = RMSNorm(cfg.d_model)
        self.to_psi = nn.Linear(cfg.d_model, self.K)          # initial wavefunction amplitude
        self.to_pi = nn.Linear(cfg.d_model, self.K)           # conjugate momentum
        self.to_curv = nn.Linear(cfg.d_model, self.K)         # per-mode curvature R*sqrt(g)

        # DeWitt supermetric G^{-1}: a LEARNABLE indefinite (Lorentzian) metric on the K
        # modes. Parameterised by a full matrix, symmetrised at use; a per-mode signature
        # (sigma, init [-1,+1,...,+1]) fixes the Lorentzian character -- mode 0 is timelike
        # (the volume / scale direction, the emergent "clock"), the rest spacelike.
        self.ginv_raw = nn.Parameter(torch.zeros(self.K, self.K))
        sig = torch.ones(self.K); sig[0] = -1.0               # Lorentzian signature
        self.register_buffer("signature", sig)
        self.log_lapse = nn.Parameter(torch.zeros(1))         # lapse N: scales the intrinsic-time step

        # readout on diffeo-INVARIANTS of the final state: [Psi, Pi, timelike comp, ||space||].
        # (never on raw coordinate values -> honours the momentum constraint's invariance.)
        hidden = cfg.readout_mult * cfg.d_model
        self.readout = nn.Sequential(
            nn.Linear(2 * self.K + 2, hidden), nn.GELU(),
            nn.Linear(hidden, cfg.d_model))
        self.drop = nn.Dropout(cfg.dropout)
        go = math.atanh(min(cfg.gate_init_open, 0.9)) if cfg.gate_init_open > 0 else 0.0
        self.gate = nn.Parameter(torch.full((1,), go))
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, std=0.02)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
        if cfg.use_fractal_phase_seed:
            self.psi_seed_gate = nn.Parameter(torch.zeros(1))  # zero-init -> no-op

        self.last_constraint = None                           # <H^2> for the trunk to collect

    def _supermetric(self):
        """The learnable Lorentzian DeWitt supermetric G^{-1} [K,K]. Symmetric, with the
        signature baked into the diagonal so it starts genuinely indefinite (timelike mode
        0). ginv_raw learns the off-diagonal mode couplings on top of that."""
        A = 0.5 * (self.ginv_raw + self.ginv_raw.t())         # symmetric part
        return A + torch.diag(self.signature)                 # Lorentzian at init (A=0)

    def forward(self, x, ring_ctl=None, phase_seed=None):
        cfg = self.cfg
        B, T, _ = x.shape
        h = self.norm(x)

        psi = self.to_psi(h)                                  # [B,T,K] wavefunction amplitude
        # NOTE: the FractalSeed trait is a NO-OP in this project. The trained Wheeler-LM
        # checkpoint left every psi_seed_gate ~0 (seed on-vs-off flips 0/64 argmax tokens,
        # ~0.1% of logit scale), so the fractal table is effectively UNUSED here. It is
        # kept only for family parity / possible future continued training -- do not treat
        # it as an active input to the Wheeler-DeWitt dynamics.
        if phase_seed is not None:                            # optional fractal seed of Psi (gated ~0)
            seed = phase_seed[..., :self.K].view(B, T, self.K)
            psi = psi + torch.tanh(self.psi_seed_gate) * seed
        pi = self.to_pi(h)                                    # [B,T,K] conjugate momentum
        r = self.to_curv(h)                                   # [B,T,K] curvature potential R*sqrt(g)

        Ginv = self._supermetric()                            # [K,K] Lorentzian
        dt = cfg.wdw_dt * F.softplus(self.log_lapse).clamp(max=4.0)  # lapse-scaled intrinsic time

        # leapfrog integration of the Wheeler-DeWitt wave equation on the K modes. Because
        # Ginv is INDEFINITE, the mode-0 (volume) direction is a genuine wave/"time" while
        # the spacelike modes oscillate -- a Klein-Gordon eq on minisuperspace.
        pi = pi - 0.5 * dt * (r * psi)                        # half kick
        for _ in range(cfg.wdw_steps):
            psi = psi + dt * torch.matmul(pi, Ginv.t())       # drift: Psi_dot = G^{-1} Pi
            last = (_ == cfg.wdw_steps - 1)
            pi = pi - (0.5 if last else 1.0) * dt * (r * psi)  # kick: Pi_dot = -(R sqrt g) Psi

        # Hamiltonian constraint H = 1/2 Pi^T G^{-1} Pi + 1/2 sum_k r_k Psi_k^2. Physical
        # states satisfy H = 0; expose <H^2> so the trunk can pressure the block onto that
        # constraint surface (the closest thing to "solving" H Psi = 0 in a feed-forward net).
        kinetic = 0.5 * (torch.matmul(pi, Ginv.t()) * pi).sum(-1)   # [B,T]
        potential = 0.5 * (r * psi * psi).sum(-1)                    # [B,T]
        H = kinetic + potential
        self.last_constraint = (H ** 2).mean() if self.training else None

        # diffeo-INVARIANT readout features: the timelike component and the norm of the
        # spacelike part are invariant under the residual coordinate symmetry.
        t_comp = psi[..., :1]                                 # timelike (volume/scale) mode
        space_norm = psi[..., 1:].norm(dim=-1, keepdim=True)  # ||spacelike Psi||
        feat = torch.cat([psi, pi, t_comp, space_norm], dim=-1)
        out = self.drop(soft_clamp(self.readout(feat) * torch.tanh(self.gate), cfg.osc_bound))

        rec = _viz.get_rec()
        if rec is not None and rec.enabled:                  # live-viz: emergent-time clock
            rec.log_ring([float(H[0, -1].abs())], [float(t_comp[0, -1, 0])],
                         psi[0, -1].tolist())
            rec.log_quaz_norm(out[0, -1].norm().item())
            rec.flush_spec()
        return out