"""
merge_datasets.py
=================
Merges all cleaned datasets into a unified format with 8-class labels.

Label schema:
  0 - safe
  1 - constructive_criticism
  2 - personal_attack
  3 - hate_speech
  4 - threat
  5 - sexual_harassment
  6 - profanity
  7 - bullying

Usage:
    python preprocessing/merge_datasets.py
"""

import logging
from pathlib import Path

import pandas as pd
import yaml

from preprocessing.clean import clean_text
from preprocessing.normalize import normalize_hinglish

# ---------------------------------------------------------------------------
# Setup
# ---------------------------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
)
logger = logging.getLogger(__name__)

PROJECT_ROOT = Path(__file__).resolve().parent.parent
RAW_DIR = PROJECT_ROOT / "data" / "raw"
PROCESSED_DIR = PROJECT_ROOT / "data" / "processed"
PROCESSED_DIR.mkdir(parents=True, exist_ok=True)

# Load config for label list
with open(PROJECT_ROOT / "configs" / "config.yaml", "r") as f:
    CONFIG = yaml.safe_load(f)

LABELS = CONFIG["labels"]
LABEL_TO_ID = {label: idx for idx, label in enumerate(LABELS)}


# ---------------------------------------------------------------------------
# Per-dataset mapping functions
# ---------------------------------------------------------------------------

def map_civil_comments(raw_path: Path) -> pd.DataFrame:
    """
    Map Civil Comments toxicity scores to 8-class labels.

    Strategy:
      - toxicity < 0.5 → safe
      - threat >= 0.5 → threat
      - identity_attack >= 0.5 → hate_speech
      - sexual_explicit >= 0.5 → sexual_harassment
      - insult >= 0.5 → personal_attack
      - obscene >= 0.5 → profanity
      - severe_toxicity >= 0.5 → bullying (fallback for severe cases)
      - toxicity >= 0.5 (no specific sub-label) → profanity
    """
    logger.info("Mapping Civil Comments...")
    df = pd.read_csv(raw_path)

    def _classify(row):
        if row.get("toxicity", 0) < 0.5:
            return "safe"
        if row.get("threat", 0) >= 0.5:
            return "threat"
        if row.get("identity_attack", 0) >= 0.5:
            return "hate_speech"
        if row.get("sexual_explicit", 0) >= 0.5:
            return "sexual_harassment"
        if row.get("insult", 0) >= 0.5:
            return "personal_attack"
        if row.get("severe_toxicity", 0) >= 0.5:
            return "bullying"
        if row.get("obscene", 0) >= 0.5:
            return "profanity"
        # Generic toxic
        return "profanity"

    df["label"] = df.apply(_classify, axis=1)
    df = df[["text", "label"]]

    logger.info(f"  Civil Comments: {len(df):,} rows")
    logger.info(f"  Distribution:\n{df['label'].value_counts().to_string()}")
    return df


def map_hatexplain(raw_path: Path) -> pd.DataFrame:
    """
    Map HateXplain labels to 8-class schema.

    Mapping:
      - normal → safe
      - hatespeech → hate_speech
      - offensive → profanity
    """
    logger.info("Mapping HateXplain...")
    df = pd.read_csv(raw_path)

    hatex_map = {
        "normal": "safe",
        "hatespeech": "hate_speech",
        "offensive": "profanity",
    }
    df["label"] = df["label"].map(hatex_map)
    df = df.dropna(subset=["label"])
    df = df[["text", "label"]]

    logger.info(f"  HateXplain: {len(df):,} rows")
    logger.info(f"  Distribution:\n{df['label'].value_counts().to_string()}")
    return df


def map_olid(raw_path: Path) -> pd.DataFrame:
    """
    Map OLID labels to 8-class schema.

    Strategy:
      - NOT → safe
      - OFF with subtask_c=IND → personal_attack
      - OFF with subtask_c=GRP → hate_speech
      - OFF (generic) → profanity
    """
    logger.info("Mapping OLID...")
    df = pd.read_csv(raw_path)

    def _classify(row):
        label_a = str(row.get("label_a", "")).strip().upper()
        label_c = str(row.get("label_c", "")).strip().upper()

        if label_a == "NOT":
            return "safe"
        if label_c == "IND":
            return "personal_attack"
        if label_c == "GRP":
            return "hate_speech"
        return "profanity"

    df["label"] = df.apply(_classify, axis=1)
    df = df[["text", "label"]]

    logger.info(f"  OLID: {len(df):,} rows")
    logger.info(f"  Distribution:\n{df['label'].value_counts().to_string()}")
    return df


def map_jigsaw_toxic(raw_path: Path) -> pd.DataFrame:
    """
    Map Jigsaw Toxic multi-label to 8-class schema.

    Strategy (priority order):
      - threat=1 → threat
      - identity_hate=1 → hate_speech
      - insult=1 → personal_attack
      - severe_toxic=1 → bullying
      - obscene=1 → profanity
      - toxic=1 (no specific sub) → profanity
      - all zeros → safe
    """
    logger.info("Mapping Jigsaw Toxic...")
    df = pd.read_csv(raw_path)

    def _classify(row):
        if row.get("threat", 0) == 1:
            return "threat"
        if row.get("identity_hate", 0) == 1:
            return "hate_speech"
        if row.get("insult", 0) == 1:
            return "personal_attack"
        if row.get("severe_toxic", 0) == 1:
            return "bullying"
        if row.get("obscene", 0) == 1:
            return "profanity"
        if row.get("toxic", 0) == 1:
            return "profanity"
        return "safe"

    df["label"] = df.apply(_classify, axis=1)
    df = df[["text", "label"]]

    logger.info(f"  Jigsaw Toxic: {len(df):,} rows")
    logger.info(f"  Distribution:\n{df['label'].value_counts().to_string()}")
    return df


def load_hinglish_seed(seed_path: Path) -> pd.DataFrame:
    """Load the hand-crafted Hinglish seed dataset."""
    logger.info("Loading Hinglish seed data...")
    df = pd.read_csv(seed_path)
    df = df[["text", "label"]]
    logger.info(f"  Hinglish seed: {len(df)} rows")
    return df


def map_125words(raw_path: Path) -> pd.DataFrame:
    """
    Map 125words.csv to 8-class label schema.

    The file uses:
      - "offensive"     → profanity
      - "not offensive" → safe
    """
    logger.info("Mapping 125words dataset...")
    df = pd.read_csv(raw_path)

    label_map = {
        "offensive": "profanity",
        "not offensive": "safe",
    }

    # Normalise column names (strip whitespace)
    df.columns = df.columns.str.strip().str.lower()
    df["label"] = df["label"].str.strip().str.lower().map(label_map)
    df = df.dropna(subset=["label"])
    df = df[["text", "label"]]

    logger.info(f"  125words: {len(df):,} rows")
    logger.info(f"  Distribution:\n{df['label'].value_counts().to_string()}")
    return df


def load_hinglish_augmented(augmented_path: Path) -> pd.DataFrame:
    """Load the pre-generated Hinglish augmented toxic dataset."""
    logger.info("Loading Hinglish augmented dataset...")
    df = pd.read_csv(augmented_path)
    df = df[["text", "label"]]
    logger.info(f"  Hinglish augmented: {len(df):,} rows")
    logger.info(f"  Distribution:\n{df['label'].value_counts().to_string()}")
    return df


# ---------------------------------------------------------------------------
# Merge all datasets
# ---------------------------------------------------------------------------

def merge_all() -> pd.DataFrame:
    """
    Load, map, clean, normalize, and merge all datasets.

    Returns:
        Unified DataFrame with columns: text, label, label_id
    """
    dfs = []

    # --- Civil Comments ---
    civil_path = RAW_DIR / "civil_comments.csv"
    if civil_path.exists():
        dfs.append(map_civil_comments(civil_path))
    else:
        logger.warning("civil_comments.csv not found, skipping.")

    # --- HateXplain ---
    hatex_path = RAW_DIR / "hatexplain.csv"
    if hatex_path.exists():
        dfs.append(map_hatexplain(hatex_path))
    else:
        logger.warning("hatexplain.csv not found, skipping.")

    # --- OLID ---
    olid_path = RAW_DIR / "olid.csv"
    if olid_path.exists():
        dfs.append(map_olid(olid_path))
    else:
        logger.warning("olid.csv not found, skipping.")

    # --- Jigsaw Toxic ---
    jigsaw_path = RAW_DIR / "jigsaw_toxic.csv"
    if jigsaw_path.exists():
        dfs.append(map_jigsaw_toxic(jigsaw_path))
    else:
        logger.warning("jigsaw_toxic.csv not found, skipping.")

    # --- Hinglish Seed ---
    hinglish_path = PROJECT_ROOT / "data" / "hinglish_seed.csv"
    if hinglish_path.exists():
        dfs.append(load_hinglish_seed(hinglish_path))
    else:
        logger.warning("hinglish_seed.csv not found, skipping.")

    # --- Augmented Seed (constructive_criticism + bullying examples) ---
    augmented_path = PROJECT_ROOT / "data" / "augmented_seed.csv"
    if augmented_path.exists():
        logger.info("Loading augmented seed data...")
        aug_df = pd.read_csv(augmented_path)[["text", "label"]]
        logger.info(f"  Augmented seed: {len(aug_df)} rows")
        logger.info(f"  Distribution:\n{aug_df['label'].value_counts().to_string()}")
        dfs.append(aug_df)
    else:
        logger.warning("augmented_seed.csv not found, skipping.")

    # --- 125words Hinglish dataset ---
    words125_path = RAW_DIR / "125words.csv"
    if words125_path.exists():
        dfs.append(map_125words(words125_path))
    else:
        logger.warning("125words.csv not found, skipping.")

    # --- Hinglish Augmented (generated from hindi150swearwords) ---
    hinglish_aug_path = PROCESSED_DIR / "hinglish_augmented.csv"
    if hinglish_aug_path.exists():
        dfs.append(load_hinglish_augmented(hinglish_aug_path))
    else:
        logger.warning(
            "hinglish_augmented.csv not found. "
            "Run: python preprocessing/generate_hinglish_augmented.py"
        )

    if not dfs:
        logger.error("No datasets found! Run download_datasets.py first.")
        return pd.DataFrame(columns=["text", "label", "label_id"])

    # --- Concatenate ---
    logger.info("Merging all datasets...")
    merged = pd.concat(dfs, ignore_index=True)

    # --- Clean & Normalize ---
    logger.info("Applying text cleaning...")
    merged["text"] = merged["text"].apply(clean_text)

    logger.info("Applying Hinglish normalization...")
    merged["text"] = merged["text"].apply(normalize_hinglish)

    # --- Remove empty/short texts ---
    merged = merged[merged["text"].str.len() >= 3]

    # --- Deduplicate ---
    before_dedup = len(merged)
    merged = merged.drop_duplicates(subset=["text"], keep="first")
    logger.info(f"Deduplication: {before_dedup:,} → {len(merged):,} ({before_dedup - len(merged):,} removed)")

    # --- Validate labels ---
    valid_labels = set(LABELS)
    invalid = merged[~merged["label"].isin(valid_labels)]
    if len(invalid) > 0:
        logger.warning(f"Dropping {len(invalid)} rows with invalid labels: {invalid['label'].unique()}")
        merged = merged[merged["label"].isin(valid_labels)]

    # --- Add numeric label ID ---
    merged["label_id"] = merged["label"].map(LABEL_TO_ID)

    # --- Final stats ---
    logger.info(f"\n{'=' * 60}")
    logger.info(f"Final merged dataset: {len(merged):,} rows")
    logger.info(f"Label distribution:")
    for label, count in merged["label"].value_counts().items():
        pct = count / len(merged) * 100
        logger.info(f"  {label:25s}: {count:>8,} ({pct:5.1f}%)")
    logger.info(f"{'=' * 60}")

    return merged


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    merged = merge_all()

    output_path = PROCESSED_DIR / "merged.csv"
    merged.to_csv(output_path, index=False)
    logger.info(f"Saved merged dataset to {output_path}")


if __name__ == "__main__":
    main()