"""
Data Preprocessing Module
Handles text cleaning, tokenization, lemmatization, and class balancing.
"""

import re
import string
import pandas as pd
import numpy as np
from collections import Counter
import nltk
from nltk.corpus import stopwords
from nltk.stem import WordNetLemmatizer
from textblob import TextBlob

# Download required NLTK data
def download_nltk_data():
    resources = ['stopwords', 'wordnet', 'omw-1.4', 'punkt', 'averaged_perceptron_tagger']
    for r in resources:
        try:
            nltk.download(r, quiet=True)
        except Exception:
            pass

download_nltk_data()

# ─────────────────────────────────────────────
# Label Mapping
# ─────────────────────────────────────────────

# Map dataset labels → unified 7-class schema
LABEL_MAP = {
    "Normal":               "Normal",
    "Anxiety":              "Anxiety",
    "Depression":           "Depression",
    "Suicidal":             "Suicidal",
    "Bipolar":              "Bipolar",
    "Stress":               "Stress",
    "Personality disorder": "Personality Disorder",
    # Aliases (in case of variations)
    "PTSD":                 "Stress",
    "BPD":                  "Personality Disorder",
    "Borderline":           "Personality Disorder",
}

LABEL2ID = {
    "Normal":               0,
    "Anxiety":              1,
    "Depression":           2,
    "Suicidal":             3,
    "Bipolar":              4,
    "Stress":               5,
    "Personality Disorder": 6,
}

ID2LABEL = {v: k for k, v in LABEL2ID.items()}

NUM_LABELS = len(LABEL2ID)

# ─────────────────────────────────────────────
# Text Cleaning
# ─────────────────────────────────────────────

def clean_text(text: str) -> str:
    """
    Thorough text cleaning pipeline:
    - Lowercase
    - Remove URLs, mentions, hashtags
    - Remove special characters / digits
    - Normalize whitespace
    """
    if not isinstance(text, str):
        return ""

    # Lowercase
    text = text.lower()

    # Remove URLs
    text = re.sub(r"http\S+|www\S+|https\S+", "", text, flags=re.MULTILINE)

    # Remove Reddit-style mentions and subreddits
    text = re.sub(r"@\w+|r/\w+|u/\w+", "", text)

    # Remove HTML entities
    text = re.sub(r"&[a-z]+;", " ", text)

    # Remove digits
    text = re.sub(r"\d+", "", text)

    # Remove punctuation (keep apostrophes for contractions)
    text = text.translate(str.maketrans("", "", string.punctuation.replace("'", "")))

    # Remove extra whitespace
    text = re.sub(r"\s+", " ", text).strip()

    return text


def lemmatize_text(text: str) -> str:
    """Lemmatize tokens and remove stopwords."""
    lemmatizer = WordNetLemmatizer()
    stop_words = set(stopwords.words("english"))

    tokens = text.split()
    tokens = [lemmatizer.lemmatize(t) for t in tokens if t not in stop_words and len(t) > 2]
    return " ".join(tokens)


def get_sentiment_features(text: str) -> dict:
    """
    Extract TextBlob sentiment features:
    - polarity (-1 to 1)
    - subjectivity (0 to 1)
    - sentiment_label: Positive / Neutral / Negative
    """
    blob = TextBlob(text)
    polarity = blob.sentiment.polarity
    subjectivity = blob.sentiment.subjectivity

    if polarity > 0.05:
        label = "Positive"
    elif polarity < -0.05:
        label = "Negative"
    else:
        label = "Neutral"

    return {
        "polarity": polarity,
        "subjectivity": subjectivity,
        "sentiment_label": label,
    }


def full_preprocess(text: str, lemmatize: bool = True) -> str:
    """Full preprocessing pipeline for inference."""
    text = clean_text(text)
    if lemmatize:
        text = lemmatize_text(text)
    return text


# ─────────────────────────────────────────────
# Dataset Loading
# ─────────────────────────────────────────────

def load_dataset(csv_path: str, max_len: int = 512) -> pd.DataFrame:
    """
    Load and preprocess the CSV dataset.

    Expected columns: (index), statement, status
    Returns a clean DataFrame with columns: text, label, label_id, sentiment_*
    """
    print(f"[DATA] Loading dataset from: {csv_path}")
    df = pd.read_csv(csv_path, index_col=0)

    # Rename columns
    df.columns = [c.strip().lower() for c in df.columns]
    if "statement" in df.columns:
        df.rename(columns={"statement": "text", "status": "label"}, inplace=True)

    # Drop nulls
    df.dropna(subset=["text", "label"], inplace=True)

    # Normalize labels
    df["label"] = df["label"].str.strip().map(LABEL_MAP)
    df.dropna(subset=["label"], inplace=True)

    # Map to numeric IDs
    df["label_id"] = df["label"].map(LABEL2ID)

    # Clean text
    print("[DATA] Cleaning text...")
    df["text_clean"] = df["text"].apply(clean_text)

    # Filter too-short texts
    df = df[df["text_clean"].str.len() > 5].reset_index(drop=True)

    # Sentiment features
    print("[DATA] Extracting sentiment features...")
    sentiment = df["text"].apply(get_sentiment_features).apply(pd.Series)
    df = pd.concat([df, sentiment], axis=1)

    print(f"[DATA] Dataset loaded: {len(df)} samples")
    print(f"[DATA] Label distribution:\n{df['label'].value_counts()}")

    return df


def compute_class_weights(label_ids: np.ndarray) -> np.ndarray:
    """
    Compute inverse-frequency class weights for weighted loss.
    Returns array of shape [num_classes].
    """
    counts = Counter(label_ids)
    total = sum(counts.values())
    weights = np.zeros(NUM_LABELS, dtype=np.float32)
    for cls_id, count in counts.items():
        weights[cls_id] = total / (NUM_LABELS * count)
    return weights