src/preprocess.py

"""
MuRIL preprocessing: Unicode normalize, mask URLs/phones/PII, ASCII-only lowercasing,
tokenize with google/muril-base-cased, stratified 70/15/15 split → HuggingFace DatasetDict.
"""

from __future__ import annotations

import json
import logging
import re
import unicodedata
from pathlib import Path

import numpy as np
import pandas as pd
from datasets import Dataset, DatasetDict
from sklearn.model_selection import train_test_split
from sklearn.utils.class_weight import compute_class_weight
from transformers import AutoTokenizer, PreTrainedTokenizerBase

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
log = logging.getLogger(__name__)

ROOT = Path(__file__).resolve().parent.parent
INPUT_CSV = ROOT / "data" / "processed" / "combined.csv"
OUTPUT_DIR = ROOT / "data" / "processed" / "tokenized"
TOKENIZER_SAVE_DIR = OUTPUT_DIR / "muril_tokenizer"

MODEL_NAME = "google/muril-base-cased"
MAX_LENGTH = 128
TRAIN_RATIO = 0.70
VAL_RATIO = 0.15
RANDOM_SEED = 42

CUSTOM_TOKENS = ["[URL]", "[PHONE]", "[EMAIL]", "[AMOUNT]", "[CODE]", "[AADHAAR]", "[PAN]"]

# ASCII lowercasing would mangle placeholders ([URL] -> [url]); restore for tokenizer add_tokens.
_PLACEHOLDER_LOWER_TO_CANON = {
    "[url]": "[URL]",
    "[phone]": "[PHONE]",
    "[email]": "[EMAIL]",
    "[amount]": "[AMOUNT]",
    "[code]": "[CODE]",
    "[aadhaar]": "[AADHAAR]",
    "[pan]": "[PAN]",
}

INDIC_RANGE_RE = re.compile(r"[\u0900-\u0D7F]+")

_URL_RE = re.compile(r"https?://\S+|www\.\S+", re.IGNORECASE)
_PHONE_RE = re.compile(
    r"(\+91[\s\-]?)?[6-9]\d{9}"
    r"|\b\d{10}\b"
    r"|\b\d{5}[\s\-]\d{5}\b"
    r"|\+\d{1,3}[\s\-]\d{6,14}",
)
_EMAIL_RE = re.compile(r"\S+@\S+\.\S+")
_AMOUNT_RE = re.compile(r"₹\s?\d[\d,]*(\.\d+)?|\brs\.?\s?\d[\d,]*", re.IGNORECASE)
_OTP_RE = re.compile(r"\b\d{4,8}\b")
_AADHAAR_RE = re.compile(r"\b\d{4}\s\d{4}\s\d{4}\b")
_PAN_RE = re.compile(r"\b[A-Z]{5}[0-9]{4}[A-Z]\b")
_WHITESPACE = re.compile(r"\s{2,}")

_FLAG_PATTERNS = {
    "urgency": re.compile(
        r"\b(urgent|immediately|right now|abhi|turant|jaldi|do not disconnect"
        r"|tatkaal|asap|within \d+ (minutes?|hours?|seconds?))\b",
        re.IGNORECASE,
    ),
    "authority_impersonation": re.compile(
        r"\b(cbi|ed|ncb|nia|police|officer|inspector|narcotics|cyber cell"
        r"|income tax|enforcement|ib|raw|customs|interpol|court|judge"
        r"|sarkar|government|adhikari)\b",
        re.IGNORECASE,
    ),
    "threat": re.compile(
        r"\b(arrest|warrant|case register|jail|prison|legal action|sue|fir"
        r"|criminal|giriaftari|case darj|pakad|band|block|freeze)\b",
        re.IGNORECASE,
    ),
    "payment_demand": re.compile(
        r"\b(pay|transfer|send money|deposit|upi|neft|rtgs|wire|bhej"
        r"|paisa|rupee|amount|fee|fine|penalty|bail|zamaanat)\b",
        re.IGNORECASE,
    ),
    "secrecy": re.compile(
        r"\b(do not tell|kisi ko mat batao|secret|confidential|family ko mat"
        r"|don.t inform|keep quiet|chup raho|private)\b",
        re.IGNORECASE,
    ),
}


def normalize_unicode(text: str) -> str:
    return unicodedata.normalize("NFC", text)


def lowercase_ascii_only(text: str) -> str:
    return "".join(ch.lower() if ord(ch) < 128 else ch for ch in text)


def replace_sensitive_tokens(text: str) -> str:
    text = _AADHAAR_RE.sub("[AADHAAR]", text)
    text = _PAN_RE.sub("[PAN]", text)
    text = _URL_RE.sub("[URL]", text)
    text = _EMAIL_RE.sub("[EMAIL]", text)
    text = _PHONE_RE.sub("[PHONE]", text)
    text = _AMOUNT_RE.sub("[AMOUNT]", text)
    text = _OTP_RE.sub("[CODE]", text)
    return text


def clean_whitespace(text: str) -> str:
    return _WHITESPACE.sub(" ", text).strip()


def restore_placeholder_tokens(text: str) -> str:
    for lower, canon in _PLACEHOLDER_LOWER_TO_CANON.items():
        text = text.replace(lower, canon)
    return text


def full_normalize(text: str) -> str:
    if not isinstance(text, str) or not text.strip():
        return ""
    text = normalize_unicode(text)
    text = replace_sensitive_tokens(text)
    text = lowercase_ascii_only(text)
    text = restore_placeholder_tokens(text)
    text = clean_whitespace(text)
    return text


def extract_flags(text: str) -> list[str]:
    return [flag for flag, pattern in _FLAG_PATTERNS.items() if pattern.search(text)]


def has_indic_script(text: str) -> bool:
    return bool(INDIC_RANGE_RE.search(text))


def load_and_preprocess(path: Path = INPUT_CSV) -> pd.DataFrame:
    log.info("Loading combined dataset from %s …", path)
    if not path.exists():
        raise FileNotFoundError(
            f"Missing {path}. Run: python src/data_loader.py",
        )
    df = pd.read_csv(path)
    required = {"text", "label", "category"}
    if not required.issubset(df.columns):
        raise ValueError(f"CSV missing columns: {required - set(df.columns)}")

    log.info("  Raw rows: %d", len(df))

    log.info("Extracting scam flags …")
    df = df.copy()
    df["flags"] = df["text"].apply(extract_flags)
    df["flag_count"] = df["flags"].apply(len)
    df["has_indic"] = df["text"].apply(has_indic_script)

    log.info("Normalizing text (Unicode NFC, [URL]/[PHONE]/…, ASCII lower only) …")
    df["text"] = df["text"].apply(full_normalize)

    empty_mask = df["text"].str.len() < 5
    if empty_mask.sum() > 0:
        log.warning("  Dropping %d rows with text < 5 chars after normalization", int(empty_mask.sum()))
        df = df[~empty_mask].reset_index(drop=True)

    df["label"] = df["label"].astype(int)

    log.info("  After normalization: %d rows", len(df))
    log.info("  Indic script present: %d rows (%.1f%%)", df["has_indic"].sum(), df["has_indic"].mean() * 100)
    scam_mean = df.loc[df["label"] == 1, "flag_count"].mean()
    log.info("  Avg flags per scam message: %.2f", scam_mean if np.isfinite(scam_mean) else 0.0)

    return df


def _train_val_test_split(df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    test_size = 1.0 - TRAIN_RATIO - VAL_RATIO
    val_relative = VAL_RATIO / (TRAIN_RATIO + VAL_RATIO)

    def _stratify(series: pd.Series | None):
        return series if series is not None and series.nunique() > 1 else None

    try:
        train_val, test = train_test_split(
            df,
            test_size=test_size,
            stratify=_stratify(df["label"]),
            random_state=RANDOM_SEED,
        )
        train, val = train_test_split(
            train_val,
            test_size=val_relative,
            stratify=_stratify(train_val["label"]),
            random_state=RANDOM_SEED,
        )
    except ValueError as e:
        log.warning("Stratified split failed (%s); falling back to random split.", e)
        train_val, test = train_test_split(df, test_size=test_size, random_state=RANDOM_SEED)
        train, val = train_test_split(train_val, test_size=val_relative, random_state=RANDOM_SEED)

    log.info("Split → train: %d | val: %d | test: %d", len(train), len(val), len(test))
    for name, split in [("train", train), ("val", val), ("test", test)]:
        scam_pct = split["label"].mean() * 100
        log.info("  %s: scam=%.1f%%", name, scam_pct)

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


def prepare_tokenizer() -> PreTrainedTokenizerBase:
    log.info("Loading tokenizer: %s", MODEL_NAME)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    log.info("  Vocab size: %d", tokenizer.vocab_size)
    log.info("  Max model length: %s", tokenizer.model_max_length)
    added = tokenizer.add_tokens(CUSTOM_TOKENS)
    log.info("  Added %d custom tokens: %s", added, CUSTOM_TOKENS)
    return tokenizer


def tokenize_dataset(
    tokenizer: PreTrainedTokenizerBase,
    df_dict: dict[str, pd.DataFrame],
) -> DatasetDict:
    split_name_map = {"train": "train", "val": "validation", "test": "test"}

    def tokenize_batch(batch: dict) -> dict:
        encoded = tokenizer(
            batch["text"],
            truncation=True,
            padding="max_length",
            max_length=MAX_LENGTH,
            return_token_type_ids=True,
        )
        if "token_type_ids" not in encoded:
            encoded["token_type_ids"] = [[0] * MAX_LENGTH for _ in batch["text"]]
        return {
            "input_ids": encoded["input_ids"],
            "attention_mask": encoded["attention_mask"],
            "token_type_ids": encoded["token_type_ids"],
            "labels": batch["label"],
        }

    out: dict[str, Dataset] = {}
    for key, frame in df_dict.items():
        hf_name = split_name_map[key]
        ds = Dataset.from_pandas(frame[["text", "label"]], preserve_index=False)
        ds = ds.map(
            tokenize_batch,
            batched=True,
            batch_size=256,
            desc=f"Tokenizing {hf_name}",
            remove_columns=["text", "label"],
        )
        out[hf_name] = ds

    return DatasetDict(out)


def compute_class_weights(train_df: pd.DataFrame) -> list[float]:
    classes = np.array([0, 1])
    weights = compute_class_weight(
        class_weight="balanced",
        classes=classes,
        y=train_df["label"].values,
    )
    log.info("Class weights → safe: %.3f | scam: %.3f", weights[0], weights[1])
    return weights.tolist()


def inspect_tokenization(text: str, tokenizer: PreTrainedTokenizerBase) -> None:
    normalized = full_normalize(text)
    flags = extract_flags(text)
    encoding = tokenizer(
        normalized,
        return_tensors="pt",
        max_length=MAX_LENGTH,
        truncation=True,
        padding=False,
        return_token_type_ids=True,
    )
    tokens = tokenizer.convert_ids_to_tokens(encoding["input_ids"][0])

    print(f"\nOriginal  : {text}")
    print(f"Normalized: {normalized}")
    print(f"Flags     : {flags}")
    print(f"Tokens ({len(tokens)}): {tokens}")
    print(f"Token IDs : {encoding['input_ids'][0].tolist()}")


def run_pipeline() -> tuple[DatasetDict, list[float], PreTrainedTokenizerBase]:
    df = load_and_preprocess()
    train_df, val_df, test_df = _train_val_test_split(df)
    class_weights = compute_class_weights(train_df)

    tokenizer = prepare_tokenizer()
    dataset = tokenize_dataset(
        tokenizer,
        {"train": train_df, "val": val_df, "test": test_df},
    )

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    dataset.save_to_disk(str(OUTPUT_DIR))
    log.info("Tokenized dataset saved → %s", OUTPUT_DIR)

    weights_path = OUTPUT_DIR / "class_weights.json"
    with open(weights_path, "w", encoding="utf-8") as f:
        json.dump({"weights": class_weights, "labels": ["safe", "scam"]}, f, indent=2)
    log.info("Class weights saved → %s", weights_path)

    TOKENIZER_SAVE_DIR.mkdir(parents=True, exist_ok=True)
    tokenizer.save_pretrained(str(TOKENIZER_SAVE_DIR))
    log.info("Tokenizer saved → %s", TOKENIZER_SAVE_DIR)

    return dataset, class_weights, tokenizer


if __name__ == "__main__":
    import sys

    if hasattr(sys.stdout, "reconfigure"):
        try:
            sys.stdout.reconfigure(encoding="utf-8", errors="replace")
            sys.stderr.reconfigure(encoding="utf-8", errors="replace")
        except (OSError, ValueError):
            pass

    dataset, class_weights, tok = run_pipeline()

    print("\n-- Sample token inspection (Hindi / mixed scam) --")
    inspect_tokenization(
        "CBI officer bol raha hoon. Aapka Aadhaar money laundering case mein hai. "
        "Turant ₹50000 UPI se bhejo https://evil.example/pay warna giraftari.",
        tok,
    )

    print("\n-- Sample token inspection (safe message) --")
    inspect_tokenization(
        "Kal meeting 3 baje hai. Conference room B mein milte hain.",
        tok,
    )

    print("\n-- Dataset summary --")
    for split_name, ds in dataset.items():
        print(f"  {split_name:12s}: {len(ds):5d} examples | columns: {ds.column_names}")

    ratio = class_weights[1] / class_weights[0] if class_weights[0] else 0.0
    print(f"\nClass weights: safe={class_weights[0]:.3f}, scam={class_weights[1]:.3f}")
    if ratio > 2.0 or (1.0 / ratio if ratio else 0) > 2.0:
        print(
            "Note: imbalance >2:1 — use WeightedRandomSampler or class_weight in loss during training.",
        )
    print("\nNext: python src/train.py  (after you add training script)")