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

# A tiny parameterized Neural Network that acts as our ODE continuous vector field
class ODEVectorField(nn.Module):
    def __init__(self, dim):
        super().__init__()
        # Simulates how the hidden state flows over abstract execution time
        self.net = nn.Sequential(
            nn.Linear(dim, dim),
            nn.Tanh(),
            nn.Linear(dim, dim)
        )
    def forward(self, t, h):
        # h comes from the ODE solver as a numpy array
        h_tensor = torch.from_numpy(h).float()
        with torch.no_grad():
            dhdt = self.net(h_tensor).numpy()
        return dhdt

class ODERoutedSLMSystem:
    def __init__(self, base_model_id, specialty_ids):
        # 1. Initialize the Core Base SLM which generates the Initial Condition C
        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 # e.g., 576 for SmolLM2
        
        # 2. Setup our continuous trajectory router (Neural ODE field)
        self.ode_func = ODEVectorField(self.hidden_dim)
        
        # 3. Store our Specialized Experts profiles with anchor target coordinates in space
        self.specialists = specialty_ids
        # Pre-assign coordinates in the latent space for each expert specialization
        # In a trained environment, these vectors represent the average signature of the domain
        np.random.seed(42)
        self.expert_anchors = {
            name: np.random.randn(self.hidden_dim) for name in specialty_ids.keys()
        }

    def compute_initial_condition(self, prompt):
        """Processes text through the base model to output the h(0) vector."""
        inputs = self.tokenizer(prompt, return_tensors="pt")
        with torch.no_grad():
            outputs = self.core_model(**inputs, output_hidden_states=True)
            # Take the mean of the final hidden layer across tokens
            # FIX: Added .float() to safely cast BFloat16 down to Float32 for NumPy compatibility
            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):
        """Integrates h(0) through the vector field to find the target expert coordinates."""
        # Solve the Initial Value Problem (IVP): dh/dt = f(h, t) from t=0 to t=integration_time
        sol = solve_ivp(self.ode_func, [0, integration_time], h_0, method='RK45')
        h_final = sol.y[:, -1] # The terminal trajectory coordinate
        
        # Determine closest expert using cosine distance or negative euclidean distance
        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, h_final

    def generate(self, prompt):
        # Step 1: Extract Initial Condition C from text
        h_0 = self.compute_initial_condition(prompt)
        
        # Step 2: Continuous ODE trajectory tracking
        expert_name, final_state = self.route_via_ode(h_0)
        print(f"\n[ODE Trajectory Route] Target destination selected: {expert_name}")
        
        # Step 3: Lazy-load or call the targeted high-accuracy specialization model
        # (For simulation, we fetch the model path string to execute generation)
        target_model_path = self.specialists[expert_name]
        print(f"Executing generation with specialized weights from: {target_model_path}")
        
        # In deployment, load/swap target layer weights or call the specific pipeline
        return expert_name

if __name__ == "__main__":
    # Define our dictionary of specialized expert models
    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"
    }
    
    # Initialize system with the base model providing the initial condition
    system = ODERoutedSLMSystem(
        base_model_id="HuggingFaceTB/SmolLM2-135M-Instruct", 
        specialty_ids=specialists
    )
    
    # Test Prompts
    prompt_math = "Calculate the derivative of x^2 log(x) using integration steps."
    prompt_code = "Write a fast matrix multiplication script in pure C language."
    
    # Run text through ODE routing
    system.generate(prompt_math)
    system.generate(prompt_code)