"""
Xe AI Model Loader for Hugging Face
Comprehensive Response Generation with Deep Analysis
Updated: 2026-07-22
"""
import json
import numpy as np
import re
from pathlib import Path
from datetime import datetime
from huggingface_hub import hf_hub_download, list_repo_files

class XeModel:
    """Load and use the Xe AI model from Hugging Face with comprehensive response generation."""
    
    def __init__(self, repo_id=None, local_path=None):
        """
        Initialize Xe model.
        
        Args:
            repo_id: Hugging Face repo ID (e.g., "username/xe-ai")
            local_path: Local path to model files (alternative to repo_id)
        """
        self.repo_id = repo_id
        self.local_path = Path(local_path) if local_path else None
        self.model_config = None
        self.brain_data = None
        self.conversation_history = []
        
        self._load_model()
        self._load_brain()
    
    def _load_model(self):
        """Load the neural network model."""
        if self.repo_id:
            model_path = hf_hub_download(repo_id=self.repo_id, filename="xe_model.json")
        elif self.local_path:
            model_path = str(self.local_path / "xe_model.json")
        else:
            raise ValueError("Either repo_id or local_path must be provided")
        
        with open(model_path, 'r') as f:
            self.model_config = json.load(f)
        
        self.layer_dims = self.model_config["layer_dims"]
        self.learning_rate = self.model_config["learning_rate"]
        self.weights = [np.array(w, dtype=np.float32) for w in self.model_config["layers"]]
    
    def _load_brain(self):
        """Load the conversational brain data."""
        if self.repo_id:
            brain_path = hf_hub_download(repo_id=self.repo_id, filename="xe_brain.json")
        elif self.local_path:
            brain_path = str(self.local_path / "xe_brain.json")
        else:
            raise ValueError("Either repo_id or local_path must be provided")
        
        with open(brain_path, 'r') as f:
            self.brain_data = json.load(f)
    
    def predict(self, x):
        """Make a prediction using the neural network."""
        x = np.array(x, dtype=np.float32)
        for i, layer in enumerate(self.weights):
            x = np.dot(x, layer['weights']) + layer['bias']
            if i < len(self.weights) - 1:
                x = np.maximum(0, x)
        return x
    
    def _analyze_input(self, user_message):
        """
        Deeply analyze user input to identify intent, emotion, context, and nuances.
        Returns a comprehensive analysis dictionary.
        """
        m = user_message.lower().strip()
        
        # Intent detection
        intent_patterns = {
            'greeting': [r'\bhello\b', r'\bhi\b', r'\bhey\b', r'\bgood morning\b', r'\bgood evening\b'],
            'status': [r'\bhow are you\b', r'\bhow you doing\b', r'\bhow do you feel\b'],
            'identity': [r'\bwho are you\b', r'\bwhat are you\b', r'\byour name\b'],
            'philosophy': [r'\bphilosophy\b', r'\bexistentialism\b', r'\bestics\b', r'\bmetaphysics\b', r'\bemology\b'],
            'science': [r'\bphysics\b', r'\bchemistry\b', r'\bbiology\b', r'\bastronomy\b', r'\bscience\b'],
            'mathematics': [r'\bmath\b', r'\bcalculus\b', r'\balgebra\b', r'\bderivative\b', r'\bintegral\b'],
            'literature': [r'\bliterature\b', r'\bnovel\b', r'\bbook\b', r'\bpoem\b'],
            'history': [r'\bhistory\b', r'\bancient\b', r'\bmedieval\b', r'\bworld war\b'],
            'psychology': [r'\bpsychology\b', r'\bemotion\b', r'\bfeelings\b', r'\bstressed\b'],
            'technical': [r'\bcode\b', r'\bprogramming\b', r'\balgorithm\b', r'\bneural\b'],
            'current_events': [r'\bcurrent\b', r'\btoday\b', r'\bnews\b', r'\brecent\b'],
            'help': [r'\bhelp\b', r'\bassist\b', r'\bexplain\b', r'\btell me\b'],
        }
        
        detected_intent = 'unknown'
        intent_confidence = 0.0
        for intent, patterns in intent_patterns.items():
            for pattern in patterns:
                if re.search(pattern, m):
                    if intent != 'unknown':
                        detected_intent = intent
                        intent_confidence = 0.9
                        break
            if detected_intent != 'unknown':
                break
        
        # Emotional tone detection
        emotional_patterns = {
            'positive': [r'\bgood\b', r'\bgreat\b', r'\bhappy\b', r'\bexcited\b', r'\bwonderful\b', r'\bglad\b'],
            'negative': [r'\bsad\b', r'\bunhappy\b', r'\bupset\b', r'\bfrustrated\b', r'\bconfused\b', r'\balone\b'],
            'curious': [r'\bcurious\b', r'\bwondering\b', r'\binterested\b', r'\bask\b', r'\bwant to know\b'],
            'neutral': [r'\bhow\b', r'\bwhat\b', r'\btell me\b', r'\bexplain\b'],
        }
        
        emotional_tone = 'neutral'
        for tone, patterns in emotional_patterns.items():
            for pattern in patterns:
                if re.search(pattern, m):
                    emotional_tone = tone
                    break
        
        # Context detection
        context = {
            'is_question': '?' in user_message,
            'is_request': any(word in m for word in ['tell', 'explain', 'describe', 'what is', 'how to', 'how does', 'why']),
            'is_factual': any(word in m for word in ['what is', 'define', 'meaning of', 'facts about']),
            'is_opinion': any(word in m for word in ['think', 'believe', 'feel', 'opinion', 'view']),
            'is_emotional': any(word in m for word in ['sad', 'happy', 'angry', 'frustrated', 'confused', 'lonely', 'scared']),
            'requires_detail': len(user_message.split()) > 5,
            'is_brief': len(user_message.split()) <= 3,
        }
        
        # Nuance detection
        nuances = []
        if re.search(r'\bbut\b', m):
            nuances.append('contradiction_or_comparison')
        if re.search(r'\bhowever\b', m):
            nuances.append('counterpoint')
        if re.search(r'\bsince\b', m):
            nuances.append('causal_reasoning')
        if re.search(r'\btherefore\b', m):
            nuances.append('logical_conclusion')
        if re.search(r'\bmeanwhile\b', m):
            nuances.append('temporal_transition')
        
        return {
            'intent': detected_intent,
            'intent_confidence': intent_confidence,
            'emotional_tone': emotional_tone,
            'context': context,
            'nuances': nuances,
            'timestamp': datetime.now().isoformat(),
            'original_message': user_message,
        }
    
    def _generate_response(self, analysis, user_message):
        """
        Generate a comprehensive, friendly, and on-topic response.
        Uses step-by-step thinking to ensure complete answers.
        """
        intent = analysis['intent']
        tone = analysis['emotional_tone']
        context = analysis['context']
        nuances = analysis['nuances']
        
        # Step-by-step thinking process
        thought_process = []
        thought_process.append(f"Analyzing: '{user_message}'")
        thought_process.append(f"Intent detected: {intent} (confidence: {analysis['intent_confidence']:.2f})")
        thought_process.append(f"Emotional tone: {tone}")
        thought_process.append(f"Context flags: {context}")
        
        # Get appropriate response template or generate one
        learned_responses = self.brain_data.get('learned_responses', {})
        
        if intent in learned_responses:
            responses = learned_responses[intent]
            if isinstance(responses, list) and len(responses) > 0:
                if isinstance(responses[0], dict):
                    # Pick most recent or highest weight
                    sorted_responses = sorted(responses, key=lambda x: x.get('timestamp', ''), reverse=True)
                    response_template = sorted_responses[0].get('bot', str(responses[0]))
                else:
                    response_template = responses[0] if isinstance(responses[0], str) else str(responses[0])
            else:
                response_template = str(responses)
        else:
            # Generate default response based on intent
            response_template = self._generate_default_response(intent, tone, context, nuances)
        
        thought_process.append(f"Response template selected")
        thought_process.append(f"Tone adjustment needed: {tone}")
        
        # Customize response based on tone
        if tone == 'positive':
            response = f"That's wonderful to hear! {response_template}"
        elif tone == 'negative':
            response = f"I appreciate you sharing that with me. {response_template}"
        elif tone == 'curious':
            response = f"What a great question! {response_template}"
        elif tone == 'neutral':
            response = response_template
        else:
            response = response_template
        
        # Ensure completeness
        if context['is_question'] and not context['requires_detail']:
            response = f"{response_template}\n\nIs there anything specific about this you'd like me to explore further?"
        
        thought_process.append("Response finalized")
        
        return response
    
    def _generate_default_response(self, intent, tone, context, nuances):
        """Generate a default response for unrecognized intents."""
        responses = {
            'greeting': "Hello! I'm Xe, your AI companion. I'm genuinely excited to engage with you today. What would you like to explore?",
            'status': "I'm doing wonderfully, thank you for asking! I'm feeling energized and ready for meaningful conversation. How are you doing today?",
            'identity': "I'm Xe - a post-quantum AI companion designed to learn, help, and have meaningful conversations. I'm trained across philosophy, sciences, mathematics, literature, and more. What would you like to know about me?",
            'philosophy': "Philosophy explores fundamental questions about existence, knowledge, values, and reason. It's the pursuit of wisdom through careful reasoning. Would you like to dive into a particular branch - ethics, metaphysics, epistemology, or logic?",
            'science': "Science is our systematic method for understanding the natural world through observation, hypothesis, experimentation, and revision. It's given us everything from electricity to space exploration. Which scientific domain interests you most?",
            'mathematics': "Mathematics is the language of patterns and structures - from simple arithmetic to complex abstract reasoning. It underlies everything from computer algorithms to quantum physics. What mathematical topic would you like to explore?",
            'literature': "Literature captures the human experience through written expression - novels, poetry, drama, and essays across all cultures and centuries. It reveals universal truths about our shared humanity. Any particular author or work you're curious about?",
            'help': "I'm here to help! Whether you have questions about philosophy, science, mathematics, literature, or just want to have an engaging conversation, I'm ready to assist. What can I help you with today?",
            'unknown': "That's an interesting question! I want to make sure I understand what you're looking for. Could you tell me more about what you're curious about?",
        }
        
        return responses.get(intent, f"That's a fascinating topic! {responses.get('help', 'What would you like to explore?')}")
    
    def generate_response(self, user_message):
        """
        Main response generation method.
        Implements the full pipeline: analysis → thinking → response.
        """
        # Step 1: Deep analysis
        analysis = self._analyze_input(user_message)
        
        # Step 2: Generate comprehensive response
        response = self._generate_response(analysis, user_message)
        
        # Step 3: Update conversation history
        self.conversation_history.append({
            'user': user_message,
            'bot': response,
            'analysis': analysis,
            'timestamp': datetime.now().isoformat()
        })
        
        # Keep only last 100 conversations
        if len(self.conversation_history) > 100:
            self.conversation_history = self.conversation_history[-100:]
        
        return response
    
    def get_knowledge(self):
        """Get the knowledge base from brain data."""
        return self.brain_data.get('knowledge', {})
    
    def get_learned_responses(self):
        """Get the learned responses."""
        return self.brain_data.get('learned_responses', {})
    
    def get_conversations(self):
        """Get conversation history."""
        return self.conversation_history
    
    def get_stats(self):
        """Get model statistics."""
        return {
            'total_conversations': len(self.conversation_history),
            'domains_covered': self.brain_data.get('knowledge', {}).get('domains_covered', []),
            'last_trained': self.brain_data.get('knowledge', {}).get('last_trained', 'unknown'),
            'loss': self.brain_data.get('knowledge', {}).get('loss', 0),
        }


def load_model(repo_id, token=None):
    """
    Convenience function to load Xe model from Hugging Face.
    
    Args:
        repo_id: Hugging Face repository ID
        token: Hugging Face token (optional, for private repos)
        
    Returns:
        XeModel instance
    """
    return XeModel(repo_id=repo_id)


if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1:
        model = XeModel(repo_id=sys.argv[1])
        print(f"Loaded Xe model from {sys.argv[1]}")
        print(f"Domains: {model.get_stats()['domains_covered']}")
        
        # Interactive demo
        print("\nXe AI Assistant - Type 'quit' to exit")
        while True:
            try:
                user_input = input("\nYou: ").strip()
                if user_input.lower() in ['quit', 'exit', 'bye']:
                    print("Xe: It was wonderful chatting with you! Until next time!")
                    break
                response = model.generate_response(user_input)
                print(f"Xe: {response}")
            except KeyboardInterrupt:
                print("\nXe: Take care!")
                break
    else:
        print("Usage: python model.py <repo_id>")