correction_engine.py

"""
correction_engine.py — Pre- and Post-translation correction pipeline
Handles: misheard words, slang, short forms, spoken English,
         confused words, and domain-specific corrections.
"""

import json
import os
import re
from typing import Dict, List, Optional, Any

from config import CORRECTIONS_FILE, DATASETS_DIR


class CorrectionEngine:
    """
    Two-stage correction pipeline:
    1. PRE-TRANSLATION:  Fix speech recognition errors before translation
    2. POST-TRANSLATION: Fix known translation errors after NLLB output
    """

    def __init__(self):
        self.pre_corrections: Dict[str, Dict] = {}   # lang → {pattern: replacement}
        self.post_corrections: Dict[str, Dict] = {}   # lang_pair → {pattern: replacement}
        self.slang_map: Dict[str, str] = {}
        self.short_forms: Dict[str, str] = {}
        self.confused_words: Dict[str, str] = {}
        self.domain_terms: Dict[str, Dict[str, str]] = {}
        self.loaded = False

    def load(self):
        """Load all correction datasets."""
        os.makedirs(DATASETS_DIR, exist_ok=True)

        # Create default corrections file if it doesn't exist
        if not os.path.exists(CORRECTIONS_FILE):
            self._create_default_corrections()

        try:
            with open(CORRECTIONS_FILE, "r", encoding="utf-8") as f:
                data = json.load(f)

            self.slang_map = data.get("slang", {})
            self.short_forms = data.get("short_forms", {})
            self.confused_words = data.get("confused_words", {})
            self.domain_terms = data.get("domain_terms", {})
            self.pre_corrections = data.get("pre_corrections", {})
            self.post_corrections = data.get("post_corrections", {})

            total = (len(self.slang_map) + len(self.short_forms) +
                     len(self.confused_words) + sum(len(v) for v in self.domain_terms.values()))
            print(f"[Correction] Loaded {total} correction entries.")
            self.loaded = True

        except Exception as e:
            print(f"[Correction] Error loading corrections: {e}")
            self._create_default_corrections()
            self.loaded = True

        # Load any additional custom correction files
        self._load_custom_corrections()

    def pre_translate(self, text: str, source_lang: str = "en") -> Dict[str, Any]:
        """
        Apply pre-translation corrections to recognized speech text.

        Pipeline:
            Raw Whisper text → Slang fix → Short forms → Confused words
            → Domain terms → Custom pre-corrections → Cleaned text

        Returns:
            dict with 'text' (corrected), 'corrections' (list of changes made)
        """
        corrections = []
        original = text

        # 1. Normalize whitespace
        text = re.sub(r'\s+', ' ', text).strip()

        # 2. Apply slang corrections (case-insensitive)
        for slang, replacement in self.slang_map.items():
            pattern = re.compile(r'\b' + re.escape(slang) + r'\b', re.IGNORECASE)
            if pattern.search(text):
                text = pattern.sub(replacement, text)
                corrections.append({"type": "slang", "from": slang, "to": replacement})

        # 3. Apply short form expansions
        for short, full in self.short_forms.items():
            pattern = re.compile(r'\b' + re.escape(short) + r'\b', re.IGNORECASE)
            if pattern.search(text):
                text = pattern.sub(full, text)
                corrections.append({"type": "short_form", "from": short, "to": full})

        # 4. Fix commonly confused words
        for wrong, correct in self.confused_words.items():
            pattern = re.compile(r'\b' + re.escape(wrong) + r'\b', re.IGNORECASE)
            if pattern.search(text):
                text = pattern.sub(correct, text)
                corrections.append({"type": "confused", "from": wrong, "to": correct})

        # 5. Apply domain-specific corrections
        for domain, terms in self.domain_terms.items():
            for wrong, correct in terms.items():
                pattern = re.compile(r'\b' + re.escape(wrong) + r'\b', re.IGNORECASE)
                if pattern.search(text):
                    text = pattern.sub(correct, text)
                    corrections.append({
                        "type": "domain",
                        "domain": domain,
                        "from": wrong,
                        "to": correct
                    })

        # 6. Apply language-specific pre-corrections
        lang_corrections = self.pre_corrections.get(source_lang, {})
        for wrong, correct in lang_corrections.items():
            pattern = re.compile(re.escape(wrong), re.IGNORECASE)
            if pattern.search(text):
                text = pattern.sub(correct, text)
                corrections.append({"type": "pre_lang", "from": wrong, "to": correct})

        return {
            "text": text,
            "original": original,
            "corrected": text != original,
            "corrections": corrections,
            "correction_count": len(corrections)
        }

    def post_translate(self, text: str, source_lang: str, target_lang: str) -> Dict[str, Any]:
        """
        Apply post-translation corrections to NLLB output.

        Returns:
            dict with 'text' (corrected), 'corrections' (list of changes)
        """
        corrections = []
        original = text
        lang_pair = f"{source_lang}_{target_lang}"

        # Apply language-pair-specific post-corrections
        pair_corrections = self.post_corrections.get(lang_pair, {})
        for wrong, correct in pair_corrections.items():
            if wrong in text:
                text = text.replace(wrong, correct)
                corrections.append({"type": "post_lang", "from": wrong, "to": correct})

        # Target-language-specific corrections
        target_corrections = self.post_corrections.get(target_lang, {})
        for wrong, correct in target_corrections.items():
            if wrong in text:
                text = text.replace(wrong, correct)
                corrections.append({"type": "post_target", "from": wrong, "to": correct})

        return {
            "text": text,
            "original": original,
            "corrected": text != original,
            "corrections": corrections
        }

    def add_correction(self, category: str, wrong: str, correct: str,
                       domain: Optional[str] = None) -> bool:
        """
        Add a new correction entry and persist to file.

        Args:
            category: 'slang', 'short_forms', 'confused_words', 'domain_terms'
            wrong: The incorrect/misheard form
            correct: The correct replacement
            domain: Domain name (only for 'domain_terms' category)
        """
        try:
            if category == "slang":
                self.slang_map[wrong] = correct
            elif category == "short_forms":
                self.short_forms[wrong] = correct
            elif category == "confused_words":
                self.confused_words[wrong] = correct
            elif category == "domain_terms" and domain:
                if domain not in self.domain_terms:
                    self.domain_terms[domain] = {}
                self.domain_terms[domain][wrong] = correct
            else:
                return False

            self._save_corrections()
            return True
        except Exception as e:
            print(f"[Correction] Error adding correction: {e}")
            return False

    def get_stats(self) -> Dict[str, int]:
        """Return statistics about loaded corrections."""
        return {
            "slang_entries": len(self.slang_map),
            "short_form_entries": len(self.short_forms),
            "confused_word_entries": len(self.confused_words),
            "domain_entries": sum(len(v) for v in self.domain_terms.values()),
            "pre_correction_langs": len(self.pre_corrections),
            "post_correction_pairs": len(self.post_corrections),
            "total": (len(self.slang_map) + len(self.short_forms) +
                      len(self.confused_words) +
                      sum(len(v) for v in self.domain_terms.values()))
        }

    def _load_custom_corrections(self):
        """Load additional correction files from datasets/ folder."""
        custom_dir = os.path.join(DATASETS_DIR, "custom_corrections")
        if not os.path.exists(custom_dir):
            os.makedirs(custom_dir, exist_ok=True)
            return

        for filename in os.listdir(custom_dir):
            if filename.endswith(".json"):
                try:
                    filepath = os.path.join(custom_dir, filename)
                    with open(filepath, "r", encoding="utf-8") as f:
                        data = json.load(f)

                    # Merge custom corrections
                    if "slang" in data:
                        self.slang_map.update(data["slang"])
                    if "short_forms" in data:
                        self.short_forms.update(data["short_forms"])
                    if "confused_words" in data:
                        self.confused_words.update(data["confused_words"])
                    if "domain_terms" in data:
                        for domain, terms in data["domain_terms"].items():
                            if domain not in self.domain_terms:
                                self.domain_terms[domain] = {}
                            self.domain_terms[domain].update(terms)

                    print(f"[Correction] Loaded custom: {filename}")
                except Exception as e:
                    print(f"[Correction] Error loading {filename}: {e}")

    def _save_corrections(self):
        """Persist corrections to disk."""
        data = {
            "slang": self.slang_map,
            "short_forms": self.short_forms,
            "confused_words": self.confused_words,
            "domain_terms": self.domain_terms,
            "pre_corrections": self.pre_corrections,
            "post_corrections": self.post_corrections,
        }
        os.makedirs(os.path.dirname(CORRECTIONS_FILE), exist_ok=True)
        with open(CORRECTIONS_FILE, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

    def _create_default_corrections(self):
        """Create the default corrections dataset."""
        data = {
            "slang": {
                "gonna": "going to",
                "wanna": "want to",
                "gotta": "got to",
                "lemme": "let me",
                "gimme": "give me",
                "kinda": "kind of",
                "sorta": "sort of",
                "dunno": "don't know",
                "ain't": "is not",
                "y'all": "you all",
                "bro": "brother",
                "sis": "sister",
                "dude": "friend",
                "nah": "no",
                "yep": "yes",
                "yup": "yes",
                "nope": "no",
                "ok": "okay",
                "coz": "because",
                "cuz": "because",
                "cos": "because",
                "sup": "what is up",
                "wassup": "what is up",
                "whatcha": "what are you",
                "imma": "I am going to",
                "fam": "family",
                "lit": "exciting",
                "lowkey": "slightly",
                "highkey": "very much",
                "salty": "upset",
                "shook": "shocked",
                "vibe": "feeling",
                "flex": "show off",
                "ghost": "disappear",
                "slay": "do excellently",
                "bet": "agreed",
                "cap": "lie",
                "no cap": "no lie",
                "fire": "excellent",
                "sus": "suspicious",
                "bruh": "brother",
                "hella": "very",
                "bougie": "fancy",
                "clap back": "respond sharply",
                "shade": "disrespect",
                "tea": "gossip",
                "mood": "relatable",
                "stan": "admire greatly",
                "simp": "overly devoted person",
                "yeet": "throw",
                "oof": "expressing discomfort",
                "periodt": "period",
                "deadass": "seriously",
                "idk": "I don't know",
                "tbh": "to be honest",
                "imo": "in my opinion",
                "smh": "shaking my head",
                "lol": "laughing out loud",
                "omg": "oh my god",
                "brb": "be right back",
                "btw": "by the way",
                "fyi": "for your information",
                "asap": "as soon as possible"
            },

            "short_forms": {
                "govt": "government",
                "dept": "department",
                "mgr": "manager",
                "prof": "professor",
                "doc": "doctor",
                "info": "information",
                "tech": "technology",
                "app": "application",
                "bio": "biography",
                "exam": "examination",
                "lab": "laboratory",
                "math": "mathematics",
                "stats": "statistics",
                "ref": "reference",
                "temp": "temperature",
                "approx": "approximately",
                "etc": "etcetera",
                "vs": "versus",
                "min": "minute",
                "max": "maximum",
                "avg": "average",
                "qty": "quantity",
                "yr": "year",
                "mo": "month",
                "wk": "week",
                "hr": "hour",
                "sec": "second",
                "Mon": "Monday",
                "Tue": "Tuesday",
                "Wed": "Wednesday",
                "Thu": "Thursday",
                "Fri": "Friday",
                "Sat": "Saturday",
                "Sun": "Sunday",
                "Jan": "January",
                "Feb": "February",
                "Mar": "March",
                "Apr": "April",
                "Jun": "June",
                "Jul": "July",
                "Aug": "August",
                "Sep": "September",
                "Oct": "October",
                "Nov": "November",
                "Dec": "December"
            },

            "confused_words": {
                "their there": "there",
                "your you're": "you are",
                "its it's": "it is",
                "affect effect": "effect",
                "then than": "than",
                "loose lose": "lose",
                "weather whether": "whether",
                "to too": "too",
                "accept except": "except",
                "advise advice": "advice",
                "principal principle": "principle",
                "stationary stationery": "stationery",
                "complement compliment": "compliment",
                "desert dessert": "dessert",
                "alot": "a lot",
                "definately": "definitely",
                "seperate": "separate",
                "occured": "occurred",
                "recieve": "receive",
                "untill": "until",
                "writting": "writing",
                "calender": "calendar",
                "goverment": "government",
                "enviroment": "environment",
                "tommorrow": "tomorrow",
                "accomodate": "accommodate",
                "concious": "conscious",
                "necesary": "necessary",
                "occassion": "occasion",
                "priviledge": "privilege",
                "pronounciation": "pronunciation",
                "recomend": "recommend",
                "refered": "referred",
                "succesful": "successful",
                "wierd": "weird"
            },

            "domain_terms": {
                "medical": {
                    "BP": "blood pressure",
                    "HR": "heart rate",
                    "temp": "temperature",
                    "RX": "prescription",
                    "OPD": "outpatient department",
                    "ICU": "intensive care unit",
                    "ER": "emergency room",
                    "MRI": "magnetic resonance imaging",
                    "CT scan": "computed tomography scan",
                    "ECG": "electrocardiogram"
                },
                "education": {
                    "GPA": "grade point average",
                    "CGPA": "cumulative grade point average",
                    "HOD": "head of department",
                    "PhD": "Doctor of Philosophy",
                    "MSc": "Master of Science",
                    "BSc": "Bachelor of Science",
                    "sem": "semester",
                    "assign": "assignment",
                    "ppt": "presentation"
                },
                "technology": {
                    "API": "application programming interface",
                    "UI": "user interface",
                    "UX": "user experience",
                    "DB": "database",
                    "OS": "operating system",
                    "CPU": "central processing unit",
                    "GPU": "graphics processing unit",
                    "RAM": "random access memory",
                    "SSD": "solid state drive",
                    "URL": "uniform resource locator"
                },
                "business": {
                    "CEO": "Chief Executive Officer",
                    "CFO": "Chief Financial Officer",
                    "CTO": "Chief Technology Officer",
                    "HR": "human resources",
                    "KPI": "key performance indicator",
                    "ROI": "return on investment",
                    "B2B": "business to business",
                    "B2C": "business to consumer",
                    "ETA": "estimated time of arrival",
                    "RSVP": "please respond"
                }
            },

            "pre_corrections": {
                "en": {
                    "hii": "hi",
                    "hiii": "hi",
                    "hellooo": "hello",
                    "byeee": "bye",
                    "yesss": "yes",
                    "nooo": "no",
                    "pleaseee": "please",
                    "thanksss": "thanks",
                    "sorryyy": "sorry",
                    "okayyyy": "okay",
                    "wowww": "wow",
                    "ummm": "",
                    "uhhh": "",
                    "ahhh": "",
                    "hmmmm": "",
                    "errr": "",
                    "like like": "like",
                    "you know you know": "you know",
                    "I mean I mean": "I mean"
                }
            },

            "post_corrections": {}
        }

        os.makedirs(os.path.dirname(CORRECTIONS_FILE), exist_ok=True)
        with open(CORRECTIONS_FILE, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

        # Load into memory
        self.slang_map = data["slang"]
        self.short_forms = data["short_forms"]
        self.confused_words = data["confused_words"]
        self.domain_terms = data["domain_terms"]
        self.pre_corrections = data["pre_corrections"]
        self.post_corrections = data["post_corrections"]

        print(f"[Correction] Created default corrections dataset.")