import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

class MultiExpertRoutedSystem:
    def __init__(self, base_model_id):
        print(f"Loading Initial Condition SLM Core: {base_model_id}...")
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        
        # 1. Load the core routing model and its tokenizer
        self.tokenizer_C = AutoTokenizer.from_pretrained(base_model_id)
        self.core_model = AutoModelForCausalLM.from_pretrained(base_model_id).to(self.device)
        
        # 2. Map routing choices to the core model's vocabulary token IDs
        self.math_token_id = self.tokenizer_C.encode("math", add_special_tokens=False)[0]
        self.code_token_id = self.tokenizer_C.encode("code", add_special_tokens=False)[0]
        
        # 3. Define the real Hugging Face model strings for the experts
        self.expert_models = {
            "Math_Expert": "Qwen/Qwen2.5-Math-1.5B-Instruct",
            "Code_Expert": "bigcode/starcoder2-3b"
        }

    def route_and_generate(self, user_prompt):
        # --- STAGE 1: AI Routing Decision ---
        inputs = self.tokenizer_C(user_prompt, return_tensors="pt").to(self.device)
        
        with torch.no_grad():
            outputs = self.core_model(**inputs, output_hidden_states=True)
            final_hidden = outputs.hidden_states[-1].mean(dim=1)
            logits = self.core_model.lm_head(final_hidden).squeeze()
            
            math_score = logits[self.math_token_id].float().item()
            code_score = logits[self.code_token_id].float().item()
            
        expert_name = "Math_Expert" if math_score > code_score else "Code_Expert"
        target_model_id = self.expert_models[expert_name]
        
        print(f"\nPrompt: \"{user_prompt[:50]}...\"")
        print(f"[AI Choice Logits] Math: {math_score:.2f} | Code: {code_score:.2f}")
        print(f"[Decision Route] Selected: \033[92m{expert_name}\033[0m -> {target_model_id}")
        
        # --- STAGE 2: Lazy Loading & Execution of the Target Expert ---
        print(f"Loading weights for {expert_name}...")
        expert_tokenizer = AutoTokenizer.from_pretrained(target_model_id)
        # Using device_map="auto" to handle memory efficiency for the larger models
        expert_model = AutoModelForCausalLM.from_pretrained(
            target_model_id, 
            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
            device_map="auto"
        )
        
        # Format prompt according to whether it's an Instruct chat template or raw code completion
        if "Instruct" in target_model_id:
            messages = [{"role": "user", "content": user_prompt}]
            expert_inputs = expert_tokenizer.apply_chat_template(
                messages, 
                add_generation_prompt=True, 
                return_tensors="pt"
            ).to(expert_model.device)
        else:
            # For base/completion models like StarCoder2, feed raw prompt or basic code framing
            expert_inputs = expert_tokenizer(user_prompt, return_tensors="pt").input_ids.to(expert_model.device)
            
        # Clean up container formats if necessary
        if isinstance(expert_inputs, dict) or hasattr(expert_inputs, "data"):
            input_ids = expert_inputs["input_ids"]
        else:
            input_ids = expert_inputs

        # Generate response using the specialized model
        outputs = expert_model.generate(
            input_ids, 
            max_new_tokens=100, 
            temperature=0.2, 
            do_sample=True
        )
        
        response = expert_tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
        print(f"\033[1;32m[{expert_name} Response]:\033[0m\n{response.strip()}\n")
        
        # Clean up GPU/VRAM memory to prepare for the next dynamic swap
        del expert_model
        del expert_tokenizer
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

if __name__ == "__main__":
    # Core SLM acts as our entrypoint gatekeeper
    system = MultiExpertRoutedSystem("HuggingFaceTB/SmolLM2-135M-Instruct")
    
    # Process both prompts into their native models
    system.route_and_generate("Calculate the derivative of x^2.")
    system.route_and_generate("Write a fast matrix multiplication script in C.")