# -*- coding: utf-8 -*-
# ######### COPYRIGHT #########
#
# Copyright(c) 2018
# -----------------
#
# * Laboratoire d'Informatique et Systèmes <http://www.lis-lab.fr/>
# * Université d'Aix-Marseille <http://www.univ-amu.fr/>
# * Centre National de la Recherche Scientifique <http://www.cnrs.fr/>
# * Université de Toulon <http://www.univ-tln.fr/>
#
# Contributors
# ------------
#
# * Ronan Hamon <firstname.lastname_AT_lis-lab.fr>
# * Valentin Emiya <firstname.lastname_AT_lis-lab.fr>
# * Florent Jaillet <firstname.lastname_AT_lis-lab.fr>
#
# Description
# -----------
#
# yafe: Yet Another Framework for Experiments.
#
# Licence
# -------
# This file is part of yafe.
#
# yafe 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.
#
# This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.
#
# ######### COPYRIGHT #########
"""Utils classes and functions for yafe.

.. moduleauthor:: Ronan Hamon
.. moduleauthor:: Valentin Emiya
.. moduleauthor:: Florent Jaillet
"""
import configparser
import importlib
import logging
import os
from pathlib import Path
import stat
import sys


class ConfigParser(configparser.ConfigParser):
    """Configuration file parser for yafe.

    This class inherits from ConfigParser in the configparser module.

    It enables reading the yafe configuration file ``$HOME/.config/yafe.conf``
    at initialization of an experiment. It also provides a method to properly
    read a path in the configuration file, and a static method to generate a
    basic empty configuration file.
    """
    _config_path = Path(os.path.expanduser('~')) / '.config' / 'yafe.conf'

    def __init__(self):
        super().__init__()
        loaded_files = super().read(self._config_path)
        if not len(loaded_files):
            print('Generation of the config file.')
            ConfigParser.generate_config()
            print('Please update config file at {}'.format(self._config_path))
            super().read(self._config_path)

    def get_path(self, section, option, check_exists=True):
        """Get the path filled in a given option of a given section.

        Parameters
        ----------
        section : str
            Name of the section.
        option : str
            Name of the option.
        check_exists : bool, optional
            Indicates if the existence of the path is checked.

        Returns
        -------
        pathlib.Path or None
            Path if the option is defined, None otherwise.

        Raises
        ------
        IOError
            If the parameter ``check_exists`` is set to ``True`` and the path
            does not exist.
        """
        path = super().get(section, option)

        if path == '':
            return None

        path = Path(path)

        if check_exists and not path.exists():
            errmsg = 'Path {} does not exist.'
            raise IOError(errmsg.format(path))

        return path

    @staticmethod
    def generate_config():
        """Generate an empty configuration file.

        The generated configuration file is stored in
        ``$HOME/.config/yafe.conf``.
        """

        config = configparser.ConfigParser(allow_no_value=True)

        config.add_section('USER')
        config.set('USER', '# user directory (data generated by experiments)')
        config.set('USER', 'data_path', '')

        config.add_section('LOGGER')
        config.set('LOGGER', '# path to the log file')
        config.set('LOGGER', 'path', '')

        with open(str(ConfigParser._config_path), 'w') as config_file:
            config.write(config_file)


def get_logger(name='', to_file=True, to_console=True, logger_path=None):
    """Return a yafe logger with the given name.

    Parameters
    ----------
    name : str
        Name of the logger.
    to_file : bool
        Indicates if the logger writes processing and debug information in a
        file.
    to_console : bool
        Indicates if the logger displays processing information in the console.
    logger_path : None or str or pathlib.Path
        Full file path where the data must be written by the logger when
        ``to_file == True``.
        If ``logger_path`` is ``None`` the logger path given in the yafe
        configuration file is used.

    Returns
    -------
    logging.Logger
        A logger with the given name prefixed by ``yafe.``.

    Notes
    -----
    The name of the logger is automatically prefixed by ``yafe.``.
    """
    logger = logging.getLogger('yafe.{}'.format(name))
    logger.handlers = []
    logger.setLevel(logging.DEBUG)

    if to_file:
        if logger_path is None:
            config = ConfigParser()
            logger_path = config.get_path('LOGGER', 'PATH', check_exists=False)
        else:
            logger_path = Path(logger_path)
        if not logger_path.exists():
            logger_path.touch()
        file_handler = logging.FileHandler(str(logger_path))
        file_handler.setLevel(logging.DEBUG)
        file_handler.setFormatter(
            logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'))

        logger.addHandler(file_handler)

    if to_console:
        console_handler = logging.StreamHandler(sys.stdout)
        console_handler.setLevel(logging.INFO)
        console_handler.setFormatter(
            logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'))

        logger.addHandler(console_handler)

    return logger


def generate_oar_script(script_file_path, xp_var_name, task_ids=None,
                        batch_size=1, oar_walltime='02:00:00',
                        activate_env_command=None, use_gpu=False):
    """Generate a script to launch an experiment using OAR.

    Tasks are divided into batches that are executed by oar jobs.

    The resulting script is written in the experiment folder, and the command
    to launch the jobs with OAR is displayed in the terminal.

    An example script illustrating how to use
    :func:`yafe.utils.generate_oar_script` is available in the corresponding
    :ref:`tutorial <tutorial_oar>`.

    Parameters
    ----------
    script_file_path : str
        File path to the script that defines the experiment.
    xp_var_name : str
        Name of the variable containing the experiment in the script.
    task_ids : list
        List of tasks ids to run.
        If ``task_ids`` is ``None``, the list of pending tasks of the
        experiment is used.
    batch_size : int
        Number of tasks run in each batch.
    oar_walltime : str
        Wall time for each OAR job ('HH:MM:SS').
    activate_env_command : str or None
        Optional command that must be run to activate a Python virtual
        environment before launching the experiment.
        Typically, this is a command of the form
        ``source some_virtual_env/bin/activate`` when using virtualenv and
        ``source activate some_conda_env`` when using conda.
        If ``activate_env_command`` is ``None``, no virtual environment is
        activated.
    use_gpu : bool
        Flag specifying if a gpu ressource is needed when running the
        experiment.
    """
    script_file_path = Path(script_file_path)
    script_dir = script_file_path.parent
    script_name = script_file_path.stem

    sys.path.append(str(script_dir))
    mod = importlib.import_module(script_name)
    xp = getattr(mod, xp_var_name)

    if task_ids is None:
        task_ids = xp.get_pending_task_ids()

    # split and save the tasks
    task_ids = list(map(str, task_ids))
    batches = [
        task_ids[i:(i + batch_size)]
        for i in range(0, len(task_ids), batch_size)
    ]
    file_path = xp.xp_path / 'listoftasks.txt'

    with open(str(file_path), 'wt') as fout:
        fout.write('\n'.join(map(lambda batch: ','.join(batch), batches)))

    # generate and save script
    script_path = Path(os.path.abspath(script_file_path))
    script_dir = script_path.parent
    script_name = script_path.stem

    script = '#!/bin/sh\n'

    # define the OAR parameters
    script += '#OAR -n {}\n'.format(xp.name)
    script += '#OAR --array-param-file {}\n'.format(str(file_path))
    script += '#OAR -O {}/%jobid%.out\n'.format(xp.xp_path)
    script += '#OAR -E {}/%jobid%.err\n'.format(xp.xp_path)
    script += '#OAR -l walltime={}\n'.format(oar_walltime)
    if use_gpu:
        script += '#OAR -p gpu IS NOT NULL\n'
    else:
        script += '#OAR -p gpu IS NULL\n'

    script += 'echo "OAR_JOB_ID: $OAR_JOB_ID"\n'
    script += 'echo "OAR_ARRAY_ID: $OAR_ARRAY_ID"\n'
    script += 'echo "OAR_ARRAY_INDEX: $OAR_ARRAY_INDEX"\n'

    # activate the virtual env
    if activate_env_command is not None and len(activate_env_command) > 0:
        script += '{}\n'.format(activate_env_command)

    # python command
    script += 'echo "Running {}.launch_experiment(task_ids=\'$1\')"\n'.format(
        xp_var_name)
    script += 'python -c "import sys; sys.path.append(\'{0}\'); ' \
        'from {1} import {2}; ' \
        '{2}.launch_experiment(task_ids=\'$1\')"\n'.format(
            script_dir, script_name, xp_var_name)
    script += 'exit $?'

    script_path = xp.xp_path / 'script_oar.sh'
    with script_path.open('w') as file:
        file.write(script)
    status = os.stat(script_path)
    os.chmod(script_path, status.st_mode | stat.S_IXUSR)

    print('oarsub -S {}'.format(str(script_path)))
