data/preprocess.py

"""Preprocessing: COVID-debiasing, text normalization, merge, balance, and split datasets."""

import re
import random
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split

from config import CONFIG, COVID_TERMS, REPLACEMENTS
from data.load_datasets import load_all_datasets


def normalize_text(text):
    """Normalize tweet-style text to remove stylistic shortcuts.

    This prevents the model from learning formatting patterns (lowercase,
    missing punctuation, slang) instead of actual conspiratorial meaning.
    """
    if not isinstance(text, str):
        return ""

    # Remove URLs
    text = re.sub(r"https?://\S+", "", text)
    # Remove @mentions
    text = re.sub(r"@\w+", "", text)
    # Remove hashtag symbols but keep the word
    text = re.sub(r"#(\w+)", r"\1", text)
    # Remove RT prefix
    text = re.sub(r"^RT\s+", "", text)
    # Normalize multiple spaces
    text = re.sub(r"\s+", " ", text).strip()

    # Capitalize first letter of each sentence
    def capitalize_sentences(t):
        sentences = re.split(r"([.!?]+\s*)", t)
        result = []
        for i, part in enumerate(sentences):
            if i == 0 or (i > 0 and re.match(r"[.!?]+\s*", sentences[i - 1])):
                part = part.lstrip()
                if part:
                    part = part[0].upper() + part[1:]
            result.append(part)
        return "".join(result)

    text = capitalize_sentences(text)

    # Add period at end if no sentence-ending punctuation
    if text and text[-1] not in ".!?":
        text += "."

    return text


def build_debias_pattern():
    """Build a compiled regex pattern for all COVID terms.

    Returns a list of (compiled_pattern, category) tuples sorted by term length
    (longest first) to avoid partial matches.
    """
    patterns = []
    for category, terms in COVID_TERMS.items():
        for term in sorted(terms, key=len, reverse=True):
            pattern = re.compile(re.escape(term), re.IGNORECASE)
            patterns.append((pattern, category))
    return patterns


DEBIAS_PATTERNS = build_debias_pattern()


def debias_text(text, rng=None):
    """Replace COVID-specific terms with generic equivalents.

    Applied stochastically during training to prevent topic-shortcutting.
    """
    if rng is None:
        rng = random.Random()

    for pattern, category in DEBIAS_PATTERNS:
        replacements = REPLACEMENTS[category]
        replacement = rng.choice(replacements)
        text = pattern.sub(replacement, text)
    return text


def cap_and_balance(df, max_samples, seed=42):
    """Cap a DataFrame to max_samples, maintaining class balance."""
    if len(df) <= max_samples:
        return df

    label_counts = df["label"].value_counts()
    per_class = max_samples // len(label_counts)

    balanced = []
    for label in label_counts.index:
        subset = df[df["label"] == label]
        n = min(len(subset), per_class)
        balanced.append(subset.sample(n=n, random_state=seed))

    return pd.concat(balanced, ignore_index=True)


def merge_datasets(datasets_dict, max_supplement=None, seed=42):
    """Merge all datasets with optional capping of supplementary data.

    The primary COVID dataset and hard negatives are kept in full;
    other supplements are capped.
    """
    if max_supplement is None:
        max_supplement = CONFIG["max_supplement_samples"]

    # Always keep these in full (small, curated, critical)
    keep_full = {"covid_conspiracy", "hard_negatives"}

    kept = []
    supplements = []
    for name, df in datasets_dict.items():
        if len(df) == 0:
            continue
        if name in keep_full:
            kept.append(df)
        else:
            supplements.append(df)

    if supplements:
        all_supplements = pd.concat(supplements, ignore_index=True)
        all_supplements = cap_and_balance(all_supplements, max_supplement, seed)
    else:
        all_supplements = pd.DataFrame(columns=["text", "label", "source"])

    merged = pd.concat(kept + [all_supplements], ignore_index=True)
    merged = merged.sample(frac=1, random_state=seed).reset_index(drop=True)

    return merged


def split_data(merged_df, seed=42):
    """Stratified split into train/val/test + separate COVID-only test set.

    Returns: train_df, val_df, test_df, covid_test_df
    """
    val_ratio = CONFIG["val_ratio"]
    test_ratio = CONFIG["test_ratio"]
    covid_test_size = CONFIG["covid_test_size"]

    # First, carve out a COVID-only test set
    covid_only = merged_df[merged_df["source"] == "covid_conspiracy"].copy()
    non_covid = merged_df[merged_df["source"] != "covid_conspiracy"].copy()

    if len(covid_only) > covid_test_size:
        covid_test = covid_only.sample(n=covid_test_size, random_state=seed)
        covid_remaining = covid_only.drop(covid_test.index)
    else:
        # If not enough COVID samples, use 20% for COVID test
        covid_test = covid_only.sample(frac=0.2, random_state=seed)
        covid_remaining = covid_only.drop(covid_test.index)

    # Combine remaining COVID + non-COVID for main splits
    main_data = pd.concat([covid_remaining, non_covid], ignore_index=True)

    # Stratified split: first split off test, then split remaining into train/val
    train_val, test = train_test_split(
        main_data,
        test_size=test_ratio,
        stratify=main_data["label"],
        random_state=seed,
    )

    relative_val_ratio = val_ratio / (1 - test_ratio)
    train, val = train_test_split(
        train_val,
        test_size=relative_val_ratio,
        stratify=train_val["label"],
        random_state=seed,
    )

    return (
        train.reset_index(drop=True),
        val.reset_index(drop=True),
        test.reset_index(drop=True),
        covid_test.reset_index(drop=True),
    )


def compute_class_weights(train_df):
    """Compute inverse-frequency class weights for balanced loss."""
    counts = train_df["label"].value_counts().sort_index()
    total = len(train_df)
    weights = total / (len(counts) * counts)
    return weights.values.astype(np.float32)


def preprocess_pipeline():
    """Run the full preprocessing pipeline.

    Returns: train_df, val_df, test_df, covid_test_df, class_weights
    """
    datasets_dict = load_all_datasets()

    # Normalize all text to remove stylistic shortcuts (capitalization, punctuation)
    # This prevents the model from learning tweet-formatting as a conspiracy signal
    print("\nNormalizing text across all datasets...")
    for name, df in datasets_dict.items():
        if len(df) > 0:
            df["text"] = df["text"].apply(normalize_text)
            # Remove any empty texts after normalization
            datasets_dict[name] = df[df["text"].str.strip().str.len() > 0].reset_index(drop=True)
            print(f"  {name}: {len(datasets_dict[name])} samples after normalization")

    print("\nMerging and balancing datasets...")
    merged = merge_datasets(datasets_dict)
    print(f"Merged dataset: {len(merged)} samples")
    print(f"Label distribution:\n{merged['label'].value_counts()}")
    print(f"Source distribution:\n{merged['source'].value_counts()}")

    print("\nSplitting data...")
    train_df, val_df, test_df, covid_test_df = split_data(merged)
    print(f"Train: {len(train_df)}, Val: {len(val_df)}, Test: {len(test_df)}, COVID-test: {len(covid_test_df)}")

    class_weights = compute_class_weights(train_df)
    print(f"Class weights: {class_weights}")

    # Show debiasing example
    covid_samples = train_df[train_df["source"] == "covid_conspiracy"].head(3)
    if len(covid_samples) > 0:
        print("\nDebiasing examples:")
        rng = random.Random(42)
        for _, row in covid_samples.iterrows():
            original = row["text"]
            debiased = debias_text(original, rng)
            print(f"  Original: {original[:100]}...")
            print(f"  Debiased: {debiased[:100]}...")
            print()

    return train_df, val_df, test_df, covid_test_df, class_weights


if __name__ == "__main__":
    train_df, val_df, test_df, covid_test_df, class_weights = preprocess_pipeline()