"""Word2Vec-based normalizer with automatic model downloading."""

import re
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union

import numpy as np
from gensim.models import KeyedVectors

from yakit.normalizers.base import Normalizer
from yakit.utils.download import download_model, download_training_data, get_model_path


class Word2VecNormalizer(Normalizer):
    """
    Word2Vec-based normalizer for Yakut language text.
    
    This normalizer combines:
    1. Word-level context using word2vec embeddings
    2. Character-level patterns from training data
    3. Similar word matching to find best replacements
    
    Example:
        >>> from yakit.normalizers import Word2VecNormalizer
        >>> 
        >>> # Auto-download model and training data
        >>> normalizer = Word2VecNormalizer()
        >>> 
        >>> # Or provide custom paths
        >>> normalizer = Word2VecNormalizer(
        ...     word2vec_path="/path/to/model.bin",
        ...     training_data_path="/path/to/train_pairs.txt"
        ... )
        >>> 
        >>> text = "Мин сахалыы билэбин"
        >>> normalized = normalizer.normalize(text)
    """

    # Character mappings for Yakut normalization
    CHAR_MAPPINGS = {
        'h': 'һ', 'H': 'Һ',
        'ь': 'һ', 'Ь': 'Һ',
        'г': 'ҕ', 'Г': 'Ҕ',
        'н': 'ҥ', 'Н': 'Ҥ',
        'о': 'ө', 'О': 'Ө',
        'у': 'ү', 'У': 'Ү',
    }

    def __init__(
        self,
        word2vec_path: Optional[Union[str, Path]] = None,
        training_data_path: Optional[Union[str, Path]] = None,
        training_pairs: Optional[List[Tuple[str, str]]] = None,
        *,
        # Model parameters
        context_window: int = 3,
        min_pattern_count: int = 1,
        word_similarity_threshold: float = 0.5,
        max_similar_words: int = 10,
        # Confidence thresholds
        replace_confidence_threshold: float = 1.2,
        no_replace_confidence_threshold: float = 1.5,
        rare_pattern_threshold: int = 5,
        rare_pattern_confidence_threshold: float = 2.0,
        word_similarity_high_threshold: float = 0.7,
        # Download options
        auto_download: bool = True,
        verbose: bool = True,
    ):
        """
        Initialize Word2Vec normalizer.
        
        Args:
            word2vec_path: Path to word2vec model file. If None, auto-downloads.
            training_data_path: Path to training data file. If None, auto-downloads.
            training_pairs: List of (denormalized, normalized) pairs for training.
                           If provided, training_data_path is ignored.
            context_window: Size of character context window for pattern matching.
            min_pattern_count: Minimum occurrences of pattern to use it.
            word_similarity_threshold: Minimum similarity for word matching.
            max_similar_words: Maximum number of similar words to consider.
            replace_confidence_threshold: Threshold for replacing characters.
            no_replace_confidence_threshold: Threshold for not replacing characters.
            rare_pattern_threshold: Minimum count to consider pattern as common.
            rare_pattern_confidence_threshold: Confidence threshold for rare patterns.
            word_similarity_high_threshold: High similarity threshold for word matching.
            auto_download: Whether to auto-download models if not found.
            verbose: Whether to print progress messages.
        """
        self.context_window = context_window
        self.min_pattern_count = min_pattern_count
        self.word_similarity_threshold = word_similarity_threshold
        self.max_similar_words = max_similar_words
        self.replace_confidence_threshold = replace_confidence_threshold
        self.no_replace_confidence_threshold = no_replace_confidence_threshold
        self.rare_pattern_threshold = rare_pattern_threshold
        self.rare_pattern_confidence_threshold = rare_pattern_confidence_threshold
        self.word_similarity_high_threshold = word_similarity_high_threshold
        self.verbose = verbose
        
        # Word-level patterns: word -> normalized_word
        self.word_patterns: Dict[str, str] = {}
        
        # Character-level patterns: (context_before, char, context_after) -> {replacement: count}
        self.char_patterns: Dict[Tuple[str, str, str], Dict[str, int]] = defaultdict(
            lambda: defaultdict(int)
        )
        # Patterns where character should NOT be replaced
        self.char_patterns_no_replace: Dict[Tuple[str, str, str], int] = defaultdict(int)
        
        # Load Word2Vec model
        self.word2vec_model = self._load_word2vec(word2vec_path, auto_download)
        
        # Load training data
        if training_pairs is None:
            training_pairs = self._load_training_data(training_data_path, auto_download)
        
        if training_pairs:
            self.train(training_pairs)
        else:
            if self.verbose:
                print("Warning: No training data provided. Normalizer will use basic rules only.")

    def _load_word2vec(
        self,
        path: Optional[Union[str, Path]],
        auto_download: bool,
    ) -> KeyedVectors:
        """Load Word2Vec model from path or download it."""
        if path is not None:
            path = Path(path)
            if not path.exists():
                raise FileNotFoundError(f"Word2Vec model not found at {path}")
        else:
            # Check cache first
            path = get_model_path("yak_word2vec.bin")
            if not path.exists():
                if auto_download:
                    if self.verbose:
                        print("Word2Vec model not found. Downloading...")
                    path = download_model()
                else:
                    raise FileNotFoundError(
                        "Word2Vec model not found. Either provide word2vec_path "
                        "or set auto_download=True to download automatically."
                    )
        
        if self.verbose:
            print(f"Loading Word2Vec model from {path}...")
        
        try:
            model = KeyedVectors.load_word2vec_format(str(path), binary=True)
            if self.verbose:
                print(f"Loaded Word2Vec model with {len(model.key_to_index)} words")
            return model
        except Exception as e:
            raise RuntimeError(f"Could not load Word2Vec model from {path}: {e}") from e

    def _load_training_data(
        self,
        path: Optional[Union[str, Path]],
        auto_download: bool,
    ) -> List[Tuple[str, str]]:
        """Load training data from path or download it."""
        if path is not None:
            path = Path(path)
            if not path.exists():
                raise FileNotFoundError(f"Training data not found at {path}")
        else:
            # Check cache first
            path = get_model_path("train_pairs.txt")
            if not path.exists():
                if auto_download:
                    if self.verbose:
                        print("Training data not found. Downloading...")
                    path = download_training_data()
                else:
                    return []
        
        if self.verbose:
            print(f"Loading training data from {path}...")
        
        pairs = []
        with open(path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                parts = line.split("\t")
                if len(parts) >= 2:
                    pairs.append((parts[0], parts[1]))
        
        if self.verbose:
            print(f"Loaded {len(pairs)} training pairs")
        
        return pairs

    def train(self, training_pairs: List[Tuple[str, str]]) -> None:
        """
        Train the normalizer on pairs of (denormalized, normalized) text.
        
        Args:
            training_pairs: List of (denormalized, normalized) pairs
        """
        # Build word-level patterns with frequency counting
        word_pattern_counts: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int))
        
        for denorm, norm in training_pairs:
            denorm_words = self._split_words(denorm)
            norm_words = self._split_words(norm)
            
            # Store word mappings with counts
            for d_word, n_word in zip(denorm_words, norm_words):
                if d_word and n_word:
                    word_pattern_counts[d_word][n_word] += 1
            
            # Also learn character-level patterns as fallback
            self._learn_char_patterns(denorm, norm)
        
        # Store most frequent mapping for each word
        for d_word, mappings in word_pattern_counts.items():
            best_n_word = max(mappings.items(), key=lambda x: x[1])[0]
            self.word_patterns[d_word] = best_n_word
        
        if self.verbose:
            print(f"Learned {len(self.word_patterns)} word patterns")
            print(f"Learned {len(self.char_patterns)} character patterns")
            print(f"Learned {len(self.char_patterns_no_replace)} no-replace patterns")

    def _split_words(self, text: str) -> List[str]:
        """Split text into words, preserving word boundaries."""
        return re.findall(r'\b\w+\b', text)

    def _learn_char_patterns(self, denorm: str, norm: str) -> None:
        """Learn character-level replacement patterns from a single pair."""
        i = 0
        j = 0
        
        while i < len(denorm) and j < len(norm):
            if denorm[i] == norm[j]:
                # Characters match - check if this is a mappable character that wasn't replaced
                denorm_char = denorm[i]
                if denorm_char in self.CHAR_MAPPINGS:
                    context_before = self._get_context(denorm, i, before=True)
                    context_after = self._get_context(denorm, i, before=False)
                    pattern_key = (context_before, denorm_char, context_after)
                    self.char_patterns_no_replace[pattern_key] += 1
                
                i += 1
                j += 1
                continue
            
            # Character differs
            denorm_char = denorm[i]
            norm_char = norm[j]
            
            # Check if this is a known mapping
            if denorm_char in self.CHAR_MAPPINGS:
                expected_norm = self.CHAR_MAPPINGS[denorm_char]
                
                if norm_char == expected_norm:
                    context_before = self._get_context(denorm, i, before=True)
                    context_after = self._get_context(denorm, i, before=False)
                    pattern_key = (context_before, denorm_char, context_after)
                    self.char_patterns[pattern_key][norm_char] += 1
                    
                    i += 1
                    j += 1
                else:
                    if i + 1 < len(denorm) and denorm[i + 1] == norm[j]:
                        i += 1
                    elif j + 1 < len(norm) and denorm[i] == norm[j + 1]:
                        j += 1
                    else:
                        i += 1
                        j += 1
            else:
                i += 1
                j += 1

    def _get_context(self, text: str, pos: int, before: bool) -> str:
        """Get context around position."""
        if before:
            start = max(0, pos - self.context_window)
            return text[start:pos]
        else:
            end = min(len(text), pos + 1 + self.context_window)
            return text[pos + 1:end]

    def _find_similar_words(self, word: str) -> List[Tuple[str, float]]:
        """Find similar words using word2vec."""
        if word not in self.word2vec_model.key_to_index:
            word_lower = word.lower()
            if word_lower not in self.word2vec_model.key_to_index:
                return []
            word = word_lower
        
        try:
            similar_words = self.word2vec_model.most_similar(
                word, 
                topn=self.max_similar_words
            )
            return [
                (w, score) for w, score in similar_words 
                if score >= self.word_similarity_threshold
            ]
        except Exception:
            return []

    def _get_context_embedding(self, context_words: List[str]) -> Optional[np.ndarray]:
        """Get average embedding vector for context words."""
        embeddings = []
        for word in context_words:
            if word in self.word2vec_model.key_to_index:
                embeddings.append(self.word2vec_model[word])
            else:
                word_lower = word.lower()
                if word_lower in self.word2vec_model.key_to_index:
                    embeddings.append(self.word2vec_model[word_lower])
        
        if embeddings:
            return np.mean(embeddings, axis=0)
        return None

    def _normalize_word_with_context(self, word: str, context_words: List[str]) -> str:
        """Normalize a word using word2vec context and learned patterns."""
        # First, check if we have exact word mapping
        if word in self.word_patterns:
            return self.word_patterns[word]
        
        # Try to find similar words in training data using word2vec
        similar_words = self._find_similar_words(word)
        
        for similar_word, similarity in similar_words:
            if similar_word in self.word_patterns:
                normalized_similar = self.word_patterns[similar_word]
                if len(normalized_similar) == len(word):
                    if similarity >= self.word_similarity_high_threshold:
                        return normalized_similar
        
        # Use context embedding to find words with similar context
        if context_words:
            context_vec = self._get_context_embedding(context_words)
            if context_vec is not None:
                try:
                    similar_to_context = self.word2vec_model.similar_by_vector(
                        context_vec, topn=self.max_similar_words
                    )
                    
                    for similar_word, _ in similar_to_context:
                        if similar_word in self.word_patterns:
                            if word in self.word2vec_model.key_to_index:
                                word_similarity = self.word2vec_model.similarity(
                                    word, similar_word
                                )
                                if word_similarity >= self.word_similarity_threshold:
                                    normalized_similar = self.word_patterns[similar_word]
                                    if len(normalized_similar) == len(word):
                                        return normalized_similar
                            else:
                                word_lower = word.lower()
                                if word_lower in self.word2vec_model.key_to_index:
                                    word_similarity = self.word2vec_model.similarity(
                                        word_lower, similar_word
                                    )
                                    if word_similarity >= self.word_similarity_threshold:
                                        normalized_similar = self.word_patterns[similar_word]
                                        if len(normalized_similar) == len(word):
                                            return normalized_similar
                except Exception:
                    pass
        
        return word

    def _apply_char_pattern(
        self,
        char: str,
        context_before: str,
        context_after: str,
    ) -> Optional[str]:
        """Apply character pattern to determine replacement."""
        pattern_key = (context_before, char, context_after)
        no_replace_count = self.char_patterns_no_replace.get(pattern_key, 0)
        
        if pattern_key in self.char_patterns:
            replacements = self.char_patterns[pattern_key]
            if replacements:
                best_replacement = max(replacements.items(), key=lambda x: x[1])[0]
                replace_count = sum(replacements.values())
                total_count = replace_count + no_replace_count
                
                if total_count >= self.min_pattern_count:
                    if total_count < self.rare_pattern_threshold:
                        if replace_count > no_replace_count * self.rare_pattern_confidence_threshold:
                            return best_replacement
                    else:
                        if replace_count > no_replace_count * self.replace_confidence_threshold:
                            return best_replacement
                        elif no_replace_count > replace_count * self.no_replace_confidence_threshold:
                            return None
                        elif replace_count > no_replace_count:
                            return best_replacement
        
        return None

    def _try_shorter_contexts(
        self,
        char: str,
        context_before: str,
        context_after: str,
    ) -> Optional[str]:
        """Try shorter context patterns for fallback."""
        # Try shorter before context
        if len(context_before) > 0:
            for shorter_before_len in range(len(context_before) - 1, -1, -1):
                shorter_before = context_before[-shorter_before_len:] if shorter_before_len > 0 else ""
                result = self._apply_char_pattern(char, shorter_before, context_after)
                if result is not None:
                    return result
        
        # Try shorter after context
        if len(context_after) > 0:
            for shorter_after_len in range(len(context_after) - 1, -1, -1):
                shorter_after = context_after[:shorter_after_len] if shorter_after_len > 0 else ""
                result = self._apply_char_pattern(char, context_before, shorter_after)
                if result is not None:
                    return result
        
        return None

    def normalize(self, text: str) -> str:
        """
        Normalize text using word2vec embeddings and learned patterns.
        
        Args:
            text: Input text without Yakut characters
            
        Returns:
            Normalized text with Yakut characters
        """
        # First, always replace 'h' -> 'һ' (rule: always)
        result = list(text)
        for i in range(len(result)):
            if result[i] == 'h':
                result[i] = 'һ'
            elif result[i] == 'H':
                result[i] = 'Һ'
        
        # Split into words for word-level processing
        words = self._split_words(text)
        word_positions = []
        pos = 0
        for word in words:
            start = text.find(word, pos)
            if start != -1:
                word_positions.append((start, start + len(word), word))
                pos = start + len(word)
        
        # Normalize words using word2vec context
        normalized_words = {}
        for idx, (start, end, word) in enumerate(word_positions):
            context_words = []
            if idx > 0:
                context_words.append(word_positions[idx - 1][2])
            if idx < len(word_positions) - 1:
                context_words.append(word_positions[idx + 1][2])
            
            normalized_word = self._normalize_word_with_context(word, context_words)
            if normalized_word != word:
                normalized_words[(start, end)] = normalized_word
        
        # Apply word-level normalizations (reverse order to avoid index shifting)
        for (start, end), norm_word in sorted(normalized_words.items(), reverse=True):
            result[start:end] = list(norm_word)
        
        # For remaining characters, use character-level patterns
        text_after_words = ''.join(result)
        for i in range(len(result)):
            char = result[i]
            
            # Skip if already replaced or not a mappable character
            if char not in self.CHAR_MAPPINGS or char in ['h', 'H']:
                continue
            
            context_before = self._get_context(text_after_words, i, before=True)
            context_after = self._get_context(text_after_words, i, before=False)
            
            # Try exact match first
            replacement = self._apply_char_pattern(char, context_before, context_after)
            
            # Try shorter contexts as fallback
            if replacement is None:
                replacement = self._try_shorter_contexts(char, context_before, context_after)
            
            if replacement:
                result[i] = replacement
        
        return ''.join(result)
