train_model.py

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split, cross_val_score, GridSearchCV
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.linear_model import LogisticRegression, SGDClassifier
from sklearn.ensemble import (
    RandomForestClassifier,
    GradientBoostingClassifier,
    AdaBoostClassifier,
    VotingClassifier,
)
from sklearn.svm import SVC, LinearSVC
from sklearn.tree import DecisionTreeClassifier
from sklearn.neighbors import KNeighborsClassifier
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from sklearn.preprocessing import LabelEncoder
from sklearn.cluster import KMeans
from collections import Counter
import seaborn as sns
import matplotlib.pyplot as plt
import joblib
import warnings
import os


# ================================================================================
# This file implements a complete pipeline for training and evaluating
# Vietnamese news classification models using multiple algorithms and ensemble methods.
# Output: Trained model artifacts and evaluation reports in 'train_model_assets' folder.
# ================================================================================


warnings.filterwarnings("ignore")

# Support Vietnamese characters in matplotlib
plt.rcParams["font.family"] = "DejaVu Sans"

# ==================== ASSET FOLDER SETUP ====================
ASSET_FOLDER = "train_model_assets"
os.makedirs(ASSET_FOLDER, exist_ok=True)
print(f"Asset folder created: {ASSET_FOLDER}")


# ==================== DATA LOADING ====================
def load_cleaned_data(file_path):
    """
    Load preprocessed data from CSV file.

    Args:
        file_path: Path to cleaned data CSV file

    Returns:
        DataFrame containing preprocessed data
    """
    print(f"Loading data from {file_path}...")
    df = pd.read_csv(file_path, encoding="utf-8")
    print(f"Total articles: {len(df)}")
    print(f"Number of categories: {df['category'].nunique()}")

    print(f"\nDistribution by category:")
    print(df["category"].value_counts())

    # Display source statistics if available
    if "source" in df.columns:
        print(f"\nDistribution by source:")
        print(df["source"].value_counts())

        print(f"\nCross-tabulation (category x source):")
        print(pd.crosstab(df["category"], df["source"]))

        # Analyze source imbalance
        print(f"\nAverage articles per source:")
        source_stats = df.groupby("source")["category"].count()
        print(source_stats)
        print(
            f"\nImbalance: Min={source_stats.min()}, Max={source_stats.max()}, Gap={source_stats.max() - source_stats.min()}"
        )

    return df


# ==================== DATA SPLITTING ====================
def split_data(df, test_size=0.2, val_size=0.1, random_state=42):
    """
    Split data into train, validation, and test sets with stratification.

    Args:
        df: Input DataFrame
        test_size: Proportion of test set
        val_size: Proportion of validation set
        random_state: Random seed for reproducibility

    Returns:
        Tuple of (X_train, X_val, X_test, y_train, y_val, y_test, label_encoder)
    """
    print(f"\n{'='*70}")
    print("DATA SPLITTING (TRAIN/VALIDATION/TEST)")
    print(f"{'='*70}")

    X = df["content"]
    y = df["category"]

    # Encode labels for XGBoost compatibility
    label_encoder = LabelEncoder()
    y_encoded = label_encoder.fit_transform(y)

    print(f"\nLabel encoding mapping:")
    for i, label in enumerate(label_encoder.classes_):
        print(f"  {i} -> {label}")

    # Split train and temp (test + validation)
    X_train, X_temp, y_train, y_temp = train_test_split(
        X,
        y_encoded,
        test_size=(test_size + val_size),
        random_state=random_state,
        stratify=y_encoded,
    )

    # Split temp into test and validation
    test_ratio = test_size / (test_size + val_size)
    X_val, X_test, y_val, y_test = train_test_split(
        X_temp, y_temp, test_size=test_ratio, random_state=random_state, stratify=y_temp
    )

    print(f"\nTrain set: {len(X_train)} articles ({len(X_train)/len(df)*100:.1f}%)")
    print(f"Validation set: {len(X_val)} articles ({len(X_val)/len(df)*100:.1f}%)")
    print(f"Test set: {len(X_test)} articles ({len(X_test)/len(df)*100:.1f}%)")

    return X_train, X_val, X_test, y_train, y_val, y_test, label_encoder


# ==================== TEXT VECTORIZATION ====================
def vectorize_text(X_train, X_val, X_test, max_features=5000):
    """
    Convert text to TF-IDF feature vectors.

    Args:
        X_train: Training text data
        X_val: Validation text data
        X_test: Test text data
        max_features: Maximum number of features

    Returns:
        Tuple of (X_train_tfidf, X_val_tfidf, X_test_tfidf, vectorizer)
    """
    print(f"\n{'='*70}")
    print("TEXT VECTORIZATION (TF-IDF)")
    print(f"{'='*70}")

    vectorizer = TfidfVectorizer(
        max_features=max_features,
        ngram_range=(1, 2),
        min_df=2,
        max_df=0.8,
        sublinear_tf=True,
    )

    X_train_tfidf = vectorizer.fit_transform(X_train)
    X_val_tfidf = vectorizer.transform(X_val)
    X_test_tfidf = vectorizer.transform(X_test)

    print(f"Number of features: {X_train_tfidf.shape[1]}")
    print(f"Train shape: {X_train_tfidf.shape}")
    print(f"Validation shape: {X_val_tfidf.shape}")
    print(f"Test shape: {X_test_tfidf.shape}")

    return X_train_tfidf, X_val_tfidf, X_test_tfidf, vectorizer


# ==================== MODEL TRAINING ====================
def train_models(X_train, y_train):
    """
    Train multiple classification models.

    Args:
        X_train: Training features
        y_train: Training labels

    Returns:
        Dictionary of trained models
    """
    print(f"\n{'='*70}")
    print("MODEL TRAINING (10+ ALGORITHMS)")
    print(f"{'='*70}")

    models = {
        "Naive Bayes": MultinomialNB(alpha=0.1),
        "Logistic Regression": LogisticRegression(
            max_iter=1000, C=1.0, random_state=42, solver="saga", n_jobs=-1
        ),
        "SGD Classifier": SGDClassifier(
            loss="hinge",
            penalty="l2",
            alpha=0.0001,
            random_state=42,
            max_iter=1000,
            n_jobs=-1,
        ),
        "Linear SVM": LinearSVC(C=1.0, max_iter=1000, random_state=42),
        "Decision Tree": DecisionTreeClassifier(
            max_depth=50, min_samples_split=10, random_state=42
        ),
        "Random Forest": RandomForestClassifier(
            n_estimators=200,
            max_depth=50,
            min_samples_split=5,
            random_state=42,
            n_jobs=-1,
        ),
        "Gradient Boosting": GradientBoostingClassifier(
            n_estimators=100, learning_rate=0.1, max_depth=5, random_state=42
        ),
        "XGBoost": XGBClassifier(
            n_estimators=100,
            learning_rate=0.1,
            max_depth=5,
            random_state=42,
            n_jobs=-1,
            eval_metric="mlogloss",
        ),
        "AdaBoost": AdaBoostClassifier(
            n_estimators=100, learning_rate=1.0, random_state=42, algorithm="SAMME"
        ),
        "SVM (RBF Kernel)": SVC(kernel="rbf", C=1.0, gamma="scale", random_state=42),
        "K-Nearest Neighbors": KNeighborsClassifier(
            n_neighbors=5, weights="distance", n_jobs=-1
        ),
    }

    trained_models = {}

    for name, model in models.items():
        print(f"\n[{name}] Training model...")
        try:
            model.fit(X_train, y_train)
            trained_models[name] = model
            print(f"[{name}] ✓ Training completed")
        except Exception as e:
            print(f"[{name}] ✗ Error: {e}")

    return trained_models


# ==================== ENSEMBLE MODEL ====================
def train_ensemble_model(trained_models, X_train, y_train):
    """
    Train ensemble model using voting classifier.

    Args:
        trained_models: Dictionary of trained base models
        X_train: Training features
        y_train: Training labels

    Returns:
        Trained ensemble model or None
    """
    print(f"\n{'='*70}")
    print("ENSEMBLE MODEL TRAINING (VOTING CLASSIFIER)")
    print(f"{'='*70}")

    estimators = [
        ("lr", trained_models.get("Logistic Regression")),
        ("rf", trained_models.get("Random Forest")),
        ("xgb", trained_models.get("XGBoost")),
    ]

    estimators = [(name, model) for name, model in estimators if model is not None]

    if len(estimators) >= 2:
        ensemble = VotingClassifier(
            estimators=estimators,
            voting="soft",
            n_jobs=-1,
        )

        print(f"Training ensemble model with {len(estimators)} base models...")
        ensemble.fit(X_train, y_train)
        print(f"✓ Ensemble model training completed")

        return ensemble
    else:
        print("Warning: Not enough models to create ensemble")
        return None


# ==================== VALIDATION EVALUATION ====================
def evaluate_on_validation(models, X_val, y_val):
    """
    Evaluate models on validation set.

    Args:
        models: Dictionary of trained models
        X_val: Validation features
        y_val: Validation labels

    Returns:
        Dictionary of validation accuracies
    """
    print(f"\n{'='*70}")
    print("VALIDATION SET EVALUATION")
    print(f"{'='*70}")

    val_results = {}

    for name, model in models.items():
        try:
            y_pred = model.predict(X_val)
            accuracy = accuracy_score(y_val, y_pred)
            val_results[name] = accuracy
            print(f"[{name}] Validation Accuracy: {accuracy:.4f} ({accuracy*100:.2f}%)")
        except Exception as e:
            print(f"[{name}] Error during evaluation: {e}")
            val_results[name] = 0.0

    return val_results


# ==================== TEST SET EVALUATION ====================
def evaluate_on_test(models, X_test, y_test, df_test=None, label_encoder=None):
    """
    Evaluate models on test set with detailed metrics.

    Args:
        models: Dictionary of trained models
        X_test: Test features
        y_test: Test labels
        df_test: Test DataFrame (optional, for source analysis)
        label_encoder: Label encoder for decoding predictions

    Returns:
        Dictionary of test accuracies
    """
    print(f"\n{'='*70}")
    print("TEST SET EVALUATION (FINAL)")
    print(f"{'='*70}")

    test_results = {}

    for name, model in models.items():
        try:
            y_pred = model.predict(X_test)
            accuracy = accuracy_score(y_test, y_pred)
            test_results[name] = accuracy

            print(f"\n{'='*70}")
            print(f"Model: {name}")
            print(f"{'='*70}")
            print(f"Test Accuracy: {accuracy:.4f} ({accuracy*100:.2f}%)")
            print(f"\nClassification Report:")

            if label_encoder:
                target_names = label_encoder.classes_
                print(
                    classification_report(
                        y_test, y_pred, target_names=target_names, zero_division=0
                    )
                )
            else:
                print(classification_report(y_test, y_pred, zero_division=0))

            # Source-specific accuracy if available
            if df_test is not None and "source" in df_test.columns:
                print(f"\nAccuracy by source:")

                for source in df_test["source"].unique():
                    source_mask = (df_test["source"] == source).values
                    y_test_source = y_test[source_mask]
                    y_pred_source = y_pred[source_mask]

                    if len(y_test_source) > 0:
                        source_acc = accuracy_score(y_test_source, y_pred_source)
                        print(
                            f"  {source:15s}: {source_acc:.4f} ({source_acc*100:.2f}%) - {len(y_test_source)} samples"
                        )

        except Exception as e:
            print(f"[{name}] Error during evaluation: {e}")
            test_results[name] = 0.0

    return test_results


# ==================== CROSS-VALIDATION ====================
def perform_cross_validation(models, X_train, y_train, cv=5):
    """
    Perform k-fold cross-validation on models.

    Args:
        models: Dictionary of models to validate
        X_train: Training features
        y_train: Training labels
        cv: Number of folds

    Returns:
        Dictionary of cross-validation scores
    """
    print(f"\n{'='*70}")
    print(f"CROSS-VALIDATION ({cv}-FOLD)")
    print(f"{'='*70}")

    cv_results = {}

    for name, model in models.items():
        try:
            print(f"\n[{name}] Running cross-validation...")
            cv_scores = cross_val_score(
                model, X_train, y_train, cv=cv, scoring="accuracy", n_jobs=-1
            )
            cv_results[name] = cv_scores

            print(f"[{name}] CV Scores: {[f'{score:.4f}' for score in cv_scores]}")
            print(
                f"[{name}] Mean: {cv_scores.mean():.4f} (+/- {cv_scores.std() * 2:.4f})"
            )
        except Exception as e:
            print(f"[{name}] Error during CV: {e}")
            cv_results[name] = np.array([0.0] * cv)

    return cv_results


# ==================== RESULTS COMPARISON ====================
def compare_results(val_results, test_results, cv_results):
    """
    Compare and rank all models based on multiple metrics.

    Args:
        val_results: Validation accuracies
        test_results: Test accuracies
        cv_results: Cross-validation scores

    Returns:
        Tuple of (best_model_name, comparison_dataframe)
    """
    print(f"\n{'='*70}")
    print("MODEL COMPARISON AND RANKING")
    print(f"{'='*70}")

    comparison_df = pd.DataFrame(
        {
            "Validation Acc": val_results,
            "Test Acc": test_results,
            "CV Mean": {name: scores.mean() for name, scores in cv_results.items()},
            "CV Std": {name: scores.std() for name, scores in cv_results.items()},
        }
    )

    comparison_df["Val-Test Gap"] = abs(
        comparison_df["Validation Acc"] - comparison_df["Test Acc"]
    )
    comparison_df["Stability Score"] = 1 - comparison_df["CV Std"]

    # Calculate weighted final score
    comparison_df["Final Score"] = (
        comparison_df["Test Acc"] * 0.5
        + comparison_df["Validation Acc"] * 0.2
        + comparison_df["CV Mean"] * 0.2
        + comparison_df["Stability Score"] * 0.1
    )

    comparison_df = comparison_df.sort_values("Final Score", ascending=False)

    print("\nModel Ranking (Top 10):")
    print(comparison_df.head(10).to_string())

    best_model_name = comparison_df["Final Score"].idxmax()
    best_row = comparison_df.loc[best_model_name]

    print(f"\n{'='*70}")
    print(f"BEST MODEL: {best_model_name}")
    print(f"{'='*70}")
    print(
        f"Test Accuracy: {best_row['Test Acc']:.4f} ({best_row['Test Acc']*100:.2f}%)"
    )
    print(f"Validation Accuracy: {best_row['Validation Acc']:.4f}")
    print(f"CV Mean ± Std: {best_row['CV Mean']:.4f} ± {best_row['CV Std']:.4f}")
    print(f"Final Score: {best_row['Final Score']:.4f}")
    print(f"{'='*70}")

    gap = best_row["Val-Test Gap"]

    print(f"\nOverfitting Analysis:")
    print(f"Validation-Test Gap: {gap:.4f}")

    if gap < 0.02:
        print(f"✓ Model is stable (gap < 2%)")
    elif gap < 0.05:
        print(f"⚠ Model shows slight overfitting (gap < 5%)")
        print(f"\nAlternative stable models (gap < 2%):")
        stable_models = comparison_df[comparison_df["Val-Test Gap"] < 0.02].head(3)
        if len(stable_models) > 0:
            print(
                stable_models[["Test Acc", "Val-Test Gap", "Final Score"]].to_string()
            )
    else:
        print(f"✗ Model shows significant overfitting (gap >= 5%)")
        print(f"\nSearching for alternative models (gap < 5%)...")
        stable_models = comparison_df[comparison_df["Val-Test Gap"] < 0.05]

        if len(stable_models) > 0:
            alternative_best = stable_models["Final Score"].idxmax()
            print(f"✓ Best alternative model: {alternative_best}")
            print(stable_models.loc[alternative_best].to_string())
            best_model_name = alternative_best
            print(f"\nSwitched to model: {best_model_name}")

    return best_model_name, comparison_df


# ==================== CONFUSION MATRIX VISUALIZATION ====================
def plot_confusion_matrix(model, X_test, y_test, model_name, label_encoder=None):
    """
    Generate and save confusion matrix heatmap.

    Args:
        model: Trained model
        X_test: Test features
        y_test: Test labels
        model_name: Name of the model
        label_encoder: Label encoder for class names
    """
    y_pred = model.predict(X_test)
    cm = confusion_matrix(y_test, y_pred)

    plt.figure(figsize=(14, 10))

    if label_encoder:
        labels = label_encoder.classes_
    else:
        labels = sorted(np.unique(y_test))

    sns.heatmap(
        cm,
        annot=True,
        fmt="d",
        cmap="Blues",
        xticklabels=labels,
        yticklabels=labels,
        cbar_kws={"label": "Count"},
    )
    plt.title(f"Confusion Matrix - {model_name}", fontsize=16, fontweight="bold")
    plt.ylabel("True Label", fontsize=12)
    plt.xlabel("Predicted Label", fontsize=12)
    plt.xticks(rotation=45, ha="right")
    plt.yticks(rotation=0)
    plt.tight_layout()

    # Use consistent filename - overwrites previous file
    filename = os.path.join(ASSET_FOLDER, "confusion_matrix.png")
    plt.savefig(filename, dpi=300, bbox_inches="tight")
    print(f"\n✓ Confusion matrix saved: {filename}")
    plt.close()


# ==================== COMPARISON CHART ====================
def plot_comparison(comparison_df):
    """
    Generate and save model comparison bar chart.

    Args:
        comparison_df: DataFrame containing comparison metrics
    """
    top_df = comparison_df.head(10)

    fig, ax = plt.subplots(figsize=(14, 8))

    x = np.arange(len(top_df))
    width = 0.25

    ax.bar(
        x - width, top_df["Validation Acc"], width, label="Validation", color="skyblue"
    )
    ax.bar(x, top_df["Test Acc"], width, label="Test", color="orange")
    ax.bar(x + width, top_df["CV Mean"], width, label="CV Mean", color="green")

    ax.set_xlabel("Models", fontsize=12)
    ax.set_ylabel("Accuracy", fontsize=12)
    ax.set_title("Top 10 Model Performance Comparison", fontsize=14, fontweight="bold")
    ax.set_xticks(x)
    ax.set_xticklabels(top_df.index, rotation=45, ha="right")
    ax.legend()
    ax.grid(axis="y", alpha=0.3)

    plt.tight_layout()

    # Use consistent filename - overwrites previous file
    filename = os.path.join(ASSET_FOLDER, "model_comparison.png")
    plt.savefig(filename, dpi=300, bbox_inches="tight")
    print(f"✓ Comparison chart saved: {filename}")
    plt.close()


# ==================== MODEL PERSISTENCE ====================
def save_model(model, vectorizer, model_name, label_encoder):
    """
    Save trained model, vectorizer, and label encoder to disk.

    Args:
        model: Trained model
        vectorizer: TF-IDF vectorizer
        model_name: Name of the model
        label_encoder: Label encoder
    """
    # Use consistent filenames - overwrites previous files
    model_file = os.path.join(ASSET_FOLDER, "best_model.pkl")
    vectorizer_file = os.path.join(ASSET_FOLDER, "tfidf_vectorizer.pkl")
    label_encoder_file = os.path.join(ASSET_FOLDER, "label_encoder.pkl")
    model_info_file = os.path.join(ASSET_FOLDER, "model_info.txt")

    joblib.dump(model, model_file)
    joblib.dump(vectorizer, vectorizer_file)
    joblib.dump(label_encoder, label_encoder_file)

    # Save model information to text file
    with open(model_info_file, "w", encoding="utf-8") as f:
        f.write(f"Best Model Information\n")
        f.write(f"=" * 70 + "\n\n")
        f.write(f"Model Type: {model_name}\n")
        f.write(f"Model Class: {type(model).__name__}\n")
        f.write(f"Saved on: {pd.Timestamp.now()}\n")

    print(f"\nModel artifacts saved:")
    print(f"✓ Model: {model_file} ({model_name})")
    print(f"✓ Vectorizer: {vectorizer_file}")
    print(f"✓ Label Encoder: {label_encoder_file}")
    print(f"✓ Model Info: {model_info_file}")


# ==================== PREDICTION TEST ====================
def test_prediction(model, vectorizer, sample_texts, label_encoder):
    """
    Test model predictions on sample texts.

    Args:
        model: Trained model
        vectorizer: TF-IDF vectorizer
        sample_texts: List of sample texts to predict
        label_encoder: Label encoder for decoding predictions
    """
    print(f"\n{'='*70}")
    print("SAMPLE PREDICTION TEST")
    print(f"{'='*70}")

    for i, text in enumerate(sample_texts, 1):
        text_tfidf = vectorizer.transform([text])
        prediction_encoded = model.predict(text_tfidf)[0]
        prediction = label_encoder.inverse_transform([prediction_encoded])[0]

        print(f"\nSample {i}:")
        print(f"Text: {text[:100]}...")
        print(f"Predicted Category: {prediction}")


# ==================== SOURCE PERFORMANCE ANALYSIS ====================
def analyze_source_performance(
    model, X_test, y_test, df_test, model_name, label_encoder
):
    """
    Detailed performance analysis by news source.

    Args:
        model: Trained model
        X_test: Test features
        y_test: Test labels
        df_test: Test DataFrame with source information
        model_name: Name of the model
        label_encoder: Label encoder
    """
    print(f"\n{'='*70}")
    print(f"DETAILED SOURCE ANALYSIS ({model_name})")
    print(f"{'='*70}\n")

    y_pred = model.predict(X_test)

    y_test_decoded = label_encoder.inverse_transform(y_test)
    y_pred_decoded = label_encoder.inverse_transform(y_pred)

    results_df = pd.DataFrame(
        {
            "true_label": y_test_decoded,
            "predicted_label": y_pred_decoded,
            "source": df_test["source"].values,
        }
    )

    for source in df_test["source"].unique():
        source_data = results_df[results_df["source"] == source]
        accuracy = accuracy_score(
            source_data["true_label"], source_data["predicted_label"]
        )

        print(f"Source: {source.upper()}")
        print(f"   Sample count: {len(source_data)} articles")
        print(f"   Accuracy: {accuracy:.4f} ({accuracy*100:.2f}%)")
        print(f"   Classification Report:")
        print(
            classification_report(
                source_data["true_label"],
                source_data["predicted_label"],
                zero_division=0,
            )
        )
        print()

    # Calculate source accuracies
    source_accuracies = {}
    for source in df_test["source"].unique():
        source_data = results_df[results_df["source"] == source]
        acc = accuracy_score(source_data["true_label"], source_data["predicted_label"])
        source_accuracies[source] = acc

    print(f"Accuracy comparison across sources:")
    for source, acc in sorted(
        source_accuracies.items(), key=lambda x: x[1], reverse=True
    ):
        print(f"   {source:15s}: {acc:.4f} ({acc*100:.2f}%)")

    fig, ax = plt.subplots(figsize=(12, 7))
    sources = list(source_accuracies.keys())
    accuracies = list(source_accuracies.values())

    colors = plt.cm.Set3(range(len(sources)))
    bars = ax.bar(sources, accuracies, color=colors, edgecolor="black", linewidth=1.2)

    ax.set_xlabel("News Source", fontsize=13, fontweight="bold")
    ax.set_ylabel("Accuracy", fontsize=13, fontweight="bold")
    ax.set_title(
        f"Accuracy by News Source - {model_name}",
        fontsize=15,
        fontweight="bold",
        pad=20,
    )

    ax.set_ylim([0, 1.08])

    ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda y, _: f"{y:.0%}"))

    for i, (bar, v) in enumerate(zip(bars, accuracies)):
        if v >= 0.98:
            label_y = v - 0.04
            va = "top"
            color = "white"
            weight = "bold"
        else:
            label_y = v + 0.02
            va = "bottom"
            color = "black"
            weight = "bold"

        ax.text(
            i,
            label_y,
            f"{v:.2%}",
            ha="center",
            va=va,
            fontsize=11,
            fontweight=weight,
            color=color,
        )

    ax.grid(axis="y", alpha=0.3, linestyle="--", linewidth=0.7)
    ax.set_axisbelow(True)

    if len(sources) > 6:
        plt.xticks(rotation=45, ha="right")

    plt.tight_layout()

    filename = os.path.join(ASSET_FOLDER, "accuracy_by_source.png")
    plt.savefig(filename, dpi=300, bbox_inches="tight")
    print(f"\n✓ Source analysis chart saved: {filename}")
    plt.close()


# ==================== K-MEANS CLUSTERING ANALYSIS ====================
def perform_kmeans_clustering(
    X_tfidf, vectorizer, n_clusters, label_encoder=None, y_true=None
):
    """
    Perform K-Means clustering to discover natural data structure.

    This unsupervised analysis helps understand if topics naturally separate
    without using category labels.

    Args:
        X_tfidf: TF-IDF feature matrix
        vectorizer: TF-IDF vectorizer (for feature names)
        n_clusters: Number of clusters (typically equal to number of categories)
        label_encoder: Label encoder (optional, for cluster-category mapping)
        y_true: True labels (optional, for purity analysis)

    Returns:
        Tuple of (kmeans_model, cluster_labels, cluster_keywords)
    """
    print(f"\n{'='*70}")
    print(f"K-MEANS CLUSTERING ANALYSIS (k={n_clusters})")
    print(f"{'='*70}")
    print("Discovering natural topic structure in the data...")

    # Train K-Means model
    print(f"\nTraining K-Means with {n_clusters} clusters...")
    kmeans = KMeans(
        n_clusters=n_clusters, random_state=42, max_iter=300, n_init=10, verbose=0
    )
    cluster_labels = kmeans.fit_predict(X_tfidf)
    print(f"✓ K-Means training completed")

    # Get cluster statistics
    print(f"\nCluster distribution:")
    cluster_counts = Counter(cluster_labels)
    for cluster_id in sorted(cluster_counts.keys()):
        count = cluster_counts[cluster_id]
        percentage = (count / len(cluster_labels)) * 100
        print(f"  Cluster {cluster_id}: {count} samples ({percentage:.2f}%)")

    # Extract top keywords for each cluster
    print(f"\n{'='*70}")
    print("TOP KEYWORDS FOR EACH CLUSTER")
    print(f"{'='*70}")

    feature_names = vectorizer.get_feature_names_out()
    cluster_centers = kmeans.cluster_centers_
    cluster_keywords = {}

    for cluster_id in range(n_clusters):
        # Get indices of top features for this cluster
        center = cluster_centers[cluster_id]
        top_indices = center.argsort()[-15:][::-1]  # Top 15 keywords
        top_keywords = [feature_names[i] for i in top_indices]
        cluster_keywords[cluster_id] = top_keywords

        print(f"\nCluster {cluster_id} (n={cluster_counts[cluster_id]}):")
        print(f"  Keywords: {', '.join(top_keywords)}")

    # Analyze cluster purity if true labels are provided
    if y_true is not None and label_encoder is not None:
        print(f"\n{'='*70}")
        print("CLUSTER PURITY ANALYSIS")
        print(f"{'='*70}")
        print("Analyzing how well clusters correspond to actual categories...\n")

        category_names = label_encoder.classes_

        for cluster_id in range(n_clusters):
            # Get samples in this cluster
            cluster_mask = cluster_labels == cluster_id
            cluster_true_labels = y_true[cluster_mask]

            # Count category distribution in this cluster
            label_counts = Counter(cluster_true_labels)
            total_in_cluster = len(cluster_true_labels)

            print(f"Cluster {cluster_id}:")
            print(f"  Total samples: {total_in_cluster}")
            print(f"  Category distribution:")

            # Sort by count (descending)
            sorted_labels = sorted(
                label_counts.items(), key=lambda x: x[1], reverse=True
            )

            for label_encoded, count in sorted_labels[:5]:  # Top 5 categories
                label_name = category_names[label_encoded]
                percentage = (count / total_in_cluster) * 100
                bar_length = int(percentage / 2)  # Scale to 50 chars max
                bar = "█" * bar_length
                print(f"    {label_name:15s}: {count:4d} ({percentage:5.1f}%) {bar}")

            # Calculate purity (percentage of most common category)
            if sorted_labels:
                most_common_label, most_common_count = sorted_labels[0]
                purity = (most_common_count / total_in_cluster) * 100
                dominant_category = category_names[most_common_label]
                print(f"  Purity: {purity:.2f}% (dominant: {dominant_category})")
            print()

        # Calculate overall clustering quality
        print(f"{'='*70}")
        print("CLUSTERING QUALITY METRICS")
        print(f"{'='*70}")

        # Calculate average purity
        total_purity = 0
        for cluster_id in range(n_clusters):
            cluster_mask = cluster_labels == cluster_id
            cluster_true_labels = y_true[cluster_mask]
            if len(cluster_true_labels) > 0:
                label_counts = Counter(cluster_true_labels)
                most_common_count = max(label_counts.values())
                purity = most_common_count / len(cluster_true_labels)
                total_purity += purity * len(cluster_true_labels)

        avg_purity = total_purity / len(y_true)
        print(f"Average cluster purity: {avg_purity:.4f} ({avg_purity*100:.2f}%)")

        # Calculate inertia (within-cluster sum of squares)
        inertia = kmeans.inertia_
        print(f"Inertia (within-cluster sum of squares): {inertia:.2f}")

        # Interpretation
        print(f"\nInterpretation:")
        if avg_purity >= 0.7:
            print(
                f"✓ High purity ({avg_purity*100:.1f}%) - Clusters strongly correspond to categories"
            )
            print(f"  Data has naturally separable topic structure")
        elif avg_purity >= 0.5:
            print(
                f"⚠ Moderate purity ({avg_purity*100:.1f}%) - Clusters partially correspond to categories"
            )
            print(f"  Some topic overlap exists in the data")
        else:
            print(
                f"✗ Low purity ({avg_purity*100:.1f}%) - Clusters don't align well with categories"
            )
            print(f"  Topics may be mixed or need better feature engineering")

    return kmeans, cluster_labels, cluster_keywords


def visualize_clustering_results(X_tfidf, cluster_labels, y_true, label_encoder):
    """
    Visualize clustering results using dimensionality reduction.

    Args:
        X_tfidf: TF-IDF feature matrix
        cluster_labels: Cluster assignments from K-Means
        y_true: True category labels
        label_encoder: Label encoder for category names
    """
    print(f"\n{'='*70}")
    print("CLUSTERING VISUALIZATION")
    print(f"{'='*70}")

    try:
        from sklearn.decomposition import PCA
        from sklearn.manifold import TSNE

        print("Reducing dimensions for visualization (this may take a moment)...")

        # Use PCA first to reduce to 50 dimensions (faster for t-SNE)
        if X_tfidf.shape[1] > 50:
            pca = PCA(n_components=50, random_state=42)
            X_reduced = pca.fit_transform(X_tfidf.toarray())
            print(f"✓ PCA: {X_tfidf.shape[1]} → 50 dimensions")
        else:
            X_reduced = X_tfidf.toarray()

        # Apply t-SNE for 2D visualization
        tsne = TSNE(n_components=2, random_state=42, perplexity=30, n_iter=1000)
        X_2d = tsne.fit_transform(X_reduced)
        print(f"✓ t-SNE: {X_reduced.shape[1]} → 2 dimensions")

        # Create visualization
        fig, axes = plt.subplots(1, 2, figsize=(18, 7))

        # Plot 1: Colored by K-Means clusters
        scatter1 = axes[0].scatter(
            X_2d[:, 0], X_2d[:, 1], c=cluster_labels, cmap="tab10", alpha=0.6, s=30
        )
        axes[0].set_title(
            "K-Means Clusters (Unsupervised)", fontsize=14, fontweight="bold"
        )
        axes[0].set_xlabel("t-SNE Component 1", fontsize=12)
        axes[0].set_ylabel("t-SNE Component 2", fontsize=12)
        plt.colorbar(scatter1, ax=axes[0], label="Cluster ID")

        # Plot 2: Colored by true categories
        scatter2 = axes[1].scatter(
            X_2d[:, 0], X_2d[:, 1], c=y_true, cmap="tab10", alpha=0.6, s=30
        )
        axes[1].set_title(
            "True Categories (Supervised)", fontsize=14, fontweight="bold"
        )
        axes[1].set_xlabel("t-SNE Component 1", fontsize=12)
        axes[1].set_ylabel("t-SNE Component 2", fontsize=12)

        # Create legend for true categories
        category_names = label_encoder.classes_
        handles = [
            plt.Line2D(
                [0],
                [0],
                marker="o",
                color="w",
                markerfacecolor=plt.cm.tab10(i / len(category_names)),
                markersize=8,
                label=cat,
            )
            for i, cat in enumerate(category_names)
        ]
        axes[1].legend(handles=handles, loc="center left", bbox_to_anchor=(1, 0.5))

        plt.tight_layout()

        # Use consistent filename - overwrites previous file
        filename = os.path.join(ASSET_FOLDER, "kmeans_clustering_visualization.png")
        plt.savefig(filename, dpi=300, bbox_inches="tight")
        print(f"\n✓ Clustering visualization saved: {filename}")
        plt.close()

    except ImportError as e:
        print(f"⚠ Visualization skipped: {e}")
        print("  Install scikit-learn with: pip install scikit-learn")


# ==================== MAIN EXECUTION ====================
def main():
    """Execute complete model training and evaluation pipeline."""
    print(f"\n{'#'*70}")
    print("VIETNAMESE NEWS CLASSIFICATION - MODEL TRAINING")
    print("(10+ ALGORITHMS + ENSEMBLE + K-MEANS ANALYSIS)")
    print(f"{'#'*70}\n")

    # Step 1: Load preprocessed data
    df = load_cleaned_data("cleaned_data.csv")

    # Step 2: Split data into train/val/test
    X_train, X_val, X_test, y_train, y_val, y_test, label_encoder = split_data(df)

    # Preserve test DataFrame for source analysis
    test_indices = X_test.index.tolist()
    if "source" in df.columns:
        df_test = df.loc[test_indices].reset_index(drop=True)
    else:
        df_test = None

    # Step 3: Vectorize text using TF-IDF
    X_train_tfidf, X_val_tfidf, X_test_tfidf, vectorizer = vectorize_text(
        X_train, X_val, X_test
    )

    # Step 4: K-MEANS CLUSTERING ANALYSIS
    print(f"\n{'#'*70}")
    print("UNSUPERVISED LEARNING: K-MEANS CLUSTERING")
    print(f"{'#'*70}")

    n_categories = len(label_encoder.classes_)

    # Perform K-Means on training data
    kmeans_model, cluster_labels_train, cluster_keywords = perform_kmeans_clustering(
        X_train_tfidf,
        vectorizer,
        n_clusters=n_categories,
        label_encoder=label_encoder,
        y_true=y_train,
    )

    # Visualize clustering results
    visualize_clustering_results(
        X_train_tfidf, cluster_labels_train, y_train, label_encoder
    )

    # Save K-Means model - use consistent filename
    kmeans_file = os.path.join(ASSET_FOLDER, "kmeans_model.pkl")
    joblib.dump(kmeans_model, kmeans_file)
    print(f"\n✓ K-Means model saved: {kmeans_file}")

    # Save cluster keywords for reference - use consistent filename
    keywords_file = os.path.join(ASSET_FOLDER, "cluster_keywords.txt")
    with open(keywords_file, "w", encoding="utf-8") as f:
        f.write("K-MEANS CLUSTER KEYWORDS\n")
        f.write("=" * 70 + "\n\n")
        for cluster_id, keywords in cluster_keywords.items():
            f.write(f"Cluster {cluster_id}:\n")
            f.write(f"  {', '.join(keywords)}\n\n")
    print(f"✓ Cluster keywords saved: {keywords_file}")

    # Step 5: Train multiple models (supervised learning)
    print(f"\n{'#'*70}")
    print("SUPERVISED LEARNING: CLASSIFICATION MODELS")
    print(f"{'#'*70}")

    models = train_models(X_train_tfidf, y_train)

    # Step 6: Train ensemble model
    ensemble_model = train_ensemble_model(models, X_train_tfidf, y_train)
    if ensemble_model:
        models["Ensemble (Voting)"] = ensemble_model

    # Step 7: Cross-validation on fast models
    fast_models = {
        name: model
        for name, model in models.items()
        if name
        in ["Naive Bayes", "Logistic Regression", "Linear SVM", "SGD Classifier"]
    }
    cv_results = perform_cross_validation(fast_models, X_train_tfidf, y_train, cv=5)

    # Fill CV results for models not cross-validated
    for name in models.keys():
        if name not in cv_results:
            cv_results[name] = np.array([0.0] * 5)

    # Step 8: Evaluate on validation set
    val_results = evaluate_on_validation(models, X_val_tfidf, y_val)

    # Step 9: Evaluate on test set
    test_results = evaluate_on_test(
        models, X_test_tfidf, y_test, df_test, label_encoder
    )

    # Step 10: Compare and select best model
    best_model_name, comparison_df = compare_results(
        val_results, test_results, cv_results
    )

    # Step 11: Visualize confusion matrix for best model
    best_model = models[best_model_name]
    plot_confusion_matrix(
        best_model, X_test_tfidf, y_test, best_model_name, label_encoder
    )

    # Step 12: Generate comparison chart
    plot_comparison(comparison_df)

    # Step 13: Save best model and artifacts
    save_model(best_model, vectorizer, best_model_name, label_encoder)

    # Step 14: Analyze performance by source
    if df_test is not None:
        analyze_source_performance(
            best_model, X_test_tfidf, y_test, df_test, best_model_name, label_encoder
        )

    # Step 15: Test predictions on sample texts
    sample_texts = [
        "HLV Park Hang-seo dẫn dắt đội tuyển Việt Nam giành chiến thắng",
        "Giá vàng hôm nay tăng mạnh, chứng khoán biến động",
        "Ca sĩ Sơn Tùng MTP ra mắt MV mới thu hút hàng triệu lượt xem",
    ]
    test_prediction(best_model, vectorizer, sample_texts, label_encoder)

    print(f"\n{'#'*70}")
    print("TRAINING COMPLETED SUCCESSFULLY")
    print(f"{'#'*70}\n")

    print("Summary of artifacts generated:")
    print(f"  • Best model: best_model.pkl ({best_model_name})")
    print(f"  • K-Means model: kmeans_model.pkl")
    print(f"  • Visualizations:")
    print(f"    - model_comparison.png")
    print(f"    - confusion_matrix.png")
    print(f"    - accuracy_by_source.png")
    print(f"    - kmeans_clustering_visualization.png")
    print(f"  • Analysis reports:")
    print(f"    - model_info.txt")
    print(f"    - cluster_keywords.txt")


if __name__ == "__main__":
    main()