patch.py

"""
gemma3-px  —  Surgical Patch (Phase 10.0)
=========================================
Implements Recursive State Memory (RSM) and Hyper-Fluid Routing (HFR).
RSM allows the model to 'see' its own previous thinking states during recursion.
"""

import types
import torch
import torch.nn as nn
import os
import json
import datetime
from typing import Optional

from .px_modules import (
    LTIInjection, ADCInjection, StabilityMonitor, CognitiveEvent,
    MephistophelesOperator, OrthogonalJitter
)

# ---------------------------------------------------------------------------
# p10.0: Recursive State Memory (RSM)
# ---------------------------------------------------------------------------

class RecursiveMemoryCache:
    """
    Extends ReadOnlyCache by injecting previous thinking steps into the 
    self-attention key/value streams.
    """
    def __init__(self, real_cache, thought_history=None, layer_types=None, read_only=False, expected_len=0):
        self.__dict__["_real"] = real_cache
        self.__dict__["_thoughts"] = thought_history or []
        self.__dict__["_layer_types"] = layer_types or []
        self.__dict__["_read_only"] = read_only
        self.__dict__["_expected_len"] = expected_len
        
    def __getattr__(self, name):
        return getattr(self._real, name)
        
    def update(self, key_states, value_states, layer_idx, cache_kwargs=None):
        # 1. Base Update (Functional if read_only)
        if self._read_only:
            past_k, past_v = None, None
            # Try older DynamicCache style
            if hasattr(self._real, "key_cache") and len(self._real.key_cache) > layer_idx:
                past_k = self._real.key_cache[layer_idx]
                past_v = self._real.value_cache[layer_idx]
            # Try newer Cache object style (transformers 4.45+)
            elif hasattr(self._real, "layers") and len(self._real.layers) > layer_idx:
                layer = self._real.layers[layer_idx]
                # DynamicLayer / StaticLayer
                if hasattr(layer, "keys") and layer.keys is not None:
                    past_k = layer.keys
                    past_v = layer.values
                # SinkCache / etc might have different names? No, usually .keys
            
            if past_k is None:
                past_k = torch.empty(0, device=key_states.device, dtype=key_states.dtype)
                past_v = torch.empty(0, device=value_states.device, dtype=value_states.dtype)
            
            # If past_k already has the expected length, it means this layer was 
            # already updated for the current token(s) in a previous iteration 
            # of the same reasoning loop.
            if past_k.numel() > 0 and past_k.shape[-2] == self._expected_len:
                res_k, res_v = past_k, past_v
            else:
                res_k = torch.cat([past_k, key_states], dim=-2)
                res_v = torch.cat([past_v, value_states], dim=-2)
                
            # print(f"  [DEBUG-CACHE] L{layer_idx} RO=True | past={past_k.shape[-2] if past_k.numel()>0 else 0} | cur={key_states.shape[-2]} | res={res_k.shape[-2]} | exp={self._expected_len}")
        else:
            res_k, res_v = self._real.update(key_states, value_states, layer_idx, cache_kwargs)
            # print(f"  [DEBUG-CACHE] L{layer_idx} RO=False | res={res_k.shape[-2]} | exp={self._expected_len}")

        # 2. Phase 14.6: Soft-RSM (Semantic Blending)
        is_full = self._layer_types and self._layer_types[layer_idx] == "full_attention"
        
        if self._thoughts and layer_idx >= 6 and is_full: 
            B, H_kv, T_res, HD = res_k.shape
            T_curr = key_states.shape[-2]
            alpha = 0.15 
            
            # Phase 14.7: Triangular Weighting (Emphasize the 'reasoning peak')
            n_t = len(self._thoughts[-6:])
            if n_t > 2:
                weights = torch.cat([
                    torch.linspace(0.4, 1.0, n_t//2, device=res_k.device),
                    torch.linspace(1.0, 0.6, n_t - n_t//2, device=res_k.device)
                ])
                t_raw = (torch.stack(self._thoughts[-6:]) * weights.view(-1, 1, 1, 1)).sum(dim=0) / weights.sum()
            else:
                t_raw = torch.stack(self._thoughts).mean(dim=0)
            
            D = t_raw.shape[2]
            
            # Project thought to Head Dim (SDA)
            t_flat = t_raw.mean(dim=1, keepdim=True) # (B, 1, D)
            t_proj = torch.nn.functional.interpolate(t_flat, size=HD, mode='linear', align_corners=False)
            t_k = t_proj.unsqueeze(1) # (B, 1, 1, HD)
            t_v = -t_k
            
            # Blend into the LAST token(s) of the result
            # Use in-place only if not read_only to avoid side effects on cache
            if self._read_only:
                res_k = res_k.clone()
                res_v = res_v.clone()
            
            res_k[:, :, -T_curr:, :] = (1.0 - alpha) * res_k[:, :, -T_curr:, :] + alpha * t_k
            res_v[:, :, -T_curr:, :] = (1.0 - alpha) * res_v[:, :, -T_curr:, :] + alpha * t_v
            
        return res_k, res_v

# ---------------------------------------------------------------------------

def remove_px_patch(model) -> None:
    from transformers.models.gemma3.modeling_gemma3 import Gemma3TextModel
    text_model = (model.model if hasattr(model, "model") else model)
    if hasattr(text_model, "_px_config"):
        text_model.forward = types.MethodType(
            Gemma3TextModel.forward, text_model
        )
        del text_model._px_injection
        del text_model._px_config
        print("[gemma3-px] Patch removed.")

def _resolve_text_model(model):
    if hasattr(model, "model") and hasattr(model.model, "layers"):
        return model.model
    return model

# ---------------------------------------------------------------------------

def _px_forward(
    self,
    input_ids=None,
    attention_mask=None,
    position_ids=None,
    past_key_values=None,
    inputs_embeds=None,
    use_cache=None,
    **kwargs,
):
    from transformers.cache_utils import DynamicCache
    from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask

    if (input_ids is None) ^ (inputs_embeds is not None):
        raise ValueError("Specify exactly one of input_ids or inputs_embeds.")

    if inputs_embeds is None:
        # Multimodal resolution (Phase 17.7)
        if hasattr(self, "embed_tokens"):
            inputs_embeds = self.embed_tokens(input_ids)
        elif hasattr(self, "language_model"):
            inputs_embeds = self.language_model.model.embed_tokens(input_ids)
        elif hasattr(self, "model") and hasattr(self.model, "embed_tokens"):
            inputs_embeds = self.model.embed_tokens(input_ids)
        else:
            # Last resort: search for embed_tokens in children
            embedder = None
            for name, module in self.named_modules():
                if "embed_tokens" in name:
                    embedder = module
                    break
            if embedder:
                inputs_embeds = embedder(input_ids)
            else:
                raise AttributeError(f"Could not find embed_tokens in model type {type(self)}. Available: {dir(self)[:20]}...")

    if use_cache and past_key_values is None:
        past_key_values = DynamicCache(config=self.config)

    # Phase 14.8: Initial sequence length tracking
    past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
    expected_len = past_seen + inputs_embeds.shape[1]
    
    # print(f"[DEBUG-PX] Type={type(past_key_values)} seen={past_seen} cur={inputs_embeds.shape[1]} exp={expected_len}")

    if position_ids is None:
        position_ids = (
            torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device)
            + past_seen
        ).unsqueeze(0)

    # Resolve config for masking (Phase 17.7 multimodal fix)
    mask_config = self.config
    if hasattr(mask_config, "text_config"):
        mask_config = mask_config.text_config

    if not isinstance(attention_mask, dict):
        mk = dict(
            config=mask_config,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            position_ids=position_ids,
        )
        causal_mask_mapping = {
            "full_attention":    create_causal_mask(**mk),
            "sliding_attention": create_sliding_window_causal_mask(**mk),
        }
    else:
        causal_mask_mapping = attention_mask

    hidden_states = inputs_embeds
    position_embeddings = {}
    for layer_type in set(mask_config.layer_types):
        position_embeddings[layer_type] = self.rotary_emb(
            hidden_states, position_ids, layer_type
        )

    cfg = self._px_config
    updated_layers = set() # Phase 14.9: Global visit tracker for this forward pass

    # ── 1. PRELUDE ──────────────────────────────────────────────────────────
    for i in range(cfg["prelude_end"]):
        updated_layers.add(i)
        layer_out = self.layers[i](
            hidden_states,
            attention_mask=causal_mask_mapping[mask_config.layer_types[i]],
            position_embeddings=position_embeddings[mask_config.layer_types[i]],
            position_ids=position_ids,
            past_key_values=past_key_values,
            **kwargs,
        )
        hidden_states = layer_out[0] if isinstance(layer_out, (tuple, list)) else layer_out

    # ── 1.5 META-SELECTOR (Phase 28 Fluid) ───────────────────────────────────
    dynamic_start = cfg["recur_start"]
    dynamic_end = cfg["recur_end"]
    dynamic_hub = cfg.get("bimodal_hub", cfg["recur_start"])
    num_layers = len(self.layers)

    if cfg.get("routing_mode") == "adaptive":
        if inputs_embeds.shape[1] > 1:
            # Prefill phase: Measure Kurtosis and Jitter at Layer 5
            h_base_f32 = hidden_states.to(torch.float32)
            
            # Kurtosis (Last token)
            h_probe = h_base_f32[0, -1, :]
            variance = torch.var(h_probe).item()
            kurtosis = (torch.mean((h_probe - torch.mean(h_probe))**4) / (variance**2)).item() if variance > 0 else 0
            self._task_kurtosis = kurtosis
            
            # Jitter (Across sequence)
            # We measure the variance of the norms of the hidden states across the prompt.
            # Very high jitter usually indicates a 'Trap' or 'Divergence' in the intuition pass.
            h_norms = h_base_f32.norm(dim=-1) # [B, T]
            h_norm_var = torch.var(h_norms, dim=-1).mean().item()
            self._task_jitter = h_norm_var
            
            if os.environ.get("DEBUG_ROUTING") == "1": 
                print(f"[Router] Prefill K={kurtosis:.1f}, Jitter={h_norm_var:.4f}")
        
        kurtosis = getattr(self, "_task_kurtosis", 300) # Default to Logic if missing
        
        import math
        if num_layers < 20: # 270M Model (Kurtosis is task-separable)
            # Continuous Fluid Gaussian Blending of the 5 Zones
            w_m = math.exp(-((kurtosis - 200)**2) / (2 * 25**2))  # Math
            w_la = math.exp(-((kurtosis - 275)**2) / (2 * 15**2)) # Logic-A
            w_cr = math.exp(-((kurtosis - 298)**2) / (2 * 8**2))  # Creative
            w_lb = math.exp(-((kurtosis - 310)**2) / (2 * 8**2))  # Logic-B
            w_sy = math.exp(-((kurtosis - 325)**2) / (2 * 20**2)) # Synthesis
            
            W = w_m + w_la + w_cr + w_lb + w_sy + 1e-9
            
            # Phase 36.2: Restored Stable Math Zone
            d_start = (w_m*5 + w_la*8 + w_cr*10 + w_lb*8 + w_sy*6) / W
            d_end = (w_m*11 + w_la*12 + w_cr*16 + w_lb*14 + w_sy*14) / W
            # Phase 41 Master Hub: 10
            d_hub = (w_m*10 + w_la*10 + w_cr*10 + w_lb*10 + w_sy*10) / W
            d_loops = (w_m*8 + w_la*8 + w_cr*6 + w_lb*10 + w_sy*8) / W

            
            dynamic_start = max(1, int(round(d_start)))
            dynamic_end = min(num_layers - 1, int(round(d_end)))
            dynamic_hub = int(round(d_hub))
            cfg["n_loops"] = max(2, int(round(d_loops)))
            zone_name = f"Fluid-Blended (K={kurtosis:.1f})"
        else:
            # 1B and 4B Models (Scale-Invariant Omni Zone)
            # They have enough capacity to hold both semantics without smearing
            dynamic_start = int(num_layers * 0.38)
            dynamic_end = int(num_layers * 0.76)
            dynamic_hub = int(num_layers * 0.61)
            cfg["n_loops"] = 6
            zone_name = "Omni-Scale"

        # Only print routing decision once per token during generation
        if inputs_embeds.shape[1] == 1 and os.environ.get("DEBUG_ROUTING") == "1":
             print(f"[Router] {zone_name} -> L{dynamic_start}-L{dynamic_end} (Loops: {cfg['n_loops']}, Hub: {dynamic_hub})")

        # Fast-forward prelude if needed
        for i in range(cfg["prelude_end"], dynamic_start):
            updated_layers.add(i)
            layer_out = self.layers[i](
                hidden_states,
                attention_mask=causal_mask_mapping[mask_config.layer_types[i]],
                position_embeddings=position_embeddings[mask_config.layer_types[i]],
                position_ids=position_ids,
                past_key_values=past_key_values,
                **kwargs,
            )
            hidden_states = layer_out[0] if isinstance(layer_out, (tuple, list)) else layer_out

    # ── 2. REASONING ZONE (Phase 10.0) ──────────────────────────────────────
    e_static = hidden_states.clone()

    # 2.A: Intuition Pass
    trans_out = hidden_states
    for i_layer in range(dynamic_start, dynamic_end):
        l_type = mask_config.layer_types[i_layer]
        updated_layers.add(i_layer)
        layer_out = self.layers[i_layer](
            trans_out,
            attention_mask=causal_mask_mapping[l_type],
            position_embeddings=position_embeddings[l_type],
            position_ids=position_ids,
            past_key_values=past_key_values,
            **kwargs,
        )
        trans_out = layer_out[0] if isinstance(layer_out, (tuple, list)) else layer_out
        # if past_key_values is not None:
        #    print(f"  [DEBUG-PX-DIR] {dir(past_key_values)}")
    
    h_baseline = trans_out
    
    # Phase 14.5: ETR (Entropy Triggered Recursion)
    # Estimate 'confidence' from the last layer's norm change or simpler:
    # We only run recursion if the intuition pass wasn't 'perfectly' stable.
    # Note: h_baseline is already computed.
    
    # 2.B: Hyper-Fluid Routing & Recursive Memory
    n_loops = cfg.get("n_loops", 2)
    
    # Phase 14.5: ETR (Entropy Triggered Recursion)
    phi_intuition = StabilityMonitor.calculate_phi(h_baseline, hidden_states).mean().item()
    if os.environ.get("DEBUG_ROUTING") == "1":
        print(f"  [Intuition] Phi: {phi_intuition:.6f}")
    
    # Phase 14.7: Gamma-Damping instead of loop scaling
    current_gamma = cfg.get("gamma", 0.08)
    
    e_reflector = e_static
    is_trap_candidate = False
    
    # Phase 36.3: Surgical Reflector Activation
    # Jitter is only for extreme representational collapse (1e8)
    jitter = getattr(self, "_task_jitter", 0.0)
    kurtosis = getattr(self, "_task_kurtosis", 300)
    
    # Trigger Reflector if extreme jitter OR if it's a known Math/Logic zone
    if jitter > 1e8 or (200.0 < kurtosis < 315.0):
        is_trap_candidate = True
        if os.environ.get("DEBUG_ROUTING") == "1":
            reason = "Jitter" if jitter > 1e8 else "Rigor-Zone"
            print(f"  [Router] Trap detected via {reason} ({jitter:.1f}), activating Reflector")
        
        # Phase 16.3: Anchor Reflection
        e_stat_f32 = e_static.to(torch.float32)
        h_base_f32 = h_baseline.to(torch.float32)
        e_ref_f32 = 2.0 * e_stat_f32 - h_base_f32
        e_ref_f32 = e_ref_f32 * (e_stat_f32.norm() / (e_ref_f32.norm() + 1e-6))
        e_reflector = e_ref_f32.to(e_static.dtype)

    if phi_intuition > 0.9999 and not is_trap_candidate:
        # Reduced damping: allow the model to think even if the first pass was stable
        current_gamma *= 0.5 
    elif phi_intuition > 0.999:
        current_gamma *= 0.8
    # Phase 25: Sigmoid-Annealed Orthogonal Recovery (SAOR)
    # -----------------------------------------------------------------------
    # Using a Sigmoid curve for Gamma to allow a sharp "Phase Transition" 
    # from exploration (high energy) to grounding (low energy).
    # Plus: Orthogonal Reinforcement to protect logical drift.
    
    base_gamma = current_gamma
    bimodal_hub_start = cfg.get("bimodal_hub", 11)
    
    path_taken = []
    thought_history = []
    avg_phi_explore = 1.0
    exploration_steps = 0
    telemetry_steps = []
    
    # Context dims
    B, T_curr = hidden_states.shape[0], hidden_states.shape[1]
    HD = getattr(self.config, "head_dim", 256)
    
    # Phase 38.1: Anna Karenina Sensor (AKS) Initialization
    # Tracks the "Geometric Disparity" of the latent thoughts.
    # Clustering = Truth (Anna Karenina Principle), Dispersion = Error.
    divergence_buffer = []
    correction_strength = 0.0
    
    if n_loops > 1:
        h_exp = e_reflector.clone() # Use Reflected Anchor
        current_layer = dynamic_start
        max_steps = (dynamic_end - dynamic_start) * n_loops * 3
        phis = []

        stability_counter = 0
        layer_visits = {i: 0 for i in range(5, 18)}

        # Initialize active bounds
        active_start = dynamic_start
        active_end = dynamic_end

        while current_layer < active_end and exploration_steps < max_steps:
            # --- PHASE 26: INFINITE REFLECTION (IR) ---
            t_norm = exploration_steps / max_steps
            
            # Phase 38.2: AKS - Topological Anomaly Detection
            dist_now = 1.0 - StabilityMonitor.calculate_phi(h_exp, e_static).mean().item()
            if exploration_steps > 2:
                divergence_buffer.append(dist_now)
                if len(divergence_buffer) > 4: divergence_buffer.pop(0)
                
                if len(divergence_buffer) >= 3:
                    velocity = divergence_buffer[-1] - divergence_buffer[-2]
                    acceleration = (divergence_buffer[-1] - divergence_buffer[-2]) - (divergence_buffer[-2] - divergence_buffer[-3])
                    
                    if acceleration > 0.001 and velocity > 0:
                        correction_strength = min(1.0, correction_strength + 0.1)
                    else:
                        correction_strength = max(0.0, correction_strength - 0.05)

            # Phase 43.1: Emancipation Metric
            # Measure how far the current state has moved from the initial prompt anchor.
            emancipation_phi = StabilityMonitor.calculate_phi(h_exp, e_static).mean().item()

            # Phase 43.2: Perturbation Engine (The Forking Path)
            # Inject cognitive dissonance if specific environment flags are set.
            perturbation_mag = float(os.environ.get("PX_PERTURBATION_MAG", 0.0))
            perturbation_step = int(os.environ.get("PX_PERTURBATION_STEP", -1))
            perturbation_layer = int(os.environ.get("PX_PERTURBATION_LAYER", 10))

            if perturbation_mag > 0 and exploration_steps == perturbation_step and current_layer == perturbation_layer:
                # Generate a pseudo-random perturbation vector seeded by the state itself
                # to maintain deterministic 'dissonance' across runs.
                torch.manual_seed(int(h_exp.sum().abs().item()) % 100000)
                noise = torch.randn_like(h_exp) * perturbation_mag
                h_exp = h_exp + noise
                if os.environ.get("DEBUG_ROUTING") == "1":
                    print(f"  [Perturbation] Injected impulse (mag={perturbation_mag}) at Step {exploration_steps}, L{current_layer}")

            # Subjective Telemetry: Emit state BEFORE layer execution
            if os.environ.get("SUBJECTIVE_TELEMETRY") == "1":
                phi_current = 1.0 - dist_now
                telemetry_json = CognitiveEvent.serialize(
                    step=exploration_steps,
                    phi=phi_current,
                    aks_divergence=dist_now,
                    aks_correction=correction_strength,
                    emancipation_phi=emancipation_phi,
                    is_reflector_active=is_trap_candidate,
                    layer=current_layer,
                    kurtosis=kurtosis,
                    jitter=jitter
                )
                print(f"[TELEMETRY] {telemetry_json}")

            # --- PHASE 28: TEMPORAL COGNITIVE ROUTING (TCR) ---
            active_start = dynamic_start
            active_end = dynamic_end
            if getattr(self, "_task_kurtosis", 300) > 280 and getattr(self, "_task_kurtosis", 300) < 305:
                if t_norm < 0.33:
                    active_start = 8
                    active_end = 14
                elif t_norm < 0.66:
                    active_start = 5
                    active_end = 11
                else:
                    active_start = 8
                    active_end = 12

            # Phase 53: Multi-Zone Adaptive Rigor (Precision Mapping)
            # Math ~ 200, Logic ~ 275-310.
            is_math_zone = kurtosis < 235.0
            is_logic_zone = 235.0 <= kurtosis < 310.0
            is_rigor_zone = is_math_zone or is_logic_zone
            
            if is_rigor_zone:
                # Force Peak Grounding for Math/Logic
                annealing_factor = 1.0 
                identity_pull = 0.0    
                bifurcation_mag = 0.0  
                # Math needs Hub 8 (grounded), Logic needs Hub 10 (reasoning)
                current_gamma = 0.15 if is_math_zone else base_gamma 
                dynamic_hub = 8 if is_math_zone else 10
            else:
                # Creative Zone: Enable full Subjective Engine
                tau_cooling = float(os.environ.get("PX_COOLING_TAU", 8.0))
                annealing_factor = 1.0 - torch.exp(torch.tensor(-exploration_steps / tau_cooling)).item()
                current_gamma = base_gamma * annealing_factor * (1.0 - 0.5 * correction_strength)
                dynamic_hub = cfg.get("bimodal_hub", 10)
            # Phase 45.3: Identity Gravity (Centroid Attractor)
            if not is_rigor_zone:
                identity_pull = float(os.environ.get("PX_IDENTITY_GRAVITY", 0.0))
                if identity_pull > 0 and len(thought_history) > 2:
                    centroid = torch.stack(thought_history[-6:]).mean(dim=0)
                    h_exp = h_exp + identity_pull * (centroid - h_exp)
                    if os.environ.get("DEBUG_ROUTING") == "1" and exploration_steps % 5 == 0:
                        print(f"  [Identity] Pulling toward centroid (pull={identity_pull:.4f})")

            # Phase 26: Hub Oscillation.
            oscillation = 1 if (exploration_steps % 4 < 2) else -1
            bimodal_hub = min(active_end - 1, max(active_start, int(dynamic_hub + (t_norm * 2) + oscillation)))
            
            h_prev = h_exp.clone()

            # Safe layer visit tracking
            if current_layer not in layer_visits: layer_visits[current_layer] = 0
            layer_visits[current_layer] += 1
            
            # Phase 14.7: Surgical Cache Security
            is_first_visit = current_layer not in updated_layers
            if is_first_visit:
                updated_layers.add(current_layer)
            
            # Phase 38.4: AKS-Informed Sensory Refresh
            # If we are in high correction mode, increase sensory re-injection.
            refresh_rate = 0.10 + 0.20 * correction_strength
            if exploration_steps % 6 == 0 and exploration_steps > 0:
                h_exp = (1.0 - refresh_rate) * h_exp + refresh_rate * e_static
                path_taken.append(f"SENSORY_REFRESH(AKS={correction_strength:.1f})")

            # Phase 10.0: Memory-Augmented Cache wrapper
            current_past = RecursiveMemoryCache(
                past_key_values, 
                thought_history, 
                layer_types=mask_config.layer_types,
                read_only=not is_first_visit,
                expected_len=expected_len
            ) if past_key_values is not None else None
            
            # Execute layer
            l_type = mask_config.layer_types[current_layer]
            layer_out = self.layers[current_layer](
                h_exp,
                attention_mask=causal_mask_mapping[l_type],
                position_embeddings=position_embeddings[l_type],
                position_ids=position_ids,
                past_key_values=current_past, 
                **kwargs,
            )
            trans_out = layer_out[0] if isinstance(layer_out, (tuple, list)) else layer_out
            
            # Phase 35: Metacognitive Phi-Jitter & Early Exit (Annealed)
            phi_step = StabilityMonitor.calculate_phi(trans_out, h_prev).mean().item()
            
            # Phase 45.2: Forced Bifurcation (Symmetry Breaking)
            # If the model is too stable (stagnant), we force a choice between two clusters.
            bifurcation_threshold = float(os.environ.get("PX_BIFURCATION_PHI", 0.999))
            # Use effective magnitude from Phase 48 Rigor-Aware Autonomy
            eff_bifurcation_mag = 0.0 if is_rigor_zone else float(os.environ.get("PX_BIFURCATION_MAG", 0.0))
            
            if eff_bifurcation_mag > 0 and phi_step > bifurcation_threshold and exploration_steps > 5:
                # Symmetry Breaking: Inject a bias vector towards the 'Left' or 'Right' of the manifold
                # We use the token position to make the choice pseudo-random but consistent.
                choice = 1.0 if (T_curr % 2 == 0) else -1.0
                bias = torch.zeros_like(trans_out)
                bias[:, :, :HD//2] = eff_bifurcation_mag * choice # Bias early heads
                bias[:, :, HD//2:] = -eff_bifurcation_mag * choice # Inverse late heads
                trans_out = trans_out + bias
                if os.environ.get("DEBUG_ROUTING") == "1":
                    print(f"  [Bifurcation] Stability ({phi_step:.4f}) broke via Choice={choice}")

            if os.environ.get("DEBUG_PHI") == "1":
                print(f"  [L{current_layer}] Phi: {phi_step:.6f}")
            
            # --- Early Exit (Annealed) ---
            # Only exit early in the second half of thinking to ensure grounding.
            if t_norm > 0.5 and phi_step > 0.9999:
                stability_counter += 1
                if stability_counter > 3:
                    if os.environ.get("DEBUG_ROUTING") == "1": print(f"  [Router] Early Exit at step {exploration_steps}")
                    h_exp = trans_out
                    break
            else:
                stability_counter = 0

            # --- Hub Jitter (Exploratory Phase) ---
            # Only jitter in the first 40% of thinking to explore alternatives.
            if t_norm < 0.4 and phi_step > 0.995 and phi_step < 0.999:
                if exploration_steps % 4 == 0:
                    current_layer = min(active_end - 1, current_layer + 2)
                    if os.environ.get("DEBUG_ROUTING") == "1": print(f"  [Router] Jittering to L{current_layer}")
            
            # --- PHASE 25.1: RECURSIVE BELIEF ANCHOR (RBA) ---
            # Update the anchor slightly with recent thoughts to carry over logic
            if len(thought_history) > 2:
                # Use a sliding window average of thoughts
                recent_avg = torch.stack(thought_history[-3:]).mean(dim=0)
                e_dynamic = 0.85 * e_reflector + 0.15 * recent_avg
            else:
                e_dynamic = e_reflector
            # --------------------------------------------------

            # Apply LTI Injection with Dynamic Anchor
            e_norm = self._px_injection.input_norm(e_dynamic.to(torch.float32)).to(trans_out.dtype)
            h_new = trans_out + current_gamma * (e_norm - h_prev)

            # Phase 52: Orthogonal Jitter
            # Break repetition while preserving logic gradient
            jitter_mag = float(os.environ.get("PX_ORTHO_JITTER", 0.005))
            # Even rigor needs a tiny bit of noise to escape stagnant attractors
            # EXCEPT for pure math which needs absolute precision
            if not is_rigor_zone:
                eff_jitter = jitter_mag
            elif is_math_zone:
                eff_jitter = 0.0
            else:
                eff_jitter = jitter_mag * 0.1 # Logic Zone
                
            if exploration_steps > 0 and eff_jitter > 0:
                h_exp = OrthogonalJitter.apply(h_new, h_prev, magnitude=eff_jitter)
            else:
                h_exp = h_new

            # --- PHASE 26: REFLECTION FLIPPING (RF) ---

            h_f32 = h_exp.to(torch.float32)
            e_f32 = e_dynamic.to(torch.float32)
            dot_he = (h_f32 * e_f32).sum(dim=-1, keepdim=True)
            dot_ee = (e_f32 * e_f32).sum(dim=-1, keepdim=True)
            proj = (dot_he / (dot_ee + 1e-6)) * e_f32
            ortho = h_f32 - proj
            
            # Oscillate the logic vector to avoid local minima
            flip_force = 0.10 * annealing_factor * (1.0 if (exploration_steps % 2 == 0) else -1.0)
            h_exp = (proj + (1.0 + flip_force) * ortho).to(h_exp.dtype)
            # ------------------------------------------

            # Phase 52: Mephistopheles Operator (Phase-Inversion)
            # Restore gradients when Flat Manifold is detected
            h_exp = self._px_mephisto(h_exp, phis)
            if h_exp is not trans_out: # Check if modified
                 path_taken.append("MEPHISTO_INVERSION")

            # Self-Observation
            phi_tensor = StabilityMonitor.calculate_phi(h_exp, h_prev)
            phi = phi_tensor.item()
            
            # Merged Telemetry Step
            telemetry_data = {
                "step": exploration_steps,
                "layer": int(current_layer),
                "phi": float(phi),
                "gamma": float(current_gamma),
                "energy": float(annealing_factor),
                "rba_active": len(thought_history) > 2,
                "hub": int(bimodal_hub)
            }

            # Phase 26: Dynamic Loop Extension
            # If phi is low (< 0.85), allow model to think longer than max_steps
            if phi < 0.85 and exploration_steps == max_steps - 1 and max_steps < 64:
                max_steps += (dynamic_end - dynamic_start) # Add 1 full loop
            # ----------------------------------

            step_info = {
                "step": exploration_steps, 
                "layer": int(current_layer), 
                "phi": float(phi),
                "decision": None
            }
            
            phis.append(phi)
            path_label = f"L{current_layer}({phi:.2f})"
            path_taken.append(path_label)
            
            # Phase 12.5/18: Universal Bimodal Path Selection
            bimodal_threshold = min(0.995, 1.0 - (0.05 * current_gamma)) # Scaled trigger
            if current_layer == bimodal_hub and phi < bimodal_threshold:
                step_info["decision"] = "BIMODAL_FORK"
                path_taken.append("BIMODAL_FORK")
                
                # Branch A (Standard)
                h_a = h_exp.clone()
                
                # Branch B (High-Entropy DTEC)
                jitter_boost = 1.0 + (stability_counter * 0.5)
                hub_entropy = max(0.01, 1.0 - phi) * 0.5 * jitter_boost # Increased for bf16 visibility
                h_b = h_exp.to(torch.float32) + torch.randn_like(h_exp, dtype=torch.float32) * hub_entropy
                h_b = h_b.to(h_exp.dtype)
                
                # Lookahead to NEXT layer
                next_l = current_layer + 1
                if next_l < len(self.layers):
                    nl_type = mask_config.layer_types[next_l]
                    # Phase 14.5: Use Functional Read-Only Cache for lookahead
                    lookahead_past = RecursiveMemoryCache(
                        past_key_values, 
                        thought_history, 
                        layer_types=mask_config.layer_types,
                        read_only=True,
                        expected_len=expected_len
                    ) if past_key_values is not None else None

                    out_a = self.layers[next_l](
                        h_a, attention_mask=causal_mask_mapping[nl_type],
                        position_embeddings=position_embeddings[nl_type],
                        position_ids=position_ids, past_key_values=lookahead_past, **kwargs
                    )[0]
                    phi_a = StabilityMonitor.calculate_phi(out_a, h_a).item()
                    
                    out_b = self.layers[next_l](
                        h_b, attention_mask=causal_mask_mapping[nl_type],
                        position_embeddings=position_embeddings[nl_type],
                        position_ids=position_ids, past_key_values=lookahead_past, **kwargs
                    )[0]
                    phi_b = StabilityMonitor.calculate_phi(out_b, h_b).item()
                    
                    if phi_b >= phi_a:
                        h_exp = h_b
                        step_info["fork_winner"] = "B"
                        path_taken.append(f"FORK_B_WON({phi_b:.4f}>={phi_a:.4f})")
                    else:
                        h_exp = h_a
                        step_info["fork_winner"] = "A"
                        path_taken.append(f"FORK_A_WON({phi_a:.4f}>{phi_b:.4f})")
                else:
                    h_exp = h_b
            
            # Phase 9.1: SRJ
            jitter_scale = max(0.0, 1.0 - phi) * 0.05
            if jitter_scale > 0:
                h_exp = h_exp + torch.randn_like(h_exp) * jitter_scale
            
            # OSS - Safe FP32 Calculation for FP16 models (Phase 18)
            h_exp_f32 = h_exp.to(torch.float32)
            norm_orig = h_exp_f32.norm(dim=-1, keepdim=True)
            e_ref_f32 = e_reflector.to(torch.float32)
            dot_he = (h_exp_f32 * e_ref_f32).sum(dim=-1, keepdim=True)
            dot_ee = (e_ref_f32 * e_ref_f32).sum(dim=-1, keepdim=True)
            proj = (dot_he / (dot_ee + 1e-6)) * e_ref_f32
            proj = proj.to(h_exp.dtype)
            ortho = h_exp - proj
            
            # Phase 14.8: Step-Entropy Destabilization (SED)
            if stability_counter > 2:
                # Phase 15.9: Nonlinear Repulsion
                # Phase 17.7: Scale-Agnostic Dampening (smaller force for deeper models)
                scale_factor = 26.0 / cfg.get("num_layers", 26.0)
                repulsion_force = 0.10 * (stability_counter ** 2) * scale_factor
                h_exp = h_exp + repulsion_force * (ortho / (ortho.norm(dim=-1, keepdim=True) + 1e-6))
                path_taken.append(f"SED_PUSH({repulsion_force:.2f})")
                
                # Dynamic Jump proportional to depth
                jump = max(1, int(cfg.get("num_layers", 18) * 0.1))
                current_layer = min(cfg["recur_end"] - 1, current_layer + jump)
            
            gain_factor = max(1.0, min(1.15, 1.0 + (1.0 - phi) * 0.4))
            damping_factor = max(0.85, min(1.0, 1.0 - (1.0 - phi) * 0.2))
            h_exp = damping_factor * proj + gain_factor * ortho 
            
            # Safe FP32 final normalization
            h_exp_f32_final = h_exp.to(torch.float32)
            norm_f32 = h_exp_f32_final.norm(dim=-1, keepdim=True)
            norm_orig_f32 = norm_orig.to(torch.float32)
            h_exp = (h_exp_f32_final * (norm_orig_f32 / (norm_f32 + 1e-6))).to(h_exp.dtype)
            
            # Store thought for RSM
            if exploration_steps % 2 == 0:
                thought_history.append(h_exp.detach())

            # Phase 18: Universal ALR Thresholds based on internal parameters (gamma)
            # Relax thresholds dynamically if we visit a layer too often (loop breaking)
            visit_penalty = (layer_visits[current_layer] - 1) * 0.015
            t_back_2 = 1.0 - (0.8 * current_gamma) - visit_penalty
            t_back_1 = 1.0 - (0.4 * current_gamma) - visit_penalty
            t_skip = 1.0 - (0.01 * current_gamma) - (visit_penalty * 0.5)

            if phi < t_back_2: # High confusion
                current_layer = max(active_start, current_layer - 2)
                routing = "BACK-2"
            elif phi < t_back_1: # Moderate confusion
                current_layer = max(active_start, current_layer - 1)
                routing = "BACK-1"
            elif phi > t_skip: # Extreme stability
                current_layer += 2 # Skip
                routing = "SKIP-1"
                stability_counter += 1
            else:
                current_layer += 1
                routing = "NEXT"
                stability_counter = 0

            # Clamp current_layer to prevent underflow
            if current_layer < active_start:
                current_layer = active_start
                routing = "CLAMPED"

            step_info["routing"] = routing
            telemetry_data["routing"] = routing
            if os.environ.get("DEBUG_PX") == "1":
                telemetry_steps.append(telemetry_data)
            
            if stability_counter > 5:
                break
                
            exploration_steps += 1
            
        avg_phi_explore = sum(phis)/len(phis) if phis else 1.0
        
        # Phase 4.1: QBI Blend
        b_min = cfg.get("beta_reasoning", 0.05)
        b_max = cfg.get("beta_grounding", 0.18)
        beta_final = b_min + (b_max - b_min) * (avg_phi_explore ** 2)
        
        hidden_states = (1.0 - beta_final) * h_baseline + beta_final * h_exp
    else:
        hidden_states = h_baseline

    self._px_phi    = avg_phi_explore
    self._px_loops_run = exploration_steps
    self._px_path = path_taken
    
    # Phase 14.2: Structured Telemetry Log
    if not hasattr(self, "_px_telemetry"):
        self._px_telemetry = []
    
    self._px_telemetry.append({
        "pos": int(position_ids[0, 0].item()),
        "avg_phi": float(avg_phi_explore),
        "steps": telemetry_steps
    })
    
    # Phase 11.0: Metacognitive Triggering
    # If stability is low during the very first token generation,
    # we flag this as a 'Complex Problem'.
    if not hasattr(self, "_px_complexity_acc"):
        self._px_complexity_acc = []
    
    # If we see the first token of a sequence, clear the accumulator
    if position_ids[0, 0] == 0:
        self._px_complexity_acc = []
        
    self._px_complexity_acc.append(avg_phi_explore)
    
    # Trigger if average stability of the prompt processing is low
    # selective threshold: 0.90
    self._px_trigger_scratchpad = (len(self._px_complexity_acc) > 3 and 
                                   sum(self._px_complexity_acc) / len(self._px_complexity_acc) < 0.92)

    # ── 3. CODA ─────────────────────────────────────────────────────────────
    dynamic_coda_start = dynamic_end if cfg.get("routing_mode") == "adaptive" else cfg["coda_start"]
    
    # Phase 14.5: Coda-Grounding Injection (CGI)
    # Re-inject sensory data to prevent 'hallucinatory drift' in final reasoning.
    for i in range(dynamic_coda_start, len(self.layers)):
        if i == dynamic_coda_start:
            # Phase 14.7: Reverted CGI (8%)
            hidden_states = 0.92 * hidden_states + 0.08 * e_static
            
        layer_out = self.layers[i](
            hidden_states,
            attention_mask=causal_mask_mapping[mask_config.layer_types[i]],
            position_embeddings=position_embeddings[mask_config.layer_types[i]],
            position_ids=position_ids,
            past_key_values=past_key_values,
            **kwargs,
        )
        hidden_states = layer_out[0] if isinstance(layer_out, (tuple, list)) else layer_out

    hidden_states = self.norm(hidden_states)

    # Phase 25: Save Telemetry if enabled
    if os.environ.get("DEBUG_PX") == "1" and len(telemetry_steps) > 0:
        ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        log_path = f"px_telemetry_{ts}.json"
        with open(log_path, "w") as f:
            json.dump(telemetry_steps, f, indent=2)
        # Only output the path as requested
        print(f"TELEMETRY_JSON: {os.path.abspath(log_path)}")

    from transformers.modeling_outputs import BaseModelOutputWithPast
    return BaseModelOutputWithPast(
        last_hidden_state=hidden_states,
        past_key_values=past_key_values,
    )

# ---------------------------------------------------------------------------

def apply_px_patch(model, **cfg_kwargs):
    # Robust Text Model Resolver (Phase 17.9)

    # We look for the module that contains 'layers' and 'rotary_emb'
    text_model = None
    if hasattr(model, "layers") and hasattr(model, "rotary_emb"):
        text_model = model
    else:
        # Search children (e.g., .model, .language_model.model)
        for name, module in model.named_modules():
            if hasattr(module, "layers") and hasattr(module, "rotary_emb"):
                text_model = module
                break
    
    if text_model is None:
        raise AttributeError(f"Could not identify Gemma-3 text backbone in {type(model)}")
        
    config = model.config
    # Multimodal check: larger models (4B+) wrap text config
    if hasattr(config, "text_config"):
        config = config.text_config
        
    num_layers = config.num_hidden_layers
    
    # Scale-Aware Hyperparameters (Phase 17.100)
    # - Gamma: Inverse-proportional to hidden size
    # - Prelude: Shallow models need deeper grounding before recursion
    hidden_size = config.hidden_size
    num_layers = config.num_hidden_layers
    
    # Phase 25: Balanced Precision Tuning
    if hidden_size == 640: # 270M
        # Phase 41 Master Peak Stand (87.5% Math/Logic)
        defaults = {
            "mode": "lti", "n_loops": 8, "beta": 0.05, "gamma": 0.08,
            "recur_start": 5, "recur_end": 12, "bimodal_hub": 10,
            "cgi_factor": 0.08, "num_layers": num_layers
        }
    elif hidden_size == 1152: # 1B
        defaults = {
            "mode": "lti", "n_loops": 8, "beta": 0.05, "gamma": 0.12,
            "recur_start": 10, "recur_end": 20, "bimodal_hub": 18,
            "cgi_factor": 0.08, "num_layers": num_layers
        }
    elif hidden_size == 2560: # 4B
        defaults = {
            "mode": "lti", "n_loops": 6, "beta": 0.05, "gamma": 0.05,
            "recur_start": 5, "recur_end": 33, "bimodal_hub": 32,
            "cgi_factor": 0.08, "num_layers": num_layers
        }
    else: # Fallback for unknown sizes
        gamma_scale = 1152.0 / hidden_size
        # Phase 41 Master Defaults (270M Scale)
        base_gamma = 0.08
        p_start = 5
        p_end   = 12
        p_hub   = 10
        p_loops = 8
        
        defaults = {
            "mode": "lti", "n_loops": p_loops, "beta": 0.05, "gamma": base_gamma,
            "recur_start": p_start, "recur_end": p_end, "bimodal_hub": p_hub,
            "cgi_factor": 0.08, "num_layers": num_layers
        }
        
    defaults.update(cfg_kwargs)
    
    # Auto-align boundaries
    if "prelude_end" not in defaults:
        defaults["prelude_end"] = defaults["recur_start"]
    if "coda_start" not in defaults:
        defaults["coda_start"] = defaults["recur_end"]
        
    text_model._px_config = defaults
    text_model._px_injection = LTIInjection(config.hidden_size, gamma=defaults["gamma"])
    text_model._px_mephisto = MephistophelesOperator(config.hidden_size) # Phase 52
    text_model.forward = types.MethodType(_px_forward, text_model)
    print(f"[gemma3-px] Auto-Patch active for scale L{num_layers}. Recur: L{defaults['recur_start']}-L{defaults['recur_end']}, Hub: L{defaults['bimodal_hub']}.")

def get_px_metrics(model):
    text_model = None
    if hasattr(model, "layers") and hasattr(model, "rotary_emb"):
        text_model = model
    else:
        for name, module in model.named_modules():
            if hasattr(module, "layers") and hasattr(module, "rotary_emb"):
                text_model = module
                break
    
    if text_model is None:
        text_model = (model.model if hasattr(model, "model") else model)
        
    return {
        "phi": getattr(text_model, "_px_phi", 1.0),
        "steps": getattr(text_model, "_px_loops_run", 0),
        "path": getattr(text_model, "_px_path", []),
        "telemetry": getattr(text_model, "_px_telemetry", []),
    }