from pandas import np
from sklearn.metrics import pairwise_distances
from sklearn.cluster import KMeans
from simsearch import SimSearch
import matplotlib.pyplot as plt
import numpy as np


(ksearch, ssearch) = SimSearch.load(save_dir='mhc_corpus/')
X=ssearch.index.index
for k in range(1,6):
    kmeans= KMeans(n_clusters=k)
    a = kmeans.fit_predict(X)
    
def compute_inertia(a, X):
    W = [np.mean(pairwise_distances(X[a == c, :])) for c in np.unique(a)]
    return np.mean(W)


def compute_gap(clustering, data, k_max=5, n_references=5):
    if len(data.shape) == 1:
        data = data.reshape(-1, 1)
    reference = np.random.rand(*data.shape)
    reference_inertia = []
    for k in range(1, k_max + 1):
        local_inertia = []
        for _ in range(n_references):
            clustering.n_clusters = k
            assignments = clustering.fit_predict(reference)
            local_inertia.append(compute_inertia(assignments, reference))
        reference_inertia.append(np.mean(local_inertia))

    ondata_inertia = []
    for k in range(1, k_max + 1):
        clustering.n_clusters = k
        assignments = clustering.fit_predict(data)
        ondata_inertia.append(compute_inertia(assignments, data))

    gap = np.log(reference_inertia) - np.log(ondata_inertia)
    return gap, np.log(reference_inertia), np.log(ondata_inertia)

k_max = 5
gap, reference_inertia, ondata_inertia = compute_gap(KMeans(), X, k_max)

plt.plot(range(1, k_max+1), reference_inertia,         '-o', label='reference')
plt.plot(range(1, k_max+1), ondata_inertia,         '-o', label='data')
plt.xlabel('k')
plt.ylabel('log(inertia)')
plt.show()


plt.plot(range(1, k_max+1), gap, '-o')
plt.ylabel('gap')
plt.xlabel('k')
plt.show()