import numpy as np
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt

def squad_sample_continuum_template_main(inp_dict):
    inp_dict        = read_continuum_template_from_library(inp_dict)
    inp_dict_cont   = inp_dict["inp_dict_cont"]
    Flux_cont_arr   = inp_dict_cont["Flux_cont_arr"]
    lambda_arr_rest = inp_dict_cont["lambda_arr_rest"]

    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 <= lambda_arr_cont) & (lambda_arr_cont <= lambda_max)
    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)
    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 read_continuum_template_from_library(inp_dict):
    inp_path           = inp_dict["inp_path_aux_files"] + "/CONTINUUM_TEMPLATES/"
    flag_cont_template = inp_dict["flag_cont_template"]
    flag_cont_random   = inp_dict["flag_cont_random"]

    if flag_cont_template == "SQUAD_DR1":
        inp_path_cont = inp_path + "Template_1_SQUAD_DR1/"
        inp_file_cont = "SQUAD_DR1_Continuum_Templates.npy"

    npy_filename       = inp_path_cont + inp_file_cont
    cont_template_dict = np.load(npy_filename,allow_pickle = True,encoding = 'latin1').item()

    if flag_cont_random == "y":
        N_qso             = len(cont_template_dict)
        cont_template_num = np.random.choice(N_qso,size=1)[0] + 1 # +1 is intentional
    else:
        cont_template_num = int(inp_dict["cont_template_num"])

    inp_dict_cont = cont_template_dict[cont_template_num-1]
    inp_dict.update({"inp_dict_cont":inp_dict_cont})
    return inp_dict


