import numpy as np
import os,sys
import matplotlib.pyplot as plt

plt.style.use('classic')
py_path = os.path.dirname(os.path.realpath(__file__)) + "/mnras.py"
plt.style.use(py_path)

#-------------------------- Custom Module Starts ------------------------------

from WEAVEify_routines.step_01_weave_mock_utils.util_03_plot_spectra.save_plot_routines import \
CREATE_PDF_OBJECT,SAVE_FIGURE_IN_PDF_FILE

from WEAVEify_routines.step_02_common_routines.dictionary_functions import READ_DICTIONARY_FROM_FILE


#--------------------------- Custom Module Ends -------------------------------



def make_quick_plots_main(inp_dict):
    flag_mpi             = inp_dict["flag_mpi"]
    flag_make_quick_plot = inp_dict["flag_make_quick_plot"]

    if (flag_make_quick_plot.upper() == "Y") & (flag_mpi.upper() == "N"):
        out_path      = inp_dict["out_path"]
        #out_file      = inp_dict["out_file_flux"]
        #npy_filename  = out_path + out_file
        #out_dict      = READ_DICTIONARY_FROM_FILE(npy_filename)
        #Flux_arr      = out_dict["Flux_arr"]

        out_dict      = inp_dict["out_dict_final"]
        wavelength_arr = out_dict["wavelength_arr"]
        noise_arr      = out_dict["snr_per_pixel_arr"]
        flux_arr       = out_dict["flux_arr"]
        cont_arr       = out_dict["cont_arr"]
        Flux_arr       = np.array([wavelength_arr,flux_arr,noise_arr,cont_arr]).T

        plt_dict      = generate_plt_dict()
        out_file_plot = inp_dict["out_file_plot"]
        pdf_filename  = out_path + out_file_plot
        pdf_obj       = CREATE_PDF_OBJECT(pdf_filename)


        fig,ax = PLOT_TOTAL_SPECTRUM(inp_dict,Flux_arr,plt_dict,flag_case="flux_final")
        SAVE_FIGURE_IN_PDF_FILE(pdf_obj,fig,"n")
        pdf_obj.close()
    return inp_dict



def generate_plt_dict():
    plt_dict = {"Flux":     {"color":"b","ls":"-","lw":0.5,"alpha":1.0},\
                "Continuum":{"color":"r","ls":"-","lw":4.5,"alpha":1.0},\
                 }
    return plt_dict

def PLOT_TOTAL_SPECTRUM(inp_dict,Flux_arr,plt_dict,flag_case="flux_final"):

    def plot_data(ax,xdata,ydata,plt_sub_dict,label=None):
        ls    = plt_sub_dict["ls"]
        lw    = plt_sub_dict["lw"]
        color = plt_sub_dict["color"]
        alpha = plt_sub_dict["alpha"]
        ax.plot(xdata,ydata,ls=ls,lw=lw,color=color,label=label)

    fig,ax        = GENERATE_PLOT_AXIS()
    plt_dict_flux = plt_dict["Flux"]
    plt_dict_cont = plt_dict["Continuum"]

    if flag_case == "flux_final":
        ydata_flux = Flux_arr[:,1]
        ydata_cont = Flux_arr[:,3]
    elif flag_case == "flux_without_noise":
        ydata_flux = Flux_arr[:,2]
        ydata_cont = Flux_arr[:,3]

    xdata = Flux_arr[:,0]
    plot_data(ax,xdata,ydata_flux,plt_dict["Flux"],label="Absorbed Flux (Mock)")
    plot_data(ax,xdata,ydata_cont,plt_dict["Continuum"],label="Continuum")

    SET_AX_PROPERTY_TOTAL_SPECTRUM(ax,inp_dict,Flux_arr)
    ax.legend(loc="upper right",prop={"size":28})

    return fig,ax

def SET_AX_PROPERTY_TOTAL_SPECTRUM(ax,inp_dict,Flux_arr):
    #x_lb = np.amin(Flux_arr[:,0])
    #x_ub = np.amax(Flux_arr[:,0])
    z_em  = inp_dict["z_em"]
    x_lb  = (1.0 + z_em) * 912.0

    #x_ub  = (1.0 + z_em) * 1625.0
    if inp_dict["resolution_mode"] == "LR":
        x_ub  = (1.0 + z_em) * 1600.0
    elif inp_dict["resolution_mode"] == "HR":
        #x_ub  = (1.0 + z_em) * 1530.0
        x_ub  = (1.0 + z_em) * 1600.0

    y_max = np.nanmax(Flux_arr[:,1])
    y_ub  = y_max * 1.15
    y_lb  = y_max * -0.08
    ax.set_xlim(x_lb,x_ub)
    ax.set_ylim(y_lb,y_ub)
    ax.axhline(y=0.0,ls="--",color="k",lw=4.0)
    ax.set_xlabel(r"$\lambda (\rm {\AA})$",size=30)
    ax.set_ylabel(r"$F$",size=30)

    xticks_arr_all = np.arange(3000.0,10001.0,500.0)
    bool_arr       = (x_lb <= xticks_arr_all) & (xticks_arr_all <= x_ub)
    xticks_arr     = xticks_arr_all[bool_arr]
    ax.set_xticks(xticks_arr.tolist())


def GENERATE_PLOT_AXIS():
    fig,ax       = plt.subplots(1)
    fig.set_figwidth(24.0)
    fig.set_figheight(8.5)
    return fig,ax

def step_0_get_file_dirname():
    out_path            = "./"
    out_file            = "WEAVE_Mock_Steps.pdf"
    inp_path            = "/data/emergence11/pgaikwad/Projects/data_files/Project-016_WEAVE_QSO_Analysis/WEAVE_MOCKS/"
    inp_file_dtv        = "out_dict_field.npy"
    inp_file_tau        = "out_dict_tau_total.npy"
    inp_file_flux_samp  = "out_dict_flux_resamp.npy"
    inp_file_flux_cont  = "out_dict_flux_cont.npy"
    inp_file_flux_final = "out_dict_flux_final.npy"
    inp_dict            = vars()
    return inp_dict

def step_1_plot_DLA_cases(pdf_obj,inp_dict):
    fig,ax                 = generate_three_subplot_axis()
    plot_single_model(ax[0],inp_dict,"Analysis-04_With_low_NHI_DLA_No_Metal",flag_cont="y")
    plot_single_model(ax[1],inp_dict,"Analysis-08_With_DLA_No_Metal",flag_cont="y")
    plot_single_model(ax[2],inp_dict,"Analysis-05_With_high_NHI_DLA_No_Metal",flag_cont="y")

    text_list     = [r"Low ${\rm N_{HI}}$ DLA",r"Intermediate ${\rm N_{HI}}$ DLA",r"High ${\rm N_{HI}}$ DLA"]
    text_dict_def = {"xy_loc":(0.8,0.7),"color":"m","fontsize":28,"weight":"normal"}

    for indx in range(len(ax)):
        text_dict = text_dict_def.copy()
        text_dict.update({"text":text_list[indx]})
        AX_ANNOTATE(ax[indx],text_dict)

    HIDE_AX_PROPERTY(ax[0],'x')
    HIDE_AX_PROPERTY(ax[1],'x')
    SAVE_FIGURE_IN_PDF_FILE(pdf_obj,fig,"n")

def step_2_plot_metal_cases(pdf_obj,inp_dict):
    fig,ax                 = generate_three_subplot_axis()
    plot_single_model(ax[0],inp_dict,"Analysis-06_No_DLA_With_Low_Metal",flag_cont="y")
    plot_single_model(ax[1],inp_dict,"Analysis-03_No_DLA_With_Metal",flag_cont="y")
    plot_single_model(ax[2],inp_dict,"Analysis-07_No_DLA_With_High_Metal",flag_cont="y")

    text_list     = [\
                     r"Low : $Z = 0.01 \: Z_{\odot}$",\
                     r"Default : $Z = 0.1 \: Z_{\odot}$",\
                     r"High : $Z = 1.0 \: Z_{\odot}$",\
                    ]
    text_dict_def = {"xy_loc":(0.8,0.7),"color":"m","fontsize":28,"weight":"normal"}

    for indx in range(len(ax)):
        text_dict = text_dict_def.copy()
        text_dict.update({"text":text_list[indx]})
        AX_ANNOTATE(ax[indx],text_dict)

    HIDE_AX_PROPERTY(ax[0],'x')
    HIDE_AX_PROPERTY(ax[1],'x')
    SAVE_FIGURE_IN_PDF_FILE(pdf_obj,fig,"n")

def step_3_plot_continuum_cases(pdf_obj,inp_dict):
    fig,ax                 = generate_three_subplot_axis()
    plot_single_model(ax[0],inp_dict,"Analysis-09_Continuum-1",flag_cont="y")
    plot_single_model(ax[1],inp_dict,"Analysis-10_Continuum-2",flag_cont="y")
    plot_single_model(ax[2],inp_dict,"Analysis-11_Continuum-3",flag_cont="y")
    ax[1].set_ylim(-80.0,3020.0)
    ax[2].set_ylim(-80.0,820.0)

    text_list     = [\
                     r"Continuum Template : 1",\
                     r"Continuum Template : 2",\
                     r"Continuum Template : 3",\
                    ]
    text_dict_def = {"xy_loc":(0.8,0.7),"color":"m","fontsize":28,"weight":"normal"}

    for indx in range(len(ax)):
        text_dict = text_dict_def.copy()
        text_dict.update({"text":text_list[indx]})
        AX_ANNOTATE(ax[indx],text_dict)

    HIDE_AX_PROPERTY(ax[0],'x')
    HIDE_AX_PROPERTY(ax[1],'x')
    SAVE_FIGURE_IN_PDF_FILE(pdf_obj,fig,"n")

def step_4_plot_introduction_figure_no_metal_no_dla(pdf_obj,inp_dict):
    fig,ax                 = generate_two_subplot_axis()
    plot_single_model(ax[0],inp_dict,"Analysis-02_No_DLA_No_Metal",flag_cont="n")
    plot_single_model(ax[1],inp_dict,"Analysis-02_No_DLA_No_Metal",flag_cont="y")

    text_list     = [\
                     r"Simulated spectrum ($z_{\rm em} \sim 3$)",\
                     r"WEAVEfied Spectrum ($z_{\rm em} \sim 3$)",\
                    ]
    text_dict_def = {"xy_loc":(0.8,0.7),"color":"m","fontsize":24,"weight":"normal"}

    for indx in range(len(ax)):
        text_dict = text_dict_def.copy()
        text_dict.update({"text":text_list[indx]})
        AX_ANNOTATE(ax[indx],text_dict)

    text_dict = {"text":r"Ly$(\alpha + \beta + \gamma)$ forest","xy_loc":(0.2,0.8),"color":"m","fontsize":24,"weight":"normal"}
    AX_ANNOTATE(ax[1],text_dict)

    HIDE_AX_PROPERTY(ax[0],'x')
    HIDE_AX_PROPERTY(ax[1],'x')
    SAVE_FIGURE_IN_PDF_FILE(pdf_obj,fig,"n")

def step_5_plot_introduction_figure_with_metal_no_dla(pdf_obj,inp_dict):
    fig,ax                 = generate_two_subplot_axis()
    plot_single_model(ax[0],inp_dict,"Analysis-03_No_DLA_With_Metal",flag_cont="n")
    plot_single_model(ax[1],inp_dict,"Analysis-03_No_DLA_With_Metal",flag_cont="y")

    text_list     = [\
                     r"Simulated spectrum ($z_{\rm em} \sim 3$)",\
                     r"WEAVEfied Spectrum ($z_{\rm em} \sim 3$)",\
                    ]
    text_dict_def = {"xy_loc":(0.8,0.7),"color":"m","fontsize":24,"weight":"normal"}

    for indx in range(len(ax)):
        text_dict = text_dict_def.copy()
        text_dict.update({"text":text_list[indx]})
        AX_ANNOTATE(ax[indx],text_dict)

    text_dict = {"text":r"Ly$(\alpha + \beta + \gamma)$ forest + Metals","xy_loc":(0.2,0.8),"color":"m","fontsize":24,"weight":"normal"}
    AX_ANNOTATE(ax[1],text_dict)

    HIDE_AX_PROPERTY(ax[0],'x')
    HIDE_AX_PROPERTY(ax[1],'x')
    SAVE_FIGURE_IN_PDF_FILE(pdf_obj,fig,"n")

def step_6_plot_introduction_figure_with_metal_with_dla(pdf_obj,inp_dict):
    fig,ax                 = generate_two_subplot_axis()
    plot_single_model(ax[0],inp_dict,"Analysis-01_With_DLA_With_Metal",flag_cont="n")
    plot_single_model(ax[1],inp_dict,"Analysis-01_With_DLA_With_Metal",flag_cont="y")

    text_list     = [\
                     r"Simulated spectrum ($z_{\rm em} \sim 3$)",\
                     r"WEAVEfied Spectrum ($z_{\rm em} \sim 3$)",\
                    ]
    text_dict_def = {"xy_loc":(0.8,0.7),"color":"m","fontsize":24,"weight":"normal"}

    for indx in range(len(ax)):
        text_dict = text_dict_def.copy()
        text_dict.update({"text":text_list[indx]})
        AX_ANNOTATE(ax[indx],text_dict)

    text_dict = {"text":r"Ly$(\alpha + \beta + \gamma)$ forest + Metals + DLA","xy_loc":(0.2,0.8),"color":"m","fontsize":24,"weight":"normal"}
    AX_ANNOTATE(ax[1],text_dict)

    HIDE_AX_PROPERTY(ax[0],'x')
    HIDE_AX_PROPERTY(ax[1],'x')
    SAVE_FIGURE_IN_PDF_FILE(pdf_obj,fig,"n")





def step_7_plot_step_by_step_illustration(pdf_obj,inp_dict):
    fig,ax                 = generate_single_subplot_axis()
    plot_single_model(ax,inp_dict,"Analysis-02_No_DLA_No_Metal",flag_cont="n")
    text      = r"Ly$(\alpha + \beta + \gamma)$ forest $(z<3)$"
    text_dict = {"xy_loc":(0.8,0.25),"color":"m","fontsize":24,"weight":"normal"}
    text_dict.update({"text":text})
    AX_ANNOTATE(ax,text_dict)
    SAVE_FIGURE_IN_PDF_FILE(pdf_obj,fig,"n")

    fig,ax                 = generate_single_subplot_axis()
    plot_single_model(ax,inp_dict,"Analysis-08_With_DLA_No_Metal",flag_cont="n")
    text      = r"Ly$(\alpha + \beta + \gamma)$ forest + DLA"
    text_dict = {"xy_loc":(0.8,0.25),"color":"m","fontsize":24,"weight":"normal"}
    text_dict.update({"text":text})
    AX_ANNOTATE(ax,text_dict)
    SAVE_FIGURE_IN_PDF_FILE(pdf_obj,fig,"n")

    fig,ax                 = generate_single_subplot_axis()
    plot_single_model(ax,inp_dict,"Analysis-01_With_DLA_With_Metal",flag_cont="n")
    text      = r"Ly$(\alpha + \beta + \gamma)$ forest + DLA + Metal"
    text_dict = {"xy_loc":(0.8,0.25),"color":"m","fontsize":24,"weight":"normal"}
    text_dict.update({"text":text})
    AX_ANNOTATE(ax,text_dict)
    SAVE_FIGURE_IN_PDF_FILE(pdf_obj,fig,"n")

    fig,ax                 = generate_single_subplot_axis()
    plot_single_model(ax,inp_dict,"Analysis-01_With_DLA_With_Metal",flag_cont="y")
    text      = r"Ly$(\alpha + \beta + \gamma)$ forest + DLA + Metal"
    text_dict = {"xy_loc":(0.8,0.7),"color":"m","fontsize":24,"weight":"normal"}
    text_dict.update({"text":text})
    AX_ANNOTATE(ax,text_dict)
    SAVE_FIGURE_IN_PDF_FILE(pdf_obj,fig,"n")























def plot_single_model(ax,inp_dict,dir_name,flag_cont="y"):

    inp_path_dir = inp_dict["inp_path"] + "/" + dir_name + "/"

    if flag_cont == "n":
        inp_file_flux_samp     = inp_dict["inp_file_flux_samp"]
        npy_filename_flux_samp = inp_path_dir + inp_file_flux_samp
        data_dict_flux_samp    = READ_DICTIONARY_FROM_FILE(npy_filename_flux_samp)
        Flux_arr_samp          = data_dict_flux_samp["Flux_resamp_arr"]
        lambda_arr_obs         = data_dict_flux_samp["lambda_arr_obs"]
        ax.plot(lambda_arr_obs,Flux_arr_samp,ls='-',lw=0.2)

    elif flag_cont == "y":
        inp_file_flux_final     = inp_dict["inp_file_flux_final"]
        npy_filename_flux_final = inp_path_dir + inp_file_flux_final
        data_dict_flux_final    = READ_DICTIONARY_FROM_FILE(npy_filename_flux_final)
        Flux_arr_final          = data_dict_flux_final["Flux_arr_final"]
        ax.plot(Flux_arr_final[:,0],Flux_arr_final[:,1],ls='-',lw=0.2)
        ax.plot(Flux_arr_final[:,0],Flux_arr_final[:,2],ls='-',color="r",lw=4.0)

    set_ax_property(ax,flag_cont)





def generate_single_subplot_axis():
    fig,ax = plt.subplots(1)
    fig.set_figwidth(20.0)
    fig.set_figheight(6.0)
    return fig,ax

def generate_two_subplot_axis():
    fig,ax = plt.subplots(nrows=2)
    fig.set_figwidth(20.0)
    fig.set_figheight(12.0)
    fig.subplots_adjust(hspace=0.05)
    return fig,ax

def generate_three_subplot_axis():
    fig,ax = plt.subplots(nrows=3)
    fig.set_figwidth(20.0)
    fig.set_figheight(16.0)
    fig.subplots_adjust(hspace=0.05)
    return fig,ax

def set_ax_property(ax,flag_cont="y",flag_noise="n"):
    ax.set_xlim(3200.0,6200.0)

    if flag_cont == "y":
        if flag_noise == "y":
            ax.set_ylim(-1.0,12.0)
        else:
            ax.set_ylim(-80.0,1620.0)

    elif flag_cont == "n":
        ax.set_ylim(-0.15,1.15)
        ax.axhline(y=0.0,ls="--",color="k",lw=3.0)
        ax.axhline(y=1.0,ls="--",color="k",lw=3.0)

    ax.set_xlabel(r"$\lambda \: (\rm \AA)$")
    ax.set_ylabel(r"$F$")

def step_8_plot_introduction_figure_metal_dla_combination(pdf_obj,inp_dict):
    fig,ax                 = generate_two_subplot_axis()
    plot_single_model(ax[0],inp_dict,"Analysis-02_No_DLA_No_Metal",flag_cont="n")
    plot_single_model(ax[1],inp_dict,"Analysis-01_With_DLA_With_Metal",flag_cont="y")

    text_list     = [\
                     r"Simulated spectrum ($z_{\rm em} \sim 3$)",\
                     r"WEAVEfied Spectrum ($z_{\rm em} \sim 3$)",\
                    ]
    text_dict_def = {"xy_loc":(0.8,0.7),"color":"m","fontsize":24,"weight":"normal"}

    for indx in range(len(ax)):
        text_dict = text_dict_def.copy()
        text_dict.update({"text":text_list[indx]})
        AX_ANNOTATE(ax[indx],text_dict)

    text_dict = {"text":r"Ly$(\alpha + \beta + \gamma)$ forest + Metals + DLA","xy_loc":(0.2,0.8),"color":"m","fontsize":24,"weight":"normal"}
    AX_ANNOTATE(ax[1],text_dict)

    HIDE_AX_PROPERTY(ax[0],'x')
    HIDE_AX_PROPERTY(ax[1],'x')
    SAVE_FIGURE_IN_PDF_FILE(pdf_obj,fig,"n")

def step_9_plot_noise_properties(pdf_obj,inp_dict):
    noise_type_list = [\
                        "Source_Shot_Noise_only",\
                        "Source_Shot_Noise_only",\
                        "Sky_Shot_Noise_only",\
                        "Source_And_Sky_Shot_Noise_only",\
                        "All_Noise_Contribution",\
                      ]
    label_list     = ["Mock Spectrum : No Noise", "Source Shot noise Only","Sky Shot Noise Only",\
                      "Source + Sky Shot Noise","Source + Sky\nDark + Read-out Noise"]

    inp_base_path  = "/data/emergence11/pgaikwad/Projects/data_files/Project-016_WEAVE_QSO_Analysis/WEAVE_MOCKS/Analysis-12_Noise_Analaysis-1/"
    inp_file       = "out_dict_flux_final.npy"

    indx           = 0
    for indx in range(len(noise_type_list)):
        fig,ax         = generate_single_subplot_axis()
        noise_type     = noise_type_list[indx]
        label          = label_list[indx]
        if indx == 0:
            inp_sub_dir    = "/Case-" + str("%02i"%(indx+1)) + "_" + noise_type + "/"
        else:
            inp_sub_dir    = "/Case-" + str("%02i"%indx) + "_" + noise_type + "/"
        inp_path       = inp_base_path + inp_sub_dir
        inp_dict       = np.load(inp_path + inp_file,allow_pickle=True,encoding='latin1').item()
        Flux_arr_final = inp_dict['Flux_arr_final']

        if indx == 0:
            ax.plot(Flux_arr_final[:,0],Flux_arr_final[:,2],ls='-',lw=0.2)
        else:
            ax.plot(Flux_arr_final[:,0],Flux_arr_final[:,1],ls='-',lw=0.2)

        #if indx == 0:
        #    ax.plot(Flux_arr_final[:,0],Flux_arr_final[:,3],ls='-',color="r",lw=4.0)
        ax.plot(Flux_arr_final[:,0],Flux_arr_final[:,3],ls='-',color="r",lw=4.0)


        set_ax_property(ax,flag_cont="y",flag_noise="y")

        text_dict = {"text":label,"xy_loc":(0.8,0.7),"color":"m","fontsize":24,"weight":"normal"}
        ax.axhline(y=0.0,ls="--",color="k",lw=4.0)
        AX_ANNOTATE(ax,text_dict)

        SAVE_FIGURE_IN_PDF_FILE(pdf_obj,fig,"n")

    #for indx in range(len(noise_type_list)):

