import numpy as np
import matplotlib.pyplot as plt
#from common_lib.miscellaneous.numpy_special_functions import 
#slice_array_between_two_limits

def noise_addition_routines_main(inp_dict):
    inp_dict       = step_b_generate_photon_count_common_term(inp_dict)
    inp_dict       = step_c_calculate_photon_count_for_noise_component(inp_dict)
    inp_dict       = step_d_generate_noise_realization(inp_dict)
    inp_dict       = step_e_convert_N_photon_to_flux(inp_dict)
    return inp_dict

#def step_a_convert_magintude_to_flux_density(inp_dict):

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

def step_b_generate_photon_count_common_term(inp_dict):
    inp_dict_Flux         = inp_dict["inp_dict_Flux"]
    Flux_arr_cont_added   = inp_dict_Flux["Flux_arr_cont_added"]
    lambda_arr            = Flux_arr_cont_added[:,0]
    Delta_lambda_arr      = np.ones(Flux_arr_cont_added.shape[0])
    Delta_lambda_arr[:-1] = lambda_arr[1:] - lambda_arr[:-1]
    Delta_lambda_arr[-1]  = lambda_arr[-1] - lambda_arr[-2]
    Delta_lambda_arr[:]   = np.median(lambda_arr[1:] - lambda_arr[:-1])
    N_pixel_per_ccd       = inp_dict["N_pixel_per_ccd"]
    A_eff_per_ccd         = inp_dict["A_eff_per_ccd"]
    t_exp                 = inp_dict["t_exp"]
    E_arr_in_erg          = lambda_angstroms_to_energy_erg(lambda_arr)
    photon_count_term_arr = N_pixel_per_ccd * A_eff_per_ccd * t_exp * Delta_lambda_arr / E_arr_in_erg

    inp_dict_Flux.update({"photon_count_term_arr":photon_count_term_arr})
    inp_dict.update({"inp_dict_Flux":inp_dict_Flux})
    return inp_dict

def lambda_angstroms_to_energy_erg(lambda_arr):
    light_speed_cgs = 2.99792458e+10 # in cm / s
    nu_arr          = light_speed_cgs / (lambda_arr * 1e-8)
    h_planck        = 6.62607004e-27 # in  erg s
    E_arr_in_erg    = h_planck * nu_arr
    return E_arr_in_erg

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

def step_c_calculate_photon_count_for_noise_component(inp_dict):
    inp_dict_Flux         = inp_dict["inp_dict_Flux"]
    Flux_arr_cont_added   = inp_dict_Flux["Flux_arr_cont_added"]
    lambda_arr            = Flux_arr_cont_added[:,0]

    N_photon_source_arr   = photon_count_type_1_source_shot_noise(inp_dict)
    N_photon_sky_arr      = photon_count_type_2_sky_shot_noise(inp_dict)
    N_photon_dark_arr     = photon_count_type_3_dark_noise(inp_dict)
    N_photon_read_out_arr = photon_count_type_4_read_out_noise(inp_dict)
    N_photon_total_arr    = N_photon_source_arr + N_photon_sky_arr + N_photon_dark_arr + N_photon_read_out_arr
    #N_photon_total_arr    = N_photon_sky_arr 

    inp_dict_Flux.update({"N_photon_source_arr":N_photon_source_arr,\
                          "N_photon_sky_arr":N_photon_sky_arr,\
                          "N_photon_dark_arr":N_photon_dark_arr,\
                          "N_photon_read_out_arr":N_photon_read_out_arr,\
                          "N_photon_total_arr":N_photon_total_arr})
    inp_dict.update({"inp_dict_Flux":inp_dict_Flux})
    return inp_dict

def photon_count_type_1_source_shot_noise(inp_dict):
    inp_dict_Flux         = inp_dict["inp_dict_Flux"]
    Flux_arr_cont_added   = inp_dict_Flux["Flux_arr_cont_added"]
    photon_count_term_arr = inp_dict_Flux["photon_count_term_arr"]
    f_1500                = inp_dict["f_1500"]
    Flux_arr              = Flux_arr_cont_added[:,1]
    N_photon_source_arr   = Flux_arr * f_1500 * photon_count_term_arr
    return N_photon_source_arr

def photon_count_type_2_sky_shot_noise(inp_dict):
    inp_dict_Flux         = inp_dict["inp_dict_Flux"]
    Flux_arr_cont_added   = inp_dict_Flux["Flux_arr_cont_added"]
    photon_count_term_arr = inp_dict_Flux["photon_count_term_arr"]
    lambda_arr            = Flux_arr_cont_added[:,0]
    f_lambda_sky          = inp_dict["f_lambda_sky"]
    Flux_arr_sky          = np.ones(Flux_arr_cont_added.shape[0])
    N_photon_sky_arr      = Flux_arr_sky * f_lambda_sky * photon_count_term_arr
    return N_photon_sky_arr

def photon_count_type_3_dark_noise(inp_dict):
    N_photon_dark     = inp_dict["N_photon_dark"]
    t_exp             = inp_dict["t_exp"]
    N_pixel_per_ccd   = inp_dict["N_pixel_per_ccd"]
    N_photon_dark_arr = N_photon_dark * t_exp * N_pixel_per_ccd
    return N_photon_dark_arr

def photon_count_type_4_read_out_noise(inp_dict):
    N_pixel_per_ccd       = inp_dict["N_pixel_per_ccd"]
    N_read_out_epoch      = inp_dict["N_read_out_epoch"]
    N_photon_read_out     = inp_dict["N_photon_read_out"]
    N_photon_read_out_arr = (N_pixel_per_ccd * N_read_out_epoch) * N_photon_read_out**2.0
    return N_photon_read_out_arr


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

def step_d_generate_noise_realization(inp_dict):
    inp_dict_Flux         = inp_dict["inp_dict_Flux"]
    N_photon_source_arr   = inp_dict_Flux["N_photon_source_arr"]
    N_photon_total_arr    = inp_dict_Flux["N_photon_total_arr"]
    lambda_arr            = inp_dict_Flux["Flux_arr_cont_added"][:,0]
    N_photon_poisson_arr  = np.random.poisson(N_photon_total_arr)
    inp_dict_Flux.update({"N_photon_poisson_arr":N_photon_poisson_arr})
    inp_dict.update({"inp_dict_Flux":inp_dict_Flux})
    return inp_dict

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

def step_e_convert_N_photon_to_flux(inp_dict):
    inp_dict_Flux          = inp_dict["inp_dict_Flux"]
    N_photon_poisson_arr   = inp_dict_Flux["N_photon_poisson_arr"]
    Flux_arr_cont_added    = inp_dict_Flux["Flux_arr_cont_added"]
    photon_count_term_arr  = inp_dict_Flux["photon_count_term_arr"]
    f_1500                 = inp_dict["f_1500"]
    Flux_arr               = Flux_arr_cont_added[:,1]
    Flux_arr_with_noise    = N_photon_poisson_arr / photon_count_term_arr / f_1500
    Flux_arr_final         = np.zeros((Flux_arr.shape[0],4))
    Flux_arr_final[:,0]    = Flux_arr_cont_added[:,0]
    #Flux_arr_final[:,1]    = Flux_arr_with_noise # Ucomment this line and delete following line for noise addition
    Flux_arr_final[:,1]    = Flux_arr_cont_added[:,1]
    Flux_arr_final[:,3]    = Flux_arr_cont_added[:,2]
    inp_dict_Flux.update({"Flux_arr_final":Flux_arr_final})
    inp_dict.update({"inp_dict_Flux":inp_dict_Flux})
    return inp_dict


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

def simple_noise_addition_routines_main(inp_dict):
    inp_path        = inp_dict["inp_path_param"]
    inp_file        = inp_dict["inp_file_noise"]
    lambda_arr_obs  = inp_dict["lambda_arr_obs"]
    Flux_resamp_arr = inp_dict["Flux_resamp_arr"]
    noise_data_arr  = np.loadtxt(inp_path + inp_file)
    Flux_noise_arr  = np.zeros(Flux_resamp_arr.shape[0])
    N_sigma         = 3.0

    for noise_indx in range(noise_data_arr.shape[0]):
        lambda_min_noise = noise_data_arr[noise_indx,0]
        lambda_max_noise = noise_data_arr[noise_indx,1]
        SNR_per_pixel    = noise_data_arr[noise_indx,2]
        noise_background = noise_data_arr[noise_indx,3] * N_sigma
        bool_arr_region  = (lambda_min_noise <= lambda_arr_obs) & (lambda_arr_obs <= lambda_max_noise) 
        N_pixel          = np.sum(bool_arr_region)

        Flux_resamp_arr_region = Flux_resamp_arr[bool_arr_region]
        bool_arr               = Flux_resamp_arr_region > noise_background
        noise_arr              = np.random.normal(0.0, 1.0 / SNR_per_pixel,N_pixel)
        Flux_resamp_arr_region[bool_arr]  = Flux_resamp_arr_region[bool_arr] * (1.0 + noise_arr[bool_arr])
        Flux_resamp_arr_region[~bool_arr] = Flux_resamp_arr_region[~bool_arr] + noise_arr[~bool_arr]
        Flux_noise_arr[bool_arr_region] = Flux_resamp_arr_region

    out_dict_flux_noise = {"Flux_noise_arr":Flux_noise_arr,"lambda_arr_obs":lambda_arr_obs}
    return out_dict_flux_noise
