import os
import uuid

import matplotlib.pyplot as plt
import pandas as pd
from sklearn.cluster import KMeans
from sklearn_extra.cluster import KMedoids


def get_best_cluster_cnt(X, cluster_type):
    """ Use an elbow method for figuring out most suitable number of clusters for given data """
    wcss = []

    if cluster_type == 'kmedoids':
        for i in range(3, 10):
            algo = KMedoids(n_clusters=i, init='k-medoids++', random_state=0, metric='mahalanobis')
            algo.fit(X)
            wcss.append(algo.inertia_)

    elif cluster_type == 'kmeans':
        for i in range(3, 10):
            kmeans = KMeans(n_clusters=i, init='k-means++', random_state=0)
            kmeans.fit(X)
            wcss.append(kmeans.inertia_)

    """ Eblow is calculated by first calculating delta of current and previous row,
    calculating 2nd delta on top of the first, and then calculating the difference
    of delta2 and delta1 on next row"""
    df = pd.DataFrame({'wcss': wcss, 'k': range(1, 16)})
    df['delta1'] = df['wcss'].diff(periods=1) * -1
    df['delta2'] = df['delta1'].diff(periods=1) * -1
    df['strength'] = df['delta2'].shift(-1) - df['delta1'].shift(-1)

    # Select max strenght, but make sure it's always more than 2 clusters
    # TODO! maybe sometimes just 2 clusters is ok?
    max_idx = df['strength'].loc[df['k'] > 2].idxmax()

    best_k = df['k'].iloc[max_idx]

    # TODO! ax = sns.lineplot(x="timepoint", y="signal", data=fmri)
    fig = plt.figure()
    plt.plot(range(1, 16), wcss)
    plt.axvline(best_k, color='gray', linestyle='--')
    plt.title('Elbow Method')
    plt.xlabel('Number of clusters')
    plt.ylabel('wcss')

    current_path = os.path.dirname(__file__)
    unique_id = uuid.uuid4()
    elbow_graph = f'{current_path}/../resources/KMedoids_elbow_graph_{unique_id}.png'
    fig.savefig(elbow_graph)

    return best_k, elbow_graph