"""Synonym replacement for adversarial data augmentation.

Replaces words with synonyms to create challenging training examples
that help improve model robustness.
"""

import random
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

import yaml


class MyanmarSynonymReplacer:
    """Replace words with Myanmar synonyms for data augmentation."""
    
    # Basic synonym dictionary for Myanmar
    SYNONYMS = {
        "ကျေးဇူး": ["ခန့်ညား", "ဂုဏ်ပြု", "အားထုတ်"],
        "ပါး": ["များ", "အရမ်း", "အလွန်"],
        "သိပ်": ["အရမ်း", "များစွာ", "ပါး"],
        "ကောင်း": ["မွန်", "သန့်", "စင်"],
        "ဆိုး": ["မကောင်း", "ယုတ်", "ညံ့"],
        "ပျော်": ["ရွှင်", "ပီတိ", "မင်္ဂလာ"],
        "စိတ်မကောင်း": ["စိတ်ဓာတ်ကျ", "ဝမ်းနည်း", "ဒေါသ"],
        "လာ": ["ရောက်", "သွား", "ပို့"],
        "သွား": ["သွား", "ထွက်", "မောင်း"],
        "ပေး": ["ပေးဆောင်", "မွေးစား", "အပ်"],
        "ယူ": ["ရ", "လက်ခံ", "ရှာ"],
        "မှန်": ["ဟုတ်", "ကျိ", "သင့်"],
        "မှား": ["မဟုတ်", "အမှား", "ယွင်း"],
    }
    
    def __init__(
        self,
        synonym_file: Optional[str] = None,
        seed: int = 42,
    ):
        """
        Args:
            synonym_file: Path to YAML file with custom synonyms
            seed: Random seed for reproducibility
        """
        random.seed(seed)
        
        self.synonyms = dict(self.SYNONYMS)
        
        if synonym_file and Path(synonym_file).exists():
            self._load_synonyms(synonym_file)
    
    def _load_synonyms(self, path: str) -> None:
        """Load synonyms from YAML file."""
        with open(path, "r", encoding="utf-8") as f:
            custom = yaml.safe_load(f)
        
        if custom:
            self.synonyms.update(custom)
    
    def get_synonyms(self, word: str) -> List[str]:
        """Get synonyms for a word."""
        return self.synonyms.get(word, [])
    
    def replace_word(self, word: str) -> Tuple[str, bool]:
        """Replace a word with a synonym.
        
        Args:
            word: Word to replace
        
        Returns:
            (replaced_word, was_replaced)
        """
        synonyms = self.get_synonyms(word)
        
        if synonyms:
            replacement = random.choice(synonyms)
            return replacement, True
        
        return word, False
    
    def augment_text(
        self,
        text: str,
        replace_prob: float = 0.3,
        max_replacements: int = 3,
    ) -> Tuple[str, List[Tuple[str, str]]]:
        """Augment text by replacing words with synonyms.
        
        Args:
            text: Myanmar text
            replace_prob: Probability of replacing each synonym word
            max_replacements: Maximum number of replacements
        
        Returns:
            (augmented_text, list_of_replacements)
        """
        words = text.split()
        replacements = []
        augmented_words = []
        
        num_replaced = 0
        
        for word in words:
            if num_replaced >= max_replacements:
                augmented_words.append(word)
                continue
            
            if word in self.synonyms and random.random() < replace_prob:
                new_word, replaced = self.replace_word(word)
                if replaced:
                    augmented_words.append(new_word)
                    replacements.append((word, new_word))
                    num_replaced += 1
                else:
                    augmented_words.append(word)
            else:
                augmented_words.append(word)
        
        return " ".join(augmented_words), replacements
    
    def augment_dataset(
        self,
        samples: List[Dict],
        replace_prob: float = 0.3,
        max_replacements: int = 3,
        n_augmentations: int = 2,
    ) -> List[Dict]:
        """Augment entire dataset.
        
        Args:
            samples: List of sample dictionaries
            replace_prob: Probability of replacing each word
            max_replacements: Maximum replacements per text
            n_augmentations: Number of augmentations per sample
        
        Returns:
            List of augmented samples
        """
        augmented = []
        
        for sample in samples:
            text = sample.get("text", "")
            
            for i in range(n_augmentations):
                aug_text, replacements = self.augment_text(
                    text,
                    replace_prob=replace_prob,
                    max_replacements=max_replacements,
                )
                
                if replacements:  # Only add if something was replaced
                    aug_sample = sample.copy()
                    aug_sample["text"] = aug_text
                    aug_sample["augmentation_id"] = i
                    aug_sample["replacements"] = replacements
                    aug_sample["is_augmented"] = True
                    augmented.append(aug_sample)
        
        return augmented
    
    def get_replacement_stats(self) -> Dict:
        """Get statistics about synonym coverage."""
        total_words = sum(len(syns) for syns in self.synonyms.values())
        return {
            "num_synonym_groups": len(self.synonyms),
            "total_synonyms": total_words,
            "avg_synonyms_per_word": total_words / len(self.synonyms) if self.synonyms else 0,
        }


class ContextualSynonymReplacer:
    """Synonym replacer that considers context."""
    
    def __init__(self):
        # Context-dependent synonyms
        self.CONTEXT_SYNONYMS = {
            "formal": {
                "ကျေးဇူး": ["ဂုဏ်ပြုမှတ်ရှိပါ", "အထူးပင်ကျေးဇူးတင်ပါ"],
                "သိပ်": ["အလွန်", "အထူးသဖြင့်"],
            },
            "informal": {
                "ကျေးဇူး": ["ခန့်ညား", "ကျေးဇူးလည်းပါ"],
                "ပါး": ["ပို", "အရမ်း"],
            },
        }
    
    def augment_by_context(
        self,
        text: str,
        context: str = "formal",
    ) -> str:
        """Augment text using context-specific synonyms.
        
        Args:
            text: Myanmar text
            context: Context type ("formal" or "informal")
        
        Returns:
            Augmented text
        """
        context_syns = self.CONTEXT_SYNONYMS.get(context, {})
        
        for word, synonyms in context_syns.items():
            if word in text:
                replacement = random.choice(synonyms)
                text = text.replace(word, replacement, 1)
        
        return text


def create_synonym_replacer(
    synonym_file: Optional[str] = None,
) -> MyanmarSynonymReplacer:
    """Factory function to create synonym replacer."""
    return MyanmarSynonymReplacer(synonym_file=synonym_file)


if __name__ == "__main__":
    replacer = create_synonym_replacer()
    
    test_text = "ကျေးဇူးပါးသိပ်ကောင်းတယ်"
    
    for i in range(3):
        aug, replacements = replacer.augment_text(test_text)
        print(f"Original: {test_text}")
        print(f"Augmented: {aug}")
        print(f"Replacements: {replacements}")
        print()