"""
=============================================================
TF-IDF & K-MEANS CLUSTERING
=============================================================

Tujuan:
- TF-IDF transformation untuk teks pekerjaan
- K-Means Clustering dengan K=4 (sesuai PRD)
- Evaluasi menggunakan Elbow Method dan Silhouette Score

PRD Reference: Section 14, 15, 16
=============================================================
"""

import pandas as pd
import numpy as np
import json
import os
from datetime import datetime
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

# Path configuration
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_INPUT = os.path.join(BASE_DIR, '..', 'data', 'processed', 'dataset_untuk_clustering.csv')
MODEL_OUTPUT = os.path.join(BASE_DIR, '..', 'data', 'processed')
PLOT_OUTPUT = os.path.join(BASE_DIR, '..', 'data', 'processed')

# K yang akan diuji (PRD Section 16.1)
K_CANDIDATES = [2, 3, 4, 5, 6]
FIXED_K = 4  # K=4 sesuai PRD Section 5.4, 15.2


def load_data():
    """Load dataset untuk clustering."""
    print(">> Loading dataset untuk clustering...")
    df = pd.read_csv(DATA_INPUT, encoding='utf-8-sig')
    print(f"   Total data: {len(df)} alumni")
    return df


def tfidf_vectorization(texts):
    """
    TF-IDF transformation.

    PRD Section 14:
    TF-IDF mengubah teks pekerjaan menjadi feature matrix.
    """
    print("\n>> Melakukan TF-IDF Vectorization...")

    # TF-IDF Vectorizer
    vectorizer = TfidfVectorizer(
        max_features=500,      # Maksimum 500 fitur
        min_df=2,              # Minimum document frequency 2
        max_df=0.95,          # Maximum document frequency 95%
        ngram_range=(1, 2),    # Unigrams dan bigrams
        lowercase=True,
        token_pattern=r'\b[a-zA-Z]{2,}\b'  # Minimal 2 karakter, hanya alfabet
    )

    # Fit dan transform
    tfidf_matrix = vectorizer.fit_transform(texts)
    feature_names = vectorizer.get_feature_names_out()

    print(f"   TF-IDF Matrix shape: {tfidf_matrix.shape}")
    print(f"   Number of features: {len(feature_names)}")
    print(f"   Sample features: {list(feature_names[:10])}...")

    return vectorizer, tfidf_matrix, feature_names


def elbow_method(tfidf_matrix, k_range):
    """
    Elbow Method untuk menentukan K optimal.
    Menghitung inertia/WCSS untuk setiap K.

    PRD Section 16.2
    """
    print("\n>> Elbow Method Analysis...")

    inertias = []

    for k in k_range:
        print(f"   Testing K={k}...", end=" ")
        kmeans = KMeans(
            n_clusters=k,
            init='k-means++',
            n_init=10,
            max_iter=300,
            random_state=42
        )
        kmeans.fit(tfidf_matrix)
        inertias.append(kmeans.inertia_)
        print(f"Inertia: {kmeans.inertia_:.2f}")

    return inertias


def silhouette_analysis(tfidf_matrix, k_range):
    """
    Silhouette Score analysis.

    PRD Section 16.3
    """
    print("\n>> Silhouette Score Analysis...")

    silhouettes = []

    for k in k_range:
        print(f"   Testing K={k}...", end=" ")
        kmeans = KMeans(
            n_clusters=k,
            init='k-means++',
            n_init=10,
            max_iter=300,
            random_state=42
        )
        labels = kmeans.fit_predict(tfidf_matrix)

        # Hitung silhouette score
        score = silhouette_score(tfidf_matrix, labels)
        silhouettes.append(score)
        print(f"Silhouette: {score:.4f}")

    return silhouettes


def plot_evaluation(k_range, inertias, silhouettes, output_path):
    """
    Plot Elbow Method dan Silhouette Score.
    """
    print("\n>> Generating evaluation plots...")

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    # Elbow Plot
    axes[0].plot(k_range, inertias, 'bo-', linewidth=2, markersize=8)
    axes[0].axvline(x=FIXED_K, color='r', linestyle='--', label=f'K={FIXED_K} (PRD)')
    axes[0].set_xlabel('Number of Clusters (K)', fontsize=12)
    axes[0].set_ylabel('Inertia (WCSS)', fontsize=12)
    axes[0].set_title('Elbow Method', fontsize=14, fontweight='bold')
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)

    # Silhouette Plot
    axes[1].plot(k_range, silhouettes, 'go-', linewidth=2, markersize=8)
    axes[1].axvline(x=FIXED_K, color='r', linestyle='--', label=f'K={FIXED_K} (PRD)')
    axes[1].set_xlabel('Number of Clusters (K)', fontsize=12)
    axes[1].set_ylabel('Silhouette Score', fontsize=12)
    axes[1].set_title('Silhouette Score', fontsize=14, fontweight='bold')
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(os.path.join(output_path, 'evaluation_plots.png'), dpi=150, bbox_inches='tight')
    print(f"   Plot saved: {output_path}/evaluation_plots.png")


def kmeans_clustering(tfidf_matrix, k, n_init=10):
    """
    K-Means Clustering dengan K tetap.

    PRD Section 15
    """
    print(f"\n>> Running K-Means with K={k}...")

    kmeans = KMeans(
        n_clusters=k,
        init='k-means++',
        n_init=n_init,
        max_iter=300,
        random_state=42
    )

    cluster_labels = kmeans.fit_predict(tfidf_matrix)
    inertia = kmeans.inertia_

    # Hitung silhouette score
    silhouette = silhouette_score(tfidf_matrix, cluster_labels)

    print(f"   Inertia: {inertia:.2f}")
    print(f"   Silhouette Score: {silhouette:.4f}")

    return kmeans, cluster_labels, inertia, silhouette


def get_cluster_statistics(df, cluster_labels, vectorizer, feature_names):
    """
    Get statistics untuk setiap cluster.
    """
    print("\n>> Analyzing cluster statistics...")

    df_result = df.copy()
    df_result['cluster'] = cluster_labels

    stats = {}

    for cluster_id in range(FIXED_K):
        cluster_data = df_result[df_result['cluster'] == cluster_id]
        stats[cluster_id] = {
            'jumlah': len(cluster_data),
            'persentase': len(cluster_data) / len(df_result) * 100
        }

        # Top terms dari centroid
        centroid = kmeans.cluster_centers_[cluster_id]
        top_indices = centroid.argsort()[-10:][::-1]
        top_terms = [feature_names[i] for i in top_indices]
        stats[cluster_id]['top_terms'] = top_terms

        print(f"\n   Cluster {cluster_id}:")
        print(f"      Jumlah: {stats[cluster_id]['jumlah']} ({stats[cluster_id]['persentase']:.1f}%)")
        print(f"      Top terms: {', '.join(top_terms[:5])}")

    return df_result, stats


def save_results(df_result, kmeans, vectorizer, tfidf_matrix, feature_names,
                 inertia, silhouette, inertias, silhouettes, output_path):
    """
    Simpan hasil clustering.
    """
    print("\n>> Saving results...")

    # 1. Save hasil clustering
    clustering_output = os.path.join(output_path, 'clustering_results.csv')
    df_result.to_csv(clustering_output, index=False, encoding='utf-8-sig')
    print(f"   Clustering results: {clustering_output}")

    # 2. Save model info
    model_info = {
        'model_version': datetime.now().strftime('%Y%m%d_%H%M%S'),
        'k': FIXED_K,
        'inertia': float(inertia),
        'silhouette_score': float(silhouette),
        'n_samples': len(df_result),
        'n_features': len(feature_names),
        'evaluation': {
            'k_candidates': K_CANDIDATES,
            'inertias': [float(i) for i in inertias],
            'silhouettes': [float(s) for s in silhouettes]
        },
        'created_at': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    }

    model_output = os.path.join(output_path, 'model_info.json')
    with open(model_output, 'w', encoding='utf-8') as f:
        json.dump(model_info, f, indent=2, ensure_ascii=False)
    print(f"   Model info: {model_output}")

    # 3. Save cluster statistics
    cluster_stats_output = os.path.join(output_path, 'cluster_statistics.json')
    with open(cluster_stats_output, 'w', encoding='utf-8') as f:
        json.dump(cluster_stats, f, indent=2, ensure_ascii=False)
    print(f"   Cluster statistics: {cluster_stats_output}")

    return model_info


def generate_summary(model_info, stats):
    """
    Generate summary report.
    """
    summary = []
    summary.append("=" * 60)
    summary.append("CLUSTERING SUMMARY")
    summary.append(f"Tanggal: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    summary.append("=" * 60)

    summary.append(f"\n[MODEL INFO]")
    summary.append(f"   K = {model_info['k']} (sesuai PRD)")
    summary.append(f"   Inertia: {model_info['inertia']:.2f}")
    summary.append(f"   Silhouette Score: {model_info['silhouette_score']:.4f}")
    summary.append(f"   Total Samples: {model_info['n_samples']}")
    summary.append(f"   Total Features: {model_info['n_features']}")

    summary.append(f"\n[EVALUATION RESULTS]")
    for i, k in enumerate(model_info['evaluation']['k_candidates']):
        summary.append(f"   K={k}: Inertia={model_info['evaluation']['inertias'][i]:.2f}, "
                      f"Silhouette={model_info['evaluation']['silhouettes'][i]:.4f}")

    summary.append(f"\n[CLUSTER DISTRIBUTION]")
    for cluster_id, stat in stats.items():
        summary.append(f"   Cluster {cluster_id}: {stat['jumlah']} ({stat['persentase']:.1f}%)")

    summary.append(f"\n[TOP TERMS PER CLUSTER]")
    for cluster_id, stat in stats.items():
        summary.append(f"   Cluster {cluster_id}: {', '.join(stat['top_terms'][:5])}")

    summary.append("\n" + "=" * 60)

    return "\n".join(summary)


# Global variable untuk simpan kmeans dan stats
kmeans = None
cluster_stats = None


def main():
    """
    Main function untuk TF-IDF dan K-Means Clustering.
    """
    global kmeans, cluster_stats

    print("=" * 60)
    print("TF-IDF & K-MEANS CLUSTERING")
    print(f"K = {FIXED_K} (sesuai PRD)")
    print("=" * 60)

    # 1. Load data
    df = load_data()

    # 2. TF-IDF Vectorization
    texts = df['teks_pekerjaan'].fillna('').tolist()
    vectorizer, tfidf_matrix, feature_names = tfidf_vectorization(texts)

    # 3. Evaluasi Elbow dan Silhouette
    inertias = elbow_method(tfidf_matrix, K_CANDIDATES)
    silhouettes = silhouette_analysis(tfidf_matrix, K_CANDIDATES)

    # 4. Plot evaluasi
    plot_evaluation(K_CANDIDATES, inertias, silhouettes, PLOT_OUTPUT)

    # 5. K-Means dengan K=4
    kmeans, cluster_labels, inertia, silhouette = kmeans_clustering(tfidf_matrix, FIXED_K)

    # 6. Get cluster statistics
    df_result, cluster_stats = get_cluster_statistics(df, cluster_labels, vectorizer, feature_names)

    # 7. Save results
    model_info = save_results(
        df_result, kmeans, vectorizer, tfidf_matrix, feature_names,
        inertia, silhouette, inertias, silhouettes, MODEL_OUTPUT
    )

    # 8. Generate summary
    summary = generate_summary(model_info, cluster_stats)
    print("\n" + summary)

    # Save summary
    summary_output = os.path.join(MODEL_OUTPUT, 'clustering_summary.txt')
    with open(summary_output, 'w', encoding='utf-8') as f:
        f.write(summary)
    print(f"\n>> Summary saved: {summary_output}")

    print("\n" + "=" * 60)
    print("[SUCCESS] CLUSTERING SELESAI!")
    print("=" * 60)

    return df_result, model_info, cluster_stats


if __name__ == "__main__":
    df_result, model_info, cluster_stats = main()