"""
Clustering Module
================
Handles K-means clustering operations and evaluation.
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import StandardScaler
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
import logging
from typing import Tuple, Optional, List
from pathlib import Path

logger = logging.getLogger(__name__)


def prepare_features_for_clustering(
    feature_df: pd.DataFrame, sample_size: Optional[int] = None
) -> Tuple[pd.DataFrame, np.ndarray, StandardScaler]:
    """
    Scale features and optionally sample data for clustering.

    Args:
        feature_df: DataFrame with user features
        sample_size: Optional sample size for memory efficiency

    Returns:
        Tuple of (sampled_df, scaled_features, scaler)
    """
    logger.info("Preparing features for clustering")

    # Sample data if requested
    if sample_size and len(feature_df) > sample_size:
        logger.info(f"Sampling {sample_size} users from {len(feature_df)} total users")

        # Stratified sampling based on activity level
        # Create activity bins
        feature_df["activity_bin"] = pd.qcut(
            feature_df["total_clips_created"],
            q=5,
            labels=["very_low", "low", "medium", "high", "very_high"],
        )

        # Sample proportionally from each bin
        sampled_df = feature_df.groupby("activity_bin", group_keys=False).apply(
            lambda x: x.sample(
                n=int(len(x) * sample_size / len(feature_df)), random_state=42
            )
        )

        # Drop the temporary column
        sampled_df = sampled_df.drop("activity_bin", axis=1)
        feature_df = feature_df.drop("activity_bin", axis=1)

        logger.info(f"Sampled {len(sampled_df)} users")
    else:
        sampled_df = feature_df.copy()

    # Identify numeric columns for scaling
    numeric_columns = sampled_df.select_dtypes(include=[np.number]).columns.tolist()
    if "user_id" in numeric_columns:
        numeric_columns.remove("user_id")  # Don't scale user_id

    # Handle highly skewed features with log transformation
    skewed_features = []
    for col in numeric_columns:
        skewness = float(sampled_df[col].skew())
        if abs(skewness) > 2:  # High skewness
            skewed_features.append(col)
            # Log transform (add 1 to handle zeros)
            sampled_df[f"{col}_log"] = np.log1p(sampled_df[col])
            numeric_columns.append(f"{col}_log")
            numeric_columns.remove(col)

    if skewed_features:
        logger.info(f"Applied log transformation to skewed features: {skewed_features}")

    # Remove highly correlated features
    correlation_matrix = sampled_df[numeric_columns].corr().abs()
    upper_triangle = correlation_matrix.where(
        np.triu(np.ones(correlation_matrix.shape), k=1).astype(bool)
    )

    # Find features with correlation > 0.95
    high_corr_features = [
        column
        for column in upper_triangle.columns
        if any(upper_triangle[column] > 0.95)
    ]

    if high_corr_features:
        logger.info(f"Removing highly correlated features: {high_corr_features}")
        numeric_columns = [
            col for col in numeric_columns if col not in high_corr_features
        ]

    # Scale features
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(sampled_df[numeric_columns])

    logger.info(f"Scaled {X_scaled.shape[1]} features for {X_scaled.shape[0]} users")

    # Store feature names for later use
    sampled_df._numeric_columns = numeric_columns

    return sampled_df, X_scaled, scaler


def find_optimal_clusters(
    X_scaled: np.ndarray,
    k_range: Tuple[int, int] = (4, 6),
    output_dir: Optional[str] = None,
    min_cluster_size_ratio: float = 0.01,
) -> int:
    """
    Determine optimal number of clusters using elbow and silhouette methods.

    Args:
        X_scaled: Scaled feature matrix
        k_range: Range of k values to test
        output_dir: Directory to save the optimization plot
        min_cluster_size_ratio: Minimum cluster size as ratio of total population

    Returns:
        Optimal number of clusters
    """
    logger.info(f"Finding optimal clusters in range {k_range}")
    logger.info(
        f"Minimum cluster size: {min_cluster_size_ratio:.1%} of total population"
    )

    k_values = range(k_range[0], k_range[1] + 1)
    inertias = []
    silhouette_scores = []
    valid_k_values = []

    min_cluster_size = int(len(X_scaled) * min_cluster_size_ratio)

    for k in k_values:
        logger.info(f"Testing k={k}")

        # Fit k-means
        kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
        kmeans.fit(X_scaled)

        # Check cluster sizes
        unique, counts = np.unique(kmeans.labels_, return_counts=True)
        min_size = min(counts)

        if min_size < min_cluster_size:
            logger.info(
                f"k={k}: Smallest cluster has {min_size} users (below minimum {min_cluster_size}), skipping"
            )
            continue

        # Calculate metrics
        inertias.append(kmeans.inertia_)
        silhouette = silhouette_score(X_scaled, kmeans.labels_)
        silhouette_scores.append(silhouette)
        valid_k_values.append(k)

        logger.info(
            f"k={k}: inertia={kmeans.inertia_:.2f}, silhouette={silhouette:.3f}, min_cluster_size={min_size}"
        )

    if not valid_k_values:
        logger.warning("No valid k values found with minimum cluster size constraint!")
        # Fallback to the middle of the range
        return (k_range[0] + k_range[1]) // 2

    # Plot results
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

    # Elbow plot
    ax1.plot(valid_k_values, inertias, "bo-")
    ax1.set_xlabel("Number of Clusters (k)")
    ax1.set_ylabel("Inertia")
    ax1.set_title("Elbow Method")
    ax1.grid(True)

    # Silhouette plot
    ax2.plot(valid_k_values, silhouette_scores, "ro-")
    ax2.set_xlabel("Number of Clusters (k)")
    ax2.set_ylabel("Silhouette Score")
    ax2.set_title("Silhouette Analysis")
    ax2.grid(True)

    plt.tight_layout()

    # Save to output directory if provided
    if output_dir:
        output_path = Path(output_dir) / "cluster_optimization.png"
    else:
        output_path = "cluster_optimization.png"

    plt.savefig(output_path)
    plt.close()
    logger.info(f"Saved optimization plot to {output_path}")

    # Find optimal k (highest silhouette score among valid k values)
    optimal_idx = np.argmax(silhouette_scores)
    optimal_k = valid_k_values[optimal_idx]
    logger.info(
        f"Optimal k={optimal_k} with silhouette score={silhouette_scores[optimal_idx]:.3f}"
    )

    return optimal_k


def merge_small_clusters(
    X_scaled: np.ndarray,
    cluster_labels: np.ndarray,
    kmeans_model: KMeans,
    min_cluster_size: int,
) -> Tuple[np.ndarray, KMeans]:
    """
    Merge clusters that are smaller than the minimum size with their nearest neighbors.

    Args:
        X_scaled: Scaled feature matrix
        cluster_labels: Initial cluster assignments
        kmeans_model: Fitted KMeans model
        min_cluster_size: Minimum allowed cluster size

    Returns:
        Tuple of (updated_labels, updated_model)
    """
    labels = cluster_labels.copy()
    centers = kmeans_model.cluster_centers_.copy()

    # Identify small clusters
    unique_labels, counts = np.unique(labels, return_counts=True)
    cluster_sizes = dict(zip(unique_labels, counts))
    small_clusters = [
        label for label, count in cluster_sizes.items() if count < min_cluster_size
    ]

    if not small_clusters:
        logger.info("No small clusters to merge")
        return labels, kmeans_model

    logger.info(
        f"Found {len(small_clusters)} small clusters to merge: {small_clusters}"
    )

    # Sort clusters by size (smallest first) to merge smallest ones first
    small_clusters.sort(key=lambda x: cluster_sizes[x])

    # Create mapping of old labels to new labels
    label_mapping = {i: i for i in range(len(unique_labels))}

    for small_cluster in small_clusters:
        if cluster_sizes[small_cluster] >= min_cluster_size:
            # This cluster has grown due to previous merges
            continue

        # Find nearest cluster based on centroid distance
        small_center = centers[small_cluster]

        # Calculate distances to all other clusters
        distances = []
        for other_cluster in unique_labels:
            if (
                other_cluster == small_cluster
                or cluster_sizes.get(other_cluster, 0) == 0
            ):
                continue
            other_center = centers[other_cluster]
            dist = np.linalg.norm(small_center - other_center)
            distances.append((other_cluster, dist))

        if not distances:
            logger.warning(f"No valid clusters to merge with cluster {small_cluster}")
            continue

        # Sort by distance and find nearest cluster
        distances.sort(key=lambda x: x[1])
        nearest_cluster = distances[0][0]

        logger.info(
            f"Merging cluster {small_cluster} ({cluster_sizes[small_cluster]} users) "
            f"into cluster {nearest_cluster} ({cluster_sizes[nearest_cluster]} users)"
        )

        # Update labels
        labels[labels == small_cluster] = nearest_cluster

        # Update cluster sizes
        cluster_sizes[nearest_cluster] += cluster_sizes[small_cluster]
        cluster_sizes[small_cluster] = 0

        # Update label mapping
        label_mapping[small_cluster] = nearest_cluster

    # Renumber clusters to be consecutive
    unique_remaining = np.unique(labels)
    new_label_mapping = {old: new for new, old in enumerate(sorted(unique_remaining))}
    labels = np.array([new_label_mapping[label] for label in labels])

    # Create new KMeans model with updated centers
    new_n_clusters = len(unique_remaining)
    new_kmeans = KMeans(n_clusters=new_n_clusters, random_state=42)
    new_kmeans.fit(X_scaled)
    new_kmeans.labels_ = labels

    # Log final cluster sizes
    unique_final, counts_final = np.unique(labels, return_counts=True)
    for cluster, count in zip(unique_final, counts_final):
        logger.info(
            f"Final cluster {cluster}: {count} users ({count/len(labels)*100:.1f}%)"
        )

    return labels, new_kmeans


def perform_clustering(
    X_scaled: np.ndarray, k: int, min_cluster_size_ratio: float = 0.01
) -> Tuple[KMeans, np.ndarray]:
    """
    Execute k-means clustering with minimum cluster size enforcement.

    Args:
        X_scaled: Scaled feature matrix
        k: Number of clusters
        min_cluster_size_ratio: Minimum cluster size as ratio of total population

    Returns:
        Tuple of (kmeans_model, cluster_labels)
    """
    logger.info(f"Performing k-means clustering with k={k}")

    min_cluster_size = int(len(X_scaled) * min_cluster_size_ratio)
    logger.info(
        f"Minimum cluster size: {min_cluster_size} users ({min_cluster_size_ratio:.1%})"
    )

    # Initialize and fit k-means
    kmeans = KMeans(
        n_clusters=k,
        random_state=42,
        n_init=10,
        max_iter=300,
        algorithm="elkan",  # More efficient for dense data
    )

    cluster_labels = kmeans.fit_predict(X_scaled)

    # Log initial cluster sizes
    unique, counts = np.unique(cluster_labels, return_counts=True)
    logger.info("Initial cluster sizes:")
    for cluster, count in zip(unique, counts):
        logger.info(
            f"Cluster {cluster}: {count} users ({count/len(cluster_labels)*100:.1f}%)"
        )

    # Merge small clusters
    cluster_labels, kmeans = merge_small_clusters(
        X_scaled, cluster_labels, kmeans, min_cluster_size
    )

    # Calculate final silhouette score
    final_silhouette = float(silhouette_score(X_scaled, cluster_labels))
    logger.info(f"Final silhouette score: {final_silhouette:.3f}")

    return kmeans, cluster_labels


def evaluate_cluster_stability(X_scaled: np.ndarray, k: int, n_runs: int = 5) -> float:
    """
    Evaluate cluster stability by running clustering multiple times.

    Args:
        X_scaled: Scaled feature matrix
        k: Number of clusters
        n_runs: Number of runs to perform

    Returns:
        Average stability score (0-1, higher is more stable)
    """
    logger.info(f"Evaluating cluster stability with {n_runs} runs")

    # Store labels from each run
    all_labels = []

    for i in range(n_runs):
        kmeans = KMeans(n_clusters=k, random_state=i * 42, n_init=10)
        labels = kmeans.fit_predict(X_scaled)
        all_labels.append(labels)

    # Calculate pairwise agreement between runs
    stability_scores = []

    for i in range(n_runs):
        for j in range(i + 1, n_runs):
            # Calculate agreement between run i and run j
            # Use adjusted Rand index for this
            from sklearn.metrics import adjusted_rand_score

            score = adjusted_rand_score(all_labels[i], all_labels[j])
            stability_scores.append(score)

    avg_stability = np.mean(stability_scores)
    std_stability = np.std(stability_scores)

    logger.info(f"Cluster stability: {avg_stability:.3f} ± {std_stability:.3f}")

    return avg_stability


def plot_cluster_visualization(
    X_scaled: np.ndarray,
    cluster_labels: np.ndarray,
    output_dir: Optional[str] = None,
    output_filename: str = "cluster_visualization.png",
):
    """
    Create visualization of clusters using PCA reduction.

    Args:
        X_scaled: Scaled feature matrix
        cluster_labels: Cluster assignments
        output_dir: Directory to save visualization
        output_filename: Name of the output file
    """
    from sklearn.decomposition import PCA

    logger.info("Creating cluster visualization")

    # Reduce to 2D using PCA
    pca = PCA(n_components=2)
    X_pca = pca.fit_transform(X_scaled)

    # Create scatter plot
    plt.figure(figsize=(10, 8))

    # Plot each cluster
    unique_labels = np.unique(cluster_labels)
    colors = sns.color_palette("husl", len(unique_labels))

    for label, color in zip(unique_labels, colors):
        mask = cluster_labels == label
        plt.scatter(
            X_pca[mask, 0],
            X_pca[mask, 1],
            c=[color],
            label=f"Cluster {label}",
            alpha=0.6,
            edgecolors="black",
            linewidth=0.5,
        )

    plt.xlabel(f"PC1 ({pca.explained_variance_ratio_[0]:.1%} variance)")
    plt.ylabel(f"PC2 ({pca.explained_variance_ratio_[1]:.1%} variance)")
    plt.title("User Clusters Visualization (PCA)")
    plt.legend()
    plt.grid(True, alpha=0.3)

    plt.tight_layout()

    # Save to output directory if provided
    if output_dir:
        output_path = Path(output_dir) / output_filename
    else:
        output_path = output_filename

    plt.savefig(output_path)
    plt.close()

    logger.info(f"Visualization saved to {output_path}")