"""Text normalization module for Myanmar language."""

import re
import unicodedata
from pathlib import Path
from typing import Dict, List, Optional

import pandas as pd
import yaml

logger = __import__("loguru").logger


class MyanmarTextNormalizer:
    """Normalize Myanmar (Burmese) text for consistent processing."""

    # Myanmar Unicode ranges
    MYANMAR_CHARS = re.compile(
        r"[\u1000-\u100F\u1010-\u101F\u1020-\u102A\u102C-\u1030\u1031\u1032\u1036-\u1038\u1039\u103A]"
    )
    
    # Normalization rules
    NORMALIZATION_RULES = {
        # Zero-width characters
        "\u200B": "",  # Zero-width space
        "\u200C": "",  # Zero-width non-joiner
        "\u200D": "",  # Zero-width joiner
        "\u2060": "",  # Word joiner
        # Myanmar-specific normalizations
        "\u1031\u103B": "\u103B\u1031",  # medial order
        "\u103D\u103E": "\u103E\u103D",  # stack order
    }

    def __init__(self, custom_rules_path: Optional[str] = None):
        self.custom_rules = {}
        if custom_rules_path and Path(custom_rules_path).exists():
            with open(custom_rules_path, "r", encoding="utf-8") as f:
                self.custom_rules = yaml.safe_load(f) or {}
        
        self.rules = {**self.NORMALIZATION_RULES, **self.custom_rules}

    def normalize_unicode(self, text: str) -> str:
        """Standardize Unicode representation (NFC normalization)."""
        return unicodedata.normalize("NFC", text)

    def remove_diacritics(self, text: str) -> str:
        """Remove tone marks for simplified processing."""
        diacritics = re.compile(
            r"[\u102B-\u102D\u102F-\u1032\u1034\u1036\u1037\u1039]"
        )
        return diacritics.sub("", text)

    def remove_whitespace(self, text: str) -> str:
        """Remove excessive whitespace."""
        text = re.sub(r"\s+", " ", text)
        return text.strip()

    def normalize_punctuation(self, text: str) -> str:
        """Standardize punctuation marks."""
        replacements = {
            "၊": "။",  # Myanmar comma to full stop
            "„": '"',
            "‟": '"',
            "'": "'",
            "`": "'",
            "—": "–",
            "–": "-",
        }
        for old, new in replacements.items():
            text = text.replace(old, new)
        return text

    def apply_custom_rules(self, text: str) -> str:
        """Apply user-defined normalization rules."""
        for pattern, replacement in self.rules.items():
            text = text.replace(pattern, replacement)
        return text

    def expand_abbreviations(self, text: str, abbreviations: Dict[str, str] = None) -> str:
        """Expand common abbreviations."""
        if abbreviations is None:
            abbreviations = {
                "အ.ပ.ခ": "အငြိမ်းစားပြည်ထဲရေးဝန်ကြီး",
                "ဒ.ပ.လ": "ဒုတိယသမ္မတ",
                "ပ.ရ.မှူး": "ပြည်သူ့လွှတ်တော်ဥက္ကဋ္ဌ",
            }
        
        for abbr, full in abbreviations.items():
            text = re.sub(rf"\b{re.escape(abbr)}\b", full, text)
        
        return text

    def normalize_numbers(self, text: str) -> str:
        """Convert Myanmar numerals to Arabic (0-9)."""
        myanmar_digits = "၀၁၂၃၄၅၆၇၈၉"
        arabic_digits = "0123456789"
        
        trans_table = str.maketrans(
            {myanmar_digits[i]: arabic_digits[i] for i in range(10)}
        )
        return text.translate(trans_table)

    def filter_non_myanmar(self, text: str, keep_english: bool = True) -> str:
        """Remove or keep non-Myanmar characters."""
        if keep_english:
            pattern = r"[^\u1000-\u109F\u0020-\u007E\u00A0-\u00FF]"
        else:
            pattern = r"[^\u1000-\u109F\s]"
        
        return re.sub(pattern, "", text)

    def normalize_line(self, text: str) -> str:
        """Apply all normalization steps to a single line."""
        text = self.normalize_unicode(text)
        text = self.apply_custom_rules(text)
        text = self.remove_whitespace(text)
        text = self.normalize_punctuation(text)
        return text

    def normalize_corpus(
        self,
        texts: List[str],
        remove_non_myanmar: bool = False,
    ) -> List[str]:
        """Normalize a list of texts."""
        normalized = []
        for text in texts:
            text = self.normalize_line(text)
            if remove_non_myanmar:
                text = self.filter_non_myanmar(text, keep_english=False)
            normalized.append(text)
        
        logger.info(f"Normalized {len(normalized)} texts")
        return normalized

    def normalize_dataset(
        self,
        input_path: str,
        output_path: str,
        text_column: str = "text",
    ) -> pd.DataFrame:
        """Normalize a dataset and save to file."""
        df = pd.read_csv(input_path)
        
        if text_column not in df.columns:
            raise ValueError(f"Column '{text_column}' not found in dataset")
        
        df[f"{text_column}_normalized"] = self.normalize_corpus(
            df[text_column].tolist()
        )
        
        df.to_csv(output_path, index=False)
        logger.info(f"Normalized dataset saved to {output_path}")
        
        return df


class ProsodyNormalizer:
    """Normalize prosodic features for consistent representation."""

    def normalize_pitch(self, pitch_values: List[float]) -> List[float]:
        """Normalize pitch values to semitones from mean."""
        import numpy as np
        pitch_arr = np.array(pitch_values)
        mean_pitch = np.mean(pitch_arr[pitch_arr > 0])
        if mean_pitch == 0:
            return pitch_values
        semitones = 12 * np.log2(pitch_arr / mean_pitch)
        return semitones.tolist()

    def normalize_energy(self, energy_values: List[float]) -> List[float]:
        """Normalize energy values to 0-1 range."""
        import numpy as np
        energy_arr = np.array(energy_values)
        min_e, max_e = energy_arr.min(), energy_arr.max()
        if max_e - min_e == 0:
            return [0.5] * len(energy_values)
        return ((energy_arr - min_e) / (max_e - min_e)).tolist()

    def quantize_prosody(
        self,
        prosody: dict,
        num_levels: int = 5,
    ) -> dict:
        """Quantize prosodic features for categorical representation."""
        quantized = {}
        for key, value in prosody.items():
            if isinstance(value, (int, float)) and key != "pitch_range":
                normalized = max(0, min(1, (value + 100) / 200))
                quantized[key] = int(normalized * (num_levels - 1))
            else:
                quantized[key] = value
        return quantized


def create_normalizer(config: dict = None) -> MyanmarTextNormalizer:
    """Factory function to create normalizer from config."""
    if config is None:
        config = {}
    
    return MyanmarTextNormalizer(
        custom_rules_path=config.get("custom_rules_path")
    )


if __name__ == "__main__":
    normalizer = create_normalizer()
    test_text = "  မင်္ဂလာပါ   ၊  ကျေးဇူးပါ  ပါ   သည်  "
    print(f"Original: {test_text}")
    print(f"Normalized: {normalizer.normalize_line(test_text)}")