from __future__ import print_function
from astropy.io import fits
import matplotlib.pyplot as plt
import numpy as np
import os

from WEAVEify_routines.step_11_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"]
    fits_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"]

    #plt.plot(wavelength_arr,flux_arr)
    #plt.plot(wavelength_arr,cont_arr)
    #plt.ylim(-1.0,6.0)
    #plt.xlim(3800.0,7000.0)
    #plt.show()
    #exit()

    col_name_dict = generate_extension_and_column_names(inp_dict)

    # ---------------------------------------
    # Check if file exist if not create one
    # ---------------------------------------
    check_and_create_basic_fits_file(fits_filename,overwrite=True,flag_verbose="n")
    # ---------------------------------------

    # ---------------------------------------------------------------
    # Section of the code to write Wavelength array in Common Fields
    # ---------------------------------------------------------------
    ext_name_wavelength      = col_name_dict["ext_name_wavelength"]
    col_data_list_wavelength = {0:{"col_name":col_name_dict["col_name_wavelength_flux"],"data_arr":wavelength_arr}}
    fits_overwrite_data_list(fits_filename,ext_name_wavelength,col_data_list_wavelength,format_str="float32")
    # ---------------------------------------------------------------


    # ---------------------------------------------------------------
    # Section of the code to write Flux information in Spectra
    # ---------------------------------------------------------------

    ext_name_spectra = col_name_dict["ext_name_spectra"]

    col_data_list_spectra = {\
    0:{"col_name":col_name_dict["col_name_flux"],"data_arr":flux_arr},\
    1:{"col_name":col_name_dict["col_name_flux_no_noise"],"data_arr":flux_arr_no_noise},\
    2:{"col_name":col_name_dict["col_name_cont"],"data_arr":cont_arr},\
    3:{"col_name":col_name_dict["col_name_snr_per_pixel"],"data_arr":snr_per_pixel_arr},\
    }

    fits_overwrite_data_list(fits_filename,ext_name_spectra,col_data_list_spectra,format_str="float32")

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


    # ---------------------------------------------------------------
    # Section of the code to save Header information of the spectra
    # ---------------------------------------------------------------


    fits_header_routines_main(fits_filename,ext_name_spectra,inp_dict,flag_output="flux")

    # ---------------------------------------------------------------
    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"]
    fits_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"]

    col_name_dict  = generate_extension_and_column_names(inp_dict)

    # ---------------------------------------
    # Check if file exist if not create one
    # ---------------------------------------
    check_and_create_basic_fits_file(fits_filename,overwrite=True,flag_verbose="n")
    # ---------------------------------------

    # ---------------------------------------------------------------
    # Section of the code to write Wavelength array in Common Fields
    # ---------------------------------------------------------------
    ext_name_wavelength      = col_name_dict["ext_name_wavelength"]
    col_data_list_wavelength = {0:{"col_name":col_name_dict["col_name_wavelength_tau"],"data_arr":wavelength_arr}}
    fits_overwrite_data_list(fits_filename,ext_name_wavelength,col_data_list_wavelength,format_str="float32")
    # ---------------------------------------------------------------

    # ---------------------------------------------------------------
    # Section of the code to write Flux information in Spectra
    # ---------------------------------------------------------------

    ext_name_spectra      = col_name_dict["ext_name_spectra"]
    absorber_list         = ["HI","metal","DLA"]
    kounter               = 0
    col_data_list_spectra = {}

    for flag_absorber in absorber_list:
        col_name       = col_name_dict["col_name_tau_" + flag_absorber]

        if flag_absorber.upper() in flag_tau.upper():
            data_arr   = inp_dict_tau["tau_total_" + flag_absorber + "_arr"]
        else:
            data_arr   = None

        sub_dict       = {"col_name":col_name,"data_arr":data_arr}
        col_data_list_spectra.update({kounter:sub_dict})
        kounter        = kounter + 1


    fits_overwrite_data_list(fits_filename,ext_name_spectra,col_data_list_spectra,format_str="float32")

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


    # ---------------------------------------------------------------
    # Section of the code to save Header information of the spectra
    # ---------------------------------------------------------------


    fits_header_routines_main(fits_filename,ext_name_spectra,inp_dict,flag_output="tau")

    # ---------------------------------------------------------------
    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"]
    fits_filename  = out_path_tau + out_file_tau

    col_name_dict       = generate_extension_and_column_names(inp_dict)
    ext_name_spectra    = col_name_dict["ext_name_spectra"]
    ext_name_wavelength = col_name_dict["ext_name_wavelength"]
    data_dict_spectra   = fits_read_data(fits_filename,ext_name_spectra)
    data_dict_wavelength = fits_read_data(fits_filename,ext_name_wavelength)
    tau_total_arr        = data_dict_spectra["Optical_Depth_" + flag_absorber.upper()]
    wavelength_arr       = data_dict_wavelength["Wavelength_For_Optical_Depth"]

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

    return inp_dict_tau

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

def fits_header_routines_main(fits_filename,ext_name,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])})

    modify_header(fits_filename,ext_name,header_dict,flag_verbose="n")

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

def generate_extension_and_column_names(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)

    ext_name_wavelength      = "Common_Fields"
    ext_name_spectra         = "Mock_Spectra-" + spec_id_str
    col_name_wavelength_flux = "Wavelength_For_Flux" # Previously was only Wavelength
    col_name_wavelength_tau  = "Wavelength_For_Optical_Depth"
    col_name_tau_HI          = "Optical_Depth_HI"
    col_name_tau_metal       = "Optical_Depth_METAL"
    col_name_tau_DLA         = "Optical_Depth_DLA"
    col_name_flux            = "Flux"
    col_name_flux_no_noise   = "Flux_Without_Noise"
    col_name_cont            = "Continuum_Template"
    col_name_snr_per_pixel   = "SNR_per_pixel"
    col_name_dict            = vars()

    return col_name_dict

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

def check_and_create_basic_fits_file(fits_filename,overwrite=True,flag_verbose="n"):
    if os.path.exists(fits_filename) == False:
        hdu1      = fits.PrimaryHDU()
        hdul_list = fits.HDUList([hdu1])
        hdul_list.writeto(fits_filename,overwrite=overwrite)
        if flag_verbose == "y":
            print("Fits File :",fits_filename,"is created ...")

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

def fits_read_data(fits_filename,ext_name,col_name_list=None):
    hdul     = fits.open(fits_filename)
    data_obj = hdul[ext_name].data

    if col_name_list is None:
        col_name_list = data_obj.names

    data_dict = {}
    for col_indx in range(len(col_name_list)):
        key_name        = col_name_list[col_indx]
        key_val         = data_obj[key_name].copy()
        data_dict.update({key_name:key_val})

    hdul.close()

    return data_dict

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

def fits_overwrite_data_list(fits_filename,ext_name,col_data_list,format_str="float32"):
    """ Usage Example is given below:
    data_arr_1    = np.arange(9)
    data_arr_2    = np.arange(9)*2
    ext_name      = "Mock_Spectra-0"
    col_data_list = {0:{"col_name":"a","data_arr":data_arr_1},1:{"col_name":"b","data_arr":data_arr_2}}
    fits_overwrite_data_list(fits_filename,ext_name,col_data_list,format_str="float32")
    """

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

    col_info_list   = check_if_other_data_exist_in_ext_name(fits_filename,ext_name,col_data_list,format_str)
    hdul_indx       = check_if_extension_name_exists(fits_filename,ext_name)
    for block_indx in range(len(col_data_list)):
        col_name          = col_data_list[block_indx]["col_name"]
        data_arr          = col_data_list[block_indx]["data_arr"]


        if data_arr is not None:
            col_info_indx = generate_single_column_BinHDUTable(col_name,data_arr,format_str=format_str)
            col_info_list.append(col_info_indx)

        elif hdul_indx is not None:
            try:
                data_dict     = fits_read_data(fits_filename,ext_name,col_name_list = [col_name])
                data_arr      = data_dict[col_name]
                col_info_indx = generate_single_column_BinHDUTable(col_name,data_arr,format_str = format_str)
                col_info_list.append(col_info_indx)
            except:
                pass

    hdul_table  = fits.BinTableHDU.from_columns(col_info_list,name=ext_name)

    #-------------------------------------------------------------------
    # Following piece of code checks if ext_name already exists
    # If yes, then the data is overwritten for that ext_name
    # If no, a new ext_name is created and the data is added at the end
    #-------------------------------------------------------------------
    hdul = fits.open(fits_filename,mode="update")
    if hdul_indx is not None:
        hdul[hdul_indx] = hdul_table
    else:
        hdul.append(hdul_table)

    hdul.close()

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

def check_if_other_data_exist_in_ext_name(fits_filename,ext_name,col_data_list,format_str="float32"):
    # ------------------------------------------------------------------------
    # Following code checks if there is additional BinHDU Table exists
    # other than current table. If yes, that table is appended to col_info_list
    # ------------------------------------------------------------------------

    col_name_data_list = []
    for indx in range(len(col_data_list)):
        col_name_data_list.append(col_data_list[indx]["col_name"])

    col_info_list   = []
    hdul            = fits.open(fits_filename)
    for hdul_indx in range(len(hdul)):
        if hdul[hdul_indx].name == ext_name.upper():

            # List Table names that exists under ext_name
            data_obj           = hdul[ext_name].data
            col_name_file_list = data_obj.names

            for col_name_file in col_name_file_list:
                # Check if Table names that exist under ext_name are going to
                # be overwritten. If not over written carry forward this data set
                # in col_info_list
                if col_name_file not in col_name_data_list:
                    data_dict_file = fits_read_data(fits_filename,ext_name,col_name_list = [col_name_file])
                    data_arr_file  = data_dict_file[col_name_file]
                    col_info_file  = generate_single_column_BinHDUTable(col_name_file,data_arr_file,format_str = format_str)
                    col_info_list.append(col_info_file)

    hdul.close()
    return col_info_list

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

def check_if_extension_name_exists(fits_filename,ext_name):
    hdul            = fits.open(fits_filename)
    hdul_indx_block = None
    for hdul_indx in range(len(hdul)):
        if hdul[hdul_indx].name == ext_name.upper():
            hdul_indx_block = hdul_indx
            break
    hdul.close()
    return hdul_indx_block

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

def check_if_binary_hdu_table_exists(fits_filename,ext_name,col_name):
    hdul            = fits.open(fits_filename)
    hdul_indx_block = None
    for hdul_indx in range(len(hdul)):
        if hdul[hdul_indx].name == ext_name.upper():
            data_obj          = hdul[ext_name].data
            col_name_list_org = data_obj.names
            if col_name in col_name_list_org:
                hdul_indx_block = hdul_indx
    hdul.close()
    return hdul_indx_block,col_name_list

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

def generate_single_column_BinHDUTable(col_name,data_arr,format_str="float32"):
    fits_format_dict = {"char":"A","int32":"J","int64":"K","float32":"E","float64":"D"}
    col_info         = fits.Column(name=col_name, array=data_arr, format=fits_format_dict[format_str])
    return col_info

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

def modify_header(fits_filename,ext_name,header_dict,flag_verbose="n"):

    with fits.open(fits_filename,mode = 'update') as hdul:

        hdul.verify('fix')
        hdr      = hdul[ext_name.upper()].header
        key_list = list(header_dict.keys())
        N_header = len(header_dict)

        for header_indx in range(N_header):
            key_name      = key_list[header_indx]
            key_val       = header_dict[key_name]
            hdr[key_name] = str(key_val)

        hdul.flush() # Changes are written to the file

    if flag_verbose == "y":
        print("Header data is written ...")
