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

class NeuralODEField(nn.Module):
    """
    A purely neural network-driven ODE Vector Field f_theta(h, t).
    No hardcoded rules—it relies completely on its weight matrices.
    """
    def __init__(self, dim):
        super().__init__()
        # Linear layers parameterize the continuous-depth trajectory field
        self.linear1 = nn.Linear(dim, dim, bias=False)
        self.linear2 = nn.Linear(dim, dim, bias=False)
        
        # Initialize weights with orthogonal variance to force strong trajectory divergence
        nn.init.orthogonal_(self.linear1.weight)
        nn.init.orthogonal_(self.linear2.weight)

    def forward(self, t, h):
        # h arrives from scipy solve_ivp as a numpy array
        h_tensor = torch.from_numpy(h).float()
        with torch.no_grad():
            # dh/dt = W2 * tanh(W1 * h)
            dhdt = self.linear2(torch.tanh(self.linear1(h_tensor))).numpy()
        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."
        }
        
        # 1. Instantiate the neural-driven ODE function
        self.ode_func = NeuralODEField(self.hidden_dim)
        
        # 2. Establish neural anchor targets
        # We project distinct high-dimensional orthogonal vectors as target coordinate zones
        # representing where the network trajectories should land.
        np.random.seed(1337)
        self.expert_anchors = {}
        for name in specialty_ids.keys():
            v = np.random.randn(self.hidden_dim)
            self.expert_anchors[name] = v / np.linalg.norm(v)

    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)
            # Extracted context hidden state acts as the Initial Condition h(0)
            h_0 = outputs.hidden_states[-1].mean(dim=1).squeeze().float().cpu().numpy()
        return h_0

    def route_via_neural_ode(self, h_0, integration_time=1.0):
        # Solve the IVP tracking the neural trajectory field path: dh/dt = f_theta(h, t)
        sol = solve_ivp(self.ode_func, [0, integration_time], h_0, method='RK45', rtol=1e-3)
        h_final = sol.y[:, -1]
        
        # Normalize the terminal state to check direction vectors via cosine similarity
        h_final_norm = h_final / (np.linalg.norm(h_final) + 1e-9)
        
        best_match = None
        max_similarity = -float('inf')
        
        # Match terminal continuous destination to the closest expert profile
        for name, anchor in self.expert_anchors.items():
            similarity = np.dot(h_final_norm, anchor)
            if similarity > max_similarity:
                max_similarity = similarity
                best_match = name
        return best_match

    def generate(self, user_prompt):
        # 1. Capture the initial condition h(0) from the core model
        h_0 = self.compute_initial_condition(user_prompt)
        
        # 2. Run the pure neural ODE integration function to find the routing target
        expert_name = self.route_via_neural_ode(h_0)
        
        print(f"\nPrompt: \"{user_prompt[:50]}...\"")
        print(f"[ODE Route] Neural Trajectory 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)
        
        # Safety extraction wrapper for BatchEncoding
        if isinstance(inputs, dict) or hasattr(inputs, "data"):
            input_ids = inputs["input_ids"]
        else:
            input_ids = 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
    )
    
    # Run text through the neural ODE network router
    system.generate("Calculate the derivative of x^2.")
    system.generate("Write a fast matrix multiplication script in C.")