from __future__ import print_function
import numpy as np
from math import log
import matplotlib.pyplot as plt
from scipy.ndimage import convolve


def instrumental_broadening_routines_main(inp_dict):


    lsf_tup              = generate_lsf_input_path(inp_dict)
    inp_path_lsf         = lsf_tup[0]
    inp_file_lsf         = lsf_tup[1]

    line_dict            = inp_dict["line_dict"]
    lambda_rest_HI       = line_dict["HI-1215"]["lambda_rest"]
    inp_path_aux_files   = inp_dict["inp_path_aux_files"]
    inp_dict_Flux        = inp_dict["inp_dict_Flux"]
    wavelength_total_arr = inp_dict_Flux["wavelength_total_arr"]
    lambda_min_instr     = wavelength_total_arr[0]
    lambda_max_instr     = wavelength_total_arr[-1]
    Flux_total_arr       = inp_dict_Flux["Flux_total_arr"]
    lsf_data_arr         = np.loadtxt(inp_path_lsf + inp_file_lsf)
    Flux_conv_arr        = np.zeros(Flux_total_arr.shape[0])
    #print("Using LSF data as specified in file : "+ inp_path_lsf + inp_file_lsf)

    for lsf_indx in range(lsf_data_arr.shape[0]):
        lambda_min_lsf = lsf_data_arr[lsf_indx,0] - 100.0
        lambda_max_lsf = lsf_data_arr[lsf_indx,1] + 100.0
        if lambda_min_lsf < lambda_min_instr: lambda_min_lsf = lambda_min_instr
        if lambda_max_lsf > lambda_max_instr: lambda_max_lsf = lambda_max_instr
        dv_fwhm           = lsf_data_arr[lsf_indx,2]
        dv_sigma          = dv_fwhm / (2.0*(2.0*log(2))**0.5)
        lambda_mid_lsf    = 0.5 * (lambda_min_lsf + lambda_max_lsf)
        dv_by_dlambda     = dv_by_dlambda_funct(lambda_mid_lsf)
        lambda_sigma      = dv_sigma / dv_by_dlambda

        bool_arr_lsf      = (lambda_min_lsf <= wavelength_total_arr) & (wavelength_total_arr <= lambda_max_lsf)
        if np.sum(bool_arr_lsf != 0):
            lambda_arr_lsf    = wavelength_total_arr[bool_arr_lsf]
            Flux_arr_lsf      = Flux_total_arr[bool_arr_lsf]
            gauss_lsf_arr     = gaussian_lsf(lambda_arr_lsf,lambda_sigma)
            Flux_arr_conv_lsf = convolve(Flux_arr_lsf,gauss_lsf_arr,mode="wrap")

            bool_arr_region   = (lsf_data_arr[lsf_indx,0]<= lambda_arr_lsf) & (lambda_arr_lsf <= lsf_data_arr[lsf_indx,1])
            bool_arr_total    = (lsf_data_arr[lsf_indx,0]<= wavelength_total_arr) & (wavelength_total_arr <= lsf_data_arr[lsf_indx,1])
            Flux_conv_arr[bool_arr_total] = Flux_arr_conv_lsf[bool_arr_region]
            #print("Convolution done for lsf_indx =",lsf_indx,"N_lsf =",lsf_data_arr.shape[0])

    inp_dict_Flux.update({"Flux_conv_arr":Flux_conv_arr})
    inp_dict.update({"inp_dict_Flux":inp_dict_Flux})
    return inp_dict

def generate_lsf_input_path(inp_dict):
    spectra_type       = inp_dict["spectra_type"].upper()
    instrument         = inp_dict["instrument"].upper()

    #instr_list         = ["HST_COS","JWST","SDSS_BOSS","SDSS_DESI","UVES_HIRES_MIKE",\
    #                      "WEAVE","XSHOOTER","CUSTOM"]
    #for instr_name in instr_list:
    #    if spectra_type in instr_name:
    #       instrument = instr_name 
    #       break

    #if instrument == "CUSTOM":
    #    inp_file_lsf  = inp_dict["inp_file_lsf"]
    #else:
    #    inp_file_lsf  = "instrumental_broadening_parameters.txt"

    inp_file_lsf       = "instrumental_broadening_parameters.txt"
    inp_path_aux_files = inp_dict["inp_path_aux_files"]
    inp_sub_path_lsf_1 = "RESOLUTION_AND_LSF_PROPERTIES/" + instrument
    inp_sub_path_lsf_2 = "/LSF_VARIATION_WITH_WAVELENGTH/"
    inp_path_lsf       = inp_path_aux_files + inp_sub_path_lsf_1 + inp_sub_path_lsf_2
    file_suffix        = "_" + spectra_type
    inp_file_lsf       = inp_file_lsf.replace(".txt",file_suffix + ".txt")
    return inp_path_lsf,inp_file_lsf

def dv_by_dlambda_funct(lambda_val,z_val=None,lambda_type = "redshifted"):
    c                 = 300000.0 # in km / s

    if lambda_type == "redshifted":
        dv_by_dlambda_val = c / lambda_val

    elif lambda_type  == "rest":
        dv_by_dlambda_val = c / (lambda_val * (1.0 + z))

    return dv_by_dlambda_val

def gaussian_lsf(xdata_arr,sigma_gauss):
    xdata      = xdata_arr.copy()
    N          = xdata.shape[0]
    xdata     -= xdata[int(N/2)]
    gauss      = np.exp(-0.5*xdata**2.0/sigma_gauss**2.0)
    gauss_norm = gauss / np.sum(gauss)
    return gauss_norm

