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

# [Include the SemanticDriftField from the previous script here]

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
        
        # Define the exact System Prompts to enforce expert roles during inference
        self.expert_system_prompts = {
            "Math_Expert": "You are a professional mathematician. Provide precise calculations and step-by-step mathematical proofs.",
            "Code_Expert": "You are an expert software developer. Write clean, efficient, and well-commented code blocks without unnecessary prose.",
            "Creative_Expert": "You are a creative fantasy author. Write imaginative stories with rich descriptive prose and deep lore."
        }
        
        # Build Semantic Anchors using token weights
        embed_weights = self.core_model.get_input_embeddings().weight.detach().float().cpu().numpy()
        self.expert_anchors = {
            "Math_Expert": embed_weights[self.tokenizer.encode("mathematics", add_special_tokens=False)[0]],
            "Code_Expert": embed_weights[self.tokenizer.encode("programming", add_special_tokens=False)[0]],
            "Creative_Expert": embed_weights[self.tokenizer.encode("creative", add_special_tokens=False)[0]]
        }
        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)
            h_0 = outputs.hidden_states[-1].mean(dim=1).squeeze().float().cpu().numpy()
        return h_0

    def route_via_ode(self, h_0):
        sol = solve_ivp(self.ode_func, [0, 1.0], h_0, method='RK45')
        h_final = sol.y[:, -1] 
        
        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, user_prompt):
        # 1. ODE Trajectory Routing
        h_0 = self.compute_initial_condition(user_prompt)
        expert_name = self.route_via_ode(h_0)
        
        print(f"\nPrompt: \"{user_prompt[:50]}...\"")
        print(f"[ODE Route] Routing to Expert: \033[92m{expert_name}\033[0m")
        
        # 2. Extract the specialized system prompt for the routed destination
        assigned_system_prompt = self.expert_system_prompts[expert_name]
        
        # 3. Format into the ChatML format with structural system instructions
        messages = [
            {"role": "system", "content": assigned_system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        
        # Move model to execution device (CPU/GPU)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        self.core_model.to(device)
        
        inputs = self.tokenizer.apply_chat_template(
            messages, 
            add_generation_prompt=True, 
            return_tensors="pt"
        ).to(device)
        
        # --- THE FIX ---
        # Safeguard: Extract the raw input_ids tensor from the dictionary wrapper
        if isinstance(inputs, dict) or hasattr(inputs, "data"):
            input_ids = inputs["input_ids"]
        else:
            input_ids = inputs
        
        # 4. Compute text generation using the raw tensor
        outputs = self.core_model.generate(
            input_ids,  # Use input_ids directly
            max_new_tokens=60, 
            temperature=0.3, 
            do_sample=True
        )
        
        # Use input_ids.shape[-1] safely to strip away the system and user prompt tokens
        response = self.tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
        print(f"\033[1;34m[System Role Activated]:\033[0m {assigned_system_prompt}")
        print(f"\033[1;32m[Response]:\033[0m {response.strip()}\n")
        
if __name__ == "__main__":
    # In a production environment, these values would point to independent delta weights
    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
    )
    
    # Test execution
    system.generate("Calculate the derivative of x^2.")
    system.generate("Write a fast matrix multiplication script in C.")