from __future__ import print_function
import numpy as np
import matplotlib.pyplot as plt
from math import log10

from WEAVEify_routines.step_04_optical_depth_calculation_METALS.method_1_cloudy_interpolation.\
cloudy_table_interpolation_routines import cloudy_table_interpolation_routines_main

from WEAVEify_routines.step_02_common_routines.density_calculation_routines import \
hydrogen_number_density_numpy

from WEAVEify_routines.step_02_common_routines.main_routines_tau import \
generate_od_data_main_all_transitions,extract_metal_specie_ion_list,\
generate_od_data_main_using_scaling

from WEAVEify_routines.step_02_common_routines.tau_summation_routines import \
tau_summation_routines_main


def method_1_generate_tau_metal_by_cloudy_interpolation_main(inp_dict):
    inp_dict                = generate_log_metallicity_field(inp_dict)
    inp_dict                = step_1_generate_metal_specie_fraction_from_cloudy_tables(inp_dict)
    specie_trans_list_metal = inp_dict["metal_trans_list"]
    flag_tau_each_trans     = inp_dict["flag_tau_each_trans"]

    if flag_tau_each_trans == "y":
        inp_dict_tau_metal  = generate_od_data_main_all_transitions(inp_dict,specie_trans_list_metal)
    elif flag_tau_each_trans == "n":
        inp_dict_tau_metal  = generate_od_data_main_using_scaling(inp_dict,specie_trans_list_metal)

    inp_dict.update({"inp_dict_tau_metal":inp_dict_tau_metal})

    inp_dict = tau_summation_routines_main(inp_dict,flag_absorber = "METAL")

    return inp_dict

def step_1_generate_metal_specie_fraction_from_cloudy_tables(inp_dict_org):
    zmean_cloudy_arr      = np.round(np.arange(1.0,4.1,0.1),1) # Redshifts at which Cloudy Table generated
    metal_specie_ion_list = extract_metal_specie_ion_list(inp_dict_org)
    cosmo_param           = inp_dict_org["cosmo_param"]
    phy_const             = inp_dict_org["phy_const"]
    delta_z_cloudy        = np.median(zmean_cloudy_arr[1:] - zmean_cloudy_arr[:-1]) * 0.5
    T_gas_arr             = inp_dict_org["T_gas_arr"]
    Delta_arr             = inp_dict_org["Delta_arr"]
    log_Z_arr             = inp_dict_org["log_Z_arr"]
    z_arr                 = inp_dict_org["z_arr"]
    nH_arr                = hydrogen_number_density_numpy(Delta_arr,z_arr,cosmo_param,phy_const)
    log_nH_arr            = np.log10(nH_arr)
    log_T_arr             = np.log10(T_gas_arr)
    cloudy_inp_arr        = np.array([log_nH_arr,log_T_arr,log_Z_arr]).T
    inp_dict_specie        = {}


    for metal_specie_ion in metal_specie_ion_list:
        log_f_XI_arr  = np.zeros(log_nH_arr.shape[0]) * -9999.0
        for z_indx in range(zmean_cloudy_arr.shape[0]):
            zmean_val = zmean_cloudy_arr[z_indx]
            zmin_val  = zmean_cloudy_arr[z_indx] - delta_z_cloudy
            zmax_val  = zmean_cloudy_arr[z_indx] + delta_z_cloudy
            bool_arr  = (zmin_val <= z_arr) & (z_arr <= zmax_val)

            if np.sum(bool_arr) != 0:
                inp_dict = inp_dict_org.copy()
                inp_dict.update({"metal_specie_ion":metal_specie_ion,"zmean_val":zmean_val})
                cloudy_intrp_funct     = cloudy_table_interpolation_routines_main(inp_dict)
                log_f_XI_arr[bool_arr] = cloudy_intrp_funct(cloudy_inp_arr[bool_arr,:])

        key_name = "log_f_" + metal_specie_ion + "_arr"
        inp_dict_specie.update({key_name:log_f_XI_arr})

    inp_dict_org.update({"metal_specie_ion_list":metal_specie_ion_list})
    inp_dict_org.update(inp_dict_specie)

    return inp_dict_org

def generate_log_metallicity_field(inp_dict_org):
    key_list_field = inp_dict_org.keys()
    if "log_Z_arr" not in key_list_field:
        log_Z_val_const = float(inp_dict_org["log_Z_val_const"])
        Delta_arr       = inp_dict_org["Delta_arr"]
        N_pixel         = Delta_arr.shape[0]
        log_Z_arr       = np.ones(N_pixel) * log_Z_val_const
        inp_dict_org.update({"log_Z_arr":log_Z_arr})
    return inp_dict_org





