from __future__ import print_function
import numpy as np
import h5py

from WEAVEify_routines.step_07_read_write_mock_spectra.\
common_header_routines  import variable_list_to_save_in_header,\
generate_standard_fits_file_compact_header

###########################################
##### SAVE POST-PROCESSED FLUX FIELDS #####
###########################################

def save_flux_fields(inp_dict,flag_verbose="n"):
    out_path_flux = inp_dict["out_path"]
    out_file_flux = inp_dict["out_file_flux"]
    hdf_filename  = out_path_flux + out_file_flux

    out_dict_final      = inp_dict["out_dict_final"]
    wavelength_arr      = out_dict_final["wavelength_arr"]
    flux_arr            = out_dict_final["flux_arr"]
    cont_arr            = out_dict_final["cont_arr"]
    flux_arr_no_noise   = out_dict_final["flux_arr_no_noise"]
    snr_per_pixel_arr   = out_dict_final["snr_per_pixel_arr"]
    noise_per_pixel_arr = out_dict_final["noise_per_pixel_arr"]

    block_dict            = generate_block_name_dictionary(inp_dict)
    flux_block_no_noise   = block_dict["flux_block_no_noise"]
    flux_block            = block_dict["flux_block"]
    cont_block            = block_dict["cont_block"]
    snr_per_pixel_block   = block_dict["snr_per_pixel_block"]
    noise_per_pixel_block = block_dict["noise_per_pixel_block"]
    wavelength_block      = block_dict["wavelength_block_flux"]
    header_block          = block_dict["header_block"]

    write_list       = {flux_block:flux_arr.astype(np.float32),\
                        flux_block_no_noise:flux_arr_no_noise.astype(np.float32),\
                        cont_block:cont_arr.astype(np.float32),\
                        wavelength_block:wavelength_arr.astype(np.float32),\
                        snr_per_pixel_block:snr_per_pixel_arr.astype(np.float32),\
                        noise_per_pixel_block:noise_per_pixel_arr.astype(np.float32)}

    hdf5_header_routines_main(hdf_filename,header_block,inp_dict,flag_output="flux")

    hdf_overwrite_data_list(hdf_filename,write_list)

    if flag_verbose == "y":
        print("Flux data is saved ...")

    return inp_dict

###########################################
##### SAVE OPTICAL DEPTH (tau) FIELDS #####
###########################################

def save_optical_depths(inp_dict,flag_verbose="n"):
    out_path_tau   = inp_dict["out_path"]
    out_file_tau   = inp_dict["out_file_tau"]
    hdf_filename   = out_path_tau + out_file_tau


    flag_tau       = inp_dict["flag_tau"]
    inp_dict_tau   = inp_dict["inp_dict_tau"]
    wavelength_arr = inp_dict_tau["wavelength_arr"]

    block_dict       = generate_block_name_dictionary(inp_dict)
    header_block     = block_dict["header_block"]
    wavelength_block = block_dict["wavelength_block_tau"]
    write_list       = {wavelength_block:wavelength_arr.astype(np.float32)}

    absorber_list    = ["HI","metal","DLA"]

    for flag_absorber in absorber_list:
        if flag_absorber.upper() in flag_tau.upper():
            tau_total_arr  = inp_dict_tau["tau_total_" + flag_absorber + "_arr"]
            tau_block      = block_dict["tau_block_" + flag_absorber]
            write_list.update({tau_block:tau_total_arr.astype(np.float32)})

    hdf5_header_routines_main(hdf_filename,header_block,inp_dict,flag_output="tau")
    hdf_overwrite_data_list(hdf_filename,write_list)

    if flag_verbose == "y":
        print("Optical Depth data is saved ...")

    return inp_dict


def read_optical_depths(inp_dict,flag_absorber="HI",flag_verbose="y"):
    out_path_tau   = inp_dict["out_path"]
    out_file_tau   = inp_dict["out_file_tau"]
    hdf_filename   = out_path_tau + out_file_tau

    block_dict       = generate_block_name_dictionary(inp_dict)
    header_block     = block_dict["header_block"]
    wavelength_block = block_dict["wavelength_block_tau"]
    tau_block        = block_dict["tau_block_" + flag_absorber]

    tau_total_arr  = hdf_read_data(hdf_filename,tau_block)
    wavelength_arr = hdf_read_data(hdf_filename,wavelength_block)
    header_dict    = hdf5_read_header_attribute(hdf_filename,header_block)

    inp_dict_tau   = {"tau_total_arr":tau_total_arr,\
                      "wavelength_arr":wavelength_arr,\
                      "header_dict":header_dict}

    return inp_dict_tau


################################################################################
##### COMMON HDF5 READ / WRITE ROUTINES TO BE USED THROUGHOUT THE SCRIPT #######
################################################################################


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

def generate_block_name_dictionary(inp_dict):
    spec_id      = inp_dict["spec_id"]

    if spec_id is None:
        spec_id     = 0
        spec_id_str = str("%i"%spec_id)
    else:
        if isinstance(spec_id,str):
            spec_id_str = spec_id
        else:
            spec_id_str = str("%i"%spec_id)

    common_block          = "/MOCK_SPECTRA-" + spec_id_str + "/"
    header_block          = common_block + "/Header"
    wavelength_block_flux = "/COMMON_FIELDS/Wavelength_For_Flux"
    wavelength_block_tau  = "/COMMON_FIELDS/Wavelength_For_Optical_Depth"
    tau_block_HI          = common_block + "/Optical_Depth_HI"
    tau_block_metal       = common_block + "/Optical_Depth_METAL"
    tau_block_DLA         = common_block + "/Optical_Depth_DLA"
    flux_block            = common_block + "/Flux"
    flux_block_no_noise   = common_block + "/Flux_Without_Noise"
    cont_block            = common_block + "/Continuum_Template"
    snr_per_pixel_block   = common_block + "/SNR_per_pixel"
    noise_per_pixel_block = common_block + "/Noise_per_pixel"
    block_dict            = vars()
    return block_dict


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

def hdf5_header_routines_main(hdf_filename,header_block,inp_dict,flag_output="flux"):
    std_fit_header_dict = generate_standard_fits_file_compact_header()
    header_param_list = variable_list_to_save_in_header(flag_output = flag_output)
    N_var             = len(header_param_list)
    header_dict       = {}

    for var_name in header_param_list:
        key_name = std_fit_header_dict[var_name][0]
        header_dict.update({key_name:str(inp_dict[var_name])})
        #header_dict.update({var_name:str(inp_dict[var_name])})

    hdf5_write_header_attribute(hdf_filename,header_block,header_dict)

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

def hdf5_write_header_attribute(hdf_filename,header_block,header_dict):

    hdf_write = h5py.File(hdf_filename,"a")
    flag      = header_block in hdf_write

    if flag == True:
        del hdf_write[header_block]

    hdf_header     = hdf_write.create_group(header_block)
    N_header       = len(header_dict)
    attr_name_list = list(header_dict.keys())

    for indx in range(N_header):
        attr_name = attr_name_list[indx]
        attr_val  = header_dict[attr_name]
        hdf_header.attrs.create(attr_name,attr_val)

    hdf_write.close()

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

def hdf5_read_header_attribute(hdf_filename,header_block="Header"):
    import h5py
    hdf_read 		= h5py.File(hdf_filename,"r")
    hdf_header 		= hdf_read[header_block]
    header_key_list     = list(hdf_header.attrs.keys())
    header_dict 	= {}
    for indx in range(len(header_key_list)):
        key 		= header_key_list[indx]
        attr_dict 	= {key:hdf_header.attrs[key]}
        header_dict.update(attr_dict)
    hdf_read.close()
    return header_dict

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

def hdf_overwrite_data_list(hdf_filename,block_data_list,flag_print="no"):
    hdf_write = h5py.File(hdf_filename,"a")
    for indx in range(len(block_data_list)):
        blockname = list(block_data_list.keys())[indx]
        data = block_data_list[blockname]
        if hdf_write.__contains__(blockname): # Check if blockname already exists
            tmp_data = hdf_write[blockname]
            if (tmp_data.shape == data.shape) & (tmp_data.dtype == data.dtype): 
                tmp_data[:] = data      # If data type and shape matches then overwrite
                if (flag_print == "yes") | (flag_print == "y"):
                    print("Blocks are overwritten !!!")
                    print("File :" + hdf_filename)
                    print("Block :" + blockname)
            else:
                del hdf_write[blockname]
                hdf_write.create_dataset(blockname,data=data)
                if (flag_print == "yes") | (flag_print == "y"):
                    print("Previous block deleted and new block created !!!")
                    print("File :" + hdf_filename)
                    print("Block :" + blockname)
        else:
            hdf_write.create_dataset(blockname,data=data)
            if (flag_print == "yes") | (flag_print == "y"):
                print("New block created !!!")
                print("File :" + hdf_filename)
                print("Block :" + blockname)
    hdf_write.close()

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

def hdf_read_data(hdf_filename,blockname,flag_exit="y"):
    hdf_read = h5py.File(hdf_filename,"r")
    if hdf_read.__contains__(blockname): # Check if blockname already exists
        data = np.array(hdf_read[blockname])
        hdf_read.close()
    else:
        print("Following Block name does not exists ...")
        print("File :",hdf_filename)
        print("Block :",blockname)
        hdf_read.close()
        if flag_exit == "y":
            exit()
        else:
            data = None
    return data

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