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

from scipy.integrate import trapz

from scipy.interpolate import interp1d

from WEAVEify_routines.step_02_common_routines.dictionary_functions import \
READ_DICTIONARY_FROM_FILE

from WEAVEify_routines.step_02_common_routines.python_2_3_compatibility_routines \
import genfromtxt_compatible

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

def continuum_placement_using_pca_main(inp_dict):
    inp_dict_cont = step_1_define_parameters(inp_dict)
    inp_dict_cont = step_2_read_inputs(inp_dict_cont)
    inp_dict_cont = step_3_generate_composite_spectrum_realization(inp_dict_cont)
    inp_dict_cont = step_4_PCA_reconstruction_1020_1216(inp_dict_cont)
    inp_dict_cont = step_5_create_final_artificial_continuum(inp_dict_cont)
    inp_dict      = step_6_add_continuum_to_the_spectra(inp_dict,inp_dict_cont)
    #inp_dict_cont = step_7_optional_make_analysis_plot(inp_dict_cont)
    return inp_dict

def step_6_add_continuum_to_the_spectra(inp_dict,inp_dict_cont):
    flux_arr_artificial = inp_dict_cont["flux_arr_artificial"]
    lambda_arr_rest     = flux_arr_artificial[:,0]
    Flux_cont_arr       = flux_arr_artificial[:,1]

    inp_dict_Flux   = inp_dict["inp_dict_Flux"]
    z_em            = inp_dict["z_em"]
    Flux_resamp_arr = inp_dict_Flux["Flux_resamp_arr"]
    lambda_arr_obs  = inp_dict_Flux["lambda_arr_obs"]
    lambda_arr_cont = lambda_arr_rest * (1.0 + z_em)
    lambda_min      = max(lambda_arr_cont[0],lambda_arr_obs[0])
    lambda_max      = min(lambda_arr_cont[-1],lambda_arr_obs[-1])

    bool_arr        = ((lambda_min-2.0) <= lambda_arr_cont) & (lambda_arr_cont <= (lambda_max + 2.0))
    intrp_funct     = interp1d(lambda_arr_cont[bool_arr],Flux_cont_arr[bool_arr],fill_value=0.0)
    bool_arr_obs    = (lambda_min <= lambda_arr_obs) & (lambda_arr_obs <= lambda_max)

    lambda_print_1    = lambda_arr_cont[bool_arr]
    lambda_print_2    = lambda_arr_obs[bool_arr_obs]

    Flux_cont_arr_obs = intrp_funct(lambda_arr_obs[bool_arr_obs][1:-1])
    lambda_arr_save    = lambda_arr_obs[bool_arr_obs][1:-1]
    Flux_arr_save      = Flux_cont_arr_obs*Flux_resamp_arr[bool_arr_obs][1:-1]
    Flux_arr_cont_save = Flux_cont_arr_obs
    Flux_arr_cont_added  = np.array([lambda_arr_save,Flux_arr_save,Flux_arr_cont_save]).T
    Flux_arr_cont_added  = continuum_normalization_factor(Flux_arr_cont_added,z_em)

    out_dict_flux_cont = {"Flux_arr_cont_added":Flux_arr_cont_added}
    inp_dict_Flux.update(out_dict_flux_cont)
    inp_dict.update({"inp_dict_Flux":inp_dict_Flux})
    return inp_dict

def continuum_normalization_factor(Flux_arr_cont_added,z_em):
    lambda_norm_rest_val      = 1500.0
    lambda_norm_val           = lambda_norm_rest_val * (1.0 + z_em)
    idx_min                   = np.argmin(np.abs(Flux_arr_cont_added[:,0]-lambda_norm_val))
    flux_cont_norm_val        = Flux_arr_cont_added[idx_min,2]
    Flux_arr_cont_added[:,1:] = Flux_arr_cont_added[:,1:] / flux_cont_norm_val
    return Flux_arr_cont_added

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

def step_1_define_parameters(inp_dict):
    inp_path_aux_files    = inp_dict["inp_path_aux_files"]
    random_seed_cont      = inp_dict["random_seed_cont"]
    alpha_cont            = inp_dict["alpha_cont"]
    
    inp_path              = inp_path_aux_files +"CONTINUUM_TEMPLATES/METHOD_2_PCA/"
    inp_file_composite    = "Composite_Spectrum_Fitting_Parameters_1140_3000.txt"
    inp_file_eBOSS        = "eBOSS_QSO_Composite_Spectra_Harris_2015.npy"
    inp_file_corr_mat     = "QSO_Emission_Line_Correlation_Matrix_Greig_2017.txt"
    inp_file_PCA_red      = "PCA_1216_2000_Vizier_Paris_2011.txt"
    inp_file_PCA_red_blue = "PCA_1020_2000_Vizier_Paris_2011.txt"
    inp_file_proj_matrix  = "projection_matrix_Paris_2011.txt"
    delta_lambda          = 0.5
    lambda_min            = 650.0
    lambda_max            = 4500.0
    N_pixel               = int((lambda_max - lambda_min) / delta_lambda) + 1
    lambda_arr_cont       = np.linspace(lambda_min,lambda_max,N_pixel)
    lambda_min_red        = 1216.0
    lambda_max_red        = 2000.0
    lambda_min_blue       = 1020.0
    bool_arr_min_red      = (lambda_min_red - delta_lambda*0.25) <= lambda_arr_cont
    bool_arr_max_red      = lambda_arr_cont <= (lambda_max_red + delta_lambda*0.25)
    bool_arr_min_blue     = (lambda_min_blue - delta_lambda*0.25) <= lambda_arr_cont
    bool_arr_1216_2000    = bool_arr_min_red  & bool_arr_max_red
    bool_arr_1020_2000    = bool_arr_min_blue & bool_arr_max_red
    bool_arr_1020_1216    = bool_arr_min_blue & (lambda_arr_cont <= (lambda_min_red + delta_lambda*0.25))
    idx_min_1020          = np.argmin(np.abs(lambda_min_blue - lambda_arr_cont))
    idx_min_2000          = np.argmin(np.abs(lambda_max_red - lambda_arr_cont))
    inp_dict_cont         = vars()
    inp_dict_cont.pop("inp_dict")
    return inp_dict_cont

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

def step_2_read_inputs(inp_dict_cont):
    lambda_arr_cont       = inp_dict_cont["lambda_arr_cont"]
    inp_path              = inp_dict_cont["inp_path"]
    inp_file_eBOSS        = inp_dict_cont["inp_file_eBOSS"]
    inp_file_PCA_red      = inp_dict_cont["inp_file_PCA_red"]
    inp_file_PCA_red_blue = inp_dict_cont["inp_file_PCA_red_blue"]
    inp_file_proj_matrix  = inp_dict_cont["inp_file_proj_matrix"]

    PCA_arr_red           = np.loadtxt(inp_path + inp_file_PCA_red)
    PCA_arr_blue_red      = np.loadtxt(inp_path + inp_file_PCA_red_blue)[:-1,:]
    PCA_proj_matrix       = np.loadtxt(inp_path + inp_file_proj_matrix)


    npy_filename       = inp_path + inp_file_eBOSS
    inp_dict_cont_composite = READ_DICTIONARY_FROM_FILE(npy_filename)
    flux_arr_comp_org  = inp_dict_cont_composite["flux_arr_composite"].astype(np.float64)

    flux_arr_composite = interpolate_continuum(flux_arr_comp_org[:,0],\
                         flux_arr_comp_org[:,1],lambda_arr_cont,flag_extrapolate="y")

    inp_path           = inp_dict_cont["inp_path"]
    inp_file_composite = inp_dict_cont["inp_file_composite"]
    txt_filename       = inp_path + inp_file_composite
    param_data_arr     = genfromtxt_compatible(inp_path,inp_file_composite,object)
    param_arr_best     = param_data_arr[:,2].astype(np.float32)
    param_arr_sigma    = param_data_arr[:,3].astype(np.float32)

    inp_dict_cont.update({\
                     "PCA_arr_red":PCA_arr_red,\
                     "PCA_arr_blue_red":PCA_arr_blue_red,\
                     "PCA_proj_matrix":PCA_proj_matrix,\
                     "param_arr_best":param_arr_best,\
                     "param_arr_sigma":param_arr_sigma,\
                     "flux_arr_composite":flux_arr_composite,\
                     "flux_arr_composite_org":flux_arr_comp_org,\
                    })

    return inp_dict_cont

def interpolate_continuum(xdata,ydata,lambda_arr_cont,flag_extrapolate="n",flag_out="2d"):

    if flag_extrapolate == "y":
        intrp_funct  = interp1d(xdata,ydata,bounds_error=False,fill_value=0.0)
    elif flag_extrapolate == "y":
        intrp_funct  = interp1d(xdata,ydata)

    ydata_intrp_1d = intrp_funct(lambda_arr_cont)

    if flag_out == "2d":
        flux_intrp_arr = np.array([lambda_arr_cont,ydata_intrp_1d]).T

    elif flag_out == "1d":
        flux_intrp_arr = ydata_intrp_1d.copy()

    return flux_intrp_arr

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

def step_3_generate_composite_spectrum_realization(inp_dict_cont):

    lambda_arr_cont      = inp_dict_cont["lambda_arr_cont"]
    param_arr_best       = inp_dict_cont["param_arr_best"]
    param_arr_sigma      = inp_dict_cont["param_arr_sigma"]

    random_seed_cont    = inp_dict_cont["random_seed_cont"]
    alpha_cont          = inp_dict_cont["alpha_cont"]

    if random_seed_cont is not None: np.random.seed(random_seed_cont)
    param_arr            = np.random.normal(param_arr_best,param_arr_sigma*5.0)

    if random_seed_cont is not None: np.random.seed(random_seed_cont)

    if alpha_cont is None:
        param_arr[1]     = np.random.uniform(-0.0,-3.0,1)
    else:
        param_arr[1]     = alpha_cont
    #print(param_arr[1])

    flux_arr_realization = gaussian_fitting_function(param_arr,lambda_arr_cont)

    inp_dict_cont.update({"flux_arr_realization":flux_arr_realization})

    return inp_dict_cont



def gaussian_fitting_function(param_flattened,xdata_arr):
    F_1450          = param_flattened[0]
    alpha           = param_flattened[1]
    N_component     = int((param_flattened.shape[0] - 2.0) / 3.0)
    param_indx      = 2
    ydata_fit_arr   = np.zeros(xdata_arr.shape[0])

    for line_indx in range(N_component):
        xc_val          = param_flattened[param_indx + 0]
        A_val           = param_flattened[param_indx + 1]
        sigma_gauss_val = param_flattened[param_indx + 2]
        ydata_indx      = A_val * np.exp(-1.0 * (xdata_arr-xc_val)**2.0 / (2.0*sigma_gauss_val**2.0))
        ydata_fit_arr   = ydata_fit_arr + ydata_indx
        param_indx      = param_indx + 3

    ydata_fit_arr        = (ydata_fit_arr + 1.0) * F_1450 * (xdata_arr / 1450.0)**alpha
    flux_arr_realization = np.array([xdata_arr,ydata_fit_arr]).T

    return flux_arr_realization


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

def step_4_PCA_reconstruction_1020_1216(inp_dict_cont):
    lambda_arr_cont    = inp_dict_cont["lambda_arr_cont"]
    bool_arr_1020_2000 = inp_dict_cont["bool_arr_1020_2000"]
    PCA_proj_matrix    = inp_dict_cont["PCA_proj_matrix"]
    b_arr              = step_4a_PCA_find_array_b_1216_2000(inp_dict_cont)
    a_arr              = np.matmul(b_arr,PCA_proj_matrix)
    flux_arr_1d        = step_4b_PCA_extrapolate_continuum_1020_2000(inp_dict_cont,a_arr)
    flux_arr_1020_2000 = np.array([lambda_arr_cont[bool_arr_1020_2000],flux_arr_1d]).T
    inp_dict_cont.update({"flux_arr_1020_2000":flux_arr_1020_2000})
    return inp_dict_cont

#------------------------------------------------

def step_4a_PCA_find_array_b_1216_2000(inp_dict_cont):
    lambda_arr_cont        = inp_dict_cont["lambda_arr_cont"]
    PCA_arr_red            = inp_dict_cont["PCA_arr_red"]
    flux_arr_composite     = inp_dict_cont["flux_arr_composite"]
    flux_arr_realization   = inp_dict_cont["flux_arr_realization"]
    bool_arr_1216_2000     = inp_dict_cont["bool_arr_1216_2000"]
    flux_arr_realization_red = flux_arr_realization[bool_arr_1216_2000,1]
    flux_arr_composite_red   = flux_arr_composite[bool_arr_1216_2000,1]

    N_PCA         = PCA_arr_red.shape[1] - 2
    N_pixel       = flux_arr_composite_red.shape[0]
    b_arr         = np.zeros(N_PCA)
    cont_diff_arr = (flux_arr_realization_red - flux_arr_composite_red)

    for pca_indx in range(N_PCA):
        b_matrix        = cont_diff_arr * PCA_arr_red[:,pca_indx+2]
        b_arr[pca_indx] = trapz(b_matrix,lambda_arr_cont[bool_arr_1216_2000])

    return b_arr

#------------------------------------------------

def step_4b_PCA_extrapolate_continuum_1020_2000(inp_dict_cont,a_arr):
    flux_arr_composite = inp_dict_cont["flux_arr_composite"]
    PCA_arr_blue_red   = inp_dict_cont["PCA_arr_blue_red"]
    bool_arr_1020_2000 = inp_dict_cont["bool_arr_1020_2000"]
    N_blue_red         = np.sum(bool_arr_1020_2000)
    cont_scale_arr     = np.zeros(N_blue_red)
    N_PCA              = PCA_arr_blue_red.shape[1] - 2

    for pca_indx in range(N_PCA):
        cont_scale_arr  = cont_scale_arr +  a_arr[pca_indx] * PCA_arr_blue_red[:,pca_indx+2]

    flux_arr_1020_1216 = flux_arr_composite[bool_arr_1020_2000,1] + cont_scale_arr

    return flux_arr_1020_1216

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

def step_5_create_final_artificial_continuum(inp_dict_cont):
    lambda_arr_cont      = inp_dict_cont["lambda_arr_cont"]
    flux_arr_realization = inp_dict_cont["flux_arr_realization"]
    flux_arr_composite   = inp_dict_cont["flux_arr_composite"]
    bool_arr_1020_2000   = inp_dict_cont["bool_arr_1020_2000"]
    bool_arr_1216_2000   = inp_dict_cont["bool_arr_1216_2000"]
    bool_arr_1020_1216   = inp_dict_cont["bool_arr_1020_1216"]
    flux_arr_1020_2000   = inp_dict_cont["flux_arr_1020_2000"]


    flux_arr_artificial       = np.zeros((lambda_arr_cont.shape[0],2))
    flux_arr_artificial[:,0]  = lambda_arr_cont

    # All regions are  initilize to realization
    flux_arr_artificial[:,1]  = flux_arr_realization[:,1].copy()

    # Region -->> 1020 <= lambda <= 2000
    flux_arr_artificial[bool_arr_1020_2000,1] = flux_arr_1020_2000[:,1]

    # Region -->> lambda <= 1020
    idx_min_1020      = inp_dict_cont["idx_min_1020"]
    flux_scale_factor = flux_arr_artificial[idx_min_1020,1] / flux_arr_composite[idx_min_1020,1]
    flux_arr_artificial[:idx_min_1020,1] = flux_arr_composite[:idx_min_1020,1] * flux_scale_factor

    # Region -->> 2000 <= lambda
    idx_min_2000      = inp_dict_cont["idx_min_2000"]
    flux_scale_factor = flux_arr_artificial[idx_min_2000,1] / flux_arr_realization[idx_min_2000,1]
    flux_arr_artificial[idx_min_2000+1:,1] = flux_arr_artificial[idx_min_2000+1:,1] * flux_scale_factor

    inp_dict_cont.update({"flux_arr_artificial":flux_arr_artificial})
    return inp_dict_cont


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

def step_7_optional_make_analysis_plot(inp_dict_cont):
    flux_arr_artificial = inp_dict_cont["flux_arr_artificial"]
    flux_arr_composite  = inp_dict_cont["flux_arr_composite"]

    fig,ax = plt.subplots(1)
    fig.set_figwidth(18.0)
    fig.set_figheight(6.0)

    ax.plot(flux_arr_composite[:,0],flux_arr_composite[:,1],lw=2.0,color="b")
    ax.plot(flux_arr_artificial[:,0],flux_arr_artificial[:,1],lw=2.0,color="r")

    ax.axhline(y=0.0,ls="--",color="k",lw=2.0)

    SET_AX_PROPERTY(ax)
    #plt.savefig("Continuum_Analysis.pdf",spi=200)
    plt.show()
    return inp_dict_cont


def SET_AX_PROPERTY(ax):
    ax.set_ylim(-1.0,8.0)
    ax.set_xlim(400.0,4400.0)
    ax.set_xlabel(r"$\lambda_{\rm rest} \;\; {\rm \AA}$")
    ax.set_ylabel(r"$F$")


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








