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

class AISLMVectorField:
    """
    An ODE vector field powered entirely by the AI model's native pre-trained intelligence.
    It passes the trajectory state through the model's own output layer (lm_head) 
    to read the AI's semantic routing vector.
    """
    def __init__(self, lm_head, token_ids, dim):
        self.lm_head = lm_head
        self.token_ids = token_ids  # Target slots mapped to expert indices
        self.dim = dim

    def __call__(self, t, h):
        # 1. Convert the continuous hidden state array back to a PyTorch tensor
        h_tensor = torch.from_numpy(h).to(device=self.lm_head.weight.device)
        
        # FIX: Dynamically cast the input tensor to match the exact dtype of the model weights (e.g. BFloat16)
        h_tensor = h_tensor.to(dtype=self.lm_head.weight.dtype)
        
        with torch.no_grad():
            # 2. Pass h through the AI's actual pre-trained classification layer
            logits = self.lm_head(h_tensor)
            
            # 3. Extract the AI's internal probability distribution for our experts
            expert_logits = torch.tensor([logits[tid] for tid in self.token_ids], dtype=torch.float32)
            probs = F.softmax(expert_logits, dim=0).cpu().numpy()
            
        # 4. The derivative dh/dt becomes a velocity vector pointing cleanly 
        # along the direction of the AI's native prediction confidence
        dhdt = np.zeros_like(h)
        dhdt[:len(probs)] = probs  # Use the first elements as our trajectory router state
        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
        
        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."
        }
        
        # Map our structural experts to equivalent semantic token IDs in the model's vocabulary
        self.expert_keys = list(specialty_ids.keys())
        target_tokens = ["math", "code", "story"]
        token_ids = [self.tokenizer.encode(t, add_special_tokens=False)[0] for t in target_tokens]
        
        # Initialize the vector field directly using the AI's internal pre-trained lm_head
        self.ode_func = AISLMVectorField(self.core_model.lm_head, token_ids, self.hidden_dim)

    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)
            # The initial condition vector h(0) extracted directly from the AI's deep layer
            h_0 = outputs.hidden_states[-1].mean(dim=1).squeeze().float().cpu().numpy()
        return h_0

    def route_via_ai_ode(self, h_0, integration_time=1.0):
        # Solve the continuous trajectory equation using the AI's internal logit flows
        sol = solve_ivp(self.ode_func, [0, integration_time], h_0, method='RK45')
        h_final = sol.y[:, -1]
        
        # Read the resulting state destination directly from the dynamic trajectory
        # The expert with the dominant probability velocity wins the route
        probs = h_final[:len(self.expert_keys)]
        best_idx = np.argmax(probs)
        return self.expert_keys[best_idx]

    def generate(self, user_prompt):
        # Get h(0) from the AI core
        h_0 = self.compute_initial_condition(user_prompt)
        
        # Route through the continuous pre-trained vector field
        expert_name = self.route_via_ai_ode(h_0)
        
        print(f"\nPrompt: \"{user_prompt[:50]}...\"")
        print(f"[ODE Route] Pure AI-Driven Destination: \033[92m{expert_name}\033[0m")
        
        assigned_system_prompt = self.expert_system_prompts[expert_name]
        
        messages = [
            {"role": "system", "content": assigned_system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        
        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)
        input_ids = inputs["input_ids"] if isinstance(inputs, dict) or hasattr(inputs, "data") else inputs
        
        outputs = self.core_model.generate(input_ids, max_new_tokens=40, temperature=0.3, do_sample=True)
        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__":
    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
    )
    
    system.generate("Calculate the derivative of x^2.")
    system.generate("Write a fast matrix multiplication script in C.")