#!/usr/bin/env python

import os
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"
import warnings
warnings.filterwarnings("ignore")
import logging
import ebf
import json
import optparse
from time import perf_counter as clock
import multiprocessing

import GalCraft.modules.SSPLoader as SSPLoader
import GalCraft.modules.Binner as Binner
import GalCraft.modules.SpecGenerator as SpecGenerator
import GalCraft.modules.CubeMaker as CubeMaker
import GalCraft.modules.Cot as Cot
import GalCraft.modules.utils as utils
from GalCraft._version import __version__


def run_GalCraft(CommandOptions):



    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    # - - - - - - - -   I N I T I A L I Z E   G A L C R A F T - - - - - - - - - - -
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    t_init = clock()

    # - - - - - - - - - - INITIALIZATION - - - - - - - - - -

    # Initialize cube name and paths
    cube_name = CommandOptions.configName
    setup_cube_name = 'setup_' + cube_name
    configDir, modelDir, templateDir, outputDir = utils.initialize_Dirs(CommandOptions)
    # Load the Configuration file
    with open(configDir + setup_cube_name + '.json', 'r') as f:
        params = json.load(f)
    # Setup the outputs folder
    filepath = outputDir + cube_name + '/'
    if os.path.exists(filepath)==False: os.mkdir(filepath[:-1])
    # Setup logger
    if params['other_params']['mode'] == 'continue': # For continue mode, change the logger mode from 'w' to 'a'
        params['other_params']['log_mode'] = 'a'
    utils.setupLogfile(logfile=filepath + 'outputs.log', __version__=__version__, mode=params['other_params']['log_mode'], welcome=True)
    logging.info('Loaded the setup file (path %s)' % (configDir + setup_cube_name + '.json'))
    logging.info('Outputs DIR: %s' % filepath)
    # Setup cpu cores to use
    if params['other_params']['ncpu'] == None:
        # n_cores = len(os.sched_getaffinity(0))
        n_cores = multiprocessing.cpu_count()
        logging.info('Change the CPU cores from %s to %s due to the CPU availability.' % (params['other_params']['ncpu'], n_cores) )
        params['other_params']['ncpu'] = n_cores
    else:
        logging.info('Run the framework using %s CPU cores.' % params['other_params']['ncpu'])



    # - - - - - - - - - - INSTRUMENT SETUP - - - - - - - - - -

    # Obtain ./instrument/ dir
    this_dir, this_filename = os.path.split(__file__)
    instrumentDir = os.path.join(this_dir, "instrument") + '/'
    # Load instrument information
    # If "instrument"=="MUSE-WFM", it will load "MUSE-WFM.json" and use the "spatial_resolution" and "spatial_bin";
    # If "instrument"=="DIY", you will design your own instrument using "spatial_resolution" and "spatial_bin" in the configuration file;
    # If "instrument"=="DEFAULT", it will generate a huge cube using all the particles and only will be used in the configuration file.
    inst = params['cube_params']['instrument']
    if inst.upper() != 'DIY' and inst.upper()!='DEFAULT':
        with open(instrumentDir + inst + '.json', 'r') as f:
            inst_params = json.load(f)
        for key in inst_params:
            params['cube_params'][key] = inst_params[key]
        logging.info('INST mode, apply the %s instrument spatial properties from the preset file.' % inst)
    elif inst.upper() == 'DEFAULT':
        logging.info('DEFAULT mode, will use all the particles and only the given spatial resolution will be used.')
    else:
        logging.info("DIY mode, apply the given spatial resolution and bin number.")



    # - - - - - - - - - - LOAD SPECTRAL TEMPLATES - - - - - - - - - -

    logging.info('Loading SSP models...')
    t = clock()
    ssp_model = SSPLoader.model(templateDir, instrumentDir, params['ssp_params'], params['other_params'])
    ssp_model.oversample()
    logging.info('SSP model has been successfully loaded, time elapsed: %.2f s' % (clock() - t))



    # - - - - - - - - - - LOAD STELLAR CATALOG - - - - - - - - - -

    logging.info('Loading the E-Galaxia model from %s' % (modelDir + params['other_params']['model_name']))
    t = clock()
    d_t = ebf.read(modelDir + params['other_params']['model_name'],'/')
    # Check the need to locate/rotate the E-Galaxia model.
    if 'vr' in d_t.keys():
        logging.info('Model has already been rotated, no location/rotation transformation is applied.')
        for key in params['oparams']:
            params['oparams'][key] = None
    else:
        logging.info('Specifing the location/rotation of the galaxy')
        Cot.observe(d_t, params['oparams'])
    # Check extinction values
    if params['cube_params']['use_extinc'] == True:
        assert 'exbv' in d_t.keys(), "No extinction in the loaded model, process terminated."
    # Check the need to apply the flux calibration due to the total Galaxy Stellar mass.
    if params['other_params']['nparticles'] != None:
        nparticles = params['other_params']['nparticles']
    else:
        if 'shuffle' in params['other_params']['model_name']:
            original_model_name = '_'.join([x for x in params['other_params']['model_name'][:-4].split('_') if 'shuffle' not in x]) + '.ebf'
            d_original = ebf.read(modelDir + original_model_name,'/')
            nparticles = len(d_original['px'])
            del d_original
        else:
            nparticles = len(d_t['vr'])
    # In case nparticles changed due to the "debug" mode.
    if params['other_params']['mode'] == 'debug':
        nparticles = 500000
    logging.info('Number of particles = %s or Galaxy Stellar Mass' % nparticles)
    # Calculate mass_cali_factor for flux calibration
    if params['other_params']['gal_mass'] != None:
        mass_cali_factor = params['other_params']['gal_mass'] / nparticles
        logging.info('Input Galaxy Stellar mass is %s, then the calibration factor is %s.' % (params['other_params']['gal_mass'], mass_cali_factor))
    else:
        logging.info('No flux calibration due to Galaxy Stellar mass is applied.')
        mass_cali_factor = 1
    logging.info('E-Galaxia model has been successfully loaded, time elapsed: %.2f s' % (clock() - t))




    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    # - - - - - - - - - -   A N A L Y S I S   M O D U L E S   - - - - - - - - - - -
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



    # - - - - - - - - - - SPATIAL BINNING - - - - - - - - - -

    if params['other_params']['mode'] != 'continue':
        logging.info('Start binning the model...')
        t = clock()
        d_t_l, statistic_count_l, x_edges_l, y_edges_l = Binner.spatial_binner(d_t, params['cube_params'], params['other_params'],
                                                                               ssp_model.age_grid, ssp_model.metal_grid, ssp_model.alpha_grid,
                                                                               filepath, configDir + setup_cube_name,
                                                                               params['oparams']['distance'], nparticles)
        logging.info('Binning process has been finished, time elapsed: %.2f s' % (clock() - t))
    else:
        logging.info('Continue the running from previous run...')
        t = clock()
        d_t_l, statistic_count_l, x_edges_l, y_edges_l = Binner.spatial_binner_continue(params['cube_params'], params['other_params'],
                                                                                        filepath, configDir + setup_cube_name)
        logging.info('Binning process has been finished, time elapsed: %.2f s' % (clock() - t))
    d_t = None



    # - - - - - - - - - - START DATACUBE LOOP  - - - - - - - - - -

    for cube_idx in range(len(statistic_count_l)):
        d_t = d_t_l[cube_idx]
        statistic_count = statistic_count_l[cube_idx]
        x_edges = x_edges_l[cube_idx]
        y_edges = y_edges_l[cube_idx]



    # - - - - - - - - - - GENERATE DATACUBE  - - - - - - - - - -

        if params['other_params']['mode'] == 'continue' and os.path.exists(filepath + 'data_cube_' + str(cube_idx) + '.fits'):
            logging.info("Data Cube No.%s exists in %s, will not be run in 'continue' mode" % (cube_idx+1, filepath))
            continue
        else:
            logging.info('Start generating the datacube No.%s...' % (cube_idx+1))
            t = clock()
            Spec_DataCube = SpecGenerator.Spec_Generator(d_t, x_edges, y_edges, statistic_count, ssp_model,
                                                          params['cube_params'], params['ssp_params'], params['other_params'],
                                                          filepath, cube_idx)
            Spec_DataCube()
            data_cube = Spec_DataCube.data_cube
            data_cube = data_cube * mass_cali_factor
            logging.info('Calibrated the flux based on the input Galaxy Stellar Mass=%s.' % params['other_params']['gal_mass'])
            # Free memory
            Spec_DataCube = None
            d_t = None
            statistic_count = None
            logging.info('Stellar spectra have been generated, time elapsed: %.2f s' % (clock() - t))


    # - - - - - WRITE DATACUBE  - - - - -

        CubeMaker.write_cube(data_cube, params, x_edges, y_edges, ssp_model.new_wave, filepath, cube_idx, ssp_model.velscale, __version__)
        logging.info('Data Cube No.%s has been written in %s, total time elapsed: %.2f s' % (cube_idx+1, filepath , clock() - t))




    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    # - - - - - - - -  F I N A L I S E   T H E   A N A L Y S I S  - - - - - - - - -
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



    # - - - - - SAVE CONFIG  - - - - -

    logging.info('Write the json file into the outputs folder.')
    with open(filepath + 'setup.json', 'w') as f:
        json.dump(params, f, indent=2)
    logging.info('All the datacubes have been generated successfully, total time elapsed: %.2f s' % (clock() - t_init))
    logging.info("GalCraft completed successfully.")







# ============================================================================ #
#                           M A I N   F U N C T I O N                          #
# ============================================================================ #
def main(args=None):

    # Capture command-line arguments
    parser = optparse.OptionParser(usage="%GalCraft [options] arg")
    parser.add_option("--config",      dest="configName", type="string", \
            help="State the name of the MasterConfig file.")
    parser.add_option("--default-dir", dest="defaultDir", type="string", \
            help="File defining default directories for input, outputs, configuration files, and spectral templates.")
    (CommandOptions, args) = parser.parse_args()

    # Check if required command-line argument is given
    assert CommandOptions.configName != None, "Please specify the path of the Configuration.json file to be used. Exit!"

    # Run the framework
    run_GalCraft(CommandOptions)


if __name__ == '__main__':
    main()
