# Paper Source: (1) Mehrjou A., Hosseini R. and Araabi B. N., 2016.
# Improved Bayesian information criterion for mixture model selection.
# Pattern Recognition Letters, 69, pp. 22-27.
# (2) H. Akaike, 1974.
# A new look at the statistical model identification.
# IEEE Trans. Autom. Control 19(6) pp. 716-723.

# Min number of clusters = 2
# The time complexity is O(nk), where n is the number of data points
# To estimate the number of clusters: Use min of AIC


import numpy as np
from scipy.spatial.distance import cdist


def aic(data, centers, labels):
    ni = np.fmax(
        np.unique(labels, return_counts=True)[1], np.finfo(float).eps
    )
    labelmask = np.zeros((centers.shape[0], data.shape[0]))
    for i in range(centers.shape[0]):
        labelmask[i, labels==i] = 1
    denom = data.shape[1] / (data.shape[0] - centers.shape[0])
    sigma = (cdist(centers, data, metric='sqeuclidean') * labelmask).sum() / denom
    '''
    return ((
        (ni * np.log(ni / data.shape[0]))
        - (0.5 * ni * data.shape[1] * np.log(2*np.pi))
        - (0.5 * ni * np.log(sigmai)) - (0.5 * (ni - centers.shape[0]))
    ).sum() - centers.shape[0])
    '''
    return (-2 *
        (ni * np.log(ni)) - (ni * data.shape[0])
        - (0.5 * ni * data.shape[1]) * np.log(2*np.pi*sigma)
        - ((ni - 1) * data.shape[1] * 0.5)
    ).sum() + (2 * centers.shape[0])
    #'''


if __name__ == '__main__':
    minK = 2
    from sklearn.datasets import load_iris
    data = load_iris().data
    maxK = int(np.ceil(data.shape[0] ** 0.5))

    from sklearn.cluster import KMeans
    index = np.zeros((maxK-minK))
    for k in range(minK,maxK):
        km1 = KMeans(
            n_clusters=k, n_init=30, max_iter=300, tol=1e-9
        ).fit(data)
        index[k-minK] = aic(
            data, km1.cluster_centers_, km1.labels_
        )
    #est_k = index.argmin() + minK
    from elbow_method import elbow_method
    est_k = elbow_method(index) + minK
    print('For Iris:\nSelected k =', est_k)

    import matplotlib.pyplot as plt
    fig, axisArray = plt.subplots(2,4)
    axisArray[0,0].scatter(data[:,2], data[:,3], marker='x', c='gray')
    axisArray[1,0].plot(np.arange(index.shape[0])+minK, index, linewidth=0.3, c='k')
    axisArray[1,0].scatter(np.arange(index.shape[0])+minK, index, marker='x', c='b')
    axisArray[1,0].scatter(est_k, index[est_k-minK], marker='x', c='r')
    axisArray[1,0].set_xticks(np.arange(index.shape[0])+minK)

    data = np.vstack(( np.vstack(( np.vstack(( np.vstack((
        np.random.normal(loc=np.array([0,0]), scale=1, size=(50,2)),
        np.random.normal(loc=np.array([20,20]), scale=1, size=(50,2)) )),
        np.random.normal(loc=np.array([0,20]), scale=1, size=(50,2)) )),
        np.random.normal(loc=np.array([20,0]), scale=1, size=(50,2)) )),
        np.random.normal(loc=np.array([10,10]), scale=1, size=(50,2))
    ))
    maxK = int(np.ceil(data.shape[0] ** 0.5))

    index = np.zeros((maxK-minK))
    for k in range(minK,maxK):
        km1 = KMeans(
            n_clusters=k, n_init=30, max_iter=300, tol=1e-9
        ).fit(data)
        index[k-minK] = aic(
            data, km1.cluster_centers_, km1.labels_
        )
    #est_k = index.argmin() + minK
    est_k = elbow_method(index) + minK
    print('For 5 Gaussians:\nSelected k =', est_k)
    axisArray[0,1].scatter(data[:,0], data[:,1], marker='x', c='gray')
    axisArray[1,1].plot(np.arange(index.shape[0])+minK, index, linewidth=0.3, c='k')
    axisArray[1,1].scatter(np.arange(index.shape[0])+minK, index, marker='x', c='b')
    axisArray[1,1].scatter(est_k, index[est_k-minK], marker='x', c='r')
    axisArray[1,1].set_xticks(np.arange(index.shape[0])+minK)

    data = np.vstack(( np.vstack(( np.vstack((
        np.random.normal(loc=np.array([0,0]), scale=1, size=(50,2)),
        np.random.normal(loc=np.array([20,20]), scale=1, size=(50,2)) )),
        np.random.normal(loc=np.array([0,20]), scale=1, size=(50,2)) )),
        np.random.normal(loc=np.array([20,0]), scale=1, size=(50,2))
    ))
    maxK = int(np.ceil(data.shape[0] ** 0.5))

    index = np.zeros((maxK-minK))
    for k in range(minK,maxK):
        km1 = KMeans(
            n_clusters=k, n_init=30, max_iter=300, tol=1e-9
        ).fit(data)
        index[k-minK] = aic(
            data, km1.cluster_centers_, km1.labels_
        )
    #est_k = index.argmin() + minK
    est_k = elbow_method(index) + minK
    print('For 4 Gaussians:\nSelected k =', est_k)
    axisArray[0,2].scatter(data[:,0], data[:,1], marker='x', c='gray')
    axisArray[1,2].plot(np.arange(index.shape[0])+minK, index, linewidth=0.3, c='k')
    axisArray[1,2].scatter(np.arange(index.shape[0])+minK, index, marker='x', c='b')
    axisArray[1,2].scatter(est_k, index[est_k-minK], marker='x', c='r')
    axisArray[1,2].set_xticks(np.arange(index.shape[0])+minK)


    data = np.vstack(( np.vstack((
        np.random.normal(loc=np.array([0,0]), scale=1, size=(50,2)),
        np.random.normal(loc=np.array([0,20]), scale=1, size=(50,2)) )),
        np.random.normal(loc=np.array([20,0]), scale=1, size=(50,2))
    ))
    maxK = int(np.ceil(data.shape[0] ** 0.5))

    index = np.zeros((maxK-minK))
    for k in range(minK,maxK):
        km1 = KMeans(
            n_clusters=k, n_init=30, max_iter=300, tol=1e-9
        ).fit(data)
        index[k-minK] = aic(
            data, km1.cluster_centers_, km1.labels_
        )
    #est_k = index.argmin() + minK
    est_k = elbow_method(index) + minK
    print('For 4 Gaussians:\nSelected k =', est_k)
    axisArray[0,3].scatter(data[:,0], data[:,1], marker='x', c='gray')
    axisArray[1,3].plot(np.arange(index.shape[0])+minK, index, linewidth=0.3, c='k')
    axisArray[1,3].scatter(np.arange(index.shape[0])+minK, index, marker='x', c='b')
    axisArray[1,3].scatter(est_k, index[est_k-minK], marker='x', c='r')
    axisArray[1,3].set_xticks(np.arange(index.shape[0])+minK)

    plt.show()
