from __future__ import print_function
import time
start = time.time()

import numpy as np
import sys,os
import matplotlib.pyplot as plt

sys.path.append(sys.path[0])

# --------------- ASTRO_MOCKS Code Module call starts ----------------

from ASTRO_MOCKS_routines.step_02_common_routines.specie_transition_dictionary import \
generate_XI_forest_line_dict

from ASTRO_MOCKS_routines.step_02_common_routines.physics_constant import physical_constants

from ASTRO_MOCKS_routines.step_02_common_routines.echo_functions import flexible_input_from_user

from ASTRO_MOCKS_routines.step_02_common_routines.directory_functions import \
update_output_filenames

# --------------- ASTRO_MOCKS Code Module call ends ----------------


def step_a_initialize_astro_mocks_parameter(inp_dict={}):

    #-----------------
    # QSO Parameters
    #-----------------

    # Emission redshift of QSO
    z_em = 3.0

    # Name of the simulation suite
    sim_name = "SHERWOOD"

    # Length of simulation box in h^-1 cMpc, (Used only when optical depths are calculated)
    L_box = 80.0

    # Number of pixels in one skewer, (Used only when optical depths are calculated)
    N_Grid = 2048

    # spec_id -->> String or None (If string, data would be stored in files identified by qso_id)
    spec_id = None

    # Set this parameter if Auxilliary files are different location than given.
    inp_path_aux_files = "./ASTRO_MOCKS_routines/ASTRO_MOCKS_AUXILIARY_FILES/"

    # Set this parameter to path of directory where output needs to be stored.
    out_path = None

    # Set this parameter to "y" if you are using the code in parallel using MPI
    flag_mpi = "n"

    # Set N_proc to 1 rank to 0 as a  default value.
    # This variable will automatically change if MPI version of the code is used.
    N_proc   = 1
    rank     = 0

    # Following parameters are useful for checking the current status of the code.
    N_los_total = 1
    los_indx    = 1
    #-----------------

    # ----------------------------------
    # Flags to add absorbers optionally
    # ----------------------------------

    # flag__tau = None or "HI" or "metal" or "DLA" or "HI_metal" or "HI_DLA" or "HI_metal_DLA"
    flag_tau  = None

    # flag_absorber = "HI" or "metal" or "DLA" or "HI_metal" or "HI_DLA" or "HI_metal_DLA"
    flag_absorber = None

    # flag_test = String, Execute the code on test examples provided with code.
    flag_test = None
    # ----------------------------------


    # ----------------------
    # HI Forest Parameters
    # ----------------------

    # HI Transition list Ly-alpha + Ly-beta
    #HI_trans_list = ["HI-1025","HI-1215","HI-972","HI-949"]
    HI_trans_list = ["HI-1215"]

    # if "y" Computes optical depth for each tranision
    # if "n" Optical depth is calculated for one transition and then
    # scaled using ratio of oscillator strength in to rest wavelength
    flag_tau_each_trans = "n"

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

    # ----------------------------
    #  Metal addition Parameters
    # ----------------------------

    # "metal_profile" or "cloudy_interpolation"
    flag_metal_method = "cloudy_interpolation"

    # Used to model metal lines. Need of flag_metal_method is "cloudy_interpolation"
    uvb_model_metal = "FG20"

    # Used if flag_metal_method is "cloudy_interpolation" and log_Z_arr field is not supplied
    log_Z_val_const = -1.5

    # The list of available transitions are given in file
    # ASTRO_MOCKS_AUXILIARY_FILES/ABSOPRTION_TRANSITION_LIST/specie_transition_properties.txt
    metal_trans_list = ["SiIV-1402","CIV-1548","CIV-1550"]
    # ----------------------------

    # -------------------------
    #  DLA addition Parameters
    # -------------------------

    # List of log of HI Column Density, (If None DLA will not be added)
    # This is useful if users want to add multiple DLA along same sightline
    log_N_HI_DLA = 21.5

    # List of Redshift of DLA, (If None DLA will not be added)
    # This is useful if users want to add multiple DLA along same sightline
    z_DLA = 2.7


    # List of b (km/s) parameter of DLA, (If None DLA will not be added)
    # This is useful if users want to add multiple DLA along same sightline
    b_DLA = 10.0
    # -------------------------

    # -------------------------------
    # Wavelength Sampling Parameters
    # -------------------------------

    # String name of the Instrument
    instrument = "WEAVE"

    # String "LR" or "HR"
    resolution_mode = "LR"

    # Name of Wavelength Sampling Parameter File
    inp_file_lambda_samp = "wavelength_sampling_parameters.txt"

    # Float (Give some value if spectra_type = None else keep it as None)
    CRVAL1 = None

    # Float (Give some value if spectra_type = None else keep it as None)
    CDELT1 = None

    # Integer (Give some value if spectra_type = None else keep it as None)
    N_pixel_obs = None

    # String "LIN" or LOG" (Give some value if spectra_type = None else keep it as None)
    lambda_scale = None

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

    # ------------------------------
    # Parameters Controlling output
    # ------------------------------

    # Format in which outpus are saved Two options fits or hdf5
    flag_output_format = "fits" #"hdf5" #"fits"

    # If "y" Makes quick plot for analysis
    flag_make_quick_plot = "n"

    out_file_tau_base  = "ASTRO_MOCKS_Spectra" # Changed to facilitate users
    out_file_plot_base = "Quick_Plot_Mock_Spectra"
    out_file_flux_base = out_file_tau_base     # Do not change this variable 
    # ------------------------------

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

    # You do not need to change following lines in function
    cosmo_param          = cosmological_parameters()
    line_dict            = generate_XI_forest_line_dict(inp_path_aux_files)
    phy_const            = physical_constants()
    inp_dict             = vars()
    return inp_dict

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

def step_b_user_input_command_line(inp_dict_org={}):
    inp_dict_user = flexible_input_from_user()
    inp_dict_org.update(inp_dict_user)
    return inp_dict_org

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

def step_c_user_input_script(inp_dict_org={},inp_dict_script={}):
    if len(inp_dict_script) != 0:
        inp_dict_org.update(inp_dict_script)
    return inp_dict_org

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

def step_d_initialize_dependent_variables(inp_dict):
    inp_path_aux_files = inp_dict["inp_path_aux_files"]
    instrument         = inp_dict["instrument"]
    resolution_mode    = inp_dict["resolution_mode"]
    flag_output_format = inp_dict["flag_output_format"]
    out_file_flux_base = inp_dict["out_file_flux_base"]
    out_file_plot_base = inp_dict["out_file_plot_base"]
    out_file_tau_base  = inp_dict["out_file_tau_base"]
    out_file_flux      = out_file_flux_base + "." + flag_output_format
    out_file_tau       = out_file_tau_base + "." + flag_output_format

    spec_id_str        = generate_spec_id_str(inp_dict)
    out_file_plot      = out_file_plot_base + "-" + spec_id_str + ".pdf"


    # String None or "WEAVE_LR" or "WEAVE_HR"
    spectra_type = instrument + "_" + resolution_mode

    # Path of Wavelength Sampling Parameter File
    inp_path_lambda_samp = inp_path_aux_files + "/RESOLUTION_AND_LSF_PROPERTIES"\
                           + "/WEAVE/WAVELENGTH_SAMPLING_PARAMETERS/"

    # Location where telescope throughputs are stored
    inp_path_throughput = inp_path_aux_files + "NOISE_ADDITION/TYPE_3_OPR3_MOCKS/"

    # Location where transmission filters are stored
    inp_path_trans_filter = inp_path_throughput + "FILTER_TRANSMISSION_CURVE/"

    inp_dict.update({"spectra_type":spectra_type,\
                     "inp_path_lambda_samp":inp_path_lambda_samp,\
                     "inp_path_throughput":inp_path_throughput,\
                     "inp_path_trans_filter":inp_path_trans_filter,\
                     "out_file_flux":out_file_flux,\
                     "out_file_tau":out_file_tau,\
                     "out_file_plot":out_file_plot,\
                     })

    # Update File names if MPI version of the code is used.
    # Spectra are written seperately for each processor
    flag_mpi = inp_dict["flag_mpi"]

    if flag_mpi.upper() == "Y":
        rank                  = inp_dict["rank"]
        N_proc                = inp_dict["N_proc"]
        file_suffix           = "." + inp_dict["flag_output_format"]
        out_file_flux         = inp_dict["out_file_flux"]
        out_file_tau          = inp_dict["out_file_tau"]
        N_digits              = len(str(N_proc))
        file_suffix_new       = "_Rank_" + str(rank).zfill(N_digits) + file_suffix
        out_file_flux         = out_file_flux.replace(file_suffix,file_suffix_new)
        out_file_tau          = out_file_tau.replace(file_suffix,file_suffix_new)
        inp_dict.update({"out_file_flux":out_file_flux,\
                         "out_file_tau":out_file_tau})
    return inp_dict

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

def generate_spec_id_str(inp_dict):
    spec_id      = inp_dict["spec_id"]

    if spec_id is None:
        spec_id     = 0
        spec_id_str = str("%i"%spec_id)
    else:
        if isinstance(spec_id,str):
            spec_id_str = spec_id
        else:
            spec_id_str = str("%i"%spec_id)

    return spec_id_str

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

def cosmological_parameters():

    #--------------------- Cosmological Parameters -------------------------
    # Cosmology in Sherwood Simulation Suite is consistent with Planck+2014
    # LCDM Planck + WP + highL + BAO
    #-----------------------------------------------------------------------


    omega_m      = 0.308
    omega_l      = 0.692
    omega_r      = 0.0
    omega_k      = 0.0
    omega_b      = 0.0482
    hubble_param = 0.678
    sigma_8      = 0.829
    Y            = 0.24
    omega_bhsq   = omega_b * hubble_param**2.0

    # in gm / cm^3  rho_c (Present day critical density) -->> (3 x H(0)^2)  / (8 x pi x G)
    rho_c_by_hsq = 1.8791e-29

    ns           = 0.961

    # You do not need to change following lines in function
    cosmo_param  = vars()
    return cosmo_param

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



