# BSD 2-Clause License
#
# Copyright (c) 2021, Hewlett Packard Enterprise
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os.path as osp
import time
from os import getcwd
from pprint import pformat

import pandas as pd
from tqdm import trange

from .control import Controller
from .entity import Ensemble, EntityList, Model, SmartSimEntity
from .error import SmartSimError
from .generation import Generator
from .utils import get_logger
from .utils.entityutils import separate_entities
from .utils.helpers import colorize, init_default

logger = get_logger(__name__)


class Experiment:
    """Experiments are the main user interface in SmartSim.

    Experiments can create instances to launch called ``Model``
    and ``Ensemble``. Through the ``Experiment`` interface, users
    can programmatically create, configure, start, stop, and
    query the instances they create.
    """

    def __init__(self, name, exp_path=None, launcher="local"):
        """Initialize an ``Experiment``

        :param name: name for the ``Experiment``
        :type name: str
        :param exp_path: path to location of ``Experiment`` directory if generated
        :type exp_path: str
        :param launcher: type of launcher, options are "slurm", "pbs",
                         "cobalt", or "local". Defaults to "local"
        :type launcher: str
        """
        self.name = name
        if exp_path:
            if not isinstance(exp_path, str):
                raise TypeError("exp_path argument was not of type str")
            if not osp.isdir(osp.abspath(exp_path)):
                raise NotADirectoryError("Experiment path provided does not exist")
            exp_path = osp.abspath(exp_path)
        self.exp_path = init_default(osp.join(getcwd(), name), exp_path, str)
        self._control = Controller(launcher=launcher)

    def start(self, *args, block=True, summary=False):
        """Launch instances passed as arguments

        Start the ``Experiment`` by turning specified instances into jobs
        for the underlying launcher and launching them.

        Instances of ``Model``, ``Ensemble`` and ``Orchestrator``
        can all be passed as arguments to the start method.

        :param block: block execution until all non-database
                      jobs are finished, defaults to True
        :type block: bool, optional
        :param summary: print a launch summary prior to launch,
                        defaults to False
        :type summary: bool, optional
        """
        try:
            if summary:
                self._launch_summary(*args)
            self._control.start(*args, block=block)
        except SmartSimError as e:
            logger.error(e)
            raise

    def stop(self, *args):
        """Stop specific instances launched by this ``Experiment``

        Instances of ``Model``, ``Ensemble`` and ``Orchestrator``
        can all be passed as arguments to the start method.

        :raises TypeError: if wrong type
        :raises SmartSimError: if stop request fails
        """
        try:
            for entity in args:
                if isinstance(entity, SmartSimEntity):
                    self._control.stop_entity(entity)
                elif isinstance(entity, EntityList):
                    self._control.stop_entity_list(entity)
                else:
                    raise TypeError(
                        f"Argument was of type {type(entity)} not SmartSimEntity or EntityList"
                    )
        except SmartSimError as e:
            logger.error(e)
            raise

    def generate(self, *args, tag=None, overwrite=False):
        """Generate the file structure for an ``Experiment``

        ``Experiment.generate`` creates directories for each instance
        passed to organize Experiments that launch many instances

        If files or directories are attached to ``Model`` objects
        using ``Model.attach_generator_files()``, those files or
        directories will be symlinked, copied, or configured and
        written into the created directory for that instance.

        Instances of ``Model``, ``Ensemble`` and ``Orchestrator``
        can all be passed as arguments to the generate method.

        :param tag: tag used in `to_configure` generator files
        :type tag: str, optional
        :param overwrite: overwrite existing folders and contents
        :type overwrite: bool, optional
        """
        try:
            generator = Generator(self.exp_path, overwrite=overwrite)
            if tag:
                generator.set_tag(tag)
            generator.generate_experiment(*args)
        except SmartSimError as e:
            logger.error(e)
            raise

    def poll(self, interval=10, verbose=True):
        """Monitor jobs through logging to stdout.

        This method should only be used if jobs were launched
        with ``Experiment.start(block=False)``

        :param interval: frequency of logging to stdout
        :type interval: int
        :param verbose: set verbosity
        :type verbose: bool
        :raises SmartSimError:
        """
        try:
            self._control.poll(interval, verbose)
        except SmartSimError as e:
            logger.error(e)
            raise

    def finished(self, entity):
        """Query if a job as completed

        A instance of ``Model``, ``Ensemble`` can be passed
        as an argument.

        :param entity: object launched by this ``Experiment``
        :type entity: SmartSimEntity | EntityList
        :returns: True if job has completed
        :rtype: bool
        """
        try:
            return self._control.finished(entity)
        except SmartSimError as e:
            logger.error(e)
            raise

    def get_status(self, *args):
        """Query the status of a job

        Instances of ``Model``, ``Ensemble`` and ``Orchestrator``
        can all be passed as arguments to ``Experiment.get_status()``

        :returns: status of the job
        :rtype: list[str]
        :raises SmartSimError: if status retrieval fails
        :raises TypeError:
        """
        try:
            statuses = []
            for entity in args:
                if isinstance(entity, SmartSimEntity):
                    statuses.append(self._control.get_entity_status(entity))
                elif isinstance(entity, EntityList):
                    statuses.extend(self._control.get_entity_list_status(entity))
                else:
                    raise TypeError(
                        f"Argument was of type {type(entity)} not SmartSimEntity or EntityList"
                    )
            return statuses
        except SmartSimError as e:
            logger.error(e)
            raise

    def create_ensemble(
        self,
        name,
        params=None,
        batch_settings=None,
        run_settings=None,
        replicas=None,
        perm_strategy="all_perm",
        **kwargs,
    ):
        """Create an ``Ensemble`` of ``Model`` instances

        Ensembles can be launched sequentially or as a batch
        if using a non-local launcher. e.g. slurm

        Ensembles require one of the following combinations
        of arguments
          - ``run_settings`` and ``params``
          - ``run_settings`` and ``replicas``
          - ``batch_settings``
          - ``batch_settings``, ``run_settings``, and ``params``
          - ``batch_settings``, ``run_settings``, and ``replicas``

        If given solely batch settings, an empty ensemble
        will be created that models can be added to manually
        through ``Ensemble.add_model()``.
        The entire ensemble will launch as one batch.

        Provided batch and run settings, either ``params``
        or ``replicas`` must be passed and the entire ensemble
        will launch as a single batch.

        Provided solely run settings, either ``params``
        or ``replicas`` must be passed and the ensemble members
        will each launch sequentially.

        :param name: name of the ensemble
        :type name: str
        :param params: parameters to expand into ``Model`` members
        :type params: dict[str, Any]
        :param batch_settings: describes settings for ``Ensemble`` as batch workload
        :type batch_settings: BatchSettings
        :param run_settings: describes how each ``Model`` should be executed
        :type run_settings: RunSettings
        :param replicas: number of replicas to create
        :type replicas: int
        :param perm_strategy: strategy for expanding ``params`` into
                             ``Model`` instances from params argument
                             options are "all_perm", "stepped", "random"
                             or a callable function
        :type perm_strategy: str
        :raises SmartSimError: if initialization fails
        :return: ``Ensemble`` instance
        :rtype: Ensemble
        """
        try:
            new_ensemble = Ensemble(
                name,
                params,
                batch_settings=batch_settings,
                run_settings=run_settings,
                perm_strat=perm_strategy,
                replicas=replicas,
                **kwargs,
            )
            return new_ensemble
        except SmartSimError as e:
            logger.error(e)
            raise

    def create_model(
        self, name, run_settings, params=None, path=None, enable_key_prefixing=False
    ):
        """Create a ``Model``

        By default, all ``Model`` instances start with the cwd
        as their path unless specified. If specified or not, upon
        user passing the instance to ``Experiment.generate()``,
        the ``Model`` path will be overwritten and replaced
        with the created directory for the ``Model``

        :param name: name of the model
        :type name: str
        :param run_settings: defines how ``Model`` should be run,
        :type run_settings: RunSettings
        :param params: model parameters for writing into configuration files
        :type params: dict, optional
        :param path: path to where the model should be executed at runtime
        :type path: str, optional
        :param enable_key_prefixing: If true, data sent to the Orchestrator
                                     using SmartRedis from this ``Model`` will
                                     be prefixed with the ``Model`` name.
        :type enable_key_prefixing: bool
        :raises SmartSimError: if initialization fails
        :return: the created ``Model``
        :rtype: Model
        """
        path = init_default(getcwd(), path, str)
        params = init_default({}, params, dict)
        try:
            new_model = Model(name, params, path, run_settings)
            if enable_key_prefixing:
                new_model.enable_key_prefixing()
            return new_model
        except SmartSimError as e:
            logger.error(e)
            raise

    def reconnect_orchestrator(self, checkpoint):
        """Reconnect to a running ``Orchestrator``

        This method can be used to connect to a Redis deployment
        that was launched by a previous ``Experiment``. This way
        users can run many experiments utilizing the same Redis
        deployment

        :param checkpoint: the `smartsim_db.dat` file created
                           when an ``Orchestrator`` is launched
        :type checkpoint: str
        """
        try:
            orc = self._control.reload_saved_db(checkpoint)
            return orc
        except SmartSimError as e:
            logger.error(e)
            raise

    def summary(self):
        """Return a summary of the ``Experiment``

        The summary will show each instance that has been
        launched and completed in this ``Experiment``

        :return: Dataframe of ``Experiment`` history
        :rtype: pd.DataFrame
        """
        index = 0
        df = pd.DataFrame(
            columns=[
                "Name",
                "Entity-Type",
                "JobID",
                "RunID",
                "Time",
                "Status",
                "Returncode",
            ]
        )
        # TODO should this include running jobs?
        for job in self._control._jobs.completed.values():
            for run in range(job.history.runs + 1):
                df.loc[index] = [
                    job.entity.name,
                    job.entity.type,
                    job.history.jids[run],
                    run,
                    job.history.job_times[run],
                    job.history.statuses[run],
                    job.history.returns[run],
                ]
                index += 1
        return df

    def _launch_summary(self, *args):
        """Experiment pre-launch summary of entities that will be launched"""

        def sprint(p):
            print(p, flush=True)

        sprint("\n")
        models, ensembles, orchestrator = separate_entities(args)

        header = colorize("=== LAUNCH SUMMARY ===", color="cyan", bold=True)
        exname = colorize("Experiment: " + self.name, color="green", bold=True)
        expath = colorize("Experiment Path: " + self.exp_path, color="green")
        launch = colorize(
            "Launching with: " + str(self._control._launcher), color="green"
        )
        numens = colorize("# of Ensembles: " + str(len(ensembles)), color="green")
        numods = colorize("# of Models: " + str(len(models)), color="green")
        has_orc = "yes" if orchestrator else "no"
        orches = colorize("Database: " + has_orc, color="green")

        sprint(f"{header}")
        sprint(f"{exname}\n{expath}\n{launch}\n{numens}\n{numods}\n{orches}\n")

        if ensembles:
            sprint(colorize("=== ENSEMBLES ===", color="cyan", bold=True))
            for ens in ensembles:
                name = colorize(ens.name, color="green", bold=True)
                num_models = colorize(
                    "# of models in ensemble: " + str(len(ens)), color="green"
                )
                batch_settings = colorize(
                    "Batch Settings: \n" + str(ens.batch_settings),
                    color="green",
                )
                run_settng = colorize(
                    "Run Settings: \n" + str(ens.run_settings),
                    color="green",
                )
                batch = colorize(f"Launching as batch: {ens.batch}", color="green")

                sprint(f"{name}")
                sprint(f"{num_models}")
                sprint(f"{batch}")
                if ens.batch:
                    print(f"{batch_settings}")
                else:
                    sprint(f"{run_settng}")
            sprint("\n")
        if models:
            sprint(colorize("=== MODELS ===", color="cyan", bold=True))
            for model in models:
                model_name = colorize(model.name, color="green", bold=True)
                parameters = colorize(
                    "Model Parameters: \n" + pformat(model.params), color="green"
                )
                run_settng = colorize(
                    "Model Run Settings: \n" + str(model.run_settings),
                    color="green",
                )
                sprint(f"{model_name}")
                sprint(f"{parameters}")
                sprint(f"{run_settng}")
            sprint("\n")
        if orchestrator:
            sprint(colorize("=== DATABASE ===", color="cyan", bold=True))
            size = colorize(
                "# of database nodes: " + str(len(orchestrator)), color="green"
            )
            batch = colorize(f"Launching as batch: {orchestrator.batch}", color="green")
            sprint(f"{batch}")
            sprint(f"{size}")

        sprint("\n")

        wait, steps = 10, 100
        prog_bar = trange(
            steps,
            desc="Launching in...",
            leave=False,
            ncols=80,
            mininterval=0.25,
            bar_format="{desc}: {bar}| {remaining} {elapsed}",
        )
        for _ in prog_bar:
            time.sleep(wait / steps)

    def __str__(self):
        return self.name

pandas.DataFrame