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

# ---------------------------------------------------------------------------
# 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
            d_hub = (w_m*8 + w_la*11 + w_cr*14 + w_lb*11 + 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 Logic-Trap zone
    if jitter > 1e8 or (280 < kurtosis < 310):
        is_trap_candidate = True
        if os.environ.get("DEBUG_ROUTING") == "1":
            reason = "Jitter" if jitter > 1e8 else "Logic-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:
        current_gamma *= 0.1 
    elif phi_intuition > 0.999:
        current_gamma *= 0.5
    
    # 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) ---
            # Center of the thinking process (0.5)
            t_norm = exploration_steps / max_steps
            
            # Phase 38.2: AKS - Topological Anomaly Detection
            # We measure the 'velocity' of the latent state across the last 3 steps.
            # If the states move AWAY from the Anchor (Disparity), we trigger Correction.
            if exploration_steps > 2:
                # Proximity to Anchor (e_static)
                dist_now = 1.0 - StabilityMonitor.calculate_phi(h_exp, e_static).mean().item()
                divergence_buffer.append(dist_now)
                if len(divergence_buffer) > 4: divergence_buffer.pop(0)
                
                # Check for "Accelerating Dispersion" (The Anna Karenina Signal)
                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 we are accelerating AWAY from the truth (dist increasing)
                    if acceleration > 0.001 and velocity > 0:
                        correction_strength = min(1.0, correction_strength + 0.1)
                        if os.environ.get("DEBUG_PHI") == "1":
                            print(f"  [AKS] Anomaly! Accel={acceleration:.4f}, Correction={correction_strength:.2f}")
                    else:
                        correction_strength = max(0.0, correction_strength - 0.05)

            # --- 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

            # Sigmoid transition: sharp drop in energy after 50% of thinking
            k_steep = 12.0
            energy_factor = 1.0 - (1.0 / (1.0 + torch.exp(torch.tensor(-k_steep * (t_norm - 0.5))))).item()
            
            # Phase 38.3: AKS-Modulated Gamma
            # High correction strength forces the model to ground HARDER (lower exploration).
            current_gamma = base_gamma * (0.5 + energy_factor) * (1.0 - 0.5 * correction_strength)

            # 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()
            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_exp = trans_out + current_gamma * (e_norm - h_prev)
            
            # --- 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 * energy_factor * (1.0 if (exploration_steps % 2 == 0) else -1.0)
            h_exp = (proj + (1.0 + flip_force) * ortho).to(h_exp.dtype)
            # ------------------------------------------

            # 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(energy_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):
    from .px_modules import LTIInjection
    
    # 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
        defaults = {
            "mode": "lti", "n_loops": 16, "beta": 0.05, "gamma": 0.12,
            "recur_start": 8, "recur_end": 12, "bimodal_hub": 11,
            "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
        base_gamma = 0.12 * gamma_scale
        prelude_ratio = 0.40 if num_layers < 20 else 0.15
        p_start = max(4, int(num_layers * prelude_ratio))
        p_end   = max(p_start + 1, int(num_layers * 0.98))
        p_hub   = max(p_start + 1, int(num_layers * 0.95))
        
        defaults = {
            "mode": "lti", "n_loops": 4, "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.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", []),
    }