import numpy as np
import sys
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
from ASTRO_MOCKS_routines.step_02_common_routines.dictionary_functions import CHECK_AND_INITIALIZE_KEY_VALUE

from ASTRO_MOCKS_routines.step_02_common_routines.python_2_3_compatibility_routines \
import genfromtxt_compatible

def wavelength_array_calculation_main(inp_dict):
    inp_dict               = read_wavelength_sampling_parameters_weave(inp_dict)
    out_dict               = compute_wavelength_array_from_coefficients(inp_dict,flag_print_range = "n")
    lambda_arr_obs_default = out_dict["lambda_arr_obs"]
    CHECK_AND_INITIALIZE_KEY_VALUE(inp_dict,"lambda_arr_obs",lambda_arr_obs_default)
    return inp_dict

def compute_wavelength_array_from_coefficients(inp_dict,flag_print_range="n"):
    """ The wavelength array is usually not given in fits file it
        needs to be calculated from two variables given in HEADER

        If wavelengths are sampled on linear scale then
        wavelength  = CRVAL1+CDELT1*np.arange(N_pixel_obs)
        Wavelength in WEAVE spectra are sampled on linear scales


        If wavelengths are sampled on logarithmic scale then
        wavelength  = 10.0**(CRVAL1+CDELT1*np.arange(N_pixel_obs))
        Wavelength in Keck / UVES spectra are usually sampled on log scales

        where N_pixel_obs is number of pixels in observed spectra
    """
    CHECK_AND_INITIALIZE_KEY_VALUE(inp_dict,"spectra_type","WEAVE_LR")
    spectra_type     = inp_dict["spectra_type"].upper()
    out_dict         = {}

    if spectra_type is None:
        key_list = ["CRVAL1","CDELT1","N_pixel_obs","lambda_scale"]
        CHECK_IF_KEY_PRESENT(key_list,inp_dict)
        CRVAL1          = float(inp_dict["CRVAL1"])
        CDELT1          = float(inp_dict["CDELT1"])
        N_pixel_obs     = int(inp_dict["N_pixel_obs"]) + 1
        lambda_scale    = inp_dict["lambda_scale"]

    elif spectra_type == "UVES":
        CRVAL1          = 3.4841826183453
        CDELT1          = 3.62161104717198e-06
        N_pixel_obs     = 142000
        lambda_scale    = "LOG"

    elif (spectra_type == "WEAVE_LR") | (spectra_type == "WEAVE_HR"):
        lambda_param_dict = inp_dict["lambda_arm_param_dict"]
        resolution_mode   = inp_dict["resolution_mode"]
        arm_type_list     = ["blue","red"]

        for arm_indx in range(len(arm_type_list)):
            arm_type              = arm_type_list[arm_indx].lower()
            lambda_param_sub_dict = lambda_param_dict[(resolution_mode,arm_type)]
            lambda_scale          = lambda_param_sub_dict["lambda_scale"]
            N_pixel_obs           = lambda_param_sub_dict["N_pixel_obs"]
            CDELT1                = lambda_param_sub_dict["CDELT1"]
            CRVAL1                = lambda_param_sub_dict["CRVAL1"]
            lambda_arr            = compute_wavelength_array(CRVAL1,CDELT1,N_pixel_obs,lambda_scale)

            if arm_type == "blue":
                lambda_min_spec = lambda_arr[0]

            elif arm_type == "red":
                lambda_max_spec = lambda_arr[-1]

        CRVAL1          = lambda_min_spec
        N_pixel_obs     = int(round((lambda_max_spec - lambda_min_spec) / CDELT1,1)) + 1


    # You do not need to change following lines in function
    lambda_arr_obs = compute_wavelength_array(CRVAL1,CDELT1,N_pixel_obs,lambda_scale)


    lambda_min_instr = lambda_arr_obs[0]    # Minimum wavelength of Mock spectra in angstroms
    lambda_max_instr = lambda_arr_obs[-1]   # Maximum wavelength of Mock spectra in angstroms

    out_dict.update({"lambda_arr_obs":lambda_arr_obs})

    return out_dict


def compute_wavelength_array(CRVAL1,CDELT1,N_pixel_obs,lambda_scale):

    if "LIN" in lambda_scale:
        lambda_arr_obs   = CRVAL1+CDELT1*np.arange(N_pixel_obs)

    elif "LOG" in lambda_scale:
        lambda_arr_obs   = 10.0**(CRVAL1+CDELT1*np.arange(N_pixel_obs))

    return lambda_arr_obs



def read_wavelength_sampling_parameters_weave(inp_dict):
    instrument = inp_dict["instrument"]

    if instrument.upper() == "WEAVE":

        resolution_mode       = inp_dict["resolution_mode"]
        inp_path              = inp_dict["inp_path_lambda_samp"]
        inp_file              = inp_dict["inp_file_lambda_samp"]
        data_arr_all          = genfromtxt_compatible(inp_path,inp_file,object)
        lambda_arm_param_dict = {}

        for indx in range(data_arr_all.shape[0]):
            resolution_mode_indx = data_arr_all[indx,0]
            arm_type_indx        = data_arr_all[indx,1].lower()
            CRVAL1_indx          = float(data_arr_all[indx,2])
            CDELT1_indx          = float(data_arr_all[indx,3])
            N_pixel_obs_indx     = int(data_arr_all[indx,4])
            lambda_scale_indx    = data_arr_all[indx,5]
            cwave_indx           = float(data_arr_all[indx,6])
            key_name             = (resolution_mode_indx.decode(),arm_type_indx.decode())
            key_val_dict         = {"CRVAL1":CRVAL1_indx,"CDELT1":CDELT1_indx,\
                                    "N_pixel_obs":N_pixel_obs_indx,\
                                    "lambda_scale":lambda_scale_indx.decode(),\
                                    "cwave":cwave_indx}
            lambda_arm_param_dict.update({key_name:key_val_dict})
        inp_dict.update({"lambda_arm_param_dict":lambda_arm_param_dict})

    return inp_dict





