# -*- coding: utf-8 -*-
"""reV batching framework for parametric runs.

The batch framework allows users to modify key-value pairs in input jsons files
based on a batch config file. The batch module will create run directories for
all combinations of input parametrics, and run the reV pipelines for each job.

Created on Mon Jun 10 13:49:53 2019

@author: gbuster
"""
import pandas as pd
import copy
import json
import yaml
import os
import shutil
import itertools
import logging
from warnings import warn
from abc import ABC, abstractmethod

from reV.config.batch import BatchConfig
from reV.pipeline.cli_pipeline import pipeline_monitor_background
from reV.pipeline.pipeline import Pipeline
from reV.utilities.exceptions import InputError, PipelineError

from rex.utilities import safe_json_load, safe_yaml_load, parse_year
from rex.utilities.loggers import init_logger


logger = logging.getLogger(__name__)


class BatchJob:
    """Framework for building a batched job suite."""

    def __init__(self, config, verbose=False):
        """
        Parameters
        ----------
        config : str
            File path to config json or csv (str).
        verbose : bool
            Flag to turn on debug logging.
        """
        self._job_tags = None

        self._config = BatchConfig(config)
        self._base_dir = self._config.config_dir
        os.chdir(self._base_dir)

        if 'logging' in self._config:
            logging_kwargs = self._config.logging
            if verbose:
                logging_kwargs['log_level'] = 'DEBUG'

            init_logger('reV.batch', **logging_kwargs)

        x = self._parse_config(self._config)
        self._arg_combs, self._file_sets, self._set_tags = x

        logger.info('Batch job initialized with {} sub jobs.'
                    .format(len(self._arg_combs)))

    @staticmethod
    def _parse_config(config):
        """Parse batch config object for useful data.

        Parameters
        ----------
        config : from reV.config.batch.BatchConfig
            reV batch config object that emulates a dict.

        Returns
        -------
        arg_combs : list
            List of dictionaries representing the different arg/value
            combinations made available in the batch config json.
        file_sets : list
            List representing the files to manipulate for each arg comb
            (same length as arg_combs).
        set_tags : list
            List of strings of tags for each batch job set
            (same length as arg_combs).
        """

        arg_combs = []
        file_sets = []
        set_tags = []
        sets = []

        # iterate through batch sets
        for s in config['sets']:
            set_tag = s.get('set_tag', '')
            if set_tag in sets:
                msg = ('Found multiple sets with the same set_tag: "{}"'
                       .format(set_tag))
                logger.error(msg)
                raise ValueError(msg)
            else:
                sets.append(set_tag)

            # iterate through combinations of arg values
            for comb in itertools.product(*list(s['args'].values())):

                # make a dictionary representation of this combination
                comb_dict = {}
                for i, k in enumerate(s['args'].keys()):
                    comb_dict[k] = comb[i]

                # append the unique dictionary representation to the attr
                arg_combs.append(comb_dict)
                file_sets.append(s['files'])
                set_tags.append(set_tag)

        return arg_combs, file_sets, set_tags

    @staticmethod
    def _fix_tag_w_year(value):
        """If one of the tag values looks like a year, add a zero for the tag.

        Parameters
        ----------
        value : str | int | float
            Value from a batch run.

        Returns
        -------
        value : str
            Value from a batch run converted to a string tag, if the input
            value looks like a year, a zero will be added.
        """

        value = str(value).replace('.', '')

        if parse_year('_' + value, option='bool'):
            value += '0'

        return value

    def _make_job_tag(self, set_tag, arg_comb):
        """Make a job tags from a unique combination of args + values.

        Parameters
        ----------
        set_tag : str
            Optional set tag to prefix job tag.
        arg_comb : dict
            Key-value pairs for this argument combination.

        Returns
        -------
        job_tag : str
            Identifying string from the arg comb.
        """

        job_tag = []

        for arg, value in arg_comb.items():
            temp = arg.split('_')
            temp = ''.join([s[0] for s in temp])

            if isinstance(value, (int, float)):
                temp += self._fix_tag_w_year(value)

            else:
                i = 0
                for s in self._config['sets']:
                    if arg in s['args']:
                        try:
                            i = s['args'][arg].index(value)
                        except ValueError:
                            pass
                        else:
                            break
                temp += str(i)

            job_tag.append(temp)

        job_tag = '_'.join(job_tag)

        if set_tag:
            job_tag = '{}_{}'.format(set_tag, job_tag)

        if job_tag.endswith('_'):
            job_tag = job_tag.rstrip('_')

        return job_tag

    def _clean_arg_comb_tag(self, set_tag, arg_comb):
        """Clean a dictionary of arg combinations for a single job by removing
        any args that only have one value in the current set tag.

        Parameters
        ----------
        set_tag : str
            Optional set tag to prefix job tag.
        arg_comb : dict
            Key-value pairs for this argument combination.

        Returns
        -------
        tag_arg_comb : str
            Arg combinations just for making the job tags. This may not have
            all the arg combinations for the actual job setup.
        """

        ignore_tags = []

        for arg in arg_comb.keys():
            all_values = []
            for batch_set in self._config['sets']:
                if (batch_set.get('set_tag', '') == set_tag
                        and arg in batch_set['args']):
                    all_values += batch_set['args'][arg]

            if len(all_values) <= 1:
                ignore_tags.append(arg)

        tag_arg_comb = {k: v for k, v in arg_comb.items()
                        if k not in ignore_tags}

        return tag_arg_comb

    @property
    def job_table(self):
        """Get a dataframe summarizing the batch jobs."""
        table = pd.DataFrame()
        for i, job_tag in enumerate(self.job_tags):
            job_info = {k: str(v) for k, v in self.arg_combs[i].items()}
            job_info['set_tag'] = str(self._set_tags[i])
            job_info['files'] = str(self.file_sets[i])
            job_info = pd.DataFrame(job_info, index=[job_tag])
            table = table.append(job_info)

        table.index.name = 'job'

        return table

    @property
    def job_tags(self):
        """Ordered list of job tags corresponding to unique arg/value combs.

        Returns
        -------
        job_tags : list
            List of job tag strings corresponding to the unique arg/value
            combinations.
        """
        if self._job_tags is None:
            self._job_tags = []
            for i, arg_comb in enumerate(self.arg_combs):
                tag_arg_comb = self._clean_arg_comb_tag(self._set_tags[i],
                                                        arg_comb)
                self._job_tags.append(self._make_job_tag(self._set_tags[i],
                                                         tag_arg_comb))
        return self._job_tags

    @property
    def arg_combs(self):
        """List of all viable arg/value combinations.

        Returns
        -------
        _arg_combs : list
            List of dictionaries representing the different arg/value
            combinations made available in the batch config json.
        """

        return self._arg_combs

    @property
    def file_sets(self):
        """List of files to be manipulated for each arg comb batch job.

        Returns
        -------
        file_sets : list
            List of same length as arg_combs, representing the files to
            manipulate for each arg comb.
        """

        return self._file_sets

    @property
    def sub_dirs(self):
        """List of job sub directories.

        Returns
        -------
        sub_dirs : list
            List of strings of job sub directories.
        """
        sub_dirs = [os.path.join(self._base_dir, tag + '/')
                    for tag in self.job_tags]
        return sub_dirs

    @staticmethod
    def _clean_arg(arg):
        """Perform any cleaning steps required before writing an arg to a json

        Cleaning steps:
            1. Convert stringified dictionary to python dictionary object

        Parameters
        ----------
        arg : str | int | float | dict
            Value to be written to a batch job json file

        Returns
        -------
        arg : str | int | float | dict
            Cleaned value ready to be written to a batch job json file.
        """

        if isinstance(arg, str):
            if '{' in arg and '}' in arg and ("'" in arg or '"' in arg):
                arg = arg.replace("'", '"')
                try:
                    arg = json.loads(arg)
                except json.decoder.JSONDecodeError as e:
                    msg = 'Could not load json string: {}'.format(arg)
                    logger.exception(msg)
                    raise e

        return arg

    @classmethod
    def _mod_dict(cls, inp, arg_mods):
        """Recursively modify key/value pairs in a dictionary.

        Parameters
        ----------
        inp : dict | list | str | int | float
            Input to modify such as a loaded json reV config to modify with
            batch. Can also be nested data from a recursive call. Should be a
            dict first, the recusive call can input nested values as
            dict/list/str etc...
        arg_mods : dict
            Key/value pairs to insert in the inp.

        Returns
        -------
        out : dict | list | str | int | float
            Modified inp with arg_mods.
        """

        out = copy.deepcopy(inp)

        if isinstance(inp, dict):
            for k, v in inp.items():
                if k in arg_mods:
                    out[k] = cls._clean_arg(arg_mods[k])
                elif isinstance(v, (list, dict)):
                    out[k] = cls._mod_dict(v, arg_mods)

        elif isinstance(inp, list):
            for i, entry in enumerate(inp):
                out[i] = cls._mod_dict(entry, arg_mods)

        return out

    @classmethod
    def _mod_file(cls, file, fpath_out, arg_mods):
        """Import and modify the contents of a json. Dump to new file.

        Parameters
        ---------
        file : `FileToMod`
            Instance of a `FileToMod` object representing the
            input file that will be imported/modified.
        fpath_out : str
            File path to dump new modified json.
        arg_mods : dict
            Dictionary representing one set of arg/value combinations to
            implement in the json
        """

        data = file.load()
        file.data = cls._mod_dict(data, arg_mods)
        file.write(fpath_out)

    def _make_job_dirs(self):
        """Copy job files from the batch config dir into sub job dirs."""

        table = self.job_table
        table.to_csv(os.path.join(self._base_dir, 'batch_jobs.csv'))
        logger.debug('Batch jobs list: {}'
                     .format(sorted(table.index.values.tolist())))
        logger.info('Preparing batch job directories...')

        # walk through current directory getting everything to copy
        for source_dir, _, filenames in os.walk(self._base_dir):

            # do make additional copies of job sub directories.
            skip = any(job_tag in source_dir for job_tag in self.job_tags)

            if not skip:

                # For each dir level, iterate through the batch arg combos
                for i, arg_comb in enumerate(self.arg_combs):
                    # tag and files to mod corresponding to this arg combo
                    tag = self.job_tags[i]
                    mod_files = self.file_sets[i]
                    mod_fnames = [os.path.basename(fn) for fn in mod_files]

                    # Add the job tag to the directory path.
                    # This will copy config subdirs into the job subdirs
                    destination_dir = source_dir.replace(
                        self._base_dir,
                        os.path.join(self._base_dir, tag + '/'))

                    if not os.path.exists(destination_dir):
                        logger.debug('Making new job directory: {}'
                                     .format(destination_dir))
                        os.makedirs(destination_dir)

                    for fn in filenames:
                        self._copy_batch_file(fn, source_dir, destination_dir,
                                              tag, arg_comb, mod_fnames)

        logger.info('Batch job directories ready for execution.')

    def _copy_batch_file(self, fn, source_dir, destination_dir, job_name,
                         arg_comb, mod_fnames):
        """Copy a file in the batch directory into a job directory with
        appropriate batch modifications (permutations) as defined by arg_comb.

        Parameters
        ----------
        fn : str
            Filename being copied.
        source_dir : str
            Directory that the file is sourced from (usually the base batch
            project directory)
        destination_dir : str
            Directory that the file should be copied to (usually a job's
            sub directory in source_dir)
        job_name : str
            Batch job name (for logging purposes)
        arg_comb : dict
            Dictionary representing one set of arg/value combinations to
            implement in the fn if the fn is a JSON or YAML file specified
            as one of the batch permutation files to modify.
        mod_fnames : list
            List of filenames that need to be modified by the batch module.

        Raises
        ------
        InputError
            If fn does not end in ('.json', '.yml', '.yaml'), i.e.
            unrecognized input mod file format.
        """

        if fn in mod_fnames:
            logger.debug('Copying and modifying run file '
                         '"{}" to job: "{}"'.format(fn, job_name))
            fn_in = os.path.join(source_dir, fn)
            fn_out = os.path.join(destination_dir, fn)
            if fn.endswith('.json'):
                in_file = JSONFileToMod(fn_in)
            elif fn.endswith(('.yml', '.yaml')):
                in_file = YAMLFileToMod(fn_in)
            else:
                msg = ('Unknown file-to-modify type: {!r}. Please ensure '
                       'the input files to modify are in JSON or YAML format '
                       '(i.e. end in ".json", ".yml", or ".yaml").')
                e = msg.format(fn)
                logger.error(e)
                raise InputError(e)
            self._mod_file(in_file, fn_out, arg_comb)

        else:
            # straight copy of non-mod and non-json
            fp_source = os.path.join(source_dir, fn)
            fp_target = os.path.join(destination_dir, fn)

            modified = True
            if os.path.exists(fp_target):
                modified = (os.path.getmtime(fp_source)
                            > os.path.getmtime(fp_target))

            if modified:
                logger.debug('Copying run file "{}" to job: '
                             '"{}"'.format(fn, job_name))
                try:
                    shutil.copy(fp_source, fp_target)
                except Exception:
                    msg = ('Could not copy "{}" to job: "{}"'
                           .format(fn, job_name))
                    logger.warning(msg)
                    warn(msg)

    def _run_pipelines(self, monitor_background=False, verbose=False):
        """Run the reV pipeline modules for each batch job.

        Parameters
        ----------
        monitor_background : bool
            Flag to monitor all batch pipelines continuously
            in the background using the nohup command. Note that the
            stdout/stderr will not be captured, but you can set a
            pipeline "log_file" to capture logs.
        verbose : bool
            Flag to turn on debug logging for the pipelines.
        """

        for d in self.sub_dirs:
            pipeline_config = os.path.join(
                d, os.path.basename(self._config.pipeline_config))
            if not os.path.isfile(pipeline_config):
                raise PipelineError('Could not find pipeline config to run: '
                                    '"{}"'.format(pipeline_config))
            elif monitor_background:
                pipeline_monitor_background(pipeline_config, verbose=verbose)
            else:
                Pipeline.run(pipeline_config, monitor=False, verbose=verbose)

    def _cancel_all(self):
        """Cancel all reV pipeline modules for all batch jobs."""
        for d in self.sub_dirs:
            pipeline_config = os.path.join(
                d, os.path.basename(self._config.pipeline_config))
            if os.path.isfile(pipeline_config):
                Pipeline.cancel_all(pipeline_config)

    def _delete_all(self):
        """Clear all of the batch sub job folders based on the job summary
        csv file in the batch config directory."""

        fp_job_table = os.path.join(self._base_dir, 'batch_jobs.csv')
        if not os.path.exists(fp_job_table):
            msg = ('Cannot delete batch jobs without jobs summary table: {}'
                   .format(fp_job_table))
            logger.error(msg)
            raise FileNotFoundError(msg)

        job_table = pd.read_csv(fp_job_table, index_col=0)

        if job_table.index.name != 'job':
            msg = ('Cannot delete batch jobs when the batch summary table '
                   'does not have "job" as the index key')
            logger.error(msg)
            raise ValueError(msg)

        for sub_dir in job_table.index:
            job_dir = os.path.join(self._base_dir, sub_dir)
            if os.path.exists(job_dir):
                logger.info('Removing batch job directory: {}'.format(sub_dir))
                shutil.rmtree(job_dir)
            else:
                w = 'Cannot find batch job directory: {}'.format(sub_dir)
                logger.warning(w)
                warn(w)

        os.remove(fp_job_table)

    @classmethod
    def cancel_all(cls, config, verbose=False):
        """Cancel all reV pipeline modules for all batch jobs.

        Parameters
        ----------
        config : str
            File path to batch config json or csv (str).
        verbose : bool
            Flag to turn on debug logging.
        """

        b = cls(config, verbose=verbose)
        b._cancel_all()

    @classmethod
    def delete_all(cls, config, verbose=False):
        """Delete all reV batch sub job folders based on the job summary csv
        in the batch config directory.

        Parameters
        ----------
        config : str
            File path to batch config json or csv (str).
        verbose : bool
            Flag to turn on debug logging.
        """
        b = cls(config, verbose=verbose)
        b._delete_all()

    @classmethod
    def run(cls, config, dry_run=False, delete=False, monitor_background=False,
            verbose=False):
        """Run the reV batch job from a config file.

        Parameters
        ----------
        config : str
            File path to config json or csv (str).
        dry_run : bool
            Flag to make job directories without running.
        delete : bool
            Flag to delete all batch job sub directories based on the job
            summary csv in the batch config directory.
        monitor_background : bool
            Flag to monitor all batch pipelines continuously
            in the background using the nohup command. Note that the
            stdout/stderr will not be captured, but you can set a
            pipeline "log_file" to capture logs.
        verbose : bool
            Flag to turn on debug logging for the pipelines.
        """

        b = cls(config, verbose=verbose)
        if delete:
            b._delete_all()
        else:
            b._make_job_dirs()
            if not dry_run:
                b._run_pipelines(monitor_background=monitor_background,
                                 verbose=verbose)


class FileToMod(ABC):
    """Abstract base class representing a file that is modded by BatchJob. """

    def __init__(self, file_path):
        """
        Parameters
        ----------
        file_path : str
            Path to data file.
        """
        super().__init__()
        self.file_path = file_path
        self.data = None

    @abstractmethod
    def load(self):
        """Load the file contents."""

    @abstractmethod
    def write(self, file_name_out):
        """Write data out in the correct file format. """


class JSONFileToMod(FileToMod):
    """JSON file that is modded by BatchJob. """

    def load(self):
        """Load the JSON file contents."""
        self.data = safe_json_load(self.file_path)
        return self.data

    def write(self, file_name_out):
        """Write data out in the correct file format. """
        with open(file_name_out, 'w') as f:
            json.dump(self.data, f, indent=4, separators=(',', ': '))


class YAMLFileToMod(FileToMod):
    """YAML file that is modded by BatchJob. """

    def load(self):
        """Load the YAML file contents."""
        self.data = safe_yaml_load(self.file_path)
        return self.data

    def write(self, file_name_out):
        """Write data out in the correct file format. """
        with open(file_name_out, 'w') as f:
            yaml.dump(self.data, f, indent=2, sort_keys=False)
