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

from WEAVEify_routines.step_02_common_routines.echo_functions import \
find_index_of_string_character

from WEAVEify_routines.step_02_common_routines.sub_routines_tau import \
sub_routines_tau_primary

from WEAVEify_routines.step_02_common_routines.density_calculation_routines import \
hydrogen_number_density_numpy

from WEAVEify_routines.step_02_common_routines.echo_functions import \
process_completion_print_function

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

def generate_od_data_main_using_scaling(inp_dict,specie_trans_list):
    #specie_trans_list = inp_dict["specie_trans_list"]
    inp_dict_tau          = {}
    specie_group_dict     = extract_single_specie_trans_from_list(inp_dict,specie_trans_list)
    N_group               = len(specie_group_dict)
    specie_group_key_list = specie_group_dict.keys()

    for specie_group_key in specie_group_key_list:
        specie_group_sub_dict   = specie_group_dict[specie_group_key]
        specie_trans_tau        = specie_group_sub_dict["specie_trans_tau"]
        specie_trans_list_group = specie_group_sub_dict["specie_trans_list_group"]
        tau_factor_list_group   = specie_group_sub_dict["tau_factor_list"]

        idx     = specie_trans_tau.find("-")
        process_completion_print_function(inp_dict,"Optical Depth calculation started for specie = " + specie_trans_tau[:idx])
        out_dict       = generate_od_data_for_given_specie_trans(inp_dict,specie_trans_tau)
        tau_XI_arr     = out_dict["tau_XI_arr"]
        wavelength_arr = out_dict["wavelength_arr"]

        N_specie       = len(specie_trans_list_group)

        for specie_indx in range(N_specie):
            specie_trans_loop   = specie_trans_list_group[specie_indx]
            tau_factor_val_loop = tau_factor_list_group[specie_indx]
            tau_XI_arr_loop     = tau_XI_arr * tau_factor_val_loop

            inp_sub_dict_tau = {"tau_" + specie_trans_loop + "_arr":tau_XI_arr_loop.astype(np.float32),\
                                "wavelength_arr":wavelength_arr.astype(np.float32)}
            inp_dict_tau.update(inp_sub_dict_tau)
        process_completion_print_function(inp_dict,"Optical Depth calculation completed for specie = " + specie_trans_tau[:idx])

    return inp_dict_tau

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

def extract_single_specie_trans_from_list(inp_dict,specie_trans_list):
    """ Suppose specie_trans_list has following values
    specie_trans_list = ["HI-1025","HI-1215","CIV-1548","CIV-1550","SiIV-1402"]

    This routine first organizes Same species in  one group. That is 
    there will be three groups for the above mentioned specie trans_list
    HI, CIV and SiIV

    Output directory would have following form

    {"HI":{"specie_trans_tau":"HI-1215",\
            "specie_trans_list":["HI-1215","HI-1025"],\
            "tau_factor_list":[1.0,0.16]},\
    "CIV":{"specie_trans_tau":"CIV-1548",\
            "specie_trans_list":["CIV-1548","CIV-1550"],\
            "tau_factor_list":[1.0,0.5]},\
    "SiIV":{"specie_trans_tau":"SiIV-1402",\
            "specie_trans_list":["SiIV-1402"],\
            "tau_factor_list":[1.0]},\
    }
    """
    line_dict         = inp_dict["line_dict"]
    specie_list       = find_unique_specie_list(specie_trans_list)
    N_group           = len(specie_list)
    specie_group_dict = {}

    for group_indx in range(N_group):
        specie                  = specie_list[group_indx]
        specie_trans_list_group = []

        for specie_trans in specie_trans_list:
            idx = specie_trans.find("-")
            if specie == specie_trans[:idx]:
                specie_trans_list_group.append(specie_trans)

        specie_group_sub_dict = find_optical_depth_scaling_factor_for_given_specie_and_all_transition(line_dict,specie_trans_list_group)
        specie_group_dict.update({specie:specie_group_sub_dict})
    return specie_group_dict


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

def find_optical_depth_scaling_factor_for_given_specie_and_all_transition(line_dict,specie_trans_list_group):
    N_specie        = len(specie_trans_list_group)
    f_lu_lambda_arr = np.zeros(N_specie)

    for specie_indx in range(N_specie):
        specie_trans                 = specie_trans_list_group[specie_indx]
        f_lu_val                     = line_dict[specie_trans]["f_lu"]
        lambda_rest_val              = line_dict[specie_trans]["lambda_rest"]
        f_lu_lambda_arr[specie_indx] = f_lu_val * lambda_rest_val

    idx               = np.argmax(f_lu_lambda_arr)
    f_lu_lambda_max   = f_lu_lambda_arr[idx]
    specie_trans_tau  = specie_trans_list_group[idx]
    tau_factor_arr    = f_lu_lambda_arr / f_lu_lambda_max
    out_dict          = {"specie_trans_tau":specie_trans_tau,\
                         "specie_trans_list_group":specie_trans_list_group,\
                         "tau_factor_list":tau_factor_arr.tolist()}

    return out_dict


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

def find_unique_specie_list(specie_trans_list):
    specie_list = []
    for specie_trans in specie_trans_list:
        idx        = specie_trans.find("-")
        specie_list.append(specie_trans[:idx])
    return list(set(specie_list))

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

def generate_od_data_main_all_transitions(inp_dict,specie_trans_list):
    #specie_trans_list = inp_dict["specie_trans_list"]
    inp_dict_tau      = {}
    for specie_trans in specie_trans_list:
        print("")
        print("Optical Depth calculation started for specie_trans =",specie_trans)
        out_dict       = generate_od_data_for_given_specie_trans(inp_dict,specie_trans)
        tau_XI_arr     = out_dict["tau_XI_arr"]
        wavelength_arr = out_dict["wavelength_arr"]
        inp_sub_dict_tau = {"tau_" + specie_trans + "_arr":tau_XI_arr.astype(np.float32),\
                            "wavelength_arr":wavelength_arr.astype(np.float32)}
        inp_dict_tau.update(inp_sub_dict_tau)
        #print
        #print "Optical Depth calculation done for specie_trans =",specie_trans
        #print
    return inp_dict_tau

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

def generate_od_data_for_given_specie_trans(inp_dict,specie_trans):
    # Do not change variable names in this routine
    # You can change variable values
    flag_mpi       = "n"
    Delta_arr      = inp_dict["Delta_arr"]
    T_XI_arr       = inp_dict["T_gas_arr"]
    v_XI_arr       = inp_dict["v_gas_arr"]
    z_arr_sim      = inp_dict["z_arr"]

    if "HI" in specie_trans:
        f_XI_arr     = inp_dict["f_HI_arr"]
    else:
        char_indx    = specie_trans.find("-")
        specie_ion   = specie_trans[:char_indx] 
        log_f_XI_arr = inp_dict["log_f_" + specie_ion + "_arr"]
        f_XI_arr     = 10**log_f_XI_arr

    cosmo_param  = inp_dict["cosmo_param"]
    phy_const    = inp_dict["phy_const"]
    nH_arr       = hydrogen_number_density_numpy(Delta_arr,z_arr_sim,cosmo_param,phy_const)
    n_XI_arr     = nH_arr * f_XI_arr
    hubble_param = cosmo_param["hubble_param"]
    L_box        = inp_dict["L_box"]
    N_Grid       = inp_dict["N_Grid"]
    line_dict    = inp_dict["line_dict"]
    phy_const    = inp_dict["phy_const"]
    flag_mpi     = inp_dict["flag_mpi"]
    rank         = inp_dict["rank"]
    spec_id      = inp_dict["spec_id"]
    N_los_total  = inp_dict["N_los_total"]
    los_indx     = inp_dict["los_indx"]
    inp_dict_tau = vars()
    out_dict     = sub_routines_tau_primary(inp_dict_tau)
    return out_dict

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

def extract_specie_list(inp_dict,specie="HI"):
    specie_trans_list      = inp_dict["specie_trans_list"]
    specie_trans_list_extr = []

    for specie_trans in specie_trans_list:

        if (specie == "HI") & ("HI-" == specie_trans[:3]):
            specie_trans_list_extr.append(specie_trans)

        elif (specie != "HI") & ("HI-" != specie_trans[:3]):
            specie_trans_list_extr.append(specie_trans)

    return specie_trans_list_extr

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

def extract_metal_specie_ion_list(inp_dict):
    metal_trans_list      = inp_dict["metal_trans_list"]
    metal_specie_ion_list = []

    for metal_trans in metal_trans_list:
        indx = find_index_of_string_character(metal_trans)
        metal_specie_ion_list.append(metal_trans[:indx])

    return list(set(metal_specie_ion_list))

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

def extract_metal_specie_ion_list_test():
    from echo_functions import find_index_of_string_character
    metal_trans_list = ["SiIV-1402","CIV-1548","CIV-1550"]
    inp_dict         = vars()
    print(extract_metal_specie_ion_list(inp_dict))

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


#extract_metal_specie_ion_list_test()
