#!/usr/bin/env python3
"""
LEGO-MoE: Complete Hybrid MoE System
TinyLlama mother model + 23 micro-experts + ISL routing
"""

import sys
from pathlib import Path
from tinyllama_wrapper import TinyLlamaWrapper
from expert_router import ExpertRouter
import torch

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

class LegoMoE:
    """Complete LEGO-MoE inference system"""
    
    def __init__(self):
        print("="*70)
        print("🧠 LEGO-MoE System Initialization")
        print("="*70)
        
        # Load components
        print("\n1. Loading TinyLlama mother model...")
        self.tinyllama = TinyLlamaWrapper()
        
        print("\n2. Loading expert router...")
        self.router = ExpertRouter()
        
        print("\n" + "="*70)
        print("✅ LEGO-MoE System Ready!")
        print("="*70)
    
    def run_expert(self, expert_name: str, query: str):
        """Run a single expert on query"""
        expert = self.router.experts.get(expert_name)
        if not expert:
            return f"Expert {expert_name} not available"
        
        # Encode query
        embedding = self.router.encode_query(query)
        
        # Run expert
        with torch.no_grad():
            _, confidence, logits = expert(embedding)
            prediction = torch.argmax(logits, dim=1).item()
        
        # Format output based on expert type
        if expert_name == "mathematical_reasoning":
            # For math, try to extract answer
            return f"Prediction: {prediction}, Confidence: {confidence.item():.2f}"
        elif expert_name == "grammar":
            return "Correct" if prediction == 0 else "Needs correction"
        else:
            return f"Analysis: {prediction}"
    
    def generate(self, query: str, verbose: bool = True) -> str:
        """
        Generate response using LEGO-MoE system
        
        Args:
            query: User query
            verbose: Show routing decisions
        
        Returns:
            Generated response
        """
        if verbose:
            print(f"\n{'='*70}")
            print(f"Query: {query}")
            print(f"{'='*70}")
        
        # 1. Route with ISL
        routing = self.router.route(query)
        
        if verbose:
            print(f"\n🧭 ISL Decision: {routing['decision']}")
            print(f"📋 Selected Experts: {routing['experts'] if routing['experts'] else 'None (direct to LLM)'}")
        
        # 2. Run experts if needed
        expert_outputs = {}
        if routing['experts']:
            if verbose:
                print(f"\n🔬 Running {len(routing['experts'])} expert(s)...")
            
            for expert_name in routing['experts']:
                output = self.run_expert(expert_name, query)
                expert_outputs[expert_name] = output
                
                if verbose:
                    print(f"  ✓ {expert_name}: {output}")
        
        # 3. Generate with TinyLlama
        if verbose:
            print(f"\n🤖 Generating response with TinyLlama...")
        
        if expert_outputs:
            response = self.tinyllama.generate_with_context(
                query,
                expert_outputs,
                temperature=0.7
            )
        else:
            response = self.tinyllama.generate(
                query,
                temperature=0.7,
                max_tokens=256
            )
        
        if verbose:
            print(f"\n{'='*70}")
            print(f"Response:")
            print(f"{'='*70}")
            print(response)
            print(f"{'='*70}\n")
        
        return response


def main():
    """CLI interface for LEGO-MoE"""
    print("\n🧠 LEGO-MoE - Hybrid Mixture of Experts")
    print("TinyLlama + 23 Micro-Experts + ISL Routing\n")
    
    # Initialize system
    moe = LegoMoE()
    
    # Interactive mode
    if len(sys.argv) > 1:
        # Single query mode
        query = " ".join(sys.argv[1:])
        moe.generate(query)
    else:
        # Interactive mode
        print("\nInteractive Mode (type 'quit' to exit)")
        print("="*70)
        
        while True:
            try:
                query = input("\nYou: ").strip()
                
                if not query:
                    continue
                
                if query.lower() in ['quit', 'exit', 'q']:
                    print("\nGoodbye! 👋")
                    break
                
                moe.generate(query)
            
            except KeyboardInterrupt:
                print("\n\nGoodbye! 👋")
                break
            except Exception as e:
                print(f"\nError: {e}")


if __name__ == "__main__":
    main()