src/features/augmentation.py

"""
Toxic-only back-translation augmentation with cosine deduplication.

Augments only the positive (toxic) class in the training set via EN→ES→EN,
then drops synthetic samples too similar to the original training corpus.
"""

from __future__ import annotations

import time
from typing import Iterable

import numpy as np
import pandas as pd
from sklearn.metrics.pairwise import cosine_similarity

from src.utils.logger import get_logger

logger = get_logger(__name__)


def toxic_back_translation(
    texts: Iterable[str],
    labels: Iterable[int | bool],
    *,
    source_lang: str = "en",
    pivot_lang: str = "es",
    min_words: int = 3,
    max_words: int = 60,
    rate_limit_every: int = 50,
    rate_limit_sleep_sec: float = 1.0,
    seed: int = 42,
) -> tuple[list[str], list[int]]:
    """
    Back-translate toxic samples only (label == 1).

    Returns parallel lists of augmented texts and labels (all toxic).
    """
    try:
        from deep_translator import GoogleTranslator
    except ImportError as e:
        raise ImportError(
            "Install augmentation deps: uv sync --extra train"
        ) from e

    import random

    random.seed(seed)
    to_pivot = GoogleTranslator(source=source_lang, target=pivot_lang)
    to_source = GoogleTranslator(source=pivot_lang, target=source_lang)

    aug_texts: list[str] = []
    aug_labels: list[int] = []

    pairs = [
        (str(t).strip(), int(bool(l)))
        for t, l in zip(texts, labels, strict=False)
        if int(bool(l)) == 1
    ]
    logger.info(f"Back-translation: {len(pairs)} toxic samples")

    for i, (text, label) in enumerate(pairs):
        words = text.split()
        if len(words) < min_words:
            continue
        try:
            short = " ".join(words[:max_words])
            pivot = to_pivot.translate(short)
            back = to_source.translate(pivot)
            if back and back.strip() and back.strip() != short.strip():
                aug_texts.append(back.strip())
                aug_labels.append(label)
        except Exception as exc:
            logger.warning(f"Back-translation failed at index {i}: {exc}")
            continue

        if rate_limit_every > 0 and i > 0 and i % rate_limit_every == 0:
            time.sleep(rate_limit_sleep_sec)

    logger.info(f"Back-translation produced {len(aug_texts)} samples")
    return aug_texts, aug_labels


def back_translate_texts(
    texts: Iterable[str],
    *,
    source_lang: str = "en",
    pivot_lang: str = "de",
    max_words: int = 60,
    rate_limit_every: int = 50,
    rate_limit_sleep_sec: float = 1.0,
    fallback_to_original: bool = True,
) -> list[str]:
    """
    Back-translate every text (EN→pivot→EN) for test-time augmentation.

    On failure, returns the original string when ``fallback_to_original`` is True.
    """
    try:
        from deep_translator import GoogleTranslator
    except ImportError as e:
        raise ImportError(
            "Install augmentation deps: uv sync --extra train"
        ) from e

    to_pivot = GoogleTranslator(source=source_lang, target=pivot_lang)
    to_source = GoogleTranslator(source=pivot_lang, target=source_lang)

    out: list[str] = []
    for i, raw in enumerate(texts):
        text = str(raw).strip()
        if not text:
            out.append(text)
            continue
        words = text.split()
        short = " ".join(words[:max_words])
        try:
            pivot = to_pivot.translate(short)
            back = to_source.translate(pivot)
            out.append(back.strip() if back and back.strip() else text)
        except Exception as exc:
            logger.warning(f"TTA back-translation failed at index {i}: {exc}")
            out.append(text if fallback_to_original else short)

        if rate_limit_every > 0 and i > 0 and i % rate_limit_every == 0:
            time.sleep(rate_limit_sleep_sec)

    return out


def deduplicate_by_cosine(
    synthetic_texts: list[str],
    synthetic_labels: list[int],
    reference_texts: list[str],
    *,
    threshold: float = 0.95,
    embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
) -> tuple[list[str], list[int]]:
    """
    Remove synthetic samples with max cosine similarity > threshold vs reference.
    """
    if not synthetic_texts:
        return [], []

    try:
        from sentence_transformers import SentenceTransformer
    except ImportError as e:
        raise ImportError(
            "Install augmentation deps: uv sync --extra train"
        ) from e

    model = SentenceTransformer(embedding_model)
    ref_emb = model.encode(reference_texts, show_progress_bar=False, convert_to_numpy=True)
    syn_emb = model.encode(synthetic_texts, show_progress_bar=False, convert_to_numpy=True)

    sims = cosine_similarity(syn_emb, ref_emb)
    max_sim = sims.max(axis=1)

    kept_texts: list[str] = []
    kept_labels: list[int] = []
    dropped = 0
    for text, label, sim in zip(synthetic_texts, synthetic_labels, max_sim, strict=False):
        if sim <= threshold:
            kept_texts.append(text)
            kept_labels.append(label)
        else:
            dropped += 1

    logger.info(
        f"Dedup: kept {len(kept_texts)}/{len(synthetic_texts)} "
        f"(dropped {dropped} with cosine > {threshold})"
    )
    return kept_texts, kept_labels


def augment_toxic_train(
    X_train: pd.Series,
    y_train: pd.Series,
    cfg: dict,
    *,
    seed: int = 42,
) -> tuple[pd.Series, pd.Series]:
    """
    Append toxic-only back-translated samples to training data (with dedup).
    """
    aug_cfg = cfg.get("augmentation", {})
    if not aug_cfg.get("enabled", True):
        return X_train, y_train

    syn_texts, syn_labels = toxic_back_translation(
        X_train.tolist(),
        y_train.tolist(),
        source_lang=aug_cfg.get("source_lang", "en"),
        pivot_lang=aug_cfg.get("pivot_lang", "es"),
        min_words=aug_cfg.get("min_words", 3),
        max_words=aug_cfg.get("max_words", 60),
        rate_limit_every=aug_cfg.get("rate_limit_every", 50),
        rate_limit_sleep_sec=aug_cfg.get("rate_limit_sleep_sec", 1.0),
        seed=seed,
    )

    dedup_cfg = aug_cfg.get("dedup", {})
    if dedup_cfg.get("enabled", True) and syn_texts:
        syn_texts, syn_labels = deduplicate_by_cosine(
            syn_texts,
            syn_labels,
            X_train.tolist(),
            threshold=float(dedup_cfg.get("cosine_threshold", 0.95)),
            embedding_model=dedup_cfg.get(
                "embedding_model", "sentence-transformers/all-MiniLM-L6-v2"
            ),
        )

    if not syn_texts:
        return X_train, y_train

    X_aug = pd.concat(
        [X_train, pd.Series(syn_texts, name=X_train.name)],
        ignore_index=True,
    )
    y_aug = pd.concat(
        [y_train, pd.Series(syn_labels, name=y_train.name)],
        ignore_index=True,
    )
    logger.info(f"Train size after augmentation: {len(X_aug)} (+{len(syn_texts)})")
    return X_aug, y_aug