#!/usr/bin/env python3
"""
Hybrid Inference: OLMoE-0125 + 18 Micro-Experts
Phase 1: Hint Injection for immediate quality boost
"""

import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
from pathlib import Path
import json
import time
import sys

# Add paths
sys.path.append(str(Path(__file__).parent / 'training' / 'encoder'))
sys.path.append(str(Path(__file__).parent / 'training' / 'distillation'))

from model import create_encoder
from models import ClassificationExpert

class HybridMoESystem:
    """Hierarchical Hybrid MoE: OLMoE Mother + Micro-Experts"""
    
    def __init__(
        self,
        olmoe_model_path: str,
        encoder_path: Path,
        experts_dir: Path,
        device: str = "auto"
    ):
        self.device = device
        
        # Load OLMoE Mother
        print("🔄 Loading OLMoE-0125 Mother...")
        self.olmoe = AutoModelForCausalLM.from_pretrained(
            olmoe_model_path,
            trust_remote_code=True,
            torch_dtype=torch.float16,
            device_map=device
        )
        self.tokenizer = AutoTokenizer.from_pretrained(olmoe_model_path)
        print("✓ OLMoE loaded")
        
        # Load Shared Encoder
        print("\n🔄 Loading Shared Encoder...")
        self.encoder_tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
        self.encoder = create_encoder({
            'vocab_size': self.encoder_tokenizer.vocab_size,
            'embed_dim': 256,
            'num_layers': 4,
            'num_heads': 8,
            'ff_dim': 1024,
            'dropout': 0.1,
            'max_seq_len': 128
        })
        self.encoder.load_state_dict(torch.load(encoder_path, map_location='cpu'))
        self.encoder.eval()
        print("✓ Encoder loaded")
        
        # Load Micro-Experts
        print("\n🔄 Loading Micro-Experts...")
        self.experts = self._load_experts(experts_dir)
        print(f"✓ {len(self.experts)} experts loaded")
    
    def _load_experts(self, experts_dir: Path):
        """Load all trained micro-experts"""
        experts = {}
        
        # Grammar Expert
        grammar_path = experts_dir / 'grammar_expert_v2.pt'
        if grammar_path.exists():
            expert = ClassificationExpert(256, 256, 2, 3, 0.1)
            expert.load_state_dict(torch.load(grammar_path, map_location='cpu'))
            expert.eval()
            experts['grammar'] = {
                'model': expert,
                'name': 'Grammar & Syntax',
                'threshold': 0.7,
                'hint_template': "[GRAMMAR EXPERT]: {advice}"
            }
        
        # Math Expert
        math_path = experts_dir / 'mathematical_reasoning_expert.pt'
        if math_path.exists():
            expert = ClassificationExpert(256, 256, 1, 3, 0.1)
            expert.load_state_dict(torch.load(math_path, map_location='cpu'))
            expert.eval()
            experts['math'] = {
                'model': expert,
                'name': 'Mathematical Reasoning',
                'threshold': 0.8,
                'hint_template': "[MATH EXPERT]: {advice}"
            }
        
        # Sentiment Expert
        sentiment_path = experts_dir / 'sentiment_analysis_expert.pt'
        if sentiment_path.exists():
            expert = ClassificationExpert(256, 256, 4, 3, 0.1)
            expert.load_state_dict(torch.load(sentiment_path, map_location='cpu'))
            expert.eval()
            experts['sentiment'] = {
                'model': expert,
                'name': 'Sentiment Analysis',
                'threshold': 0.7,
                'hint_template': "[SENTIMENT EXPERT]: {advice}"
            }
        
        return experts
    
    def encode_prompt(self, prompt: str):
        """Encode prompt with shared encoder"""
        inputs = self.encoder_tokenizer(
            prompt,
            padding='max_length',
            truncation=True,
            max_length=128,
            return_tensors='pt'
        )
        
        with torch.no_grad():
            embedding = self.encoder(inputs['input_ids'], inputs['attention_mask'])
        
        return embedding
    
    def run_experts(self, embedding):
        """Run all experts in parallel"""
        results = []
        
        for expert_id, expert_data in self.experts.items():
            model = expert_data['model']
            
            with torch.no_grad():
                output_emb, confidence, logits = model(embedding)
            
            # Get confidence score
            probs = torch.softmax(logits, dim=-1)
            max_conf = probs.max().item()
            predicted_class = probs.argmax().item()
            
            results.append({
                'id': expert_id,
                'name': expert_data['name'],
                'confidence': max_conf,
                'prediction': predicted_class,
                'threshold': expert_data['threshold'],
                'template': expert_data['hint_template']
            })
        
        return results
    
    def generate_hints(self, expert_results, prompt):
        """Generate expert hints for high-confidence activations"""
        hints = []
        
        for result in expert_results:
            if result['confidence'] > result['threshold']:
                # Generate domain-specific advice
                if result['id'] == 'grammar':
                    advice = "Ensure perfect grammar and sentence structure."
                elif result['id'] == 'math':
                    advice = "Provide step-by-step mathematical reasoning."
                elif result['id'] == 'sentiment':
                    advice = "Consider emotional tone and sentiment."
                else:
                    advice = f"Apply {result['name']} expertise."
                
                hint = result['template'].format(advice=advice)
                hints.append(hint)
        
        return hints
    
    def hybrid_generate(
        self,
        prompt: str,
        max_new_tokens: int = 200,
        temperature: float = 0.7,
        top_k: int = 4
    ):
        """Hybrid generation with expert hint injection"""
        
        print(f"\n{'='*70}")
        print(f"🎯 Hybrid Inference")
        print(f"{'='*70}")
        print(f"Prompt: {prompt}\n")
        
        # Step 1: Encode prompt
        print("1️⃣ Encoding prompt...")
        embedding = self.encode_prompt(prompt)
        
        # Step 2: Run experts
        print("2️⃣ Running micro-experts...")
        expert_results = self.run_experts(embedding)
        
        # Step 3: Filter by confidence
        active_experts = [
            e for e in expert_results 
            if e['confidence'] > e['threshold']
        ]
        
        print(f"   Active experts: {len(active_experts)}/{len(expert_results)}")
        for exp in active_experts:
            print(f"   - {exp['name']}: {exp['confidence']:.2%}")
        
        # Step 4: Select top-k
        top_experts = sorted(
            active_experts,
            key=lambda x: x['confidence'],
            reverse=True
        )[:top_k]
        
        # Step 5: Generate hints
        print("\n3️⃣ Generating expert hints...")
        hints = self.generate_hints(top_experts, prompt)
        
        # Step 6: Inject into system prompt
        if hints:
            hint_text = "\n".join(hints)
            enriched_prompt = f"""{hint_text}

User: {prompt}
Assistant:"""
        else:
            enriched_prompt = f"User: {prompt}\nAssistant:"
        
        print(f"\n4️⃣ OLMoE generation...")
        
        # Step 7: Generate with OLMoE
        inputs = self.tokenizer(enriched_prompt, return_tensors="pt").to(self.olmoe.device)
        
        start = time.time()
        outputs = self.olmoe.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            do_sample=True
        )
        elapsed = time.time() - start
        
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract assistant response
        if "Assistant:" in response:
            response = response.split("Assistant:")[-1].strip()
        
        tokens_generated = len(outputs[0]) - len(inputs['input_ids'][0])
        tokens_per_sec = tokens_generated / elapsed
        
        print(f"\n{'='*70}")
        print(f"📊 Results")
        print(f"{'='*70}")
        print(f"Response: {response}\n")
        print(f"⚡ Performance:")
        print(f"   Tokens: {tokens_generated}")
        print(f"   Time: {elapsed:.2f}s")
        print(f"   Speed: {tokens_per_sec:.1f} t/s")
        print(f"   Active Experts: {len(top_experts)}")
        print(f"{'='*70}\n")
        
        return {
            'response': response,
            'active_experts': top_experts,
            'tokens_per_sec': tokens_per_sec,
            'hints': hints
        }


def main():
    print("="*70)
    print("🚀 Hybrid MoE Inference System")
    print("   OLMoE-0125 + Micro-Experts")
    print("="*70)
    
    # Initialize system
    system = HybridMoESystem(
        olmoe_model_path="models/olmoe-0125",
        encoder_path=Path("models/encoder_bootstrap.pt"),
        experts_dir=Path("models"),
        device="auto"
    )
    
    # Test prompts
    test_prompts = [
        "What is 15 + 27?",
        "Fix this sentence: Me and him goes to the store yesterday",
        "Write a Python function to calculate factorial",
        "This movie was absolutely terrible and I hated it",
    ]
    
    for prompt in test_prompts:
        result = system.hybrid_generate(prompt)
        time.sleep(1)  # Brief pause between tests


if __name__ == "__main__":
    main()