Mental Health Status Classification
from Social Media Posts

class TextPreprocessor:
    """
    Cleans and normalises raw social media text for TF-IDF vectorisation.

    Key steps:
      1. Lowercase conversion
      2. URL & HTML removal
      3. Punctuation removal (preserving apostrophes for contractions)
      4. Stopword removal (NLTK English list)
      5. Lemmatisation (WordNetLemmatizer)
    """

    def __init__(self):
        self.lemmatizer = WordNetLemmatizer()
        self.stop_words = set(stopwords.words('english'))
        # Preserve negation words — critical for mental health context
        negations = {'not', 'no', 'never', 'neither', 'nor',
                     "don't", "won't", "can't", "couldn't", "wouldn't",
                     "shouldn't", "isn't", "wasn't", "aren't"}
        self.stop_words -= negations

    def clean(self, text: str) -> str:
        """Apply full preprocessing pipeline to a single text string."""
        if pd.isna(text) or not str(text).strip():
            return ""

        text = str(text).lower()

        # Remove URLs
        text = re.sub(r'https?://\S+|www\.\S+', ' ', text)

        # Remove HTML tags
        text = re.sub(r'<[^>]+>', ' ', text)

        # Remove non-alphabetic characters (keep spaces)
        text = re.sub(r"[^a-z\s']", ' ', text)

        # Tokenise
        tokens = word_tokenize(text)

        # Remove stopwords and lemmatise
        tokens = [
            self.lemmatizer.lemmatize(tok)
            for tok in tokens
            if tok not in self.stop_words and len(tok) > 2
        ]

        return ' '.join(tokens)

    def fit_transform_labels(self, labels: pd.Series) -> pd.Series:
        """Map raw dataset labels to canonical four-class schema."""
        mapped = labels.str.lower().str.strip().map(LABEL_MAP)
        # Any unrecognised label → 'normal' (conservative assumption)
        unmapped = mapped.isna().sum()
        if unmapped > 0:
            logging.warning(f"{unmapped} labels not in LABEL_MAP — assigned 'normal'")
        mapped = mapped.fillna('normal')
        return mapped


def load_and_prepare_data(filepath: str) -> tuple:
    """
    Load dataset, apply preprocessing, and return cleaned (X, y) pair.

    Parameters
    ----------
    filepath : str
        Path to the CSV dataset file.

    Returns
    -------
    X : list of str
        Cleaned text for each post.
    y : pd.Series
        Canonical labels (depression | anxiety | ptsd | normal).
    """
    df = pd.read_csv(filepath)
    logging.info(f"Loaded dataset: {df.shape[0]} rows, {df.shape[1]} columns")

    # Detect text and label columns flexibly
    text_col = next((c for c in df.columns
                     if c.lower() in ['text', 'statement', 'post', 'content']), None)
    label_col = next((c for c in df.columns
                      if c.lower() in ['label', 'status', 'category', 'class']), None)

    assert text_col, "❌ Could not detect text column."
    assert label_col, "❌ Could not detect label column."

    # Drop rows with missing values in text or label
    df = df.dropna(subset=[text_col, label_col])
    df = df[df[text_col].str.strip().str.len() >= 10]

    preprocessor = TextPreprocessor()

    # Preprocess text
    logging.info("Preprocessing text (this may take a moment)...")
    X = df[text_col].apply(preprocessor.clean).tolist()

    # Map labels to canonical four-class schema
    y = preprocessor.fit_transform_labels(df[label_col])

    logging.info(f"Label distribution:\n{y.value_counts()}")
    return X, y

def build_tfidf_vectorizer() -> TfidfVectorizer:
    """
    Construct a TF-IDF vectoriser configured for mental health text classification.

    Configuration rationale
    ──────────────────────
    max_features=8000  : Captures rich vocabulary without excessive dimensionality.
    ngram_range=(1,2)  : Unigrams + bigrams. Bigrams capture phrases like
                         "can't sleep", "feeling hopeless", "panic attack".
    sublinear_tf=True  : Applies log(1 + tf). Prevents long posts from
                         dominating short ones in raw term frequency.
    min_df=3           : Ignores extremely rare terms (noise / typos).
    max_df=0.85        : Ignores near-universal terms that carry no signal.
    analyzer='word'    : Word-level tokenisation.
    """
    return TfidfVectorizer(
        max_features=8000,
        ngram_range=(1, 2),
        sublinear_tf=True,
        min_df=3,
        max_df=0.85,
        analyzer='word',
        strip_accents='unicode',
        decode_error='replace'
    )


def inspect_top_features(vectorizer: TfidfVectorizer, model, n: int = 15):
    """
    Print top-N TF-IDF features for each class as per SVM coefficients.
    Useful for academic interpretability analysis.
    """
    feature_names = np.array(vectorizer.get_feature_names_out())

    # CalibratedClassifierCV wraps the base LinearSVC
    base_model = model.calibrated_classifiers_[0].estimator

    for i, cls in enumerate(CLASSES):
        if i < len(base_model.coef_):
            coefs = base_model.coef_[i]
            top_pos = feature_names[np.argsort(coefs)[-n:]][::-1]
            print(f"\n🔍 Top {n} features for '{cls}':")
            print('  ' + ', '.join(top_pos))


# ── Example: visualise TF-IDF term importance ──────────────────────────────
def plot_top_tfidf_terms(vectorizer, model, top_n: int = 10):
    """Bar chart of top TF-IDF terms per category."""
    feature_names = np.array(vectorizer.get_feature_names_out())
    base_model = model.calibrated_classifiers_[0].estimator
    fig, axes = plt.subplots(1, len(CLASSES), figsize=(18, 5), sharey=False)

    colors = ['#6b4c9a', '#c0572a', '#8a6c2a', '#2a7a4a']

    for idx, (cls, ax, color) in enumerate(zip(CLASSES, axes, colors)):
        if idx < len(base_model.coef_):
            coefs = base_model.coef_[idx]
            top_idx = np.argsort(coefs)[-top_n:]
            top_terms = feature_names[top_idx]
            top_scores = coefs[top_idx]

            ax.barh(top_terms, top_scores, color=color, alpha=0.85)
            ax.set_title(f'{cls.upper()}', fontweight='bold', fontsize=11)
            ax.set_xlabel('SVM Coefficient', fontsize=8)
            ax.tick_params(axis='y', labelsize=8)

    plt.suptitle('Top TF-IDF Features per Mental Health Category',
                 fontsize=13, fontweight='bold', y=1.02)
    plt.tight_layout()
    plt.savefig('tfidf_top_features.png', dpi=150, bbox_inches='tight')
    plt.show()

class MentalHealthSVMClassifier:
    """
    Four-class SVM classifier for mental health post classification.

    Architecture
    ────────────
    • TF-IDF vectoriser (fitted on training data only — no leakage)
    • LinearSVC wrapped with CalibratedClassifierCV for probability output
    • class_weight='balanced' to handle label imbalance
    • Normal Guard: if max probability for any MH class < threshold → Normal

    Parameters
    ----------
    confidence_threshold : float, default=0.45
        Minimum confidence required to classify a post as a mental health
        category. Posts below this threshold are returned as 'normal'.
        Tune on the validation set to control false-positive rate.
    """

    def __init__(self, confidence_threshold: float = 0.45):
        self.confidence_threshold = confidence_threshold
        self.vectorizer = build_tfidf_vectorizer()
        self.model = None
        self.label_encoder = LabelEncoder()
        self.classes_ = None

    def fit(self, X_train: list, y_train: pd.Series):
        """Fit vectoriser and SVM on training data."""
        logging.info("Fitting TF-IDF vectoriser...")
        X_vec = self.vectorizer.fit_transform(X_train)

        # Encode string labels to integers
        y_enc = self.label_encoder.fit_transform(y_train)
        self.classes_ = self.label_encoder.classes_

        logging.info(f"Training SVM on {X_vec.shape[0]} samples, {X_vec.shape[1]} features...")

        base_svc = LinearSVC(
            class_weight='balanced',
            max_iter=2000,
            C=1.0,
            random_state=42
        )
        # CalibratedClassifierCV uses cross-validation internally
        # to produce well-calibrated probability estimates
        self.model = CalibratedClassifierCV(base_svc, cv=5, method='sigmoid')
        self.model.fit(X_vec, y_enc)

        logging.info("✅ Model trained successfully.")
        return self

    def predict_with_confidence(self, texts: list) -> list:
        """
        Predict category for each text with Normal Guard applied.

        Returns
        -------
        list of dict: [{
            'label': str,         # predicted category
            'confidence': float,  # confidence for predicted label
            'probabilities': dict # full probability vector
        }]
        """
        X_vec = self.vectorizer.transform(texts)
        proba_matrix = self.model.predict_proba(X_vec)  # shape: (n, 4)

        results = []
        normal_idx = list(self.classes_).index('normal') \
                     if 'normal' in self.classes_ else -1

        for proba in proba_matrix:
            prob_dict = {cls: float(p)
                        for cls, p in zip(self.classes_, proba)}

            # ─── NORMAL GUARD ────────────────────────────────────────────
            # Compute max confidence across non-normal MH classes only.
            # If none exceed the threshold, assign Normal.
            # This is the KEY FIX for the false-positive problem.
            mh_classes = [c for c in self.classes_ if c != 'normal']
            max_mh_conf = max(prob_dict[c] for c in mh_classes)

            if max_mh_conf < self.confidence_threshold:
                predicted_label = 'normal'
                confidence = prob_dict['normal']
            else:
                predicted_label = max(mh_classes, key=lambda c: prob_dict[c])
                confidence = prob_dict[predicted_label]
            # ─────────────────────────────────────────────────────────────

            results.append({
                'label': predicted_label,
                'confidence': confidence,
                'probabilities': prob_dict
            })

        return results

    def predict(self, texts: list) -> list:
        """Convenience method: return only predicted labels."""
        return [r['label'] for r in self.predict_with_confidence(texts)]

def evaluate_classifier(classifier: MentalHealthSVMClassifier,
                          X_test: list, y_test: pd.Series) -> dict:
    """
    Comprehensive evaluation of the 4-class SVM classifier.

    Metrics reported
    ─────────────────
    • Per-class Precision, Recall, F1-score
    • Macro-averaged Precision, Recall, F1
    • Normal class Specificity (True Negative Rate)
      → Critical: measures how well we avoid false alarms on healthy posts
    • Confusion matrix (normalised)
    """
    y_pred = classifier.predict(X_test)

    logging.info("\n══════════ CLASSIFICATION REPORT ══════════")
    print(classification_report(y_test, y_pred, target_names=CLASSES,
                                digits=4, zero_division=0))

    # ── Per-class metrics ──────────────────────────────────────────────────
    metrics = {}
    for cls in CLASSES:
        binary_y_true = (y_test == cls).astype(int)
        binary_y_pred = (pd.Series(y_pred) == cls).astype(int)

        TP = ((binary_y_true == 1) & (binary_y_pred == 1)).sum()
        TN = ((binary_y_true == 0) & (binary_y_pred == 0)).sum()
        FP = ((binary_y_true == 0) & (binary_y_pred == 1)).sum()
        FN = ((binary_y_true == 1) & (binary_y_pred == 0)).sum()

        sensitivity = TP / (TP + FN) if (TP + FN) > 0 else 0.0
        specificity = TN / (TN + FP) if (TN + FP) > 0 else 0.0

        metrics[cls] = {
            'sensitivity': sensitivity,
            'specificity': specificity,
            'precision': precision_score(binary_y_true, binary_y_pred, zero_division=0),
            'recall': recall_score(binary_y_true, binary_y_pred, zero_division=0),
            'f1': f1_score(binary_y_true, binary_y_pred, zero_division=0)
        }

    logging.info("\n══════════ SENSITIVITY & SPECIFICITY ══════════")
    for cls, m in metrics.items():
        print(f"  {cls:12s} | Sensitivity: {m['sensitivity']:.4f} | Specificity: {m['specificity']:.4f}")

    # ── Confusion matrix ────────────────────────────────────────────────────
    cm = confusion_matrix(y_test, y_pred, labels=CLASSES, normalize='true')
    fig, ax = plt.subplots(figsize=(7, 6))
    sns.heatmap(cm, annot=True, fmt='.2f', cmap='Blues',
                xticklabels=CLASSES, yticklabels=CLASSES, ax=ax)
    ax.set_xlabel('Predicted Label', fontsize=11)
    ax.set_ylabel('True Label', fontsize=11)
    ax.set_title('Normalised Confusion Matrix — TF-IDF + SVM',
                 fontsize=12, fontweight='bold')
    plt.tight_layout()
    plt.savefig('confusion_matrix.png', dpi=150)
    plt.show()

    return metrics


def cross_validate_pipeline(X: list, y: pd.Series, n_splits: int = 5) -> None:
    """
    5-fold stratified cross-validation to report stable performance estimates.
    Prevents overfitting to a single train/test split.
    """
    logging.info(f"\n5-fold Stratified Cross-Validation...")
    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)

    # Build a sklearn Pipeline for CV compatibility
    pipeline = Pipeline([
        ('tfidf', build_tfidf_vectorizer()),
        ('svm', CalibratedClassifierCV(
            LinearSVC(class_weight='balanced', max_iter=2000, random_state=42),
            cv=3
        ))
    ])

    X_arr = np.array(X)
    y_arr = y.values

    f1_scores = cross_val_score(pipeline, X_arr, y_arr,
                                cv=skf, scoring='f1_macro')

    print(f"\n  CV F1-Macro Scores: {f1_scores.round(4)}")
    print(f"  Mean F1-Macro: {f1_scores.mean():.4f} ± {f1_scores.std():.4f}")

"""
Complete end-to-end execution pipeline.
Run this file directly: python train_model.py
"""

if __name__ == '__main__':

    DATASET_PATH = 'mental_health.csv'
    CONFIDENCE_THRESHOLD = 0.45  # Tune this to balance sensitivity vs specificity
    TEST_SIZE = 0.20
    RANDOM_STATE = 42

    start_time = time.time()
    print("=" * 65)
    print("  MENTAL HEALTH STATUS CLASSIFICATION — TF-IDF + SVM")
    print("=" * 65)

    # ── Step 1: Load & Preprocess ───────────────────────────────────────
    print("\n[Step 1/5] Loading and preprocessing data...")
    X, y = load_and_prepare_data(DATASET_PATH)
    print(f"  Ready: {len(X)} samples across {y.nunique()} classes")

    # ── Step 2: Train/Test Split ─────────────────────────────────────────
    print("\n[Step 2/5] Splitting dataset (80/20 stratified)...")
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=TEST_SIZE, random_state=RANDOM_STATE, stratify=y
    )
    print(f"  Train: {len(X_train)} | Test: {len(X_test)}")

    # ── Step 3: Cross-Validation ──────────────────────────────────────────
    print("\n[Step 3/5] Running 5-fold cross-validation...")
    cross_validate_pipeline(X_train, y_train)

    # ── Step 4: Train Final Model ─────────────────────────────────────────
    print("\n[Step 4/5] Training final SVM classifier...")
    classifier = MentalHealthSVMClassifier(confidence_threshold=CONFIDENCE_THRESHOLD)
    classifier.fit(X_train, y_train)

    # ── Step 5: Evaluate ──────────────────────────────────────────────────
    print("\n[Step 5/5] Evaluating on held-out test set...")
    metrics = evaluate_classifier(classifier, X_test, y_test)

    # ── Feature Importance ────────────────────────────────────────────────
    print("\n[Bonus] Top discriminative TF-IDF features per category:")
    inspect_top_features(classifier.vectorizer, classifier.model, n=10)
    plot_top_tfidf_terms(classifier.vectorizer, classifier.model)

    elapsed = time.time() - start_time
    print(f"\n✅ Total pipeline time: {elapsed:.1f}s")
    print("=" * 65)

    # ── Inference example ────────────────────────────────────────────────
    test_posts = [
        "I had a great time at the park today, the weather was lovely.",
        "I can't stop crying and I don't even know why. Everything feels pointless.",
        "My heart was racing during the meeting, I thought I was going to collapse.",
        "Every time I hear a loud sound I jump. I keep seeing it happen again.",
    ]

    print("\n── Inference on Sample Posts ───────────────────────────────────────")
    preprocessor = TextPreprocessor()
    cleaned_posts = [preprocessor.clean(p) for p in test_posts]
    predictions = classifier.predict_with_confidence(cleaned_posts)

    for post, pred in zip(test_posts, predictions):
        print(f"\n  Post    : {post[:70]}...")
        print(f"  → Label : {pred['label'].upper()} ({pred['confidence']:.1%} confidence)")
        probs = ' | '.join(f"{k}: {v:.2f}" for k, v in pred['probabilities'].items())
        print(f"     Probs: {probs}")