turing/features.py

import ast
import hashlib
from pathlib import Path
import random
import re
from typing import List, Tuple

import nltk
from nltk.corpus import stopwords, wordnet
from nltk.stem import PorterStemmer, WordNetLemmatizer
import numpy as np
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.feature_selection import SelectKBest, chi2
import typer

from turing.config import (
    INTERIM_DATA_DIR,
    LABEL_COLUMN,
    LANGS,
)
from turing.data_validation import run_custom_deepchecks, run_targeted_nlp_checks
from turing.dataset import DatasetManager

# --- NLTK Resource Check ---
REQUIRED_NLTK_PACKAGES = [
    "stopwords",
    "wordnet",
    "omw-1.4",
    "averaged_perceptron_tagger",
    "punkt",
]
for package in REQUIRED_NLTK_PACKAGES:
    try:
        nltk.data.find(f"corpora/{package}")
    except LookupError:
        try:
            nltk.download(package, quiet=True)
        except Exception:
            pass

app = typer.Typer()


# --- CONFIGURATION CLASS ---
class FeaturePipelineConfig:
    """
    Configuration holder for the pipeline. Generates a unique ID based on parameters
    to version the output directories.
    """

    def __init__(
        self,
        use_stopwords: bool,
        use_lemmatization: bool,
        use_combo_feature: bool,
        max_features: int,
        min_comment_length: int,
        max_comment_length: int,
        enable_augmentation: bool,
        custom_tags: str = "base",
    ):
        self.use_stopwords = use_stopwords
        self.use_lemmatization = use_lemmatization
        self.use_combo_feature = use_combo_feature
        self.max_features = max_features
        self.min_comment_length = min_comment_length
        self.max_comment_length = max_comment_length
        self.enable_augmentation = enable_augmentation
        self.custom_tags = custom_tags
        self.hash_id = self._generate_readable_id()

    def _generate_readable_id(self) -> str:
        tags = ["clean"]
        if self.enable_augmentation:
            tags.append("aug-soft")
        tags.append(f"k{self.max_features}")
        if self.custom_tags != "base":
            tags.append(self.custom_tags)
        return "-".join(tags)


# --- TEXT UTILITIES ---
class TextCanonicalizer:
    """
    Reduces text to a 'canonical' form (stemmed, lowercase)
    to detect semantic duplicates.
    preserves javadoc tags to distinguish usage (@return) from summary (Returns).
    """

    def __init__(self):
        self.stemmer = PorterStemmer()
        self.stop_words = set(stopwords.words("english"))
        # Code keywords are preserved as they carry semantic weight
        self.code_keywords = {
            "return",
            "true",
            "false",
            "null",
            "if",
            "else",
            "void",
            "int",
            "boolean",
            "param", 
            "throws",
            "exception",
        }

    def to_canonical(self, text: str) -> str:
        if pd.isna(text):
            return ""
        text = str(text).lower()
        text = re.sub(r"[^a-z0-9\s@]", " ", text)
        
        words = text.split()
        canonical_words = []
        
        for w in words:
            # If the word starts with @ (e.g., @return), keep it as is
            if w.startswith("@"):
                canonical_words.append(w)
                continue

            if w in self.stop_words and w not in self.code_keywords:
                continue
            
            stemmed = self.stemmer.stem(w)
            canonical_words.append(stemmed)
            
        return " ".join(canonical_words).strip()


class TextProcessor:
    """
    Standard text cleaning logic for final feature extraction (TF-IDF).
    """

    def __init__(self, config: FeaturePipelineConfig, language: str = "english"):
        self.config = config
        self.stop_words = set(stopwords.words(language))
        self.lemmatizer = WordNetLemmatizer()

    def clean_text(self, text: str) -> str:
        if pd.isna(text):
            return ""
        text = str(text).lower()
        # Remove heavy code markers but keep text structure
        text = re.sub(r"(^\s*//+|^\s*/\*+|\*/$)", "", text)
        # Keep only alpha characters for NLP model (plus pipe for combo)
        text = re.sub(r"[^a-z\s|]", " ", text)
        tokens = text.split()
        if self.config.use_stopwords:
            tokens = [w for w in tokens if w not in self.stop_words]
        if self.config.use_lemmatization:
            tokens = [self.lemmatizer.lemmatize(w) for w in tokens]
        return " ".join(tokens)


# --- AUGMENTATION ---
class SafeAugmenter:
    """
    protects reserved keywords from synonym replacement.
    """

    def __init__(self, aug_prob=0.3):
        self.aug_prob = aug_prob
        self.protected_words = {
            "return",
            "public",
            "private",
            "void",
            "class",
            "static",
            "final",
            "if",
            "else",
            "for",
            "while",
            "try",
            "catch",
            "import",
            "package",
            "null",
            "true",
            "false",
            "self",
            "def",
            "todo",
            "fixme",
            "param",
            "throw",
        }

    def get_synonyms(self, word):
        synonyms = set()
        for syn in wordnet.synsets(word):
            for lemma in syn.lemmas():
                name = lemma.name().replace("_", " ")
                if name.isalpha() and name.lower() != word.lower():
                    synonyms.add(name)
        return list(synonyms)

    def augment(self, text: str) -> str:
        if pd.isna(text) or not text:
            return ""
        words = text.split()
        if len(words) < 2:
            return text
        new_words = []
        for word in words:
            word_lower = word.lower()

            if word_lower in self.protected_words:
                new_words.append(word)
                continue

            # Random Case Injection (Noise)
            if random.random() < 0.1:
                if word[0].isupper():
                    new_words.append(word.lower())
                else:
                    new_words.append(word.capitalize())
                continue

            # Synonym Replacement
            if random.random() < self.aug_prob and len(word) > 3:
                syns = self.get_synonyms(word_lower)
                if syns:
                    replacement = random.choice(syns)
                    if word[0].isupper():
                        replacement = replacement.capitalize()
                    new_words.append(replacement)
                else:
                    new_words.append(word)
            else:
                new_words.append(word)
        return " ".join(new_words)

    def apply_balancing(
        self, df: pd.DataFrame, min_samples: int = 100
    ) -> Tuple[pd.DataFrame, pd.DataFrame]:
        """
        Generates synthetic data for minority classes.
        Returns: (Balanced DataFrame, Report DataFrame)
        """
        df["temp_label_str"] = df[LABEL_COLUMN].astype(str)
        counts = df["temp_label_str"].value_counts()
        print(
            f"\n   [Balance Check - PRE] Min class size: {counts.min()} | Max: {counts.max()}"
        )

        existing_sentences = set(df["comment_sentence"].str.strip())
        new_rows = []
        report_rows = []

        for label_str, count in counts.items():
            if count < min_samples:
                needed = min_samples - count
                class_subset = df[df["temp_label_str"] == label_str]
                if class_subset.empty:
                    continue

                samples = class_subset["comment_sentence"].tolist()
                orig_label = class_subset[LABEL_COLUMN].iloc[0]

                # Propagate 'combo' if present
                orig_combo = None
                if "combo" in class_subset.columns:
                    orig_combo = class_subset["combo"].iloc[0]

                generated = 0
                attempts = 0
                # Cap attempts to avoid infinite loops if vocabulary is too small
                while generated < needed and attempts < needed * 5:
                    attempts += 1
                    src = random.choice(samples)
                    aug_txt = self.augment(src).strip()

                    # Ensure Global Uniqueness
                    if aug_txt and aug_txt not in existing_sentences:
                        row = {
                            "comment_sentence": aug_txt,
                            LABEL_COLUMN: orig_label,
                            "partition": "train_aug",
                            "index": -1,  # Placeholder
                        }
                        if orig_combo:
                            row["combo"] = orig_combo

                        new_rows.append(row)
                        report_rows.append(
                            {
                                "original_text": src,
                                "augmented_text": aug_txt,
                                "label": label_str,
                                "reason": f"Class has {count} samples (Target {min_samples})",
                            }
                        )
                        existing_sentences.add(aug_txt)
                        generated += 1

        df = df.drop(columns=["temp_label_str"])
        df_report = pd.DataFrame(report_rows)

        if new_rows:
            augmented_df = pd.concat([df, pd.DataFrame(new_rows)], ignore_index=True)
            augmented_df["index"] = range(len(augmented_df))

            temp_counts = augmented_df[LABEL_COLUMN].astype(str).value_counts()
            print(
                f"   [Balance Check - POST] Min class size: {temp_counts.min()} | Max: {temp_counts.max()}"
            )
            return augmented_df, df_report

        return df, df_report


# --- CLEANING LOGIC ---
def clean_training_data_smart(
    df: pd.DataFrame, min_len: int, max_len: int, language: str = "english"
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Performs 'Smart Cleaning' on the Training Set with language-specific heuristics.
    """
    canon = TextCanonicalizer()
    dropped_rows = []

    print(f"   [Clean] Computing heuristics (Language: {language})...")
    df["canon_key"] = df["comment_sentence"].apply(canon.to_canonical)

    # 1. Token Length Filter
    def count_code_tokens(text):
        return len([t for t in re.split(r"[^a-zA-Z0-9]+", str(text)) if t])

    df["temp_token_len"] = df["comment_sentence"].apply(count_code_tokens)

   
    MIN_ALPHA_CHARS = 6
    MAX_SYMBOL_RATIO = 0.50

    # 2. Heuristic Filters (Tiny/Huge/Code)
    def get_heuristics(text):
        s = str(text).strip()
        char_len = len(s)
        if char_len == 0:
            return False, False, 1.0
        
        alpha_len = sum(1 for c in s if c.isalpha())
        
        non_alnum_chars = sum(1 for c in s if not c.isalnum() and not c.isspace())
        symbol_ratio = non_alnum_chars / char_len if char_len > 0 else 0

        is_tiny = alpha_len < MIN_ALPHA_CHARS
        is_huge = char_len > 800
        is_code = symbol_ratio > MAX_SYMBOL_RATIO
        
        return is_tiny, is_huge, is_code

    heuristics = df["comment_sentence"].apply(get_heuristics)
    df["is_tiny"] = [x[0] for x in heuristics]
    df["is_huge"] = [x[1] for x in heuristics]
    df["symbol_ratio"] = [x[2] for x in heuristics] 
    
    
    df["is_code"] = df["symbol_ratio"] > 0.50

    mask_keep = (
        (df["temp_token_len"] >= min_len)
        & (df["temp_token_len"] <= max_len)
        & (~df["is_tiny"])
        & (~df["is_huge"])
        & (~df["is_code"])
    )

    df_dropped_qual = df[~mask_keep].copy()
    if not df_dropped_qual.empty:
        def reason(row):
            if row["is_tiny"]:
                return f"Too Tiny (<{MIN_ALPHA_CHARS} alpha)"
            if row["is_huge"]:
                return "Too Huge (>800 chars)"
            if row["is_code"]:
                return f"Pure Code (>{int(MAX_SYMBOL_RATIO*100)}% symbols)"
            return f"Token Count ({row['temp_token_len']})"

        df_dropped_qual["drop_reason"] = df_dropped_qual.apply(reason, axis=1)
        dropped_rows.append(df_dropped_qual)

    df = df[mask_keep].copy()

    # 3. Semantic Conflicts (Ambiguity)
    df["label_s"] = df[LABEL_COLUMN].astype(str)
    conflict_counts = df.groupby("canon_key")["label_s"].nunique()
    conflicting_keys = conflict_counts[conflict_counts > 1].index

    mask_conflicts = df["canon_key"].isin(conflicting_keys)
    df_dropped_conflicts = df[mask_conflicts].copy()
    if not df_dropped_conflicts.empty:
        df_dropped_conflicts["drop_reason"] = "Semantic Conflict"
        dropped_rows.append(df_dropped_conflicts)

    df = df[~mask_conflicts].copy()

    # 4. Exact Duplicates
    mask_dupes = df.duplicated(subset=["comment_sentence"], keep="first")
    df_dropped_dupes = df[mask_dupes].copy()
    if not df_dropped_dupes.empty:
        df_dropped_dupes["drop_reason"] = "Exact Duplicate"
        dropped_rows.append(df_dropped_dupes)

    df = df[~mask_dupes].copy()

    # Cleanup columns
    cols_to_drop = [
        "canon_key",
        "label_s",
        "temp_token_len",
        "is_tiny",
        "is_huge",
        "is_code",
        "symbol_ratio"
    ]
    df = df.drop(columns=cols_to_drop, errors="ignore")

    if dropped_rows:
        df_report = pd.concat(dropped_rows, ignore_index=True)
        cols_rep = ["index", "comment_sentence", LABEL_COLUMN, "drop_reason"]
        final_cols = [c for c in cols_rep if c in df_report.columns]
        df_report = df_report[final_cols]
    else:
        df_report = pd.DataFrame(columns=["index", "comment_sentence", "drop_reason"])

    print(f"   [Clean] Removed {len(df_report)} rows. Final: {len(df)}.")
    return df, df_report

# --- FEATURE ENGINEERING ---
class FeatureEngineer:
    def __init__(self, config: FeaturePipelineConfig):
        self.config = config
        self.processor = TextProcessor(config=config)
        self.tfidf_vectorizer = TfidfVectorizer(max_features=config.max_features)

    def extract_features_for_check(self, df: pd.DataFrame) -> pd.DataFrame:
        """Extracts metadata features for analysis."""

        def analyze(text):
            s = str(text)
            words = s.split()
            n_words = len(words)
            if n_words == 0:
                return 0, 0, 0
            first_word = words[0].lower()
            starts_verb = (
                1
                if first_word.endswith("s")
                or first_word.startswith("get")
                or first_word.startswith("set")
                else 0
            )
            return (len(s), n_words, starts_verb)

        metrics = df["comment_sentence"].apply(analyze)
        df["f_length"] = [x[0] for x in metrics]
        df["f_word_count"] = [x[1] for x in metrics]
        df["f_starts_verb"] = [x[2] for x in metrics]
        # Calculate MD5 hash for efficient exact duplicate detection in Deepchecks
        df["text_hash"] = df["comment_sentence"].apply(
            lambda x: hashlib.md5(str(x).encode()).hexdigest()
        )
        return df

    def vectorize_and_select(self, df_train, df_test):
        def clean_fn(x):
            return re.sub(r"[^a-zA-Z\s]", "", str(x).lower())

        X_train = self.tfidf_vectorizer.fit_transform(
            df_train["comment_sentence"].apply(clean_fn)
        )
        y_train = np.stack(df_train[LABEL_COLUMN].values)

        # Handling multi-label for Chi2 (using sum or max)
        y_train_sum = (
            y_train.sum(axis=1) if len(y_train.shape) > 1 else y_train
        )
        selector = SelectKBest(
            chi2, k=min(self.config.max_features, X_train.shape[1])
        )
        X_train = selector.fit_transform(X_train, y_train_sum)

        X_test = self.tfidf_vectorizer.transform(
            df_test["comment_sentence"].apply(clean_fn)
        )
        X_test = selector.transform(X_test)

        vocab = [
            self.tfidf_vectorizer.get_feature_names_out()[i]
            for i in selector.get_support(indices=True)
        ]
        return X_train, X_test, vocab


# --- MAIN EXECUTION ---
def main(
    feature_dir: Path = typer.Option(
        INTERIM_DATA_DIR / "features", help="Output dir."
    ),
    reports_root: Path = typer.Option(
        Path("reports/data"), help="Reports root."
    ),
    max_features: int = typer.Option(5000),
    min_comment_length: int = typer.Option(
        2, help="Remove comments shorter than chars."
    ),
    max_comment_length: int = typer.Option(300),
    augment: bool = typer.Option(False, "--augment", help="Enable augmentation."),
    balance_threshold: int = typer.Option(100, help="Min samples per class."),
    run_vectorization: bool = typer.Option(False, "--run-vectorization"),
    run_nlp_check: bool = typer.Option(
        True, "--run-nlp", help="Run Deepchecks NLP suite."
    ),
    custom_tags: str = typer.Option("base", help="Custom tags."),
    save_full_csv: bool = typer.Option(False, "--save-full-csv"),
    languages: List[str] = typer.Option(LANGS, show_default=False),
):

    config = FeaturePipelineConfig(
        True,
        True,
        True,
        max_features,
        min_comment_length,
        max_comment_length,
        augment,
        custom_tags,
    )
    print(f"=== Pipeline ID: {config.hash_id} ===")

    dm = DatasetManager()
    full_dataset = dm.get_dataset()
    fe = FeatureEngineer(config)
    augmenter = SafeAugmenter()

    feat_output_dir = feature_dir / config.hash_id
    feat_output_dir.mkdir(parents=True, exist_ok=True)
    report_output_dir = reports_root / config.hash_id

    for lang in languages:
        print(f"\n{'='*30}\nPROCESSING LANGUAGE: {lang.upper()}\n{'='*30}")
        df_train = full_dataset[f"{lang}_train"].to_pandas()
        df_test = full_dataset[f"{lang}_test"].to_pandas()

        # Standardize Label Format
        for df in [df_train, df_test]:
            if isinstance(df[LABEL_COLUMN].iloc[0], str):
                df[LABEL_COLUMN] = (
                    df[LABEL_COLUMN]
                    .str.replace(r"\s+", ", ", regex=True)
                    .apply(ast.literal_eval)
                )

        lang_report_dir = report_output_dir / lang

        # 1. RAW AUDIT
        print("   >>> Phase 1: Auditing RAW Data")
        df_train_raw = fe.extract_features_for_check(df_train.copy())
        df_test_raw = fe.extract_features_for_check(df_test.copy())
        run_custom_deepchecks(
            df_train_raw, df_test_raw, lang_report_dir, "raw", lang
        )
        if run_nlp_check:
            run_targeted_nlp_checks(
                df_train_raw, df_test_raw, lang_report_dir, "raw"
            )

        # 2. CLEANING & AUGMENTATION
        print("\n   >>> Phase 2: Smart Cleaning & Augmentation")
        df_train, df_dropped = clean_training_data_smart(
            df_train, min_comment_length, max_comment_length, language=lang
        )

        if not df_dropped.empty:
            dropped_path = lang_report_dir / "dropped_rows.csv"
            df_dropped.to_csv(dropped_path, index=False)
            print(f"   [Report] Dropped rows details saved to: {dropped_path}")

        if augment:
            print("   [Augment] Applying Soft Balancing...")
            df_train, df_aug_report = augmenter.apply_balancing(
                df_train, min_samples=balance_threshold
            )

            if not df_aug_report.empty:
                aug_path = lang_report_dir / "augmentation_report.csv"
                df_aug_report.to_csv(aug_path, index=False)
                print(
                    f"   [Report] Augmentation details saved to: {aug_path}"
                )

        # 3. PROCESSED AUDIT
        print("\n   >>> Phase 3: Auditing PROCESSED Data")
        df_train = fe.extract_features_for_check(df_train)
        df_test = fe.extract_features_for_check(df_test)
        run_custom_deepchecks(
            df_train, df_test, lang_report_dir, "processed", lang
        )
        if run_nlp_check:
            run_targeted_nlp_checks(
                df_train, df_test, lang_report_dir, "processed"
            )

        # 4. FINAL PROCESSING & SAVING
        print("\n   >>> Phase 4: Final Processing & Save")
        df_train["comment_clean"] = df_train["comment_sentence"].apply(
            fe.processor.clean_text
        )
        df_test["comment_clean"] = df_test["comment_sentence"].apply(
            fe.processor.clean_text
        )

        if config.use_combo_feature:
            if "combo" in df_train.columns:
                df_train["combo_clean"] = df_train["combo"].apply(
                    fe.processor.clean_text
                )
            if "combo" in df_test.columns:
                df_test["combo_clean"] = df_test["combo"].apply(
                    fe.processor.clean_text
                )

        X_train, X_test, vocab = None, None, []
        if run_vectorization:
            print("   [Vectorization] TF-IDF & Chi2...")
            X_train, X_test, vocab = fe.vectorize_and_select(df_train, df_test)
        def format_label_robust(lbl):
            if hasattr(lbl, "tolist"): # Check if numpy array
                lbl = lbl.tolist()
            return str(lbl)

        df_train[LABEL_COLUMN] = df_train[LABEL_COLUMN].apply(format_label_robust)
        df_test[LABEL_COLUMN] = df_test[LABEL_COLUMN].apply(format_label_robust)

        cols_to_save = [
            "index",
            LABEL_COLUMN,
            "comment_sentence",
            "comment_clean",
        ]
        if "combo" in df_train.columns:
            cols_to_save.append("combo")
        if "combo_clean" in df_train.columns:
            cols_to_save.append("combo_clean")
        meta_cols = [c for c in df_train.columns if c.startswith("f_")]
        cols_to_save.extend(meta_cols)

        print(f"   [Save] Columns: {cols_to_save}")
        df_train[cols_to_save].to_csv(
            feat_output_dir / f"{lang}_train.csv", index=False
        )
        df_test[cols_to_save].to_csv(
            feat_output_dir / f"{lang}_test.csv", index=False
        )

        if run_vectorization and X_train is not None:
            from scipy.sparse import save_npz

            save_npz(feat_output_dir / f"{lang}_train_tfidf.npz", X_train)
            save_npz(feat_output_dir / f"{lang}_test_tfidf.npz", X_test)
            with open(
                feat_output_dir / f"{lang}_vocab.txt", "w", encoding="utf-8"
            ) as f:
                f.write("\n".join(vocab))

    print(f"\nAll Done. Reports in: {report_output_dir}")


if __name__ == "__main__":
    typer.run(main)