# AUTOGENERATED! DO NOT EDIT! File to edit: ../00_cluster_grid_search.ipynb.

# %% auto 0
__all__ = ['assign_cluster_labels', 'calculate_silhouette_avg', 'grid_search']

# %% ../00_cluster_grid_search.ipynb 3
from sklearn import cluster
import pandas as pd
import numpy as np
from typing import Union
from sklearn.base import ClusterMixin
from sklearn import metrics
from sklearn.metrics import silhouette_score
from tqdm import tqdm
from sklearn.model_selection import ParameterGrid

# %% ../00_cluster_grid_search.ipynb 7
def assign_cluster_labels(
    data_to_cluster: Union[pd.DataFrame, np.ndarray, list], # Data points to cluster
    cluster_model: ClusterMixin, # Sk-learn clustering model. Must be specified in the format `cluster.model`, e.g `cluster.KMeans`
    model_kwargs: dict, # Keyword arguments for sk-learn clustering model
) -> np.ndarray: # Label assignments
    """Assigns cluster labels to `data_to_cluster` using specified `cluster_model`"""
    cluster_model = cluster_model(**model_kwargs)
    return cluster_model.fit_predict(data_to_cluster)

# %% ../00_cluster_grid_search.ipynb 12
def calculate_silhouette_avg(
    data: Union[pd.DataFrame, np.ndarray, list], # The data that was used for the clustering
    cluster_labels: np.ndarray, # Cluster labels
) -> float: # Average silhouette score
    """Calculate the average silhouette score"""
    return silhouette_score(data, cluster_labels)

# %% ../00_cluster_grid_search.ipynb 17
def grid_search(
    data_to_cluster: Union[pd.DataFrame, np.ndarray, list], # Data points to cluster
    cluster_models: list, # List of Sk-learn clustering models to iterate through. Each model must be specified in the format `cluster.model`, e.g `cluster.KMeans`
    model_kwargs_list: list, # List of dicts of keyword arguments for the sk-learn clustering models to iterate through.
    highlight: bool, # True to highlight highest avg_silhouette_score
    sort: bool, # True to sort by highest avg_silhouette_score
) -> pd.DataFrame: # Table showing avg_silhouette_score for each model and parameter specified
    """Perform grid search for the specified clustering models and parameters"""
    cluster_model = []
    model_parameters = []
    silhouette_score = []
    for model, model_kwargs in tqdm(zip(cluster_models, model_kwargs_list)):
        for parameters in tqdm(ParameterGrid(model_kwargs)):
            labels = assign_cluster_labels(
                data_to_cluster,
                model,
                parameters
            )
            avg_silhouette_score = calculate_silhouette_avg(
                data_to_cluster,
                labels
            )
            cluster_model.append(str(model).split(".")[-1][:-2])
            model_parameters.append(parameters)
            silhouette_score.append(avg_silhouette_score)
    results = pd.DataFrame.from_dict({
        "cluster_model": cluster_model,
        "model_params": model_parameters,
        "avg_silhouette_score": silhouette_score
    })
    if sort:
        results = results.sort_values("avg_silhouette_score", ascending=False)
    if highlight:
        results = results.style.highlight_max(
            subset = ["avg_silhouette_score"],
            color = 'lightgreen', axis = 0
        )
    return results
