from __future__ import print_function
import numpy as np
import h5py
import matplotlib.pyplot as plt

def read_and_plot_synthetic_spectra_main():
    # model_indx options
    # 12 --->> f_HI Corresponds to model passing through best fit f_HI values from Gaikwad+2023 (LATE REIONIZATION)
    # 0  --->> f_HI Corresponds to model passing through Lower  f_HI values from Gaikwad+2023 (EARLY REIONIZATION)
    # 25 --->> f_HI Corresponds to model passing through best fit f_HI values from Gaikwad+2023 (ULTRA-LATE REIONIZATION)
    model_indx = 12

    # mock_indx could be anywhere between 0 and 45 (Sightlines at different locations)
    mock_indx  = 0

    # Path to synthetic spectra file
    inp_path = "./"

    # File name containing synthetic spectra
    inp_file = "Synthetic_Spectra_Dark_Gap_Stacking.hdf5"




    # List of all QSO names in single mock
    qso_name_list = ["PSO_J323+12", "PSO_J231-20", "VDES_J0224-4711", \
    "PSO_J1212+0505", "DELS_J1535+1943", "ATLAS_J2211-3206", "PSO_J060+24", \
    "PSO_J065-26", "PSO_J359-06", "PSO_J217-07", "PSO_J217-16", "PSO_J239-07", \
    "SDSS_J0842+1218", "ATLAS_J158-14", "VDES_J0408-5632", "ATLAS_J029-36", \
    "SDSS_J2310+1855", "PSO_J007+04", "PSO_J029-29", "PSO_J108+08", \
    "PSO_J183-12", "PSO_J025-11", "PSO_J242-12", "PSO_J065+01", "PSO_J308-27", \
    "PSO_J036+03", "SDSS_J0818+1722", "SDSS_J0927+2001", "SDSS_J1030+0524", \
    "SDSS_J1306+0356", "ULAS_J0148+0600", "ULAS_J1319+0950", \
    "VST-ATLAS_J025-33", "CFHQS_J1509-1749 ", "SDSS_J0100+2802", \
    "SDSS_J0836+0054", "PSO_J159-02", "PSO_J308-21", "ULAS_J1207+0630", \
    "PSO_J340-18", "PSO_J056-16", "PSO_J215-16", "J1335-0328", "J0108+0711"]

    print("")
    # You can iterate over all qso_indx i.e., 0 to len(qso_name_list)-1
    qso_indx           = 0

    qso_name = qso_name_list[qso_indx]
    print("Name of the QSO is :",qso_name)

    # The factor by which Gamma_12 is scaled to get different neutral fraction
    # Iterate over Gamma_12_scale_list to get different neutral fraction for same mock
    Gamma_12_scale_list = [0.3,0.4,0.5,0.6,1.0,1.5,2.0,2.5,3.0]
    Gamma_12_indx       = 4
    Gamma_12_scale_val  = Gamma_12_scale_list[Gamma_12_indx]
    print("Gamma_12 factor used to scale Neutral Fraction is :",Gamma_12_scale_val)

    hdf_filename = inp_path + inp_file
    flux_arr     = read_flux_for_given_qso(hdf_filename,qso_name,model_indx,Gamma_12_scale_val,mock_indx)

    # Zeroth column is the wavelengths
    wavelength_arr   = flux_arr[:,0]
    print("Size of Wavelength Array is :",wavelength_arr.shape)

    # First column is flux without noise, instrumental or continuum effect
    flux_arr_perfect = flux_arr[:,1]
    print("Size of Perfect flux Array is :",flux_arr_perfect.shape)

    # Second  column is flux with noise, instrumental or continuum effect
    flux_arr_all_effect = flux_arr[:,2]
    print("Size of Observational systematics contaminated flux Array is :",flux_arr_all_effect.shape)

    # Third  column is continuum used to contaminate spectra
    cont_arr  = flux_arr[:,3]
    print("Size of Continuum Array is :",cont_arr.shape)

    # Normalized flux is obtained by diving flux_arr_all_effect by continuum
    bool_arr           = cont_arr == 0
    cont_arr[bool_arr] = 1.0
    flux_arr_normalize = flux_arr_all_effect / cont_arr
    print("Size of Normalized flux Array is :",flux_arr_normalize.shape)

    # Neutral Fraction along the sightline
    f_HI_arr = flux_arr[:,4]
    print("Size of Neutral fraction Array is :",f_HI_arr.shape)
    print("")


    #plt.plot(wavelength_arr,f_HI_arr)
    #plt.show()

def read_flux_for_given_qso(hdf_filename,qso_name,model_indx,Gamma_12_scale_val,mock_indx):
    block_1  = "/REION_HIST_TYPE_01_EARLY_LATE/Model-" + str("%03i"%model_indx)
    block_2  = "/Best_Fit_Gamma_12_Evolution_Scaled_By_" + str("%1.4f"%Gamma_12_scale_val)
    block_3  = "/Mock-" + str("%04i"%mock_indx) + "/" + qso_name + "/Flux_Data"
    block    = block_1 + block_2 + block_3
    flux_arr = hdf_read_data(hdf_filename,block)
    return flux_arr

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


read_and_plot_synthetic_spectra_main()