from transformers.models.gemma3.modeling_gemma3 import Gemma3ForCausalLM, Gemma3Model
from .configuration_openmythos import OpenMythosConfig
from .patch import apply_px_patch

class OpenMythosForCausalLM(Gemma3ForCausalLM):
    config_class = OpenMythosConfig
    
    def __init__(self, config):
        super().__init__(config)
        # We apply the patch immediately after initialization
        # Extract default parameters from the config if they exist
        routing_mode = getattr(config, "px_routing_mode", "adaptive")
        gamma = getattr(config, "px_gamma", 0.10)
        
        # We MUST start the initial prelude at Layer 5 so the Kurtosis probe 
        # (which is calibrated for L5) reads the correct tensor. 
        # The router will fast-forward the remaining layers if dynamic_start > 5.
        apply_px_patch(self, recur_start=5, recur_end=12, routing_mode=routing_mode, gamma=gamma)
        
# For safety, if they instantiate the base model directly:
class OpenMythosModel(Gemma3Model):
    config_class = OpenMythosConfig
    
    def __init__(self, config):
        super().__init__(config)