model.py

"""
Geometric Transformer — GeoLIP Pipeline Integration
=====================================================
Dual-stream transformer with constellation-routed attention,
quaternion composition, and per-layer Cayley alignment.

Uses REAL geolip_core components:
    core.associate.constellation  — ConstellationObserver (anchors + triangulation + patchwork)
    core.curate.gate              — AnchorGate (CM determinant validity)
    core.align.procrustes         — CayleyOrthogonal rotation in SO(d)
    pipeline.observer             — TorchComponent / BaseTower interfaces

NEW components (transformer-specific):
    ManifoldProjection            — Input stage: hidden_state → S^(d-1)
    PositionGeometricContext      — Curation: constellation output → FiLM context
    FiLMLayer                     — Feature-wise Linear Modulation (proven in Ryan Spearman)
    GeometricAttention            — Attention with FiLM on Q,K from curated constellation
    QuaternionCompose             — Hamilton product of dual-stream outputs (proven)
    CayleyOrthogonal              — SO(d) rotation via Cayley map (proven)
    DualStreamBlock               — Content + geometric streams, aligned + composed
    GeometricTransformerLayer     — Full layer: project → observe → attend → compose
    GeometricTransformer          — Stack of layers with cross-layer rotation

Architecture per layer:
    1. ManifoldProjection:  h_i → emb_i on S^(d-1) per position
    2. ConstellationObserver: emb_i → {triangulation, assignment, patchwork, bridge}
    3. PositionGeometricContext: constellation output → (B, L, context_dim)
    4. Stream A (content):  standard self-attention
    5. Stream B (geometric): attention with FiLM(Q,K | geo_ctx), V unmodulated
    6. CayleyOrthogonal:    align B → A basis
    7. QuaternionCompose:   w=content, i=aligned_geo, j=disagree, k=agree
    8. Gated residual

Design principles from Ryan Spearman (ρ=0.309, 76/84 wins):
    - FiLM on Q,K ONLY — geometry routes attention, V stays pure
    - FiLM on individual arms BEFORE composition, not after
    - Quaternion algebra as structural regularizer (non-commutative coupling)
    - Disagreement arm (j) carries the transferable signal
    - CayleyOrthogonal guarantees pure rotation (det=1 always)
    - Never global average pool — per-position geometric context

Usage:
    from geometric_transformer import GeometricTransformer

    model = GeometricTransformer('geo_xfmr', d_model=512, n_layers=4)
    out = model(hidden_states)

    # Or as a head on frozen ESM-2:
    model = GeometricTransformer('esm2_geo', d_model=1280, n_layers=6)
    out = model(esm2_hidden_states)

Dependencies:
    pip install geolip-core  (includes constellation, patchwork, gate, observer interfaces)
"""

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

# ═══════════════════════════════════════════════════════════════════════════════
# GEOLIP IMPORTS — real components, not reimplementations
# ═══════════════════════════════════════════════════════════════════════════════

try:
    from geolip_core.core.associate.constellation import (
        ConstellationObserver, ConstellationAssociation, ConstellationCuration,
        Constellation, init_anchors_repulsion,
    )
    from geolip_core.core.curate.gate import AnchorGate
    from geolip_core.pipeline.observer import (
        TorchComponent, BaseTower, Input, Curation, Distinction,
    )
    _HAS_GEOLIP = True
except ImportError:
    _HAS_GEOLIP = False

    # ── Fallback stubs ──
    class TorchComponent(nn.Module):
        def __init__(self, name=None, **kwargs):
            super().__init__()
            self._component_name = name or self.__class__.__name__

    class BaseTower(nn.Module):
        def __init__(self, name=None, **kwargs):
            super().__init__()
            self._tower_name = name or self.__class__.__name__
            self._components = nn.ModuleDict()
            self._cache = {}

        def attach(self, name, module):
            if isinstance(module, nn.Module):
                self._components[name] = module
            return self

        def has(self, name):
            return name in self._components

        def __getitem__(self, key):
            return self._components[key]

        def cache_set(self, key, value):
            self._cache[key] = value

        def cache_get(self, key, default=None):
            return self._cache.get(key, default)

        def cache_clear(self):
            self._cache.clear()

    Input = TorchComponent
    Curation = TorchComponent
    Distinction = TorchComponent

    class Constellation(nn.Module):
        """Learned anchors on S^(d-1). Triangulates input embeddings."""
        def __init__(self, n_anchors, dim, anchor_drop=0.0, anchor_init='repulsion'):
            super().__init__()
            self.n_anchors = n_anchors
            self.dim = dim
            self.anchor_drop = anchor_drop
            anchors = torch.randn(n_anchors, dim)
            # Repulsion-initialized
            anchors = F.normalize(anchors, dim=-1)
            for _ in range(200):
                sim = anchors @ anchors.T
                sim.fill_diagonal_(-2.0)
                anchors = F.normalize(anchors - 0.05 * anchors[sim.argmax(dim=1)], dim=-1)
            self.anchors = nn.Parameter(anchors)

        def triangulate(self, emb, training=False):
            anchors = F.normalize(self.anchors, dim=-1)
            cos = emb @ anchors.T
            tri = 1.0 - cos
            _, nearest = cos.max(dim=-1)
            return tri, nearest

        def forward(self, emb, training=False):
            return self.triangulate(emb, training)

    class ConstellationAssociation(TorchComponent):
        """Association through constellation anchors."""
        def __init__(self, dim=256, n_anchors=32, anchor_drop=0.0,
                     anchor_init='repulsion', assign_temp=0.1, **kwargs):
            super().__init__(**kwargs)
            self.assign_temp = assign_temp
            self.constellation = Constellation(n_anchors, dim, anchor_drop, anchor_init)

        @property
        def frame_dim(self):
            return self.constellation.n_anchors

        def associate(self, emb, **context):
            anchors_n = F.normalize(self.constellation.anchors, dim=-1)
            cos = emb @ anchors_n.T
            tri = 1.0 - cos
            _, nearest = cos.max(dim=-1)
            soft_assign = F.softmax(cos / self.assign_temp, dim=-1)
            mag = context.get('mag', None)
            distances_weighted = tri * mag if mag is not None else tri
            return {
                'distances': tri, 'distances_weighted': distances_weighted,
                'cos_to_anchors': cos, 'assignment': soft_assign,
                'nearest': nearest,
            }

        def forward(self, emb, **context):
            return self.associate(emb, **context)

    class Patchwork(nn.Module):
        """Round-robin patchwork compartments."""
        def __init__(self, n_anchors, n_comp=8, d_comp=32, activation='gelu'):
            super().__init__()
            self.n_comp = n_comp
            anchors_per = max(1, n_anchors // n_comp)
            self.compartments = nn.ModuleList([
                nn.Sequential(nn.Linear(anchors_per, d_comp), nn.GELU(), nn.Linear(d_comp, d_comp))
                for _ in range(n_comp)
            ])
            self.output_dim = n_comp * d_comp
            self.anchors_per = anchors_per

        def forward(self, distances):
            parts = []
            for i, comp in enumerate(self.compartments):
                start = i * self.anchors_per
                end = start + self.anchors_per
                chunk = distances[..., start:end]
                if chunk.shape[-1] < self.anchors_per:
                    chunk = F.pad(chunk, (0, self.anchors_per - chunk.shape[-1]))
                parts.append(comp(chunk))
            return torch.cat(parts, dim=-1)

    class ConstellationCuration(Curation):
        """Curation through patchwork compartments + bridge."""
        def __init__(self, n_anchors=32, dim=256, n_comp=8, d_comp=32,
                     activation='gelu', **kwargs):
            super().__init__(**kwargs)
            self.dim = dim
            self.n_anchors = n_anchors
            self.patchwork = Patchwork(n_anchors, n_comp, d_comp, activation)
            pw_dim = self.patchwork.output_dim
            self.bridge = nn.Linear(pw_dim, n_anchors)
            self._feature_dim = n_anchors + pw_dim + dim

        @property
        def feature_dim(self):
            return self._feature_dim

        def curate_full(self, association_output, emb=None, **context):
            distances = association_output['distances_weighted']
            assignment = association_output['assignment']
            pw = self.patchwork(distances)
            bridge = self.bridge(pw)
            parts = [assignment, pw]
            if emb is not None:
                parts.append(emb)
            features = torch.cat(parts, dim=-1)
            return {'patchwork': pw, 'bridge': bridge, 'features': features}

        def forward(self, association_output, emb=None, **context):
            return self.curate_full(association_output, emb=emb, **context)['features']

    class ConstellationObserver(nn.Module):
        """Composed association + curation."""
        def __init__(self, dim=256, n_anchors=32, n_comp=8, d_comp=32,
                     anchor_drop=0.0, anchor_init='repulsion',
                     activation='gelu', assign_temp=0.1):
            super().__init__()
            self.association = ConstellationAssociation(
                dim=dim, n_anchors=n_anchors, anchor_drop=anchor_drop,
                anchor_init=anchor_init, assign_temp=assign_temp)
            self.curation = ConstellationCuration(
                n_anchors=n_anchors, dim=dim, n_comp=n_comp,
                d_comp=d_comp, activation=activation)

        @property
        def constellation(self):
            return self.association.constellation

        @property
        def patchwork(self):
            return self.curation.patchwork

        @property
        def feature_dim(self):
            return self.curation.feature_dim

        def observe(self, emb, **context):
            a_out = self.association(emb, **context)
            c_out = self.curation.curate_full(a_out, emb=emb, **context)
            return {
                'embedding': emb, 'features': c_out['features'],
                'triangulation': a_out['distances'],
                'cos_to_anchors': a_out['cos_to_anchors'],
                'nearest': a_out['nearest'],
                'assignment': a_out['assignment'],
                'patchwork': c_out['patchwork'], 'bridge': c_out['bridge'],
            }

        def forward(self, emb, **context):
            return self.observe(emb, **context)


# ═══════════════════════════════════════════════════════════════════════════════
# PROVEN COMPONENTS — from Ryan Spearman (unchanged, tested)
# ═══════════════════════════════════════════════════════════════════════════════

class FiLMLayer(TorchComponent):
    """Feature-wise Linear Modulation. Proven in Ryan Spearman.

    Produces γ * x + β from geometric context.
    Identity-initialized: γ=1, β=0 at init.
    """
    def __init__(self, name, feature_dim, context_dim):
        super().__init__(name)
        self.to_gamma = nn.Linear(context_dim, feature_dim)
        self.to_beta = nn.Linear(context_dim, feature_dim)
        nn.init.zeros_(self.to_gamma.weight); nn.init.ones_(self.to_gamma.bias)
        nn.init.zeros_(self.to_beta.weight); nn.init.zeros_(self.to_beta.bias)

    def forward(self, x, ctx):
        """x: (B, L, D), ctx: (B, L, C) → (B, L, D)"""
        return self.to_gamma(ctx) * x + self.to_beta(ctx)


class CayleyOrthogonal(TorchComponent):
    """Guaranteed SO(d) rotation via Cayley map. Proven in Procrustes alignment.

    Q = (I - A)(I + A)^(-1) where A is skew-symmetric.
    det(Q) = 1 always. ‖R-I‖ ≈ 4.1 at convergence in SO(256).

    Caches the rotation matrix — only recomputes when A_upper changes
    (i.e. after optimizer.step()). The solve is input-independent.
    """
    def __init__(self, name, dim):
        super().__init__(name)
        self.dim = dim
        self.A_upper = nn.Parameter(torch.zeros(dim * (dim - 1) // 2) * 0.01)
        self._cached_R = None
        self._cached_A_version = None

    def _param_version(self):
        """Track parameter changes via data_ptr + requires_grad state."""
        return self.A_upper.data_ptr(), self.A_upper._version

    def get_rotation(self):
        # During training: always recompute (autograd graph needed fresh)
        # During eval: cache the rotation (params don't change)
        if self.training:
            self._cached_R = None

        version = self._param_version()
        if self._cached_R is not None and self._cached_A_version == version:
            return self._cached_R

        d = self.dim
        A = torch.zeros(d, d, device=self.A_upper.device, dtype=self.A_upper.dtype)
        idx = torch.triu_indices(d, d, offset=1, device=A.device)
        A[idx[0], idx[1]] = self.A_upper
        A = A - A.T
        I = torch.eye(d, device=A.device, dtype=A.dtype)
        R = torch.linalg.solve(I + A, I - A)

        if not self.training:
            self._cached_R = R
            self._cached_A_version = version
        return R

    def invalidate_cache(self):
        """Call after optimizer.step() if needed."""
        self._cached_R = None
        self._cached_A_version = None

    def forward(self, x):
        """(..., dim) → (..., dim) rotated."""
        return x @ self.get_rotation().T


def quaternion_multiply(q1, q2):
    """Hamilton product. q = (w, x, y, z) along dim=-2.

    Supports batched: (..., 4, D) × (..., 4, D) → (..., 4, D)
    Or scalar:        (..., 4) × (..., 4) → (..., 4)
    """
    w1, x1, y1, z1 = q1.unbind(-2) if q1.dim() >= 2 and q1.shape[-2] == 4 else q1.unbind(-1)
    w2, x2, y2, z2 = q2.unbind(-2) if q2.dim() >= 2 and q2.shape[-2] == 4 else q2.unbind(-1)
    stack_dim = -2 if q1.dim() >= 2 and q1.shape[-2] == 4 else -1
    return torch.stack([
        w1*w2 - x1*x2 - y1*y2 - z1*z2,
        w1*x2 + x1*w2 + y1*z2 - z1*y2,
        w1*y2 - x1*z2 + y1*w2 + z1*x2,
        w1*z2 + x1*y2 - y1*x2 + z1*w2,
    ], dim=stack_dim)


def quaternion_multiply_batched(q1, q2):
    """Hamilton product on (B, 4, D) tensors. Fully vectorized, no loops.

    Each of the 4 slices along dim=1 is one quaternion component.
    The D dimension is batched — all D quaternions multiplied in parallel.
    """
    w1, x1, y1, z1 = q1[:, 0], q1[:, 1], q1[:, 2], q1[:, 3]
    w2, x2, y2, z2 = q2[:, 0], q2[:, 1], q2[:, 2], q2[:, 3]
    return torch.stack([
        w1*w2 - x1*x2 - y1*y2 - z1*z2,
        w1*x2 + x1*w2 + y1*z2 - z1*y2,
        w1*y2 - x1*z2 + y1*w2 + z1*x2,
        w1*z2 + x1*y2 - y1*x2 + z1*w2,
    ], dim=1)  # (B, 4, D)


class QuaternionCompose(TorchComponent):
    """Four-arm Hamilton product composition. Proven in GeoQuat head.

    The algebra forces cross-term interactions between arms.
    Arms cannot independently memorize — the non-commutative
    product couples their outputs as structural regularizer.

    Fully vectorized: single batched Hamilton product, no Python loops.
    """
    def __init__(self, name, input_dim, quat_dim=64):
        super().__init__(name)
        self.quat_dim = quat_dim
        self.proj_w = nn.Linear(input_dim, quat_dim)
        self.proj_i = nn.Linear(input_dim, quat_dim)
        self.proj_j = nn.Linear(input_dim, quat_dim)
        self.proj_k = nn.Linear(input_dim, quat_dim)
        self.rotation = nn.Parameter(torch.randn(1, 4, quat_dim) * 0.1)

    @property
    def output_dim(self):
        return self.quat_dim * 4

    def forward(self, arm_w, arm_i, arm_j, arm_k):
        """Each arm: (B, L, D) → composed: (B, L, 4*quat_dim)"""
        shape = arm_w.shape[:-1]
        D = arm_w.shape[-1]
        flat = arm_w.dim() > 2
        if flat:
            arm_w = arm_w.reshape(-1, D); arm_i = arm_i.reshape(-1, D)
            arm_j = arm_j.reshape(-1, D); arm_k = arm_k.reshape(-1, D)

        # q: (N, 4, quat_dim) — stack 4 projected arms as quaternion components
        q = torch.stack([self.proj_w(arm_w), self.proj_i(arm_i),
                         self.proj_j(arm_j), self.proj_k(arm_k)], dim=1)
        q = q / (q.norm(dim=1, keepdim=True) + 1e-8)

        # r: (N, 4, quat_dim) — broadcast learned rotation
        r = self.rotation.expand(q.shape[0], -1, -1)
        r = r / (r.norm(dim=1, keepdim=True) + 1e-8)

        # Single batched Hamilton product over all quat_dim simultaneously
        # (N, 4, quat_dim) × (N, 4, quat_dim) → (N, 4, quat_dim)
        composed = quaternion_multiply_batched(r, q)

        # Flatten 4 × quat_dim → 4*quat_dim
        composed = composed.reshape(q.shape[0], -1)

        if flat:
            composed = composed.reshape(*shape, -1)
        return composed


# ═══════════════════════════════════════════════════════════════════════════════
# NEW COMPONENTS — transformer-specific, built for this architecture
# ═══════════════════════════════════════════════════════════════════════════════

class ManifoldProjection(TorchComponent):
    """Input stage: project transformer hidden states to S^(d-1).

    Per-position, per-layer projection from model space to the
    constellation's embedding space. L2-normalized to sit on the
    unit hypersphere.

    This is the tap — it reads the representation without modifying it.
    """
    def __init__(self, name, d_model, manifold_dim):
        super().__init__(name)
        self.proj = nn.Linear(d_model, manifold_dim)
        self.norm = nn.LayerNorm(manifold_dim)

    def forward(self, hidden_states):
        """(B, L, D) → (B, L, manifold_dim) on S^(manifold_dim - 1)"""
        h = self.norm(self.proj(hidden_states))
        return F.normalize(h, dim=-1)


class PositionGeometricContext(TorchComponent):
    """Curation stage: constellation observation → FiLM context vector.

    Takes the full observation dict from ConstellationObserver and fuses
    it into a per-position conditioning vector for FiLM layers.

    Processes: cos_to_anchors, assignment, patchwork, embedding.
    These are the same features the GeoQuat head used — validated on
    ProteinGym across 84 unseen proteins.
    """
    def __init__(self, name, n_anchors, pw_dim, manifold_dim, context_dim):
        super().__init__(name)
        # Anchor features: cos + assignment + triangulation = 3 * n_anchors
        self.anchor_mlp = nn.Sequential(
            nn.Linear(n_anchors * 3, context_dim),
            nn.GELU(),
            nn.LayerNorm(context_dim),
        )
        # Structural features: patchwork + embedding
        self.struct_mlp = nn.Sequential(
            nn.Linear(pw_dim + manifold_dim, context_dim),
            nn.GELU(),
            nn.LayerNorm(context_dim),
        )
        # Fuse anchor + structural
        self.fuse = nn.Sequential(
            nn.Linear(context_dim * 2, context_dim),
            nn.GELU(),
            nn.LayerNorm(context_dim),
        )

    def forward(self, obs_dict):
        """
        Args:
            obs_dict: from ConstellationObserver.observe(), keys:
                cos_to_anchors: (B*L, A)
                assignment: (B*L, A)
                triangulation: (B*L, A)
                patchwork: (B*L, pw_dim)
                embedding: (B*L, manifold_dim)
        Returns:
            (B*L, context_dim) geometric context
        """
        anchor_feats = torch.cat([
            obs_dict['cos_to_anchors'],
            obs_dict['assignment'],
            obs_dict['triangulation'],
        ], dim=-1)

        struct_feats = torch.cat([
            obs_dict['patchwork'],
            obs_dict['embedding'],
        ], dim=-1)

        a = self.anchor_mlp(anchor_feats)
        s = self.struct_mlp(struct_feats)
        return self.fuse(torch.cat([a, s], dim=-1))


class GeometricAttention(TorchComponent):
    """Attention with FiLM from curated constellation. Stream B.

    FiLM modulates Q and K BEFORE attention — the constellation
    position controls WHERE attention flows. V stays unmodulated.
    FiLM between FFN layers conditions the nonlinearity.

    Proven principle: context before composition, not after.
    """
    def __init__(self, name, d_model, n_heads=8, context_dim=128, dropout=0.1):
        super().__init__(name)
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.scale = self.head_dim ** -0.5

        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

        # FiLM on Q and K — geometry routes attention
        self.film_q = FiLMLayer(f'{name}_film_q', d_model, context_dim)
        self.film_k = FiLMLayer(f'{name}_film_k', d_model, context_dim)

        self.norm = nn.LayerNorm(d_model)

        # FFN with FiLM between layers
        self.ffn1 = nn.Linear(d_model, d_model * 4)
        self.film_ffn = FiLMLayer(f'{name}_film_ffn', d_model * 4, context_dim)
        self.ffn2 = nn.Linear(d_model * 4, d_model)
        self.ffn_drop = nn.Dropout(dropout)
        self.ffn_norm = nn.LayerNorm(d_model)

    def forward(self, x, geo_ctx, attn_mask=None, key_padding_mask=None):
        """
        x: (B, L, D), geo_ctx: (B, L, C) → (B, L, D)
        """
        B, L, D = x.shape
        H, HD = self.n_heads, self.head_dim

        Q = self.film_q(self.w_q(x), geo_ctx)
        K = self.film_k(self.w_k(x), geo_ctx)
        V = self.w_v(x)  # V unmodulated — content stays pure

        Q = Q.view(B, L, H, HD).transpose(1, 2)
        K = K.view(B, L, H, HD).transpose(1, 2)
        V = V.view(B, L, H, HD).transpose(1, 2)

        scores = (Q @ K.transpose(-2, -1)) * self.scale
        if attn_mask is not None:
            scores = scores + attn_mask
        if key_padding_mask is not None:
            scores = scores.masked_fill(
                key_padding_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
        attn_out = (self.dropout(F.softmax(scores, dim=-1)) @ V)
        attn_out = attn_out.transpose(1, 2).reshape(B, L, D)

        x = self.norm(x + self.w_o(attn_out))

        # FFN with geometric FiLM between layers
        h = F.gelu(self.ffn1(x))
        h = self.film_ffn(h, geo_ctx)
        x = self.ffn_norm(x + self.ffn_drop(self.ffn2(h)))

        return x


class ContentAttention(TorchComponent):
    """Standard self-attention. Stream A. No geometric conditioning."""
    def __init__(self, name, d_model, n_heads=8, dropout=0.1):
        super().__init__(name)
        self.attn = nn.MultiheadAttention(
            d_model, n_heads, dropout=dropout, batch_first=True)
        self.norm = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4), nn.GELU(),
            nn.Linear(d_model * 4, d_model), nn.Dropout(dropout))
        self.ffn_norm = nn.LayerNorm(d_model)

    def forward(self, x, attn_mask=None, key_padding_mask=None):
        a, _ = self.attn(x, x, x, attn_mask=attn_mask,
                         key_padding_mask=key_padding_mask)
        x = self.norm(x + a)
        x = self.ffn_norm(x + self.ffn(x))
        return x


# ═══════════════════════════════════════════════════════════════════════════════
# LAYER — dual-stream with constellation routing
# ═══════════════════════════════════════════════════════════════════════════════

class GeometricTransformerLayer(BaseTower):
    """One layer of the geometric transformer.

    Pipeline per layer:
        1. ManifoldProjection: h_i → emb_i on S^(manifold_dim - 1)
        2. ConstellationObserver: emb_i → {triangulation, assignment, patchwork, ...}
        3. PositionGeometricContext: observation → FiLM context (B, L, context_dim)
        4. ContentAttention (Stream A): standard MHA
        5. GeometricAttention (Stream B): FiLM(Q,K | geo_ctx), V pure
        6. CayleyOrthogonal: align B basis → A basis
        7. QuaternionCompose: w=A, i=aligned_B, j=A-B, k=A*B
        8. Decode + gated residual

    Access:
        layer['projection']  → ManifoldProjection
        layer['observer']    → ConstellationObserver
        layer['context']     → PositionGeometricContext
        layer['content']     → ContentAttention
        layer['geometric']   → GeometricAttention
        layer['rotation']    → CayleyOrthogonal
        layer['compose']     → QuaternionCompose
    """
    def __init__(self, name, d_model, n_heads=8, n_anchors=32,
                 manifold_dim=256, n_comp=8, d_comp=32,
                 context_dim=128, quat_dim=64, dropout=0.1):
        super().__init__(name)
        self.d_model = d_model

        # 1. Project to manifold
        self.attach('projection', ManifoldProjection(
            f'{name}_proj', d_model, manifold_dim))

        # 2. Constellation observer (real association + curation)
        self.attach('observer', ConstellationObserver(
            dim=manifold_dim, n_anchors=n_anchors,
            n_comp=n_comp, d_comp=d_comp))

        # 3. Fuse observation into FiLM context
        pw_dim = self['observer'].curation.patchwork.output_dim
        self.attach('context', PositionGeometricContext(
            f'{name}_ctx', n_anchors, pw_dim, manifold_dim, context_dim))

        # 4. Stream A: content
        self.attach('content', ContentAttention(
            f'{name}_content', d_model, n_heads, dropout))

        # 5. Stream B: geometric
        self.attach('geometric', GeometricAttention(
            f'{name}_geo', d_model, n_heads, context_dim, dropout))

        # 6. Cayley rotation: align B → A
        self.attach('rotation', CayleyOrthogonal(f'{name}_cayley', d_model))

        # 7. Quaternion composition
        self.attach('compose', QuaternionCompose(
            f'{name}_quat', d_model, quat_dim))

        # 8. Decode + gate
        self.attach('decode', nn.Sequential(
            nn.Linear(quat_dim * 4, d_model), nn.GELU(), nn.LayerNorm(d_model)))
        self.attach('gate', nn.Sequential(
            nn.Linear(d_model * 2, d_model), nn.Sigmoid()))

    def forward(self, x, attn_mask=None, key_padding_mask=None):
        """
        Args:
            x: (B, L, D) input hidden states

        Returns:
            x_out: (B, L, D) transformed hidden states
            geo_state: dict with full geometric residual:
                'embedding':      (B, L, manifold_dim)  position on S^(d-1)
                'geo_ctx':        (B, L, context_dim)   compressed FiLM context
                'triangulation':  (B, L, A)             cosine distances to anchors
                'cos_to_anchors': (B, L, A)             raw cosine similarities
                'assignment':     (B, L, A)             soft assignment
                'nearest':        (B, L)                nearest anchor index
                'patchwork':      (B, L, pw_dim)        compartment features
                'bridge':         (B, L, A)             patchwork's assignment estimate
                'content':        (B, L, D)             Stream A output
                'geometric':      (B, L, D)             Stream B output (pre-rotation)
                'composed':       (B, L, 4*quat_dim)    raw quaternion composition
        """
        B, L, D = x.shape

        # 1. Project to manifold: per-position embedding on S^(d-1)
        emb = self['projection'](x)  # (B, L, manifold_dim)

        # 2. Constellation observation: flatten to (B*L, manifold_dim) for observer
        emb_flat = emb.reshape(B * L, -1)
        obs = self['observer'].observe(emb_flat)

        # 3. Build FiLM context
        geo_ctx_flat = self['context'](obs)  # (B*L, context_dim)
        geo_ctx = geo_ctx_flat.reshape(B, L, -1)  # (B, L, context_dim)

        # 4. Stream A: content attention
        a_out = self['content'](x, attn_mask=attn_mask,
                                key_padding_mask=key_padding_mask)

        # 5. Stream B: geometric attention
        b_out = self['geometric'](x, geo_ctx, attn_mask=attn_mask,
                                   key_padding_mask=key_padding_mask)

        # 6. Cayley rotation: align B → A
        b_aligned = self['rotation'](b_out)

        # 7. Quaternion composition
        #    w = content (what does standard attention think?)
        #    i = aligned geometry (what does geometric attention think?)
        #    j = disagreement (where do they diverge? — the surprise signal)
        #    k = agreement (where do they converge? — the confidence signal)
        composed = self['compose'](
            arm_w=a_out, arm_i=b_aligned,
            arm_j=a_out - b_aligned, arm_k=a_out * b_aligned)

        # 8. Decode + gated residual
        decoded = self['decode'](composed)
        g = self['gate'](torch.cat([x, decoded], dim=-1))
        x_out = g * decoded + (1 - g) * x

        # 9. Build full geometric state — reshape everything back to (B, L, ...)
        def unflatten(t):
            if t is None: return None
            if t.dim() == 1: return t.reshape(B, L)        # (B*L,) → (B, L)
            return t.reshape(B, L, *t.shape[1:])            # (B*L, ...) → (B, L, ...)

        geo_state = {
            'embedding':      emb,                          # already (B, L, manifold_dim)
            'geo_ctx':        geo_ctx,                      # already (B, L, context_dim)
            'triangulation':  unflatten(obs['triangulation']),
            'cos_to_anchors': unflatten(obs['cos_to_anchors']),
            'assignment':     unflatten(obs['assignment']),
            'nearest':        unflatten(obs['nearest']),
            'patchwork':      unflatten(obs['patchwork']),
            'bridge':         unflatten(obs['bridge']),
            'content':        a_out,                        # (B, L, D)
            'geometric':      b_out,                        # (B, L, D) pre-rotation
            'composed':       composed,                     # (B, L, 4*quat_dim)
        }

        return x_out, geo_state


# ═══════════════════════════════════════════════════════════════════════════════
# FULL MODEL — stack of layers
# ═══════════════════════════════════════════════════════════════════════════════

class GeometricTransformer(BaseTower):
    """Geometric Transformer — dual-stream with constellation routing.

    Stack of GeometricTransformerLayers. Optional cross-layer Cayley
    rotation aligns each layer's output basis to the next layer's
    expected input.

    Access:
        model['layer_0']      → first layer
        model['cross_rot_0']  → cross-layer rotation 0→1
        model['final_norm']   → output normalization

    Args:
        name: tower identity
        d_model: transformer model dimension
        n_heads: attention heads per stream
        n_layers: number of geometric transformer layers
        n_anchors: constellation anchor points
        manifold_dim: dimension of S^(d-1) for constellation
        n_comp: patchwork compartments
        d_comp: hidden dim per compartment
        context_dim: FiLM conditioning dimension
        quat_dim: quaternion space dimension
        dropout: dropout rate
        cross_layer_rotation: add Cayley rotation between layers
        vocab_size: if set, adds embedding + output head
    """
    def __init__(self, name, d_model=512, n_heads=8, n_layers=4,
                 n_anchors=32, manifold_dim=256, n_comp=8, d_comp=32,
                 context_dim=128, quat_dim=64, dropout=0.1,
                 cross_layer_rotation=True, vocab_size=None, max_seq_len=2048):
        super().__init__(name)
        self.d_model = d_model
        self.n_layers = n_layers

        if vocab_size is not None:
            self.attach('embed', nn.Embedding(vocab_size, d_model))
            self.attach('pos_embed', nn.Embedding(max_seq_len, d_model))
            self.attach('head', nn.Linear(d_model, vocab_size, bias=False))

        for i in range(n_layers):
            self.attach(f'layer_{i}', GeometricTransformerLayer(
                f'{name}_L{i}', d_model, n_heads, n_anchors,
                manifold_dim, n_comp, d_comp, context_dim, quat_dim, dropout))

        if cross_layer_rotation and n_layers > 1:
            for i in range(n_layers - 1):
                self.attach(f'cross_rot_{i}', CayleyOrthogonal(
                    f'{name}_xrot_{i}', d_model))

        self.attach('final_norm', nn.LayerNorm(d_model))

        self._config = dict(
            d_model=d_model, n_heads=n_heads, n_layers=n_layers,
            n_anchors=n_anchors, manifold_dim=manifold_dim,
            n_comp=n_comp, d_comp=d_comp, context_dim=context_dim,
            quat_dim=quat_dim, dropout=dropout,
            cross_layer_rotation=cross_layer_rotation,
            vocab_size=vocab_size,
        )

    @property
    def config(self):
        return self._config.copy()

    def param_report(self):
        total = 0
        name = getattr(self, '_tower_name', getattr(self, 'name', self.__class__.__name__))
        print(f"\n  {name} — parameter report")
        print(f"  {'Component':<35s}  {'Params':>12s}")
        print(f"  {'─'*35}  {'─'*12}")
        for cname, module in self.named_children():
            n = sum(p.numel() for p in module.parameters())
            total += n
            print(f"  {cname:<35s}  {n:>12,}")
        print(f"  {'─'*35}  {'─'*12}")
        print(f"  {'TOTAL':<35s}  {total:>12,}")
        return total

    def forward(self, x, attn_mask=None, key_padding_mask=None,
                return_geo_state=False):
        """
        Args:
            x: (B, L, D) hidden states or (B, L) token ids
            return_geo_state: if True, return per-layer geometric state dicts

        Returns:
            out: (B, L, D) transformed hidden states (or logits if head attached)
            geo_states: list of per-layer geo_state dicts (if return_geo_state)
                Each dict contains: embedding, geo_ctx, triangulation,
                cos_to_anchors, assignment, nearest, patchwork, bridge,
                content, geometric, composed
        """
        if self.has('embed') and x.dtype in (torch.long, torch.int32, torch.int64):
            pos = torch.arange(x.shape[1], device=x.device)
            x = self['embed'](x) + self['pos_embed'](pos)

        geo_states = []
        has_xrot = self.has('cross_rot_0')

        for i in range(self.n_layers):
            x, geo_state = self[f'layer_{i}'](
                x, attn_mask=attn_mask, key_padding_mask=key_padding_mask)
            if return_geo_state:
                geo_states.append(geo_state)
            if has_xrot and i < self.n_layers - 1:
                x = self[f'cross_rot_{i}'](x)

        x = self['final_norm'](x)
        if self.has('head'):
            x = self['head'](x)

        return (x, geo_states) if return_geo_state else x


# ═══════════════════════════════════════════════════════════════════════════════
# FACTORIES
# ═══════════════════════════════════════════════════════════════════════════════

def geo_transformer_esm2(name='geo_esm2', n_layers=6, **kw):
    """Pre-configured for ESM-2 650M (d=1280)."""
    return GeometricTransformer(name, d_model=1280, n_heads=16,
        n_layers=n_layers, n_anchors=32, manifold_dim=256,
        n_comp=8, d_comp=32, context_dim=128, quat_dim=64, **kw)

def geo_transformer_small(name='geo_small', n_layers=4, **kw):
    """Small config for prototyping."""
    return GeometricTransformer(name, d_model=256, n_heads=8,
        n_layers=n_layers, n_anchors=16, manifold_dim=128,
        n_comp=4, d_comp=16, context_dim=64, quat_dim=32, **kw)

def geo_transformer_vision(name='geo_vit', n_layers=4, **kw):
    """For scatter/SVD vision pipeline (patches as tokens)."""
    return GeometricTransformer(name, d_model=384, n_heads=8,
        n_layers=n_layers, n_anchors=32, manifold_dim=128,
        n_comp=8, d_comp=16, context_dim=64, quat_dim=32, **kw)


# ═══════════════════════════════════════════════════════════════════════════════
# SELF-TEST
# ═══════════════════════════════════════════════════════════════════════════════

if __name__ == '__main__':
    print("Geometric Transformer — Self-Test")
    print(f"  geolip_core available: {_HAS_GEOLIP}")
    print("=" * 60)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    model = geo_transformer_small('test', n_layers=2)
    if hasattr(model, 'network_to'):
        model.network_to(device=device, strict=False)
    else:
        model = model.to(device)
    total = model.param_report()

    B, L, D = 2, 32, 256
    x = torch.randn(B, L, D, device=device)

    out, geos = model(x, return_geo_state=True)
    assert out.shape == (B, L, D), f"Expected ({B},{L},{D}), got {out.shape}"
    assert len(geos) == 2

    print(f"\n  Input:  ({B}, {L}, {D})")
    print(f"  Output: {out.shape}")
    print(f"  Geo states: {len(geos)} layers")
    print(f"  State keys: {sorted(geos[0].keys())}")
    for k, v in geos[0].items():
        if v is not None:
            shape = v.shape if hasattr(v, 'shape') else type(v).__name__
            print(f"    {k:<18s}: {shape}")

    # Verify rotations
    for name, module in model.named_modules():
        if isinstance(module, CayleyOrthogonal):
            R = module.get_rotation()
            I = torch.eye(R.shape[0], device=R.device)
            print(f"  {name}: ‖RRᵀ-I‖={((R@R.T)-I).norm():.8f}  det={torch.det(R):.4f}")

    # ESM-2 scale overhead
    print(f"\n  ESM-2 scale:")
    esm = geo_transformer_esm2('esm2', n_layers=6)
    if hasattr(esm, 'network_to'):
        esm.network_to(device=device, strict=False)
    else:
        esm = esm.to(device)
    n = esm.param_report()
    print(f"  Overhead on 650M base: {n/1e6:.1f}M ({n/650e6*100:.1f}%)")

    print(f"\n{'='*60}")
    print(f"  PASSED")
    print(f"{'='*60}")