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])

# --------------- WEAVEfy Code Module call starts ----------------

from WEAVEify_routines.step_02_common_routines.specie_transition_dictionary import \
generate_XI_forest_line_dict

from WEAVEify_routines.step_02_common_routines.physics_constant import physical_constants

from WEAVEify_routines.step_02_common_routines.echo_functions import flexible_input_from_user

from WEAVEify_routines.step_02_common_routines.directory_functions import \
update_output_filenames

from WEAVEify_routines.step_02_common_routines.dictionary_functions import  \
CHECK_AND_INITIALIZE_KEY_VALUE,CHECK_IF_KEY_PRESENT

# --------------- WEAVEfy Code Module call ends ----------------


def step_a_initialize_weavefy_parameter(inp_dict={}):

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

    # Emission redshift of QSO
    z_em = 3.0

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

    # Do not change following line and name of default_parameter_WEAVEify.py file.
    # This provision allows to call weave_mocks.py from any location
    inp_path_weave_mock = os.path.realpath(__file__).replace("default_parameters_WEAVEify.py","")

    # Set this parameter if Auxilliary files are different location than given.
    inp_path_aux_files  = inp_path_weave_mock + "WEAVEify_routines/WEAVE_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

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

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

    # -------------------------------
    # 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

    # String Functional form of Line Spread Function
    lsf_type = "GAUSSIAN"
    # -------------------------------

    # -------------------------------
    # Continuum Placement Parameters
    # -------------------------------

    # Select the QSO continuum Placement Method Currently supports ony PCA
    #flag_cont_method = "template" # "PCA"
    flag_cont_method = "PCA"

    # integer or None. If integer is specified, the continuum template is
    # generated using this random seed. In such cases continuum will be random
    # only when random_seed_cont is changed.
    random_seed_cont = None

    # Power-law slope of the continuum set between -0.0 to -3.0. If None
    # the power-law index is generated randomly
    alpha_cont = None #-0.0 -0.5,-3.0

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

    # --------------------------
    # Noise Addition Parameters
    # --------------------------

    # None or Gaussian or Poisson or OpR3
    flag_noise  = "OpR3" # None

    # Type of source "point" or "extended"
    source_type = "point"

    # Diameter of Telescope in meters
    dia_telescope = 4.2

    # Diameter of Fibre in arcsec
    dia_fibre_arcsec = 1.3045

    # Appearent magnitude of source
    m_source = 19.0

    # Wave band in which source/object magnitude is specified
    wave_band_source = "R"

    # Appearent magnitude of sky
    m_sky = 20.92

    # Wave band in which sky magnitude is specified
    wave_band_sky = "B"

    # Exposure time in seconds
    t_exp = 3060.0

    # Airmass value for the atmosphere
    airmass = 1.107

    # Profile for seeing Gaussian or Moffat
    seeing_profile = "Moffat"

    # Moffat profile index
    beta_moffat = 2.5

    # Full width at half maximum of seeing in arcsec
    fwhm_seeing = 0.75
    PRI_fwhm = 0.45

    # Read noise per electron
    read_noise = 3.3

    # Noise due to dark current
    dark_current = 0.0

    # CCD pixel size in microns
    ccd_pix_size_micron = 15.0

    # Fiber injection f/ratio (Focal Plane)
    fFP = 3.2

    # Collimator f/ratio
    fcol = 3.1

    # Camera f/ratio
    fcam = 1.8

    # Beam size in degrees
    beamsize = 180.0

    # offset of fibre from obsject center in arcsec
    r_offset = 0.1

    # Scaling factor for Uncertainty on SNR per pixel. If f_scale_snr_std = 1
    # then uncertainty on SNR per pixel will  be similar to OpR3 mocks
    # If f_scale_snr_std = 0.0 there will not be any uncertainty on SNR per pixel
    # f_scale_snr_std can take any values >= 0.
    f_scale_snr_std = 0.0

    # Name of the telescope throughput file
    inp_file_throughput = "telescope_throughput_V01.txt"

    # Which tyoe of telescope throughput to use
    throughput_type     = "optimal"


    # Name of transmission filters file
    inp_file_trans_filter = "BESSELL_Transmission_Filter.txt"

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

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

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

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

    out_file_plot_base = "Quick_Plot_Mock_Spectra"
    out_file_flux_base = "WEAVEified_Mock_Spectra"
    out_file_tau_base  = "WEAVEified_Mock_Spectra" # Changed to facilitate users
    # ------------------------------

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

    # 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 step_e_reorganize_input_for_weaveify(inp_dict):
    CHECK_IF_KEY_PRESENT(["wavelength","flux"],inp_dict)
    wavelength    = inp_dict["wavelength"]
    flux          = inp_dict["flux"]
    inp_dict_Flux = {"wavelength_total_arr":wavelength,"Flux_total_arr":flux}
    inp_dict.update({"inp_dict_Flux":inp_dict_Flux})
    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

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



