src/preprocess_text.py

import pandas as pd
from sklearn.model_selection import train_test_split

from .text_cleaning import basic_stats, clean_text


def preprocess_english(en_df: pd.DataFrame) -> pd.DataFrame:
    en_df = en_df.copy()
    en_df.columns = [
        str(c).replace("\ufeff", "").strip().lower() for c in en_df.columns
    ]

    if "text" in en_df.columns and "class" in en_df.columns:
        text_col, label_col = "text", "class"
    elif "text" in en_df.columns and "label" in en_df.columns:
        text_col, label_col = "text", "label"
    else:
        obj_cols = [c for c in en_df.columns if en_df[c].dtype == "object"]
        avg_len = {c: en_df[c].astype(str).str.len().mean() for c in obj_cols}
        text_col = max(avg_len, key=avg_len.get)
        label_candidates = [c for c in obj_cols if c != text_col]
        label_col = min(
            label_candidates, key=lambda c: en_df[c].astype(str).str.len().mean()
        )

    print(f"✅ Using EN text_col: {text_col} | label_col: {label_col}")

    out = en_df[[text_col, label_col]].copy()
    out.columns = ["text", "label_raw"]
    out["text"] = out["text"].apply(clean_text)

    out["label_raw"] = out["label_raw"].astype(str).str.strip().str.lower()

    label_map = {
        "suicide": 1,
        "non-suicide": 0,
        "non suicide": 0,
        "nonsuicide": 0,
        "not suicide": 0,
    }
    out["label"] = out["label_raw"].map(label_map)

    bad = out[out["label"].isna()]["label_raw"].value_counts().head(20)
    print("\n⚠️ Unmapped EN labels (top):")
    print(bad if len(bad) else "None ✅")

    out = out.dropna(subset=["label"]).reset_index(drop=True)
    out["label"] = out["label"].astype(int)
    out["lang"] = "en"
    return out[["text", "label", "lang"]]


def preprocess_bangla(bn_df: pd.DataFrame) -> pd.DataFrame:
    bn_df = bn_df.copy()
    print("BN columns:", list(bn_df.columns))

    bn_cols_lower = {c: str(c).lower() for c in bn_df.columns}
    intention_col = None
    for c in bn_df.columns:
        if "intention" in bn_cols_lower[c]:
            intention_col = c
            break
    if intention_col is None:
        intention_col = bn_df.columns[1]  # fallback

    text_candidates = [c for c in bn_df.columns if c != intention_col]
    text_col = text_candidates[0]

    print(f"✅ Using BN text_col: {text_col} | intention_col: {intention_col}")

    out = bn_df[[text_col, intention_col]].copy()
    out.columns = ["text", "label"]

    out["text"] = out["text"].apply(clean_text)
    out["label"] = pd.to_numeric(out["label"], errors="coerce")
    out = out.dropna(subset=["label"]).reset_index(drop=True)
    out["label"] = out["label"].astype(int)
    out = out[out["label"].isin([0, 1])].reset_index(drop=True)
    out["lang"] = "bn"

    basic_stats("BN", out, "text", "label")
    return out[["text", "label", "lang"]]


def build_text_splits(
    text_all: pd.DataFrame, test_size: float, val_size: float, seed: int
):
    text_all = text_all.copy()
    text_all = text_all[text_all["text"].str.len() > 0].reset_index(drop=True)

    train_df, temp_df = train_test_split(
        text_all, test_size=test_size, random_state=seed, stratify=text_all["label"]
    )
    val_df, test_df = train_test_split(
        temp_df, test_size=val_size, random_state=seed, stratify=temp_df["label"]
    )
    return train_df, val_df, test_df