import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
from scipy.integrate import solve_ivp
import numpy as np

class SemanticDriftField:
    """
    An unparameterized continuous linear drift field.
    It guides the trajectory h(t) along a gradient path toward the correct anchor.
    """
    def __init__(self, anchors):
        self.anchors = anchors

    def __call__(self, t, h):
        # Find which anchor has the strongest pull on the current state h
        best_anchor = None
        max_sim = -float('inf')
        for name, anchor in self.anchors.items():
            # Standard dot-product similarity
            sim = np.dot(h, anchor) / (np.linalg.norm(h) * np.linalg.norm(anchor) + 1e-9)
            if sim > max_sim:
                max_sim = sim
                best_anchor = anchor
        
        # The derivative dh/dt acts as a drift vector drawing h toward the dominant domain anchor
        dhdt = 0.5 * (best_anchor - h)
        return dhdt

class ODERoutedSLMSystem:
    def __init__(self, base_model_id, specialty_ids):
        print(f"Loading Initial Condition SLM Core: {base_model_id}...")
        self.tokenizer = AutoTokenizer.from_pretrained(base_model_id)
        self.core_model = AutoModelForCausalLM.from_pretrained(base_model_id)
        self.hidden_dim = self.core_model.config.hidden_size 
        self.specialists = specialty_ids
        
        # --- FIX: Extract Semantic Anchors from the Model's Own Embedding Weights ---
        # Instead of np.random, we use actual token indices to locate true semantic coordinates
        embed_weights = self.core_model.get_input_embeddings().weight.detach().float().cpu().numpy()
        
        # Token IDs for semantic guide posts
        math_token_id = self.tokenizer.encode("mathematics", add_special_tokens=False)[0]
        code_token_id = self.tokenizer.encode("programming", add_special_tokens=False)[0]
        story_token_id = self.tokenizer.encode("creative", add_special_tokens=False)[0]
        
        self.expert_anchors = {
            "Math_Expert": embed_weights[math_token_id],
            "Code_Expert": embed_weights[code_token_id],
            "Creative_Expert": embed_weights[story_token_id]
        }
        
        # Initialize the ODE system with our true semantic drift field
        self.ode_func = SemanticDriftField(self.expert_anchors)

    def compute_initial_condition(self, prompt):
        inputs = self.tokenizer(prompt, return_tensors="pt")
        with torch.no_grad():
            outputs = self.core_model(**inputs, output_hidden_states=True)
            # Safely cast bfloat16 to float32 for scipy
            h_0_tensor = outputs.hidden_states[-1].mean(dim=1).squeeze().float()
            h_0 = h_0_tensor.cpu().numpy()
        return h_0

    def route_via_ode(self, h_0, integration_time=1.0):
        # Solve dh/dt from t=0 to t=1
        sol = solve_ivp(self.ode_func, [0, integration_time], h_0, method='RK45')
        h_final = sol.y[:, -1] 
        
        # Match terminal trajectory state to closest anchor destination
        best_match = None
        min_dist = float('inf')
        for name, anchor in self.expert_anchors.items():
            dist = np.linalg.norm(h_final - anchor)
            if dist < min_dist:
                min_dist = dist
                best_match = name
        return best_match

    def generate(self, prompt):
        h_0 = self.compute_initial_condition(prompt)
        expert_name = self.route_via_ode(h_0)
        
        print(f"\nPrompt: \"{prompt[:45]}...\"")
        print(f"[ODE Trajectory Route] Destination reached: \033[92m{expert_name}\033[0m")
        print(f"Executing generation with: {self.specialists[expert_name]}")
        return expert_name

if __name__ == "__main__":
    specialists = {
        "Math_Expert": "HuggingFaceTB/SmolLM2-135M-Instruct-Math-Delta",
        "Code_Expert": "HuggingFaceTB/SmolLM2-135M-Instruct-Code-Delta",
        "Creative_Expert": "HuggingFaceTB/SmolLM2-135M-Instruct-Creative-Delta"
    }
    
    system = ODERoutedSLMSystem(
        base_model_id="HuggingFaceTB/SmolLM2-135M-Instruct", 
        specialty_ids=specialists
    )
    
    # Run diverse test prompts to check successful routing separation
    system.generate("Calculate the derivative of x^2 log(x) using integration steps.")
    system.generate("Write a fast matrix multiplication script in pure C language.")
    system.generate("Tell me a creative fantasy story about a clockwork dragon.")