#!/usr/bin/env python3
"""
Expert Router for LEGO-MoE
Routes queries to appropriate experts based on ISL decisions
"""

import torch
import re
from pathlib import Path
import sys

sys.path.append(str(Path(__file__).parent / 'training' / 'distillation'))
from models import ClassificationExpert

class ExpertRouter:
    """Routes queries to appropriate micro-experts"""
    
    def __init__(self):
        print("Loading expert router...")
        self.isl = self.load_isl_expert()
        self.experts = self.load_all_experts()
        print(f"✓ Router loaded with {len(self.experts)} experts")
    
    def load_isl_expert(self):
        """Load ISL meta-governor"""
        try:
            expert = ClassificationExpert(
                input_dim=256,
                hidden_dim=256,
                num_layers=3,
                num_classes=3,
                dropout=0.1
            )
            expert.load_state_dict(torch.load("models/isl_meta-governor_expert.pt", map_location='cpu'))
            expert.eval()
            return expert
        except Exception as e:
            print(f"⚠️  Could not load ISL expert: {e}")
            return None
    
    def load_all_experts(self):
        """Load all available micro-experts"""
        experts = {}
        
        expert_configs = [
            ("grammar", 2),
            ("mathematical_reasoning", 2),
            ("question_answering", 2),
            ("semantic_understanding", 2),
        ]
        
        for name, num_classes in expert_configs:
            try:
                expert = ClassificationExpert(
                    input_dim=256,
                    hidden_dim=256,
                    num_layers=3,
                    num_classes=num_classes,
                    dropout=0.1
                )
                expert.load_state_dict(torch.load(f"models/{name}_expert.pt", map_location='cpu'))
                expert.eval()
                experts[name] = expert
            except:
                continue
        
        return experts
    
    def encode_query(self, query: str) -> torch.Tensor:
        """Simple query encoding"""
        features = torch.zeros(1, 256)
        
        # Set features based on query patterns
        query_lower = query.lower()
        
        # Math indicators
        if any(word in query_lower for word in ['calculate', 'compute', 'solve', 'math']):
            features[0, 0] = 1.0
        if any(char in query for char in ['+', '-', '*', '/', '=', '%']):
            features[0, 1] = 1.0
        if re.search(r'\d+', query):
            features[0, 2] = 1.0
        
        # Grammar indicators
        if any(word in query_lower for word in ['grammar', 'correct', 'fix', 'spell']):
            features[0, 3] = 1.0
        
        # Question indicators
        if query.strip().endswith('?'):
            features[0, 4] = 1.0
        if any(word in query_lower for word in ['what', 'why', 'how', 'when', 'where', 'who']):
            features[0, 5] = 1.0
        
        # Code indicators
        if any(word in query_lower for word in ['code', 'function', 'python', 'javascript']):
            features[0, 6] = 1.0
        if any(char in query for char in ['(', ')', '{', '}', '[', ']', ';']):
            features[0, 7] = 1.0
        
        return features
    
    def isl_decision(self, query: str) -> str:
        """ISL expert decides complexity"""
        if not self.isl:
            return "ESCALATE_FULL"  # Default if ISL not available
        
        embedding = self.encode_query(query)
        
        with torch.no_grad():
            _, confidence, logits = self.isl(embedding)
            prediction = torch.argmax(logits, dim=1).item()
        
        decisions = ["TERMINATE_SAFE", "ESCALATE_FULL", "ESCALATE_PARTIAL"]
        decision = decisions[prediction] if prediction < 3 else "ESCALATE_FULL"
        
        return decision
    
    def select_experts(self, query: str) -> list:
        """Select which experts to run"""
        selected = []
        query_lower = query.lower()
        
        # Math expert
        if any(word in query_lower for word in ['calculate', 'compute', 'solve', 'math', 'multiply', 'add', 'subtract']):
            if "mathematical_reasoning" in self.experts:
                selected.append("mathematical_reasoning")
        
        if re.search(r'\d+\s*[\+\-\*\/]\s*\d+', query):
            if "mathematical_reasoning" in self.experts:
                selected.append("mathematical_reasoning")
        
        # Grammar expert
        if any(word in query_lower for word in ['grammar', 'correct', 'fix', 'spell', 'punctuation']):
            if "grammar" in self.experts:
                selected.append("grammar")
        
        # QA expert
        if query.strip().endswith('?'):
            if "question_answering" in self.experts:
                selected.append("question_answering")
        
        # Semantic understanding for complex queries
        if len(query.split()) > 10:
            if "semantic_understanding" in self.experts:
                selected.append("semantic_understanding")
        
        return list(set(selected))  # Remove duplicates
    
    def route(self, query: str) -> dict:
        """
        Route query to appropriate experts
        
        Returns:
            {
                'decision': ISL decision,
                'experts': list of expert names to run,
                'direct_to_llm': bool
            }
        """
        # ISL decision
        decision = self.isl_decision(query)
        
        # If terminate safe, go direct to LLM
        if decision == "TERMINATE_SAFE":
            return {
                'decision': decision,
                'experts': [],
                'direct_to_llm': True
            }
        
        # Select experts
        experts = self.select_experts(query)
        
        return {
            'decision': decision,
            'experts': experts,
            'direct_to_llm': len(experts) == 0
        }


def test_router():
    """Test expert router"""
    print("="*70)
    print("Testing Expert Router")
    print("="*70)
    
    router = ExpertRouter()
    
    test_queries = [
        "What is 15 * 23?",
        "Fix this sentence: I goes to school yesterday",
        "What is the capital of France?",
        "Hello, how are you?",
    ]
    
    for query in test_queries:
        print(f"\nQuery: {query}")
        result = router.route(query)
        print(f"  ISL Decision: {result['decision']}")
        print(f"  Selected Experts: {result['experts']}")
        print(f"  Direct to LLM: {result['direct_to_llm']}")
    
    print("\n" + "="*70)
    print("✓ Router working!")
    print("="*70)


if __name__ == "__main__":
    test_router()