from __future__ import print_function
import numpy as np

#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 read_weaveify_output_fields(out_path,flag_output_format="fits",flag_mpi="n",rank=None,spec_id=None,flag_print_output="n"):
    """
    out_path -->> Path of the directory containing WEAVEify output files

    flag_output_format -->> Format of the output files

    flag_mpi -->> "y" if output files were generated using MPI
                 filenames are different for serial and parallel code)

    rank    -->> Rank is an intger associated with processor number that generated the file.
                 If flag_mpi ="y", then rank must be specified. 
                 For example, if flag_mpi = "y" and rank = 5, then the data
                 would be read from the file `WEAVEified_Mock_Spectra_Rank_5.fits'
    
    spec_id -->> If None, all the data available in the files will be read to
                 memory and saved in out_dict. If spec_id is specified only the 
                 the data related to MOCK_SPECTRA-<spec_id> would be read.
                 [Note: COMMON_FIELDS arrays are always read in the script
                 irrespective to spec_id]

    flag_print_output -->> If "y", this will print the information of arrays
                 in final output dictionary

    Returns:
    out_dict -->> Output dictionary containing all the fields

    out_dict[<extension or hdf5 group name>][<field name>]

    For example, to read flux from spec_id=0, we can use
    flux_arr = out_dict['MOCK_SPECTRA-0']['Flux']

    To read wavelength from COMMON_FIELDS, we can use
    wavelength_arr out_dict['COMMON_FIELDS']['Wavelength_For_Flux']

    To read header dictionary
    header_dict = out_dict['MOCK_SPECTRA-0']['Header']

    """

    # Import either astropy or h5py modules based on flag_output_format
    grb_return = step_1_import_modules_based_on_file_format(flag_output_format)

    # Generate file name and full file path of the output file
    fits_hdf_filename = step_2_generate_full_file_path(out_path,flag_output_format,flag_mpi,rank)

    # Generate Extension Name (for fits) or Block (for hdf5) structure of the file
    ext_block_dict = step_3_generate_ext_block_names_dict(fits_hdf_filename,flag_output_format,spec_id)

    # Collect data from fits or hdf5 file and for given Extension/Block  Names
    out_dict       = step_4_collect_data(fits_hdf_filename,flag_output_format,ext_block_dict)


    # Print Array Information from fits or hdf5 file
    if flag_print_output.upper() == "Y":
        out_dict  = step_5_print_data(out_dict,flag_output_format)


    return out_dict


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


def step_1_import_modules_based_on_file_format(flag_output_format="fits"):
    if flag_output_format == "fits":
        from astropy.io import fits
        global fits
    elif flag_output_format == "hdf5":
        import h5py
        global h5py
    return None

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

def step_2_generate_full_file_path(out_path,flag_output_format="fits",flag_mpi="n",rank=None):

    out_path = reformat_dir_path(out_path)
    file_base = "WEAVEified_Mock_Spectra"


    if flag_mpi == "y":
        if rank is not None:
            file_base = file_base + "_Rank_" + str("%i"%int(rank))
        else:
            print("Trying to read MPI generated output files but rank is not specified")
            print("Exiting the code ...")
            exit()

    out_file = file_base + "." + flag_output_format
    filename = out_path + out_file

    return filename

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

def step_3_generate_ext_block_names_dict(fits_hdf_filename,flag_output_format="fits",spec_id=None):

    if flag_output_format == "fits":
        ext_block_dict = step_3a_fits_generate_extension_name_dictionary(fits_hdf_filename,spec_id=spec_id)

    elif flag_output_format == "hdf5":
        ext_block_dict = step_3b_hdf5_generate_block_name_dictionary(fits_hdf_filename,spec_id=spec_id)

    return ext_block_dict

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

def step_3a_fits_generate_extension_name_dictionary(fits_filename,spec_id=None):

    ext_name_list = ["COMMON_FIELDS"]
    if spec_id is None:
        ext_name_list_all = get_all_hdu_extension_name_list(fits_filename)
        ext_name_list_all = remove_items_from_list(ext_name_list_all,['PRIMARY','COMMON_FIELDS'])
        for spec_id_indx in range(len(ext_name_list_all)):
            ext_name_spectra = ext_name_list_all[spec_id_indx]
            ext_name_list.append(ext_name_spectra)

    else:

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

        ext_name_spectra      = "MOCK_SPECTRA-" + spec_id_str
        ext_name_list.append(ext_name_spectra)


    ext_name_dict = {"ext_name_list":ext_name_list}

    return ext_name_dict

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

def step_3b_hdf5_generate_block_name_dictionary(hdf_filename,spec_id=None):

    if spec_id is None:
        group_list = list_group_dataset(hdf_filename,"group")
    else:

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

        group_name_spec = "MOCK_SPECTRA-" + spec_id_str
        group_list      = ["COMMON_FIELDS",group_name_spec]

    block_dict = {"group_list":group_list}
    return block_dict

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

def step_4_collect_data(fits_hdf_filename,flag_output_format,ext_block_dict):

    if flag_output_format == "fits":
        out_dict = step_4a_collect_data_fits(fits_hdf_filename,ext_block_dict)

    elif flag_output_format == "hdf5":
        out_dict = step_4b_collect_data_hdf5(fits_hdf_filename,ext_block_dict)

    return out_dict

def step_4a_collect_data_fits(fits_filename,ext_block_dict):
    ext_name_list = ext_block_dict["ext_name_list"]
    out_dict      = {}
    for ext_name in ext_name_list:
        out_dict_indx   = fits_read_data(fits_filename,ext_name)
        header_dict     = fits_read_header(fits_filename,ext_name)
        out_dict_indx.update(header_dict)
        out_dict.update({ext_name:out_dict_indx})
    return out_dict

def step_4b_collect_data_hdf5(hdf_filename,ext_block_dict):
    out_dict = {}
    group_list = ext_block_dict["group_list"]
    for group_name in group_list:
        dataset_list = list_dataset_in_a_block(hdf_filename,group_name)
        out_sub_dict = {}
        for dataset in dataset_list:
            blockname = "/" + group_name + "/" + dataset
            if dataset.upper() != "HEADER":
                data_arr  = hdf_read_data(hdf_filename,blockname)
                out_sub_dict.update({dataset:data_arr})
            else:
                header_dict = hdf_read_header(hdf_filename,blockname)
                out_sub_dict.update(header_dict)
        out_dict.update({group_name:out_sub_dict})
    return out_dict



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

def step_5_print_data(out_dict,flag_output_format):
    key_list = out_dict.keys()
    N_keys   = len(key_list)
    for key_indx in range(N_keys):

        key_name     = key_list[key_indx]
        var_name_str = "out_dict['" + key_name + "']"
        sub_key_list = out_dict[key_name].keys()
        for sub_key_name in sub_key_list:
            var_name     = var_name_str + "['" + sub_key_name + "']"

            if sub_key_name.upper() == "HEADER":
                header_dict  = out_dict[key_name][sub_key_name]
                header_len   = len(header_dict)
                if len(str(header_len)) == 1:
                    print("Header Length =",header_len,",      Variable Name =",var_name)
                elif len(str(header_len)) == 2:
                    print("Header Length =",header_len,",     Variable Name =",var_name)
            else:
                data_arr     = out_dict[key_name][sub_key_name]
                data_arr_shape = data_arr.shape
                print("Array Shape =",data_arr.shape,", Variable Name =",var_name)

    return out_dict


##############################################################
############# COMMON FITS FILE READING ROUTINES ##############
##############################################################

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_read_header(fits_filename,ext_name):
    hdul            = fits.open(fits_filename)
    header_obj      = hdul[ext_name].header
    header_key_list = header_obj.keys()
    header_sub_dict = {}
    for header_indx in range(len(header_key_list)):
        key_name    = header_key_list[header_indx]
        key_val     = header_obj[key_name]
        header_sub_dict.update({key_name:key_val})

    hdul.close()
    header_dict = {"Header":header_sub_dict}

    return header_dict

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

def get_all_hdu_extension_name_list(fits_filename):
    hdul          = fits.open(fits_filename)
    ext_name_list = []
    for hdu_indx in range(len(hdul)):
        ext_name_list.append(hdul[hdu_indx].name)
    hdul.close()
    return ext_name_list

##############################################################
############# COMMON HDF5 FILE READING ROUTINES ##############
##############################################################

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

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

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

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

def list_group_dataset(hdf_filename,flag_dataset_group="group"):
    hdf_read           = h5py.File(hdf_filename,"r")
    group_dataset_list = []
    hdf_read.visit(group_dataset_list.append)

    # newly added
    group_list   = []

    for blockname in group_dataset_list:
        is_dataset = isinstance(hdf_read[blockname], h5py.Dataset)

        if is_dataset is False:
            if "/" not in blockname:
                group_list.append(blockname)

    hdf_read.close()
    return group_list

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

def list_dataset_in_a_block(hdf_filename,blockname):
    hdf_read 	=	h5py.File(hdf_filename,"r")
    hdf_dict	=	hdf_read[blockname].keys()
    hdf_read.close()
    return hdf_dict

##############################################################
################## MISCELLANEOUS ROUTINES ####################
##############################################################

def reformat_dir_path(inp_str):
    if inp_str[-1] != "/":
        inp_str = inp_str + "/"
    return inp_str

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

def remove_items_from_list(ext_name_list,item_list):
    for indx in range(len(item_list)):
        item = item_list[indx]
        if item in ext_name_list:ext_name_list.remove(item)
    return ext_name_list

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

def test_read_weaveify_output_fields():
    #out_path_test_1 = "../../WEAVE_MOCKS_AUXILIARY_FILES/CODE_TEST_EXAMPLES/WEAVE/OUTPUT/SERIAL_SINGLE_SPECTRUM/"
    #out_path_test_1 = "../../WEAVE_MOCKS_AUXILIARY_FILES/CODE_TEST_EXAMPLES/WEAVE/OUTPUT/SERIAL_MULTIPLE_SPECTRA/"
    out_path_test_1 = "../../WEAVE_MOCKS_AUXILIARY_FILES/CODE_TEST_EXAMPLES/WEAVE/OUTPUT/PARALLEL_MPI_MULTIPLE_SPECTRA/"

    out_dict = read_weaveify_output_fields(out_path_test_1,\
               flag_output_format="fits",flag_mpi="y",rank=0,spec_id=101,\
               flag_print_output="y")

#test_read_weaveify_output_fields()
