import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

class AISLMRouterSystem:
    def __init__(self, model_id):
        print(f"Loading Initial Condition SLM Core: {model_id}...")
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        
        # Load model and tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        self.core_model = AutoModelForCausalLM.from_pretrained(model_id).to(self.device)
        
        # Define system roles for the target specialists
        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."
        }
        
        # Map our choices to target tokens in the model's vocabulary
        self.math_token_id = self.tokenizer.encode("math", add_special_tokens=False)[0]
        self.code_token_id = self.tokenizer.encode("code", add_special_tokens=False)[0]

    def route_and_generate(self, user_prompt):
        # 1. Process prompt tokens to extract the hidden states
        inputs = self.tokenizer(user_prompt, return_tensors="pt").to(self.device)
        
        with torch.no_grad():
            outputs = self.core_model(**inputs, output_hidden_states=True)
            # Take the final hidden state vector across tokens
            final_hidden = outputs.hidden_states[-1].mean(dim=1)
            
            # 2. Pass the hidden state through the model's native classification head (lm_head)
            logits = self.core_model.lm_head(final_hidden).squeeze()
            
            # 3. Pull out the raw logit scores for our two target paths
            math_score = logits[self.math_token_id].float().item()
            code_score = logits[self.code_token_id].float().item()
            
        # Determine the destination based purely on the AI head weights
        if math_score > code_score:
            expert_name = "Math_Expert"
        else:
            expert_name = "Code_Expert"
            
        print(f"\nPrompt: \"{user_prompt[:50]}...\"")
        print(f"[AI Choice Logits] Math Score: {math_score:.2f} | Code Score: {code_score:.2f}")
        print(f"[Decision Route] Selected Specialist: \033[92m{expert_name}\033[0m")
        
        # 4. Wrap context into ChatML using the designated role instructions
        assigned_system_prompt = self.expert_system_prompts[expert_name]
        messages = [
            {"role": "system", "content": assigned_system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        
        templated_inputs = self.tokenizer.apply_chat_template(
            messages, 
            add_generation_prompt=True, 
            return_tensors="pt"
        ).to(self.device)
        
        # Explicit input tensor extraction fix
        if isinstance(templated_inputs, dict) or hasattr(templated_inputs, "data"):
            input_ids = templated_inputs["input_ids"]
        else:
            input_ids = templated_inputs
            
        # 5. Generate final text response
        outputs = self.core_model.generate(
            input_ids, 
            max_new_tokens=50, 
            temperature=0.2, 
            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__":
    system = AISLMRouterSystem("HuggingFaceTB/SmolLM2-135M-Instruct")
    
    # Test alternative selections
    system.route_and_generate("Calculate the derivative of x^2.")
    system.route_and_generate("Write a fast matrix multiplication script in C.")