"""
NAPLY Semantic Understanding Module v5.0
=========================================

Deep semantic understanding beyond word-level token prediction.
Implements pattern recognition, contextual meaning, and conceptual understanding.

Features:
- Pattern Recognition: Identifies linguistic patterns and structures
- Contextual Semantics: Understands meaning in context
- Conceptual Embeddings: Maps concepts to semantic spaces
- Grammar & Syntax: Built-in English grammar knowledge
- Semantic Clustering: Groups related concepts
"""

import numpy as np
from typing import Dict, List, Tuple, Optional, Set
from collections import defaultdict
import json
import pickle
import re


class SemanticPattern:
    """
    Represents a linguistic pattern with semantic meaning.
    """
    def __init__(self, pattern_type: str, template: str, meaning: str, 
                 examples: List[str] = None, confidence: float = 0.0):
        self.pattern_type = pattern_type  # 'syntax', 'semantic', 'contextual'
        self.template = template
        self.meaning = meaning
        self.examples = examples or []
        self.confidence = confidence
        self.frequency = 0
    
    def __repr__(self):
        return f"SemanticPattern({self.pattern_type}: {self.meaning})"


class ConceptEmbeddings:
    """
    Semantic concept embeddings that map words to conceptual spaces.
    Enables understanding beyond token-level prediction.
    """
    
    def __init__(self, embedding_dim: int = 256, num_concepts: int = 10000):
        self.embedding_dim = embedding_dim
        self.num_concepts = num_concepts
        
        # Concept space embeddings
        self.concept_embeddings = np.random.randn(num_concepts, embedding_dim) * 0.02
        
        # Word to concept mapping
        self.word_to_concepts: Dict[str, List[int]] = defaultdict(list)
        
        # Concept to semantic meaning
        self.concept_meanings: Dict[int, str] = {}
        
        # Semantic relationships
        self.semantic_relations: Dict[int, List[Tuple[int, float]]] = defaultdict(list)
        
        self._init_english_concepts()
    
    def _init_english_concepts(self):
        """
        Initialize pre-built English semantic concepts from Foundation.
        Covers grammar, syntax, common meanings, and linguistic patterns.
        """
        from .foundation import EnglishFoundation
        
        # Load from Foundation
        self.foundation_vocab = EnglishFoundation.get_vocabulary()
        self.grammar_rules = EnglishFoundation.get_grammar_rules()
        
        # Generate concepts from vocabulary
        concept_idx = 1000
        for word, meta in self.foundation_vocab.items():
            concept_idx += 1
            meaning = meta.get('meaning', 'common word')
            self.concept_meanings[concept_idx] = meaning
            self.concept_embeddings[concept_idx] = self._generate_concept_embedding(word, meaning)
            self.map_word_to_concepts(word, [concept_idx])
            
            # Map forms if verb
            if 'forms' in meta:
                for form in meta['forms']:
                    self.map_word_to_concepts(form, [concept_idx])
            
            # Map synonyms if adj
            if 'synonyms' in meta:
                for syn in meta['synonyms']:
                    self.map_word_to_concepts(syn, [concept_idx])

        # Core Grammar concepts (0-999) - Keep specific indices from before
        grammar_concepts = {
            0: ("subject_verb_agreement", "Ensures subject and verb match in number"),
            1: ("tense_consistency", "Maintains consistent verb tenses"),
        }
        
        for concept_id, (name, meaning) in grammar_concepts.items():
            self.concept_meanings[concept_id] = meaning
            self.concept_embeddings[concept_id] = self._generate_concept_embedding(name, meaning)
    
    def _generate_concept_embedding(self, name: str, meaning: str) -> np.ndarray:
        """Generate embedding from concept name and meaning."""
        # Simple hash-based embedding for now
        text = f"{name} {meaning}"
        hash_val = hash(text) % (2**32)
        np.random.seed(hash_val)
        return np.random.randn(self.embedding_dim) * 0.1
    
    def map_word_to_concepts(self, word: str, concept_ids: List[int]):
        """Map a word to related semantic concepts."""
        self.word_to_concepts[word.lower()].extend(concept_ids)
    
    def get_word_concepts(self, word: str) -> np.ndarray:
        """Get combined concept embedding for a word."""
        concept_ids = self.word_to_concepts.get(word.lower(), [])
        if not concept_ids:
            # Return zero embedding if word unknown
            return np.zeros(self.embedding_dim)
        
        # Average concept embeddings
        embeddings = [self.concept_embeddings[cid] for cid in concept_ids]
        return np.mean(embeddings, axis=0)
    
    def compute_semantic_similarity(self, word1: str, word2: str) -> float:
        """Compute semantic similarity between two words."""
        emb1 = self.get_word_concepts(word1)
        emb2 = self.get_word_concepts(word2)
        
        if np.linalg.norm(emb1) == 0 or np.linalg.norm(emb2) == 0:
            return 0.0
        
        # Cosine similarity
        return float(np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2)))


class PatternRecognizer:
    """
    Recognizes linguistic patterns in text.
    Goes beyond token prediction to understand structural patterns.
    """
    
    def __init__(self):
        self.patterns: Dict[str, List[SemanticPattern]] = defaultdict(list)
        self._init_english_patterns()
    
    def _init_english_patterns(self):
        """Initialize pre-built English language patterns."""
        
        # Syntactic patterns
        syntax_patterns = [
            SemanticPattern(
                pattern_type="syntax",
                template="S V O",
                meaning="Basic Subject-Verb-Object structure",
                examples=["The cat eats food", "She reads books"],
                confidence=0.95
            ),
            SemanticPattern(
                pattern_type="syntax", 
                template="S V",
                meaning="Intransitive verb structure",
                examples=["He sleeps", "Birds sing"],
                confidence=0.90
            ),
            SemanticPattern(
                pattern_type="syntax",
                template="S V O C",
                meaning="Subject-Verb-Object-Complement",
                examples=["They elected him president", "She made me happy"],
                confidence=0.85
            ),
            SemanticPattern(
                pattern_type="syntax",
                template="Wh- S V",
                meaning="Question formation with wh-words",
                examples=["What do you want?", "Where is it?"],
                confidence=0.88
            ),
        ]
        
        # Semantic patterns
        semantic_patterns = [
            SemanticPattern(
                pattern_type="semantic",
                template="If X then Y",
                meaning="Conditional relationship",
                examples=["If it rains, we stay inside", "If you study, you pass"],
                confidence=0.92
            ),
            SemanticPattern(
                pattern_type="semantic",
                template="X because Y",
                meaning="Causal explanation",
                examples=["I'm happy because you came", "It failed because of errors"],
                confidence=0.90
            ),
            SemanticPattern(
                pattern_type="semantic",
                template="X although Y",
                meaning="Contrastive concession",
                examples=["He ran although he was tired", "I went although it rained"],
                confidence=0.85
            ),
            SemanticPattern(
                pattern_type="semantic",
                template="Not only X but also Y",
                meaning="Additive emphasis",
                examples=["Not only smart but also kind", "Not only fast but accurate"],
                confidence=0.82
            ),
        ]
        
        # Contextual patterns
        contextual_patterns = [
            SemanticPattern(
                pattern_type="contextual",
                template="In conclusion, ...",
                meaning="Summarizing final point",
                examples=["In conclusion, we should act now", "In conclusion, the data shows"],
                confidence=0.88
            ),
            SemanticPattern(
                pattern_type="contextual",
                template="On the other hand, ...",
                meaning="Presenting alternative view",
                examples=["On the other hand, we could wait", "On the other hand, it's expensive"],
                confidence=0.87
            ),
            SemanticPattern(
                pattern_type="contextual",
                template="For example, ...",
                meaning="Providing illustration",
                examples=["For example, consider this case", "For example, many people"],
                confidence=0.90
            ),
        ]
        
        # Store all patterns
        for pattern in syntax_patterns + semantic_patterns + contextual_patterns:
            self.patterns[pattern.pattern_type].append(pattern)
    
    def extract_patterns(self, text: str) -> List[SemanticPattern]:
        """Extract patterns from text."""
        found_patterns = []
        
        # Simple pattern matching (in production, this would use more sophisticated NLP)
        for pattern_type, patterns in self.patterns.items():
            for pattern in patterns:
                # Check if pattern template loosely matches
                if self._pattern_matches(text, pattern.template):
                    found_patterns.append(pattern)
        
        return found_patterns
    
    def _pattern_matches(self, text: str, template: str) -> bool:
        """Check if text matches a pattern template (simplified)."""
        # Simplified matching - in real implementation would use parsing
        template_parts = template.replace('...', '').strip().lower().split()
        text_lower = text.lower()
        
        # Check if key words from template appear in text
        matches = 0
        for part in template_parts:
            if part.isalpha() and len(part) > 2 and part in text_lower:
                matches += 1
        
        return matches >= len([p for p in template_parts if p.isalpha()]) * 0.5


class ContextualUnderstanding:
    """
    Maintains and understands context across sequences.
    Tracks discourse flow and maintains coherence.
    """
    
    def __init__(self, max_context_length: int = 2048):
        self.max_context_length = max_context_length
        self.context_history: List[Dict] = []
        self.entity_tracker: Dict[str, Dict] = {}
        self.discourse_stack: List[str] = []
    
    def update_context(self, text: str, embedding: Optional[np.ndarray] = None):
        """Update context with new text."""
        context_entry = {
            'text': text,
            'embedding': embedding,
            'entities': self._extract_entities(text),
            'timestamp': len(self.context_history)
        }
        
        self.context_history.append(context_entry)
        
        # Update entity tracker
        for entity in context_entry['entities']:
            if entity not in self.entity_tracker:
                self.entity_tracker[entity] = {
                    'first_seen': len(self.context_history) - 1,
                    'mentions': [],
                    'type': self._guess_entity_type(entity)
                }
            self.entity_tracker[entity]['mentions'].append(len(self.context_history) - 1)
        
        # Trim history if too long
        if len(self.context_history) > self.max_context_length:
            self.context_history = self.context_history[-self.max_context_length:]
    
    def _extract_entities(self, text: str) -> List[str]:
        """Extract named entities from text (simplified)."""
        # Simple entity extraction - capitalize words might be entities
        words = text.split()
        entities = []
        current_entity = []
        
        for word in words:
            clean_word = word.strip('.,!?;:"()[]{}')
            if clean_word and clean_word[0].isupper() and len(clean_word) > 1:
                current_entity.append(clean_word)
            else:
                if current_entity:
                    entities.append(' '.join(current_entity))
                    current_entity = []
        
        if current_entity:
            entities.append(' '.join(current_entity))
        
        return entities
    
    def _guess_entity_type(self, entity: str) -> str:
        """Guess the type of an entity."""
        # Simple heuristic
        if entity.lower() in ['monday', 'tuesday', 'wednesday', 'thursday', 'friday', 'saturday', 'sunday',
                             'january', 'february', 'march', 'april', 'may', 'june',
                             'july', 'august', 'september', 'october', 'november', 'december']:
            return 'temporal'
        elif entity[0].isupper() and entity.endswith('s') and len(entity) < 10:
            return 'organization'
        else:
            return 'unknown'
    
    def get_relevant_context(self, query: str, top_k: int = 5) -> List[Dict]:
        """Retrieve relevant context for a query."""
        if not self.context_history:
            return []
        
        # Simple relevance scoring
        scored_contexts = []
        query_words = set(query.lower().split())
        
        for i, entry in enumerate(self.context_history):
            entry_words = set(entry['text'].lower().split())
            overlap = len(query_words & entry_words)
            recency_score = i / len(self.context_history)  # More recent = higher
            
            score = overlap * 0.7 + recency_score * 0.3
            scored_contexts.append((score, entry))
        
        # Sort by score and return top_k
        scored_contexts.sort(key=lambda x: x[0], reverse=True)
        return [entry for _, entry in scored_contexts[:top_k]]


class SemanticEngine:
    """
    Main semantic understanding engine that combines all components.
    """
    
    def __init__(self, embedding_dim: int = 256):
        self.concept_embeddings = ConceptEmbeddings(embedding_dim=embedding_dim)
        self.pattern_recognizer = PatternRecognizer()
        self.contextual_understanding = ContextualUnderstanding()
        
        # Semantic cache for fast lookups
        self.semantic_cache: Dict[str, Dict] = {}
        
        self._init_english_vocabulary()
    
    def _init_english_vocabulary(self):
        """Initialize pre-built English vocabulary with semantic mappings."""
        # Common words mapped to concepts
        vocabulary_mappings = {
            # Grammar words
            'the': [2], 'a': [2], 'an': [2],
            'is': [0, 1], 'are': [0], 'was': [0, 1], 'were': [0],
            'have': [1], 'has': [1], 'had': [1],
            'do': [6], 'does': [6], 'did': [6],
            'will': [1], 'would': [8], 'could': [8], 'should': [8],
            'can': [6], 'may': [6], 'might': [6],
            
            # Semantic words
            'because': [1001], 'since': [1001], 'therefore': [1000],
            'although': [1004], 'however': [1004], 'but': [1004],
            'like': [1005], 'similar': [1005], 'same': [1005],
            'if': [1006], 'when': [1006], 'unless': [1006],
            'before': [1001], 'after': [1001], 'during': [1001],
            'above': [1002], 'below': [1002], 'between': [1002],
            'better': [1003], 'best': [1003], 'worse': [1003],
            'why': [1007], 'how': [1009], 'what': [7],
            
            # Contextual words
            'furthermore': [2004], 'moreover': [2004], 'additionally': [2004],
            'conversely': [2004], 'nevertheless': [2004], 'nonetheless': [2004],
            'specifically': [2005], 'particularly': [2005], 'especially': [2005],
            'for': [2007], 'thus': [2000], 'hence': [2000],
        }
        
        for word, concept_ids in vocabulary_mappings.items():
            self.concept_embeddings.map_word_to_concepts(word, concept_ids)
    
    def analyze_text(self, text: str) -> Dict:
        """
        Perform comprehensive semantic analysis of text.
        
        Returns dictionary with:
        - patterns: Recognized linguistic patterns
        - concepts: Semantic concepts present
        - context: Relevant context information
        - meaning: Overall semantic interpretation
        """
        # Extract patterns
        patterns = self.pattern_recognizer.extract_patterns(text)
        
        # Identify concepts in text
        words = text.lower().split()
        concepts_found = {}
        for word in words:
            word_clean = word.strip('.,!?;:"()[]{}')
            concept_emb = self.concept_embeddings.get_word_concepts(word_clean)
            if np.linalg.norm(concept_emb) > 0:
                concepts_found[word_clean] = concept_emb
        
        # Update and retrieve context
        self.contextual_understanding.update_context(text)
        relevant_context = self.contextual_understanding.get_relevant_context(text)
        
        # Cache result
        analysis = {
            'patterns': patterns,
            'concepts': concepts_found,
            'context': relevant_context,
            'text': text,
            'semantic_complexity': len(patterns) + len(concepts_found)
        }
        
        self.semantic_cache[text] = analysis
        return analysis
    
    def get_semantic_embedding(self, text: str) -> np.ndarray:
        """Get semantic embedding for text combining all understanding."""
        # Check cache
        if text in self.semantic_cache:
            analysis = self.semantic_cache[text]
        else:
            analysis = self.analyze_text(text)
        
        # Combine concept embeddings
        if analysis['concepts']:
            concept_embs = list(analysis['concepts'].values())
            combined = np.mean(concept_embs, axis=0)
        else:
            combined = np.zeros(self.concept_embeddings.embedding_dim)
        
        return combined
    
    def understand_meaning(self, text: str) -> str:
        """
        Generate human-readable semantic interpretation.
        """
        analysis = self.analyze_text(text)
        
        interpretations = []
        
        # Add pattern interpretations
        for pattern in analysis['patterns']:
            interpretations.append(f"{pattern.pattern_type}: {pattern.meaning}")
        
        # Add concept interpretations
        for concept_word in list(analysis['concepts'].keys())[:3]:
            interpretations.append(f"Concept '{concept_word}' detected")
        
        if not interpretations:
            return "Text appears to be a simple statement without complex semantic patterns."
        
        return " | ".join(interpretations)


class SemanticMemory:
    """
    Maintains semantic memory across training and inference.
    Enables the model to build up conceptual understanding over time.
    """
    
    def __init__(self, memory_size: int = 10000):
        self.memory_size = memory_size
        self.episodic_memory: List[Dict] = []  # Specific instances
        self.semantic_memory: Dict[str, Dict] = {}  # Generalized concepts
        self.procedural_memory: Dict[str, callable] = {}  # Learned procedures
    
    def store_episode(self, text: str, context: Dict, outcome: Optional[str] = None):
        """Store a specific episode in memory."""
        episode = {
            'text': text,
            'context': context,
            'outcome': outcome,
            'timestamp': len(self.episodic_memory)
        }
        
        self.episodic_memory.append(episode)
        
        # Consolidate to semantic memory periodically
        if len(self.episodic_memory) % 100 == 0:
            self._consolidate_memory()
        
        # Trim if too large
        if len(self.episodic_memory) > self.memory_size:
            self.episodic_memory = self.episodic_memory[-self.memory_size:]
    
    def _consolidate_memory(self):
        """Consolidate episodic memories into semantic memory."""
        # Group similar episodes
        # In production, would use clustering algorithms
        pass
    
    def retrieve_similar(self, query: str, top_k: int = 5) -> List[Dict]:
        """Retrieve memories similar to query."""
        query_lower = query.lower()
        scored_memories = []
        
        for episode in self.episodic_memory:
            text_lower = episode['text'].lower()
            
            # Simple similarity score
            query_words = set(query_lower.split())
            text_words = set(text_lower.split())
            overlap = len(query_words & text_words)
            
            if overlap > 0:
                scored_memories.append((overlap, episode))
        
        # Sort and return top_k
        scored_memories.sort(key=lambda x: x[0], reverse=True)
        return [ep for _, ep in scored_memories[:top_k]]
    
    def save(self, path: str):
        """Save memory to disk."""
        data = {
            'episodic': self.episodic_memory,
            'semantic': self.semantic_memory
        }
        with open(path, 'wb') as f:
            pickle.dump(data, f)
    
    def load(self, path: str):
        """Load memory from disk."""
        with open(path, 'rb') as f:
            data = pickle.load(f)
            self.episodic_memory = data.get('episodic', [])
            self.semantic_memory = data.get('semantic', {})


# Pre-built English foundation with full grammar, vocabulary, and fluency
ENGLISH_FOUNDATION = {
    'grammar_rules': {
        'subject_verb_agreement': {
            'singular': ['is', 'was', 'has', 'does', 's'],
            'plural': ['are', 'were', 'have', 'do', '']
        },
        'tense_forms': {
            'present_simple': {'suffix': 's', 'aux': 'do/does'},
            'past_simple': {'suffix': 'ed', 'aux': 'did'},
            'present_perfect': {'aux': 'have/has', 'suffix': 'ed'},
            'future': {'aux': 'will', 'suffix': ''}
        },
        'article_rules': {
            'definite': 'the',
            'indefinite_vowel': 'an',
            'indefinite_consonant': 'a'
        }
    },
    'common_phrases': {
        'greetings': ['Hello', 'Hi', 'Good morning', 'Good afternoon', 'How are you?'],
        'farewells': ['Goodbye', 'See you', 'Take care', 'Bye', 'Have a nice day'],
        'thanks': ['Thank you', 'Thanks', 'I appreciate it', 'Many thanks'],
        'apologies': ['Sorry', 'I apologize', 'My apologies', 'Excuse me'],
        'agreement': ['Yes', 'I agree', 'Exactly', 'That is right', 'Sure'],
        'disagreement': ['No', 'I disagree', 'That is not right', 'Actually'],
    },
    'sentence_structures': {
        'declarative': 'Subject + Verb + (Object)',
        'interrogative': '(Wh-word) + Auxiliary + Subject + Verb',
        'imperative': 'Verb (+ Object)',
        'exclamatory': '(What/How) + Subject + Verb + !'
    }
}


def get_english_foundation() -> Dict:
    from .foundation import EnglishFoundation
    return {
        'grammar_rules': EnglishFoundation.get_grammar_rules(),
        'common_phrases': {}, # Kept for API compatibility, filled dynamically if needed
        'sentence_structures': EnglishFoundation.get_grammar_rules().get('tenses', {})
    }


# Export main classes
__all__ = [
    'SemanticEngine',
    'SemanticPattern',
    'ConceptEmbeddings',
    'PatternRecognizer',
    'ContextualUnderstanding',
    'SemanticMemory',
    'get_english_foundation'
]