import time
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from netCDF4 import Dataset
import pickle







# create a k-mean clustering from 2D profile data
# input data must be two-dimensional where the first
# dimension is the number of profiles and the second
# dimension is the wavelengths
def create_kmean_from_data(data_2D, k_value=False):

    # data is normalized by the first value (wavelength[0])
    # for every profile
    data_2D /= data_2D[:,0,np.newaxis]      # normalize input data
    tot_profiles = data_2D.shape[0]
    print("Total profiles:                  %i" % tot_profiles)

    if not k_value:
        k0 = 10
        k1 = 100
        k = np.linspace(k0, k1, k1-k0+1).astype("int")
        inertia = []
        k_list = []
        t0 = time.time()


        for i,kk in enumerate(k):
            t1 = time.time()
            kmean = KMeans(n_clusters=kk,   # number of clusters
                           n_init=10,       # runs with different centroids
                           max_iter=200,    # max iterations for convergence
                           tol=1e-4,        # tolerance for convergence (default = 1e-4)
                           n_jobs=-1        # -1 = use all processors
                           ).fit(data_2D)   # fit model with "data_2D" data

            print("k:                               %i" % kk)
            print("Elapsed time training:           %.1f s" % (time.time()-t1))
            print("kmean iterations:                %i" % kmean.n_iter_)
            print("kmean inertia:                   %.2e" % kmean.inertia_)

            counts = np.bincount(kmean.labels_) # counts
            c_mean = 1/kk                       # count mean
            dist = counts/np.sum(counts)        # distributed counts
            rel_err = (dist-c_mean)/c_mean      # relative error
            lim = 0.50                          # acceptable limits
            rel_err_lim = np.abs(rel_err)<lim   # mark relative error limits
            n_in_lim = len(np.where(rel_err_lim)[0])    # count acceptable errors
            print("Mean:                            %.3f" % c_mean)
            print("Bins within %i%% of mean:         %i/%i = %.3f" % (100*lim,n_in_lim,kk,n_in_lim/kk))

            inertia.append(kmean.inertia_)
            k_list.append(kk)
            if kk != k0:
                print("kmean inertia gradient:          %.2e \n" % (inertia[i-1]-inertia[i]))
                plt.figure("Inertia gradient", figsize=(10,4))
                plt.clf()
                plt.subplot(121)
                norm_grad = inertia[0]-inertia[1]
                plt.plot(k_list[1:], (np.array(inertia[:-1])-np.array(inertia[1:]))/norm_grad, label="Gradient")
                plt.plot(k_list[1:], inertia[1:]/inertia[1], label="Inertia")
                plt.xlabel("k"); plt.ylabel("Inertia (normalized)")
                plt.legend()
                plt.grid()

                plt.subplot(122)
                plt.title("Bins within %i%% of mean: %i/%i = %.3f" % (100*lim,n_in_lim,kk,n_in_lim/kk))
                plt.hist(kmean.labels_, bins=kk, range=(0,kk), align="left")
                plt.plot([0,kk-1], [c_mean*tot_profiles,c_mean*tot_profiles], label="mean", color="black")
                plt.plot([0,kk-1], [(1+lim)*c_mean*tot_profiles,(1+lim)*c_mean*tot_profiles], label=r"mean$\pm$%i%%"%(100*lim), color="red")
                plt.plot([0,kk-1], [(1-lim)*c_mean*tot_profiles,(1-lim)*c_mean*tot_profiles], color="red")
                plt.xlabel("Cluster"); plt.ylabel("Counts")
                plt.legend()

                plt.tight_layout()
                plt.pause(0.1)
            else: print("kmean inertia gradient:\n")

        plt.savefig("kmean_figures/Inertia_gradient_all_k%i-k%i.png" % (k0,k1))
        print("Total elapsed time:              %.1f s" % (time.time()-t0))
        plt.show()

    else:
        t1 = time.time()
        kmean = KMeans(n_clusters=int(k_value), # number of clusters
                       n_init=50,               # runs with different centroids
                       max_iter=1000,           # max iterations for convergence
                       verbose=3,               # prints
                       n_jobs=-1                # -1 = use all processors
                       ).fit(data_2D)           # fit model with "data_2D" data
        print("Elapsed time training:         %.1f s" % (time.time()-t1))

        # save the model in 'KMEAN_MODELS' folder
        # structure of filename:
        # profiles_clusters_lower-percentile_upper-percentile_weight-indices_type-of-trained-profile.*
        model = {"kmean_model": kmean}
        pickle.dump(model, open("KMEAN_MODELS/%i_%i.pkl"
                                % (data_2D.shape[0], k_value),"wb"))


        # saving cluster centers and labels and also
        # individual profiles (used for training)
        clab = kmean.labels_
        cent = kmean.cluster_centers_
        clab.reshape(data_2D.shape[0])
        ncfile = Dataset("KMEAN_MODELS/%i_%i.nc"
                          % (data_2D.shape[0], k_value),"w")
        ncfile.title = "labels and cluster centers"

        xs = ncfile.createDimension("xsize", data_2D.shape[0])
        wl = ncfile.createDimension("wav", data_2D.shape[1])
        nc = ncfile.createDimension("noc", k_value)

        vlab = ncfile.createVariable("labels", np.int8, ("xsize"))
        vcen = ncfile.createVariable("centers", np.float64, ("noc","wav"))
        vprof = ncfile.createVariable("profiles", np.float32, ("xsize", "wav"))

        vlab[:] = clab
        vcen[:] = cent
        vprof[:] = data_2D
        ncfile.close()


##############################################################


# test of saved k-mean model
def test_kmean_model(file, test_profiles, wavelength=False, k_nearest=False):

    # load model file with pickle format
    kmm = pickle.load(open(file+".pkl","rb"))
    kmean = kmm["kmean_model"]
    num_p = test_profiles.shape[0]          # number of testing profiles

    # predict profiles with K-mean clustering
    pred_line = test_profiles / test_profiles[:,0,np.newaxis]      # normalize data
    pred = kmean.predict(pred_line)                     # predict data by clusters
    ds = Dataset(file+".nc")                            # load model attributes from .nc file
    cent = ds["centers"]                                # centre of clusters
    prob_pred = np.bincount(pred, minlength=len(cent))/num_p # predict distribution


    # predict profile(s) by K-nearest neighbors
    if k_nearest:
        labs = ds["labels"]                     # label of training profiles
        prof = ds["profiles"]                   # profiles used for training

        neigh = KNeighborsClassifier(n_neighbors=k_nearest).fit(prof, labs)
        neigh_pred = neigh.predict(pred_line)   # predict data by neighbors

        C_KNN_diff = np.bincount(neigh_pred==pred)
        print("Cluster and KNN predict equality [False, True]: ", C_KNN_diff)


    # plot with specified wavelengths on x-axis
    if wavelength.any() != False:
        plt.figure("lines")
        if k_nearest==False or C_KNN_diff[0]==0: r = np.random.randint(num_p)  # a random profile to plot
        else: r = np.where((neigh_pred==pred)==False)[0][np.random.randint(C_KNN_diff[0])]    # random non-equal predicted profiles

        plt.title("Test line %i" % r)
        plt.plot(wavelength, test_profiles[r,:] / test_profiles[r,0,np.newaxis],label="Test line")##
        plt.plot(wavelength, cent[pred[r],:], "*", color="black", label="Cluster predicts %i" % pred[r])##
        if k_nearest:
            plt.title("Test line %i, non-equality predictions %i = %.2f %%" %
                      (r, C_KNN_diff[0], C_KNN_diff[0]/np.sum(C_KNN_diff)))
            plt.plot(wavelength, cent[neigh_pred[r],:], ".", color="red", label="KNN predicts %i" % neigh_pred[r])##

        # plot all representative profiles from model
        for i in range(cent.shape[0]):
            plt.plot(wavelength, cent[i,:], "--", lw=1, label="C-line %i (%.1f%%)"%(i,100*prob_pred[i]))
        plt.legend()
        plt.show()

    # plot without specified wavelengths on x-axis
    else:
        plt.figure("lines")
        if k_nearest==False or C_KNN_diff[0]==0: r = np.random.randint(num_p)  # a random profile to plot
        else: r = np.where((neigh_pred==pred)==False)[0][np.random.randint(C_KNN_diff[0])]    # random non-equal predicted profiles

        plt.title("Test line %i" % r)
        plt.plot(test_profiles[r,:] / test_profiles[r,0,np.newaxis],label="Test line")##
        plt.plot(cent[pred[r],:], "*", color="black", label="Cluster predicts %i" % pred[r])##
        if k_nearest:
            plt.title("Test line %i, non-equality predictions %i = %.2f %%" %
                      (r, C_KNN_diff[0], C_KNN_diff[0]/np.sum(C_KNN_diff)))
            plt.plot(cent[neigh_pred[r],:], ".", color="red", label="KNN predicts %i" % neigh_pred[r])##

        # plot all representative profiles from model
        for i in range(cent.shape[0]):
            plt.plot(cent[i,:], "--", lw=1, label="C-line %i (%.1f%%)"%(i,100*prob_pred[i]))
        plt.legend()
        plt.show()
