import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage.filters import convolve
from scipy.interpolate import interp1d
from math import log


# *************** Custom Module Starts *********************

from common_lib.file_format.hdf5_functions import hdf_overwrite_data_list,\
hdf_read_data

from common_lib.plot_setting import load_plt

# *************** Custom Module Ends *********************

def extract_continuum_squad_dr1_main():
    inp_path       = "/data/emergence12/pgaikwad/Projects/data_files/Project-6_Gamma_HI_measurements/OBSERVATION/SQUAD_DR1/INPUT_FILES/EXTRACTED_FILES/"
    qso_name_list  = ["J162116-004250","J005758-264314","J013724-422417"]
    qso_indx       = 1
    qso_name       = qso_name_list[qso_indx]
    flag_plot_data = "n"
    #pdf_obj        = CREATE_PDF_OBJECT("Test.pdf")
    inp_dict       = step_0_read_input_observations(inp_path,qso_name)
    out_dict_flux_cont       = plot_continuum(inp_dict)
    out_path       = "/data/emergence11/pgaikwad/Projects/data_files/Project-016_WEAVE_QSO_Analysis/WEAVE_MOCKS/"
    out_file       = "out_dict_flux_cont.npy"
    np.save(out_path + out_file,out_dict_flux_cont,allow_pickle=True)

def step_0_read_input_observations(inp_path,qso_name):

    from common_lib.file_format.fits_functions import READ_HEADER_INFO,\
    READ_TABLE_COLUMN,READ_SINGLE_COLUMN_DATA,COMPUTE_WAVELENGTH_ARRAY_FROM_FITS_FILE

    from common_lib.observations.SQUAD.coadd_routine import MASK_BAD_PIXELS

    inp_file         = qso_name + ".fits"
    fits_filename    = inp_path + qso_name + "/" + inp_file
    file_header      = READ_HEADER_INFO(fits_filename,"n")
    z_em             = file_header["UP_ZEM"]
    fit_data_arr     = READ_SINGLE_COLUMN_DATA(fits_filename).T
    N_pixels         = fit_data_arr.shape[0]
    wavelength_arr   = COMPUTE_WAVELENGTH_ARRAY_FROM_FITS_FILE(fits_filename,N_pixels,flag_print_range = "n")

    Flux_arr_norm    = fit_data_arr[:,0]
    Flux_err_arr     = fit_data_arr[:,1]
    Flux_cont_arr    = fit_data_arr[:,3] # Continuum fit by Michael Murphy 

    Flux_arr         = np.zeros((N_pixels,4))
    Flux_arr[:,0]    = wavelength_arr
    Flux_arr[:,1]    = Flux_arr_norm * Flux_cont_arr
    Flux_arr[:,2]    = Flux_err_arr  * Flux_cont_arr
    Flux_arr[:,3]    = Flux_cont_arr
    bool_arr         = MASK_BAD_PIXELS(Flux_arr_norm,Flux_err_arr)
    inp_dict         = {"Flux_arr":Flux_arr[bool_arr,:],"z_em":z_em}
    return inp_dict

def plot_continuum(inp_dict):
    Flux_arr   = inp_dict["Flux_arr"]
    z_em       = inp_dict["z_em"]
    lambda_arr = Flux_arr[:,0]
    lambda_arr_rest = lambda_arr / (1.0 + z_em) 
    lambda_arr_to_plot = lambda_arr_rest * (1.0 + 3.0)
    #plt.plot(lambda_arr_to_plot,Flux_arr[:,1])
    #plt.plot(lambda_arr_to_plot,Flux_arr[:,3])
    #plt.ylim(0.0,1200.0)
    #plt.show()
    out_dict = {"lambda_arr_rest":lambda_arr_rest,"Flux_arr_cont":Flux_arr}
    return out_dict

extract_continuum_squad_dr1_main()
