#!/usr/bin/env python3
"""
Enhanced Expert Router - Loads All 23 Trained Experts
"""

import torch
from pathlib import Path
import sys

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

class EnhancedExpertRouter:
    """Enhanced router with all 23 experts loaded"""
    
    def __init__(self):
        print("Loading Enhanced Expert Router...")
        self.experts = {}
        self.expert_configs = self._get_all_expert_configs()
        self._load_all_experts()
        print(f"✓ Loaded {len(self.experts)} experts")
    
    def _get_all_expert_configs(self):
        """Get configurations for all 23 experts"""
        return {
            # General Experts (15)
            "grammar": {"num_classes": 2, "category": "general"},
            "grammar_v2": {"num_classes": 2, "category": "general"},
            "mathematical_reasoning": {"num_classes": 2, "category": "general"},
            "logical_inference": {"num_classes": 2, "category": "general"},
            "common_sense": {"num_classes": 2, "category": "general"},
            "question_answering": {"num_classes": 2, "category": "general"},
            "code_generation": {"num_classes": 2, "category": "general"},
            "sentiment_analysis": {"num_classes": 3, "category": "general"},
            "toxicity_detection": {"num_classes": 2, "category": "general"},
            "semantic_understanding": {"num_classes": 2, "category": "general"},
            "summarization": {"num_classes": 2, "category": "general"},
            "ner": {"num_classes": 9, "category": "general"},
            "translation": {"num_classes": 2, "category": "general"},
            "causal_reasoning": {"num_classes": 2, "category": "general"},
            
            # Security Experts (5)
            "attack_surface_mapper": {"num_classes": 2, "category": "security"},
            "authorization_logic": {"num_classes": 2, "category": "security"},
            "data_flow_analysis": {"num_classes": 2, "category": "security"},
            "business_logic": {"num_classes": 2, "category": "security"},
            "exploit_feasibility": {"num_classes": 2, "category": "security"},
            
            # Meta-Cognitive (1)
            "isl_meta-governor": {"num_classes": 3, "category": "meta"},
        }
    
    def _load_all_experts(self):
        """Load all available expert models"""
        models_dir = Path("models")
        
        for expert_name, config in self.expert_configs.items():
            model_path = models_dir / f"{expert_name}_expert.pt"
            
            if not model_path.exists():
                print(f"  ⚠️  {expert_name}: model file not found")
                continue
            
            try:
                expert = ClassificationExpert(
                    input_dim=256,
                    hidden_dim=256,
                    num_layers=3,
                    num_classes=config["num_classes"],
                    dropout=0.1
                )
                expert.load_state_dict(torch.load(model_path, map_location='cpu'))
                expert.eval()
                
                self.experts[expert_name] = {
                    "model": expert,
                    "category": config["category"],
                    "num_classes": config["num_classes"]
                }
                
                print(f"  ✓ {expert_name} ({config['category']})")
                
            except Exception as e:
                print(f"  ✗ {expert_name}: {str(e)[:50]}")
    
    def get_expert(self, name: str):
        """Get expert by name"""
        return self.experts.get(name)
    
    def get_experts_by_category(self, category: str):
        """Get all experts in a category"""
        return {
            name: info for name, info in self.experts.items()
            if info["category"] == category
        }
    
    def get_summary(self):
        """Get summary of loaded experts"""
        by_category = {}
        for name, info in self.experts.items():
            cat = info["category"]
            if cat not in by_category:
                by_category[cat] = []
            by_category[cat].append(name)
        
        return {
            "total": len(self.experts),
            "by_category": by_category
        }


def test_enhanced_router():
    """Test enhanced router"""
    print("="*70)
    print("Testing Enhanced Expert Router")
    print("="*70)
    
    router = EnhancedExpertRouter()
    
    print("\n" + "="*70)
    print("Summary:")
    summary = router.get_summary()
    print(f"Total experts loaded: {summary['total']}")
    
    for category, experts in summary['by_category'].items():
        print(f"\n{category.upper()} ({len(experts)}):")
        for expert in sorted(experts):
            print(f"  - {expert}")
    
    print("\n" + "="*70)
    print("✓ Enhanced router working!")
    print("="*70)


if __name__ == "__main__":
    test_enhanced_router()