'''Main program:
Usage: run_simulator(file_param)  \n
If no param file is specified, the default one is exemple/params_exemple.txt \n
In the first part of the program, model coordinates are read and the
SKIM swath is computing accordingly. \n
The SKIM grid parameters are saved in netcdf files, if you don't want to
recompute them, set maksgrid (in params file) to False.\n

In the second part of the program, errors are computed on SKIM grid for
each pass, for each cycle. The error free velocity is the velocity interpolated
from the model at each timestep. Note that there is no temporal interpolation
between model files and thus if several files are used in the velocity
interpolation, some discontinuities may be visible. \n

OUTPUTS are netcdf files containing the requested errors, the error free
radial velocity and the radial velocity with errors. There is one file every
pass and every cycle.

\n
\n
Copyright (C) 2017-2021 OceanDataLab
This file is part of skimulator.

skimulator is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

skimulator is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with skimulator.  If not, see <http://www.gnu.org/licenses/>.
'''
import os
import numpy
import glob
import sys
import time
import datetime
import logging
import traceback
import pkg_resources
from . import build_swath
from . import rw_data
#import skimulator.build_error as build_error
#import skimulator.build_error2 as build_error
from .error import generator
from . import mod_tools
from . import mod_run as run
from . import mod_parallel
# Define logger level for debug purposes
logger = logging.getLogger(__name__)
#logger = multiprocessing.log_to_stderr()
#logger.setLevel(logging.DEBUG)

# - Define global variables for progress bars
istep = 0
ntot = 1
ifile = 0


def run_simulator(p, die_on_error=False):
    '''Main routine to run simulator, input is the imported parameter file,
    no outputs are returned but netcdf grid and data files are written as well
    as a skimulator.output file to store all used parameter.
    '''
    # - Initialize some parameters values
    timestart = datetime.datetime.now()
    mod_tools.initialize_parameters(p)
    mod_tools.check_path(p)
    model = p.model
    '''
    p.timeshift = getattr(p, 'timeshift', 0)
    if p.shift_time is None:
        p.timeshift = 0
    '''

    # - Read list of user model files """
    if p.file_input is not None:
        model_data, list_file = run.load_coordinate_model(p)
    # if no modelbox is specified (modelbox=None), the domain of the input
    # data is taken as a modelbox
    # coordinates from the region defined by modelbox are selected
    if p.modelbox is not None:
        modelbox = numpy.array(p.modelbox, dtype='float')
        # Use convert to 360 data
        modelbox[0] = (modelbox[0] + 360) % 360
        if modelbox[1] != 360:
            modelbox[1] = (modelbox[1] + 360) % 360
    else:
        if p.file_input is not None:
            modelbox = model_data.calc_box(p)
        else:
            logger.error('modelbox should be provided if no model file is'
                         'provided')
            sys.exit(1)
    # - Extract data on modelbox
    # TODO: do only this step if modelbox is defined? Do it later?
    if p.file_input is not None:
        model_data.read_coordinates(p)
        # Select model data in the region modelbox
        model_data.model_index_lon = {}
        model_data.model_index_lat = {}
        model_data.model_index = {}
        model_data.vloncirc = {}
        for key in model_data.vlon.keys():
            _lon = + model_data.vlon[key]
            _lat = + model_data.vlat[key]
            if p.grid == 'regular':
                if modelbox[0] < modelbox[1]:
                    _tmp = numpy.where(((modelbox[0]-1) <= _lon)
                                       & (_lon <= (modelbox[1]+1)))[0]
                    model_data.model_index_lon[key] = + _tmp
                else:
                    _tmp = numpy.where(((modelbox[0]-1) <= _lon)
                                       | (_lon <= (modelbox[1]+1)))[0]
                    model_data.model_index_lon[key] = + _tmp
                _tmp = numpy.where(((modelbox[2]-1) <= _lat)
                                   & (_lat <= (modelbox[3]+1)))[0]
                model_data.model_index_lat[key] = + _tmp
                model_data.vlon[key] = + _lon[model_data.model_index_lon[key]]
                model_data.vlat[key] = + _lat[model_data.model_index_lat[key]]
            else:
                if modelbox[0] < modelbox[1]:
                    _tmp = numpy.where(((modelbox[0]-1) <= _lon)
                                       & (_lon <= (modelbox[1]+1))
                                       & ((modelbox[2]-1) <= _lat)
                                       & (_lat <= (modelbox[3]+1)))
                    model_data.model_index[key] = + _tmp
                else:
                    _tmp = numpy.where(((modelbox[0]-1) <= _lon)
                                       | (_lon <= (modelbox[1]+1))
                                       & ((modelbox[2]-1) <= _lat)
                                       & (_lat <= (modelbox[3]+1)))
                    model_data.model_index[key] = + _tmp
            # prevent IDL issues
            _wlon = model_data.vlon[key]
            model_data.vloncirc[key] = numpy.rad2deg(numpy.unwrap(_wlon))

        model_data.model = model
    else:
        model_data = []
        list_file = []

    # avoid issue with 0=360 for global modelbox
    if modelbox[1] == 0:
        modelbox[1] = 359.99
    # - Make SKIM grid if necessary """
    if p.makesgrid is True:
        logger.info('\n Force creation of SKIM grid')
        # make nadir orbit
        orb = build_swath.makeorbit(modelbox, p, orbitfile=p.filesat)
        # build swath around this orbit
        ok = False
        try:
            ok = build_swath.orbit2swath(modelbox, p, orb, die_on_error)
        except mod_parallel.MultiprocessingError:
            logger.error('An error occurred with the multiprocessing '
                         'framework')
            traceback.print_exception(*sys.exc_info())
            sys.exit(1)
        except mod_parallel.DyingOnError:
            logger.error('An error occurred and all errors are fatal')
            sys.exit(1)

        if not ok:
            logger.error('Errors occurred while generating grid files')
            sys.exit(1)

        logger.info("\n SKIM Grids and nadir tracks have been written in"
                    "{}".format(p.outdatadir))
        logger.info("-----------------------------------------------")


    # - Compute interpolated velocity and errors for each pass, at each
    #   cycle
    logger.info('Compute interpolated velocity and errors:')
    #   load all SKIM grid files (one for each pass)
    listsgridfile = sorted(glob.glob(p.filesgrid + '_p*.nc'))
    if not listsgridfile:
        logger.error('\n There is no SKIM grid file in {}, run simulator with'
                     'option makesgrid set to true in your params'
                     'file'.format(p.outdatadir))
        sys.exit(1)
    # Build model time steps from parameter file
    modeltime = numpy.arange(0, p.nstep*p.timestep, p.timestep)
    #   Remove the grid from the list of model files
    if p.file_input and p.file_grid_model is None:
        logger.info("WARNING: the first file is not used to build data")
        list_file.remove(list_file[0])

    # - Loop on SKIM grid files
    jobs = []
    p2 = mod_tools.todict(p)
    #time_yaw = None
    #vac_yaw = None
    error_generator = generator.Generator(p)

    #if p.attitude is True and os.path.isfile(p.yaw_file):
    #    time_yaw, vac_yaw = build_error.load_yaw_aocs(p.yaw_file)
    #    # time_yaw = time_yaw / 86400
    for sgridfile in listsgridfile:
        jobs.append([sgridfile, p2, listsgridfile, list_file,
                     modelbox, model_data, modeltime, error_generator])
    ok = False
    try:
        ok = make_skim_data(p.proc_count, jobs, die_on_error, p.progress_bar)
    except mod_parallel.MultiprocessingError:
        logger.error('An error occurred with the multiprocessing framework')
        traceback.print_exception(*sys.exc_info())
        sys.exit(1)
    except mod_parallel.DyingOnError:
        logger.error('An error occurred and all errors are fatal')
        sys.exit(1)

    # - Write Selected parameters in a txt file
    timestop = datetime.datetime.now()
    timestop = timestop.strftime('%Y%m%dT%H%M%SZ')
    timestart = timestart.strftime('%Y%m%dT%H%M%SZ')
    rw_data.write_params(p, os.path.join(p.outdatadir,
                                         'skim_simulator.output'))
    if ok is True:
        if p.progress_bar is True:
            __ = mod_tools.update_progress(1, 'All passes have been processed',
                                           '')
        else:
            __ = logger.info('All passes have been processed')
        logger.info("\n Simulated skim files have been written in "
                    "{}".format(p.outdatadir))
        logger.info(''.join(['-'] * 61))
        sys.exit(0)
    logger.error('\nERROR: At least one of the outputs was not saved.')
    sys.exit(1)


def make_skim_data(_proc_count, jobs, die_on_error, progress_bar):
    """ Compute SWOT-like data for all grids and all cycle, """
    # - Set up parallelisation parameters
    proc_count = min(len(jobs), _proc_count)

    status_updater = mod_tools.update_progress_multiproc
    jobs_manager = mod_parallel.JobsManager(proc_count,
                                                       status_updater,
                                                       exc_formatter,
                                                       err_formatter)
    ok = jobs_manager.submit_jobs(worker_method_skim, jobs, die_on_error,
                                  progress_bar)

    if not ok:
        # Display errors once the processing is done
        jobs_manager.show_errors()

    return ok


def exc_formatter(exc):
    """Format exception returned by sys.exc_info() as a string so that it can
    be serialized by pickle and stored in the JobsManager."""
    error_msg = traceback.format_exception(exc[0], exc[1], exc[2])
    return error_msg


def err_formatter(pid, grid, cycle, exc):
    """Transform errors stored by the JobsManager into readable messages."""
    msg = None
    if cycle < 0:
        msg = '/!\ Error occurred while processing grid {}'.format(grid)
    else:
        _msg = '/!\ Error occurred while processing cycle {} on grid {}'
        msg = _msg.format(cycle, grid)
    return msg


def worker_method_skim(*args, **kwargs):
    msg_queue, sgridfile, p2, listsgridfile = args[:4]
    list_file, modelbox, model_data, modeltime, error_generator = args[4:]
    p = mod_tools.fromdict(p2)
    #   Load SKIM grid files (Swath and nadir)
    sgrid = run.load_sgrid(sgridfile, p)
    # duplicate SKIM grids to assure that data are not modified and are
    # saved properly
    sgrid_tmp = run.load_sgrid(sgridfile, p)

    sgrid.gridfile = sgridfile
    # Convert cycle in days
    sgrid.cycle /= 86400.
    sgrid_tmp.cycle /= 86400.
    sgrid_tmp.indi = None
    # Select model data around the swath to reduce interpolation cost in
    # griddata ### TODO comment to be removed?
    # - Generate SKIM like data
    # Compute number of cycles needed to cover all nstep model timesteps
    rcycle = (p.timestep * p.nstep)/float(sgrid.cycle)
    ncycle = int(rcycle)
    error = generator.Generator(p)
    first_time = datetime.datetime.strptime(p.first_time, '%Y-%m-%dT%H:%M:%SZ')
    #  Loop on all cycles
    for cycle in range(0, ncycle+1):
        # if ifile > (p.nstep*p.timestep + 1):
        #    break
        # Create SKIM-like data
        msg_queue.put((os.getpid(), sgridfile, cycle + 1, None))
        if p.file_input is None:
            model_data = []
        # Initialize all list of variables (each beam is appended to the
        # list)
        # Initialize velocity, indices and mask to empty lists
        output_var = {}
        #for key in p.list_output:
        #    output_var[key] = []
        list_key = list(p.list_input_var.keys())
        list_key.append('vindice')
        for key in list_key:
            output_var[key] = []
        # TODO: proof
        # sigma0 complusory
        # uwnd, vwnd compulsory for Instrument and WaveDoppler
        # hs uwnd vwnd mssu mssc mssclose for WaveDoppler
        # Loop over the beams to interpolate model data
        time_all = []
        for i in range(len(p.list_pos) + 1):
            sgrid_tmp.lon = sgrid.lon[i]
            sgrid_tmp.lat = sgrid.lat[i]
            sgrid_tmp.time = sgrid.time[i]
            time = sgrid_tmp.time / 86400. + sgrid.cycle * cycle
            # Interpolate observation on SKIM swath
            create = run.interp_on_swath(cycle, list_file,
                                             modelbox, sgrid_tmp,
                                             model_data, modeltime, p)
            output_var_i, time = create
            time_all.append(time)
            for key in output_var_i.keys():
                output_var[key].append(output_var_i[key])
        for key, value in output_var.items():
            output_var[key] = numpy.transpose(numpy.array(value))
        time_arr = numpy.transpose(numpy.array(time_all))
        radial_angle = numpy.transpose(sgrid.radial_angle)
        ac_angle = numpy.array(sgrid.angle)
        radial_angle = numpy.array(sgrid.radial_angle)
        output_var.update(run.compute_geov(output_var, ac_angle, radial_angle,
                                           p))

        dic_grid = {'lon': sgrid.lon, 'lat': sgrid.lat, 'incl': sgrid.incl,
                    'angle': sgrid.angle, 'ipass': sgrid.ipass}

        #ac_angle m
        ## TODO remove sgrid.x_al_nadir variable
        all_error = error.generate(sgrid.x_al, output_var, sgrid.angle,
                                   time_arr, sgrid.cycle, first_time,
                                   radial_angle,
                                   sgrid.x_al[:, 0], sgrid.x_ac)
        # for key, value in output_var.items():
        #    output_var[key] = numpy.transpose(value)
        for k, v in all_error.items():
            output_var[k] = v

        # Compute correction
        p.delta_azim = 15
        all_corr = error.correct(dic_grid, output_var, radial_angle,
                                 p.delta_azim)
        for k, v in all_corr.items():
            output_var[k] = v


        # Add error to true interpolated fields
        obs = error.generate_obs(output_var, all_error, all_corr)
        for k, v in obs.items():
            output_var[k] = v

        # Save outputs in a netcdf file
        if ((~numpy.isnan(numpy.array(output_var['vindice']))).any()
              or not p.file_input):
            sgrid.ncycle = cycle
            run.save_SKIM(cycle, sgrid, time_all, output_var, p)
        del time

    if p.file_input is not None:
        for key in model_data.vlon.keys():
            model_data.vlon[key] = (model_data.vlon[key] + 360) % 360

    modelbox[0] = (modelbox[0] + 360) % 360
    modelbox[1] = (modelbox[1] + 360) % 360
    del sgrid
    msg_queue.put((os.getpid(), sgridfile, None, None))

