import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, TextBox
import time
from astropy.io import fits
from astropy.wcs import WCS
from sklearn.cluster import KMeans
from sklearn.neighbors import KNeighborsClassifier
#import scipy as sp
from netCDF4 import Dataset
import pickle
from iris2model_sav_to_python_format import convert_sav_to_dict
import iris2_lists
#import dask.array as da
import warnings     # ignore tedious warnings
warnings.filterwarnings("ignore")





# plot the different variables
def plot_parameters(show=True):
    ltau_valstep = LTAU[1] - LTAU[0]
    phys_var_idx = 0
    ltau_idx = int(len(LTAU)/2)

    fig, (ax1,ax2,ax3,ax4) = plt.subplots(1,4, figsize=(16,6))

    # making a slider with respect to log(tau)
    slider = Slider(plt.axes([0.025, 0.1, 0.025, 0.8]),
                    r"log($\tau$)",
                    np.min(LTAU),
                    np.max(LTAU),
                    valstep=ltau_valstep,
                    valinit=LTAU[ltau_idx],
                    orientation="vertical",
                    valfmt="%.1f")


    dat1 = MODEL[:,:,ltau_idx,0]
    dat2 = MODEL[:,:,ltau_idx,1]
    dat3 = MODEL[:,:,ltau_idx,2]
    dat4 = MODEL[:,:,ltau_idx,3]

    ax1.set_title("%s" % PHYS_VAR[0][:PHYS_VAR[0].find("[")])
    im1 = ax1.imshow(dat1.T, cmap="gray", aspect="auto", origin="lower", vmin=3e3, vmax=1.5*np.mean(dat1))
    cbar1 = fig.colorbar(im1, ax=ax1, label=PHYS_VAR[0][PHYS_VAR[0].find("["):], format="%.0e")


    ax2.set_title("%s" % PHYS_VAR[1][:PHYS_VAR[1].find("[")])
    im2 = ax2.imshow(dat2.T, cmap="gray", aspect="auto", origin="lower")
    cbar2 = fig.colorbar(im2, ax=ax2, label=PHYS_VAR[1][PHYS_VAR[1].find("["):], format="%.0e")

    ax3.set_title("%s" % PHYS_VAR[2][:PHYS_VAR[2].find("[")])
    im3 = ax3.imshow(dat3.T, cmap="gray", aspect="auto", origin="lower", vmin=np.percentile(dat3,1), vmax=np.percentile(dat3,98))
    cbar3 = fig.colorbar(im3, ax=ax3, label=PHYS_VAR[2][PHYS_VAR[2].find("["):], format="%.0e")

    ax4.set_title("%s" % PHYS_VAR[3][:PHYS_VAR[3].find("[")])
    im4 = ax4.imshow(dat4.T, cmap="gray", aspect="auto", origin="lower", vmin=np.percentile(dat4,1), vmax=np.percentile(dat4,98))
    cbar4 = fig.colorbar(im4, ax=ax4, label=PHYS_VAR[3][PHYS_VAR[3].find("["):], format="%.0e")

    def update(X):
        idx = int(np.argmin(abs(LTAU-X)))

        dat1 = MODEL[:,:,idx,0]
        dat2 = MODEL[:,:,idx,1]
        dat3 = MODEL[:,:,idx,2]
        dat4 = MODEL[:,:,idx,3]
        im1.set_data(dat1.T)
        im2.set_data(dat2.T)
        im3.set_data(dat3.T)
        im4.set_data(dat4.T)
        im1.set_clim(vmin=3e3, vmax=1.5*np.mean(dat1))
        im2.set_clim(vmin=-0.5*np.max(dat2), vmax=0.5*np.max(dat2))
        im3.set_clim(vmin=np.percentile(dat3,1), vmax=np.percentile(dat3,98))#vmin=0.1*np.mean(dat3), vmax=3*np.mean(dat3))
        im4.set_clim(vmin=np.percentile(dat4,1), vmax=np.percentile(dat4,98))#0.1*np.mean(dat4), vmax=3*np.mean(dat4))
        cbar1.update_normal(im1)
        cbar2.update_normal(im2)
        cbar3.update_normal(im3)
        cbar4.update_normal(im3)

        #fig_par.canvas.draw_idle()
        fig.canvas.draw_idle()


    plt.tight_layout(rect=(0.05,0,1,1))
    slider.on_changed(update)
    if show: plt.show()



# plot RP line and observation line along
# with the model atmosphere of choice
def plot_RP_line(tau_idx=-1, initX=100, initY=100, show=True, show_weights=False):
    temp_model = MODEL[:,:,tau_idx,0]

    global valX, valY, init_p1, init_p2, cont_p1, cont_p2
    valX = initX
    valY = initY

    init_p1 = 100            # initial percentile
    init_p2 = 0            # initial percentile
    val_p1 = np.percentile(temp_model, init_p1)
    val_p2 = np.percentile(temp_model, init_p2)
    cont_p1 = np.zeros(temp_model.shape)
    cont_p2 = np.zeros(temp_model.shape)
    cont_p1[np.where(temp_model>val_p1)] = 1
    cont_p2[np.where(temp_model<val_p2)] = 1


    left, width     = 0.1, 0.2
    bottom, height  = 0.2, 0.65

    figx,figy = 14,6
    fig_line = plt.figure("Profiles", figsize=(figx,figy))
    axes = plt.axes([left, bottom, width, height])
    axes_line = plt.axes([left+width+0.1, bottom, 1-(left+width+0.15), height])
    plt.axes(axes)
    # make slider with respect to x and y coordinate
    sliderX = Slider(   plt.axes([left, 0.05, width, 0.003*figx]),
                        "X",
                        0,
                        RP_LINES.shape[0]-1,
                        valinit=valX,
                        valfmt="%i")
    sliderY = Slider(   plt.axes([0.025, bottom, 0.003*figy, height]),
                        "Y",
                        0,
                        RP_LINES.shape[1]-1,
                        valinit=valY,
                        orientation="vertical",
                        valfmt="%i")

    # upper and lower percentile sliders
    slider_p1 = Slider( plt.axes([left, 0.95, width, 0.002*figx]),
                        label="Upper percentile",
                        valmin=0,
                        valmax=100,
                        valinit=init_p1,
                        valfmt="%i")
    slider_p2 = Slider( plt.axes([left, 0.9, width, 0.002*figx]),
                        label="Lower percentile",
                        valmin=0,
                        valmax=100,
                        valinit=init_p2,
                        valfmt="%i")


    def updateX(X):
        global valX
        valX = int(X)
        RP_plot.set_data(WAVELENGTH, RP_LINES[valX,valY,:])
        og_plot.set_data(WAVELENGTH, OG_LINES[valX,valY,:])
        idx_plot.set_data(valX,valY)
        chi2_text.set_text(r"$\chi^2$ = %.3f" % CHI2[valX,valY])
        fig_line.canvas.draw_idle()

    def updateY(Y):
        global valY
        valY = int(Y)
        RP_plot.set_data(WAVELENGTH, RP_LINES[valX,valY,:])
        og_plot.set_data(WAVELENGTH, OG_LINES[valX,valY,:])
        idx_plot.set_data(valX,valY)
        chi2_text.set_text(r"$\chi^2$ = %.3f" % CHI2[valX,valY])
        fig_line.canvas.draw_idle()


    def update_p1(P1):
        global val_p1, cont_p1, cont_plot_p1
        val_p1 = np.percentile(temp_model, int(P1))
        try:
            for coll in cont_plot_p1.collections:
                coll.remove()
            cont_p1[:,:] = 0
            cont_p1[np.where(temp_model>val_p1)] = 1
            cont_plot_p1 = plt.contour(cont_p1.T, cmap="bwr_r", levels=[0,1], linewidths=0.7)
        except:
            pass
        fig_line.canvas.draw_idle()


    def update_p2(P2):
        global val_p2, cont_p2, cont_plot_p2
        val_p2 = np.percentile(temp_model, int(P2))
        try:
            for coll in cont_plot_p2.collections:
                coll.remove()
            cont_p2[:,:] = 0
            cont_p2[np.where(temp_model<val_p2)] = 1
            cont_plot_p2 = plt.contour(cont_p2.T, cmap="winter", levels=[0,1], linewidths=0.7)
        except:
            pass
        fig_line.canvas.draw_idle()



    def submit_chi2(val):
        chi2_val = float(val)
        close_chi2_val = np.min(np.abs(CHI2-chi2_val))
        idx_chi2 = np.where(abs(CHI2-chi2_val) == close_chi2_val)

        global valX, valY
        valX, valY = int(idx_chi2[0][0]), int(idx_chi2[1][0])

        RP_plot.set_data(WAVELENGTH, RP_LINES[valX,valY,:])
        og_plot.set_data(WAVELENGTH, OG_LINES[valX,valY,:])
        idx_plot.set_data(valX,valY)
        chi2_text.set_text(r"$\chi^2$ = %.3f" % CHI2[valX,valY])

        sliderX.set_val(valX)
        sliderY.set_val(valY)

        plt.draw()



    def submit_X(X):
        global valX
        valX = int(X)

        RP_plot.set_data(WAVELENGTH, RP_LINES[valX,valY,:])
        og_plot.set_data(WAVELENGTH, OG_LINES[valX,valY,:])
        idx_plot.set_data(valX,valY)
        chi2_text.set_text(r"$\chi^2$ = %.3f" % CHI2[valX,valY])

        sliderX.set_val(valX)
        sliderY.set_val(valY)
        plt.draw()


    def submit_Y(Y):
        global valY
        valY = int(Y)

        RP_plot.set_data(WAVELENGTH, RP_LINES[valX,valY,:])
        og_plot.set_data(WAVELENGTH, OG_LINES[valX,valY,:])
        idx_plot.set_data(valX,valY)
        chi2_text.set_text(r"$\chi^2$ = %.3f" % CHI2[valX,valY])

        sliderX.set_val(valX)
        sliderY.set_val(valY)
        plt.draw()




    # search for nearest chi2 in model atmosphere
    axbox_chi2 = plt.axes([0.9, 0.05, 0.05, 0.05])
    chi2_box = TextBox(axbox_chi2, r"Find nearest $\chi^2$ ")
    chi2_box.on_submit(submit_chi2)


    # submit X and Y
    axbox_X = plt.axes([0.025, 0.05, 0.025, 0.035])
    axbox_Y = plt.axes([0.025, 0.087, 0.025, 0.035])
    X_box = TextBox(axbox_X, "X ")
    Y_box = TextBox(axbox_Y, "Y ")
    X_box.on_submit(submit_X)
    Y_box.on_submit(submit_Y)


    # left plot
    plt.axes(axes)
    idx_plot, = plt.plot(valX,valY,"*",ms=3,color="cyan")
    im = plt.imshow(temp_model.T,
                    cmap="gray", aspect="auto", origin="lower",
                    vmax=np.mean(temp_model)*2.5)
    global cont_plot_p1, cont_plot_p2
    cont_plot_p1 = plt.contour(cont_p1.T, cmap="bwr_r", levels=[0,1], linewidths=0.7)
    cont_plot_p2 = plt.contour(cont_p2.T, cmap="winter", levels=[0,1], linewidths=0.7)
    #plt.scatter(remove_pixels[0], remove_pixels[1], color="red")
    #plt.scatter(remove_neg[0], remove_neg[1], color="red")


    # right plot
    plt.axes(axes_line)
    colors = ["purple", "blue", "red", "yellow"]
    # plots the weighted ares in colors
    if show_weights:
        for i,lst in enumerate(WEIGHTS_IDX):
            for j,idx in enumerate(lst):
                if j==0:
                    plt.plot([WAVELENGTH[idx],WAVELENGTH[idx]],[YLIM_PROFILE[0],
                             YLIM_PROFILE[1]], color=colors[i],alpha=0.2,
                             label="W%i"%(i+1), lw=3)
                else:
                    plt.plot([WAVELENGTH[idx],WAVELENGTH[idx]],[YLIM_PROFILE[0],
                             YLIM_PROFILE[1]], color=colors[i],alpha=0.2, lw=3)

    RP_plot, = plt.plot(WAVELENGTH, RP_LINES[valX,valY,:], label="RP")
    og_plot, = plt.plot(WAVELENGTH, OG_LINES[valX,valY,:], label="Observation", lw=1)
    #plt.plot(WAVELENGTH, np.mean(RP_LINES,axis=(0,1)), "--", label="mean RP", lw=1)
    #plt.plot(WAVELENGTH, np.mean(OG_LINES,axis=(0,1)), "--", label="mean obs", lw=1)
    chi2_text = plt.text(np.min(WAVELENGTH), 1.1*YLIM_PROFILE[1],
                         r"$\chi^2$ = %.3f" % CHI2[valX,valY],
                         fontsize=15)
    weights_text = plt.text(1.002*np.min(WAVELENGTH), 1.1*YLIM_PROFILE[1],
                         "Weights = [%i,%i,%i,%i]" % (WEIGHTS[0],WEIGHTS[1],WEIGHTS[2],WEIGHTS[3]),
                         fontsize=15)
    plt.ylim(YLIM_PROFILE[0], YLIM_PROFILE[1])
    plt.ylabel("Itensity " + str(LINE_UNITS)[2:-1])
    plt.xlabel("Wavelength [Å]")
    plt.legend()
    plt.grid()

    plt.axes(axes)
    sliderX.on_changed(updateX)
    sliderY.on_changed(updateY)
    slider_p1.on_changed(update_p1)
    slider_p2.on_changed(update_p2)
    if show: plt.show()



# calculate the average profile in a given region
def average_profile(box):
    if type(box) == type(0):
        msg = ""
        assert False, "Integer not a valid input! Must be of type [left, \
                        \nbottom, width, height] or 2D array of coordinates"

    elif len(box) == 4:
        RP_mean_line = np.average(RP_LINES[box[0]:box[0]+box[2],
                                           box[1]:box[1]+box[3],:],axis=(0,1))
        OG_mean_line = np.average(OG_LINES[box[0]:box[0]+box[2],
                                           box[1]:box[1]+box[3],:],axis=(0,1))
        return RP_mean_line, OG_mean_line

    elif type(box)==type(()) or type(box)==type(np.array([])):
        RP_mean_line = np.average(RP_LINES[box[0],box[1],:],axis=0)
        OG_mean_line = np.average(OG_LINES[box[0],box[1],:],axis=0)
        return RP_mean_line, OG_mean_line

    else:
        print("Entry 'box' in 'average_profile(box)' not valid! Must be 4 \
               \nelements long or an 2D array with coordinates.")
        return [], []



# visualise the distribution of chi^2
def chi2_analysis(plot=False, large_lim=False, wl_lim=False):

    mean_chi2 = np.mean(CHI2[MASK_ONDATA])
    median_chi2 = np.median(CHI2[MASK_ONDATA])
    max_chi2 = np.where(CHI2 == np.max(CHI2[MASK_ONDATA]))
    m1,m2 = np.where(MASK_ONDATA)                   # only points with valid data
    w = np.array(WEIGHTS_IDX[0]+WEIGHTS_IDX[2])     # only wavelength with corresponding W1 and W3 weight
    rp_og_diff = np.sum(np.abs(np.sum(RP_LINES[m1,m2,:],axis=1)-np.sum(OG_LINES[m1,m2,:],axis=1)))
    rp_og_diff_w1w3 = np.sum(np.abs(np.sum(RP_LINES[m1,m2][:,w],axis=1)-np.sum(OG_LINES[m1,m2][:,w],axis=1)))
    # residuals
    res = np.abs(RP_LINES[m1,m2,:] - OG_LINES[m1,m2,:])
    res_w1w3 = res[:,WEIGHTS_IDX[0]+WEIGHTS_IDX[2]]
    res_mean = np.mean(np.sum(res, axis=1))
    res_mean_w1w3 = np.mean(np.sum(res_w1w3, axis=1))

    print()
    print("Weights:                            ", WEIGHTS)
    print("Mean chi2:                           %.2f" % mean_chi2)
    print("Median chi2:                         %.2f" % median_chi2)
    print("Max chi2 at [%i,%i]:               %i" % (max_chi2[0],max_chi2[1],CHI2[max_chi2][0]))
    print("abs(RP-OG)                           %.5f" % rp_og_diff)
    print("abs(RP-OG) (w1 w3)                   %.5f" % rp_og_diff_w1w3)
    print("mean abs(residuals)                  %.5e" % res_mean)
    print("mean abs(residuals) (w1 w3)          %.5e" % res_mean_w1w3)



    if large_lim:
        large_chi2 = np.where((CHI2>large_lim[0]) & (CHI2<large_lim[1]))
        r_idx = np.random.randint(len(large_chi2[0]))
        idxX, idxY = large_chi2[0][r_idx], large_chi2[1][r_idx]
        print("Limits of large chi2:                [%i,%i]" % (large_lim[0], large_lim[1]))
        print("Random large chi2[%i,%i]:          %i" % (idxX, idxY, CHI2[idxX,idxY]))
        print("Total points in limits:              %i" % len(large_chi2[0]))

    if wl_lim:
        pass

    if plot:
        range1 = (0,50)
        range2 = (50,200)
        range3 = (200,400)
        range4 = (400,1000)

        plt.figure("Histogram", figsize=(20,5))
        plt.subplot(141)
        h1 = plt.hist(CHI2[MASK_ONDATA].flatten(), bins=75, range=(range1[0],range1[1]))
        plt.title("Total points: %i" % np.sum(h1[0]))
        plt.subplot(142)
        h2 = plt.hist(CHI2[MASK_ONDATA].flatten(), bins=75, range=(range2[0],range2[1]))
        plt.title("Total points: %i" % np.sum(h2[0]))
        plt.subplot(143)
        h3 = plt.hist(CHI2[MASK_ONDATA].flatten(), bins=75, range=(range3[0],range3[1]))
        plt.title("Total points: %i" % np.sum(h3[0]))
        plt.subplot(144)
        h4 = plt.hist(CHI2[MASK_ONDATA].flatten(), bins=75, range=(range4[0],range4[1]))
        plt.title("Total points: %i" % np.sum(h4[0]))
        plt.tight_layout()

        print("Points in section [%i-%i]:            %i" % (range1[0], range1[1], np.sum(h1[0])))
        print("Points in section [%i-%i]:          %i" % (range2[0], range2[1], np.sum(h2[0])))
        print("Points in section [%i-%i]:         %i" % (range3[0], range3[1], np.sum(h3[0])))
        print("Points in section [%i-%i]:        %i" % (range4[0], range4[1], np.sum(h4[0])))

        plt.figure("Residual", figsize=(12,5))
        res_hist = plt.hist(np.sum(res, axis=1).flatten(), bins=250, range=(0, 6e-5))
        plt.plot([res_mean,res_mean], [0,np.max(res_hist[0])], label="Mean")
        plt.legend()
        plt.xlabel("abs(residual)"); plt.ylabel("Counts")
        plt.title("Mean residual: %.2e" % res_mean)
        plt.tight_layout()

        plt.show()



# compute log(tau_0) at max(response function) for an arbitrary RP
# for temperature, v_los, v_turb and electron density
def tau0_from_RF(RP_line=None, idx_wl=False, wl_isolate=False, idx_wl_label=None, show_RF=False):

    rf_path = "/astro/local/opt/solarsoft_db/iris/iris2/latest/"
    rf_temp = "rf_temp_n07_iris2.v2.0.fits"
    rf_vlos = "rf_vlos_n04_iris2.v2.0.fits"
    rf_vturb = "rf_vturb_n04_iris2.v2.0.fits"
    rf_nne = "rf_nne_n07_iris2.v2.0.fits"
    wl = "wl_iris2.v2.0.fits"

    temp = fits.open(rf_path+rf_temp)
    vlos = fits.open(rf_path+rf_vlos)
    vturb = fits.open(rf_path+rf_vturb)
    nne = fits.open(rf_path+rf_nne)
    wl = fits.open(rf_path+wl)[0].data
    c_levels = 250

    temp_data = temp[0].data
    vlos_data = vlos[0].data
    vturb_data = vturb[0].data
    nne_data = nne[0].data
    ltau_temp = np.array(temp[0].header["NODESVAL"].replace("[","").replace("]","").split(",")).astype(np.float)
    ltau_vlos = np.array(vlos[0].header["NODESVAL"].replace("[","").replace("]","").split(",")).astype(np.float)
    ltau_vturb = np.array(vturb[0].header["NODESVAL"].replace("[","").replace("]","").split(",")).astype(np.float)
    ltau_nne = np.array(nne[0].header["NODESVAL"].replace("[","").replace("]","").split(",")).astype(np.float)

    if RP_line: r = RP_line
    else: r = np.random.randint(temp_data.shape[0])


    if idx_wl:

        # making sure to use correct index in wl
        idx_wl_old = idx_wl
        idx_wl = np.argmin(np.abs(wl-WAVELENGTH[idx_wl_old]))

        # consider only wavelength surrounding input idx_wl instead of allwavelengths
        if wl_isolate:
            d_idx = 20
            wl = wl[idx_wl-d_idx:idx_wl+d_idx]
            temp_data = temp_data[:,:,idx_wl-d_idx:idx_wl+d_idx]
            vlos_data = vlos_data[:,:,idx_wl-d_idx:idx_wl+d_idx]
            vturb_data = vturb_data[:,:,idx_wl-d_idx:idx_wl+d_idx]
            nne_data = nne_data[:,:,idx_wl-d_idx:idx_wl+d_idx]
            idx_wl = np.argmin(np.abs(wl-WAVELENGTH[idx_wl_old]))

        # log(tau) at max(response function) for a given wavelength
        rf_idx_temp = temp_data[RP_line,:,idx_wl]
        rf_idx_vlos = vlos_data[RP_line,:,idx_wl]
        rf_idx_vturb = vturb_data[RP_line,:,idx_wl]
        rf_idx_nne = nne_data[RP_line,:,idx_wl]
        ltau0_temp = ltau_temp[np.argmax(rf_idx_temp)]
        ltau0_vlos = ltau_vlos[np.argmax(rf_idx_vlos)]
        ltau0_vturb = ltau_vturb[np.argmax(rf_idx_vturb)]
        ltau0_nne = ltau_nne[np.argmax(rf_idx_nne)]


    if show_RF:
        fig, ((ax_temp,ax_vlos),(ax_vturb,ax_nne)) = plt.subplots(2,2,figsize=(10,10),
                                                                  num="Response functions",
                                                                  sharex=True, sharey=True)

        # response function of temperature
        contour_temp = ax_temp.contourf(wl, ltau_temp, temp_data[r,:,:], cmap="bwr", levels=c_levels)
        if idx_wl:
            ax_temp.plot([wl[idx_wl],wl[idx_wl]],
                         [np.min(ltau_temp), np.max(ltau_temp)],
                         "--", label=idx_wl_label, color="orange")
        ax_temp.set_title("Temperature, RP %i" % r)
        ax_temp.set_ylabel(r"log($\tau$)")

        # response function of line of sight velocity
        contour_vlos = ax_vlos.contourf(wl, ltau_vlos, vlos_data[r,:,:], cmap="bwr", levels=c_levels)
        if idx_wl: ax_vlos.plot([wl[idx_wl],wl[idx_wl]],
                                [np.min(ltau_temp), np.max(ltau_temp)],
                                "--", color="orange")
        ax_vlos.set_title("Velocity los, RP %i" % r)

        # response function of turbulence velocity
        contour_vturb = ax_vturb.contourf(wl, ltau_vturb, vturb_data[r,:,:], cmap="bwr", levels=c_levels)
        if idx_wl: ax_vturb.plot([wl[idx_wl],wl[idx_wl]],
                                [np.min(ltau_temp), np.max(ltau_temp)],
                                "--", color="orange")
        ax_vturb.set_title("Velocity turb, RP %i" % r)
        ax_vturb.set_xlabel("Wavelength [Å]")
        ax_vturb.set_ylabel(r"log($\tau$)")

        # response function of electron density
        contour_nne = ax_nne.contourf(wl, ltau_nne, nne_data[r,:,:], cmap="bwr", levels=c_levels)
        if idx_wl: ax_nne.plot([wl[idx_wl],wl[idx_wl]],
                                [np.min(ltau_temp), np.max(ltau_temp)],
                                "--", color="orange")
        ax_nne.set_title("Electron density, RP %i" % r)
        ax_nne.set_xlabel("Wavelength [Å]")

        fig.legend()
        fig.tight_layout()

        if idx_wl:
            fig_idx, ((ax_t_idx,ax_vl_idx),(ax_vt_idx,ax_n_idx)) = \
                        plt.subplots(2,2,figsize=(10,10), num="Single wl response functions")
            ax_t_idx.plot(ltau_temp, rf_idx_temp, "*-", color="red")
            ax_t_idx.plot(ltau0_temp, rf_idx_temp[np.argmax(rf_idx_temp)],
                          "o", label="log(tau)=%.1f"%ltau0_temp,color="y")
            ax_t_idx.set_title("Temperature")
            ax_t_idx.set_ylabel("Response function")
            ax_t_idx.legend()

            ax_vl_idx.plot(ltau_vlos, rf_idx_vlos, "*-", color="purple")
            ax_vl_idx.plot(ltau0_vlos, rf_idx_vlos[np.argmax(rf_idx_vlos)],
                          "o", label="log(tau)=%.1f"%ltau0_vlos,color="y")
            ax_vl_idx.set_title("Velocity los")
            ax_vl_idx.legend()

            ax_vt_idx.plot(ltau_vturb, rf_idx_vturb, "*-", color="green")
            ax_vt_idx.plot(ltau0_vturb, rf_idx_vturb[np.argmax(rf_idx_vturb)],
                          "o", label="log(tau)=%.1f"%ltau0_vturb,color="y")
            ax_vt_idx.set_title("Velocity turb")
            ax_vt_idx.set_xlabel(r"log($\tau$)")
            ax_vt_idx.set_ylabel("Response function")
            ax_vt_idx.legend()

            ax_n_idx.plot(ltau_nne, rf_idx_nne, "*-", color="black")
            ax_n_idx.plot(ltau0_nne, rf_idx_nne[np.argmax(rf_idx_nne)],
                          "o", label="log(tau)=%.1f"%ltau0_nne,color="y")
            ax_n_idx.set_title("Electron density")
            ax_n_idx.set_xlabel(r"log($\tau$)")
            ax_n_idx.legend()

            fig_idx.tight_layout()
        plt.show()

    del temp, vlos, vturb, nne, temp_data, vlos_data, vturb_data, nne_data
    return ltau0_temp, ltau0_vlos, ltau0_vturb, ltau0_nne



# find nearest peaks and depression
def find_profile_signatures(wl_init, profile):
    centre = np.argmin(np.abs(WAVELENGTH-wl_init))
    rerun = True

    while rerun:
        if profile[centre+1] < profile[centre]:
            centre += 1
        elif profile[centre-1] < profile[centre]:
            centre -= 1
        else:
            rerun = False

    blue_peak = centre
    red_peak = centre

    while profile[blue_peak-1] > profile[blue_peak]:
        blue_peak -= 1
    while profile[red_peak+1] > profile[red_peak]:
        red_peak += 1

    return centre, blue_peak, red_peak



# find RP-OG profile pair with high residual value
def locate_large_residual_profiles(lower_perc, upper_perc, wl_idx=False):
    if not wl_idx:
        wl_idx = np.linspace(0,len(WAVELENGTH)-1,len(WAVELENGTH)).astype("int")

    res = np.abs(RP_LINES[:,:,wl_idx] - OG_LINES[:,:,wl_idx]) * MASK_ONDATA[:,:,np.newaxis]
    res = np.sum(res, axis=2)
    m1,m2 = np.where(MASK_ONDATA)           # only points with valid data
    res_mask = res[m1,m2]

    idx_x, idx_y = np.where((res > np.percentile(res_mask, lower_perc)) &
                            (res < np.percentile(res_mask, upper_perc)))
    #plt.contourf(MASK_ONDATA.T, levels=1)
    #plt.scatter(idx_x,idx_y, color="red")
    #plt.show()
    return idx_x, idx_y



# create a k-mean clustering from 2D profile data
def create_kmean_from_data(data_2D, percs=[0,100], k_value=False):

    # select only profiles with residuals that match input percentile domain
    X,Y = locate_large_residual_profiles(lower_perc=percs[0],
                                         upper_perc=percs[1],
                                         wl_idx=WEIGHTS_DICT_IDX[weight_ind])
    data_2D = data_2D[X,Y,:][:,WEIGHTS_DICT_IDX[weight_ind]]
    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 = 50
        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("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_%i_%i_%s_%s.pkl"
                                % (data_2D.shape[0], k_value, percs[0],
                                   percs[1], weight_ind, train_on),"wb"))


        # saving cluster centers and labels
        clab = kmean.labels_
        cent = kmean.cluster_centers_
        clab.reshape(data_2D.shape[0])
        ncfile = Dataset("KMEAN_MODELS/%i_%i_%i_%i_%s_%s.nc"
                          % (data_2D.shape[0], k_value, percs[0],
                             percs[1], weight_ind, train_on),"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_model(file, specified_lines=False, k_nearest=False):

    # extract correct WAVELENGTH indices used when training the model
    indices = WEIGHTS_DICT_IDX[weight_ind]


    # use only profiles within fixed (trained on) percentiles
    if not specified_lines:
        # extract info from filename, remove path specification
        for l in range(len(file)):
            if file[-(l+1)] == "/":
                slice = -l
                break
        text = file[slice:].replace("_", " ").split()
        lp,up = int(text[2]), int(text[3])     # lower and upper percentile from filename
        X,Y = locate_large_residual_profiles(lower_perc=lp, upper_perc=up, wl_idx=indices)

    # to predict specified profile(s)
    else:
        X = specified_lines[0]
        Y = specified_lines[1]


    # load model file with pickle format
    kmm = pickle.load(open(file+".pkl","rb"))
    kmean = kmm["kmean_model"]

    # extract lines for testing the model
    line = np.zeros((len(X), len(WAVELENGTH)))
    for i in range(line.shape[0]):
        line[i] = OG_LINES[X[i],Y[i],:]

    # predict profiles with K-mean clustering
    pred_line = line[:,indices] / line[:,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)) / len(X) # 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)


    plt.figure("lines")
    if k_nearest==False or C_KNN_diff[0]==0: r = np.random.randint(len(X))  # 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 [x,y] = [%i,%i]" % (X[r],Y[r]))
    plt.plot(WAVELENGTH[indices], line[r,indices] / line[r,0,np.newaxis],label="Test line")##
    plt.plot(WAVELENGTH[indices], cent[pred[r],:], "*", color="black", label="Cluster predicts %i" % pred[r])##
    if k_nearest:
        plt.title("Test line [x,y] = [%i,%i], non-equality predictions %i = %.2f %%" %
                  (X[r], Y[r], C_KNN_diff[0], C_KNN_diff[0]/np.sum(C_KNN_diff)))
        plt.plot(WAVELENGTH[indices], 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[indices], cent[i,:], "--", lw=1, label="C-line %i (%.1f%%)"%(i,100*prob_pred[i]))
    plt.legend()
    plt.show()


# creates a plot of the observation data at given
# wavelength along the mapping of profiles labelled
# by clusters
def plot_observation_with_cluster_label(file, spectral_line="2796"):
    # use only wavelengths that has been
    # used (weighted) when training the model.
    # it should always be stated in the file name
    indices = WEIGHTS_DICT_IDX[file.split("_")[-2]]

    # load k-mean model
    #kmm = pickle.load(open(file+".pkl","rb"))
    #kmean = kmm["kmean_model"]



    # label observation lines using the
    # k-mean model on data






# plot observational and representative spectrum
def plot_spectrum(slit_i=0):

    print("Spectrum plot of slit index: %i " % slit_i)
    fig, (ax_og, ax_rp, ax_diff) = plt.subplots(1,3, figsize=(16,5),
                                                sharex=True, sharey=True)

    # plot spectrum of the observational spectrum
    spec_og = ax_og.contourf(WAVELENGTH,
                             SLIT[MASK_ONDATA[slit_i, :]],
                             OG_LINES[slit_i, MASK_ONDATA[slit_i,:], :],
                             cmap="inferno",
                             levels=100)
    ax_og.set_xlabel("Wavelength (Å)")
    ax_og.set_ylabel("Slit (arcsecond)")
    ax_og.set_title("Observational spectrum")
    fig.colorbar(spec_og, ax=ax_og)


    # plot spectrum of the representative spectrum
    spec_rp = ax_rp.contourf(WAVELENGTH,
    SLIT[MASK_ONDATA[slit_i, :]],
    RP_LINES[slit_i, MASK_ONDATA[slit_i,:], :],
    cmap="inferno",
    levels=100)
    ax_rp.set_xlabel("Wavelength (Å)")
    #ax_rp.set_ylabel("Slit (arcsecond)")
    ax_rp.set_title("Representative spectrum")
    fig.colorbar(spec_rp, ax=ax_rp)


    # plot difference between representative and observational spectrum
    spec_diff = ax_diff.contourf(WAVELENGTH,
                             SLIT[MASK_ONDATA[slit_i, :]],
                             np.abs(RP_LINES[slit_i, MASK_ONDATA[slit_i,:], :]-
                             OG_LINES[slit_i, MASK_ONDATA[slit_i,:], :]),
                             cmap="inferno",
                             levels=100)
    ax_diff.set_xlabel("Wavelength (Å)")
    #ax_og.set_ylabel("Slit (arcsecond)")
    ax_diff.set_title("Absolute spectrum difference")
    fig.colorbar(spec_diff, ax=ax_diff)

    fig.tight_layout()
    fig.show()







if __name__ == '__main__':

    #file = "SAV_DATA/.sav"
    files = ["SAV_DATA/10_10_10_10.sav", # 0 - best residual value
             "SAV_DATA/10_09_10_08.sav", # 1
             "SAV_DATA/10_05_10_02.sav", # 2
             "SAV_DATA/10_10_10_05.sav", # 3
             "SAV_DATA/05_02_05_01.sav", # 4
             "SAV_DATA/06_03_06_01.sav", # 5
             "SAV_DATA/10_03_10_01.sav", # 6
             "SAV_DATA/10_08_10_06.sav", # 7
             "SAV_DATA/10_08_10_05.sav", # 8
             "SAV_DATA/10_07_10_02.sav", # 9
             "SAV_DATA/10_05_10_03.sav", # 10 - best abs RP-OG (w1 and w3) value
             "SAV_DATA/05_10_05_10.sav", # 11
             "SAV_DATA/10_00_10_00.sav", # 12
             "SAV_DATA/10_08_10_02.sav", # 13
             "SAV_DATA/10_08_10_04.sav", # 14
             "SAV_DATA/10_07_10_04.sav", # 15
             "SAV_DATA/10_06_10_04.sav", # 16 - best residual (w1 and w3) value
             "SAV_DATA/10_06_10_03.sav", # 17
             "SAV_DATA/10_05_10_04.sav", # 18
             "SAV_DATA/10_01_10_01.sav", # 19
             "SAV_DATA/10_08_10_08.sav", # 20
             "SAV_DATA/10_09_10_09.sav", # 21
             "SAV_DATA/10_09_10_07.sav"] # 22



    WL_K_CENTRE = 2795.5    # [Å] approximately
    WL_H_CENTRE = 2803      # [Å] approximately
    WL_TRIPLET = 2798       # [Å] approximately
    WL_CONTINUUM = 2797.5   # [Å] approximately


    for file in [files[16]]:
        t0 = time.time()
        print("Creating IRIS2 dictionary...")
        iris2_dict = convert_sav_to_dict(file)
        t1 = time.time()
        print("Elapsed time creating IRIS2 dictionary:  %.3f s" % (t1-t0), "\n")

        """
        for param in iris2_dict:
            print('"%s"' % param)
            #print(iris2_dict[param]["info"])
            print(iris2_dict[param]["type"], iris2_dict[param]["size"], "\n")
        """

        # For more info on output of IRIS2, go to
        # https://iris.lmsal.com/iris2/iris2_chapter02_01.html
        MODEL = iris2_dict["MODEL"]["data"]
        PHYS_VAR = iris2_dict["PHYSICAL_VARIABLE"]["data"]
        LTAU = iris2_dict["LTAU"]["data"]
        WAVELENGTH = iris2_dict["WL"]["data"]
        RP_LINES = iris2_dict["IRIS2_INV"]["data"]          # representative profiles (step, slit, wl)
        OG_LINES = iris2_dict["IRIS_CAL"]["data"]           # observations (step, slit, wl)
        LINE_UNITS = iris2_dict["UNITS_IRIS2_INV"]["data"]
        HEADER = iris2_dict["HEADER"]["data"]
        SLIT = np.linspace(0,HEADER[0].header["FOVY"],RP_LINES.shape[1])  # slit array in arcseconds
        SLIT += HEADER[0].header["YCEN"] - HEADER[0].header["FOVY"]/2
        MASK_ONDATA = (iris2_dict["MASK_ONDATA"]["data"]+1).astype("bool")
        CHI2 = (iris2_dict["CHI2"]["data"]).astype("float")
        MAP_INDEX_DB = iris2_dict["MAP_INDEX_DB"]["data"].astype("int")
        WEIGHTS_IN_WINDOWS = iris2_dict["WEIGHTS_IN_WINDOWS"]["data"]
        WEIGHT2NOISE = iris2_dict["WEIGHT2NOISE"]["data"]
        del iris2_dict      # free data from memory

        # other definitions
        WEIGHTS = np.zeros(4)
        WEIGHTS[0] = WEIGHTS_IN_WINDOWS[np.argmin(abs(WAVELENGTH-WL_K_CENTRE))]
        WEIGHTS[1] = WEIGHTS_IN_WINDOWS[np.argmin(abs(WAVELENGTH-WL_TRIPLET))]
        WEIGHTS[2] = WEIGHTS_IN_WINDOWS[np.argmin(abs(WAVELENGTH-WL_H_CENTRE))]
        WEIGHTS[3] = WEIGHTS_IN_WINDOWS[np.argmin(abs(WAVELENGTH-WL_CONTINUUM))]
        WEIGHTS = (10.1*WEIGHTS).astype(int)
        WEIGHTS_DICT_IDX = iris2_lists.weight_dict_index
        YLIM_PROFILE = [np.min(OG_LINES)/4, np.max(RP_LINES)/6]
        #remove_rows = [0,1,2,237,238,239,776,777,778,1042,1043,1044,1045]
        remove_middle_rows = [237,238,239,776,777,778]
        remove_pixels = np.where((CHI2>1000) | (CHI2<10)) # remove profiles with extreme chi^2
        remove_neg = np.where(np.any(OG_LINES<0,axis=2)) # remove negative valued profiles
        #MASK_ONDATA[:,remove_rows] = False
        MASK_ONDATA[:,remove_middle_rows] = False
        MASK_ONDATA[remove_pixels[0],remove_pixels[1]] = False
        MASK_ONDATA[remove_neg[0],remove_neg[1]] = False
        NOISE = WEIGHTS_IN_WINDOWS / WEIGHT2NOISE



        plt.plot(WAVELENGTH, np.mean(OG_LINES, axis=(0,1)), ".-")
        plt.grid()
        plt.tight_layout()
        plt.show()


        initX, initY = 100, 100       # indexes used below


        # run the different functions
        #plot_parameters()
        #chi2_analysis(plot=False)#, large_lim=[1000,10000])
        k3_idx, k2v_idx, k2r_idx = find_profile_signatures(WL_K_CENTRE, RP_LINES[initX,initY,:])
        ltau0 = tau0_from_RF(RP_line=MAP_INDEX_DB[initX,initY],
                    idx_wl=k3_idx, wl_isolate=True,
                    idx_wl_label="k3", show_RF=False)
        ltau0_idx = np.argmin(np.abs(LTAU-ltau0[0]))    # using temp ase reference
        plot_RP_line(tau_idx=ltau0_idx, initX=initX, initY=initY)
        #weight_ind = "1010"
        #train_on = "OG"         # whether to train on OG or RP profiles
        #percs = [0,100]         # lower and upper percentile
        #k_value = False#30      # number of clusters (k_value=15 looks good (albow method) with percs=[98,99])
        #kmean_model = "KMEAN_MODELS/414200_15_0_100_1010_OG"
        #if train_on=="OG":   create_kmean_from_data(OG_LINES, percs=percs, k_value=k_value)
        #elif train_on=="RP": create_kmean_from_data(RP_LINES, percs=percs, k_value=k_value)
        #X,Y = locate_large_residual_profiles(lower_perc=75,
        #                             upper_perc=99,
        #                             wl_idx=WEIGHTS_DICT_IDX[weight_ind])
        #test_model(file=kmean_model, specified_lines=[X,Y], k_nearest=False)
        #plot_spectrum(slit_i=np.random.randint(RP_LINES.shape[0]))

        #plot_observation_with_cluster_label(file=kmean_model)





        """
        print(ltau0, ltau0_idx)
        test_box1 = [300,560,30,40]  # left, bottom, width, height
        test_box2 = [0,0,MODEL.shape[0],MODEL.shape[1]]  # entire image
        coord = np.where(MODEL[:,:,ltau0_idx,0] > 0.9*np.max(MODEL[:,:,ltau0_idx,0]))  # 2D coordinate array
        RPline, OGline = average_profile(coord)
        plt.plot(WAVELENGTH, RPline, label="RP mean")
        plt.plot(WAVELENGTH, OGline, label="Observation mean", lw=0.7)
        plt.legend()
        plt.grid()
        plt.show()
        """


        del MODEL, PHYS_VAR, LTAU, WAVELENGTH, RP_LINES, OG_LINES
        del LINE_UNITS, MASK_ONDATA, CHI2, MAP_INDEX_DB, WEIGHTS_IN_WINDOWS
        del WEIGHT2NOISE, NOISE
        print("\n----------------------------------------------------\n")
