# Copyright 2019 <NAME>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections import defaultdict
from typing import Any, Callable, Dict, Iterable, List

from .event import Event

Callback = Callable[[dict], None]


class Runner:
    """A neural network runner.

    A runner provides a thin abstraction of iterating over batches for several epochs,
    which is typically done in neural network training. To customize the behavior during
    a run, a runner provides a way to listen to events emitted during such run.
    To listen to an event, call `Runner.on` and provide a callback which will be called
    when the event is emitted. An event callback is a callable that accepts a `dict`
    and returns nothing. The `dict` is the state of the run. By default, the state contains:

    * ``batches`` - Iterable of batches which constitutes an epoch.
    * ``max_epoch`` - Maximum number of epochs to run.
    * ``n_iters`` - Current number of batch iterations.
    * ``running`` - A boolean which equals ``True`` if the runner is still running. Can
      be set to ``False`` to stop the runner earlier.
    * ``epoch`` - Current number of epoch.
    * ``batch`` - Current batch retrieved from ``state['batches']``.

    Attributes:
        state (dict): Runner's state that is passed to event callbacks.

    Note:
        Callbacks for an event are called in the order they are passed to `~Runner.on`.

    Caution:
        All the state contents above are required for a runner to function properly.
        You are free to change their values to suit your use cases better, but be careful.
    """

    def __init__(self) -> None:
        self.state: dict = {}
        self._callbacks: Dict[Event, List[Callback]] = defaultdict(list)

    def on(self, event: Event, callbacks=None):
        """Add single/multiple callback(s) to listen to an event.

        If ``callbacks`` is ``None``, this method returns a decorator which accepts
        a single callback for the event. If ``callbacks`` is a sequence of callbacks,
        they will all be added as listeners to the event *in order*.

        Args:
            event: Event to listen.
            callbacks: Callback(s) for the event.

        Returns:
            A decorator which accepts a callback, if ``callbacks`` is ``None``.
        """
        if callbacks is not None:
            cblist = self._callbacks[event]
            try:
                cblist.extend(callbacks)
            except TypeError:  # must be a single callback
                cblist.append(callbacks)
            return

        def decorator(cb: Callback) -> Callback:
            self._callbacks[event].append(cb)
            return cb

        return decorator

    def run(self, batches: Iterable[Any], max_epoch: int = 1) -> None:
        """Run on batches for a number of epochs.

        Args:
            batches: Batches to iterate over in an epoch.
            max_epoch: Maximum number of epochs to run.
        """
        state = self.state
        state.update(
            {
                "max_epoch": max_epoch,
                "batches": batches,
                "n_iters": 0,
                "running": True,
                "epoch": 0,
            }
        )

        self._emit(Event.STARTED, state)

        while state["running"] and state["epoch"] < state["max_epoch"]:
            state["epoch"] += 1
            self._emit(Event._ETIMER_STARTED, state)
            self._emit(Event.EPOCH_STARTED, state)
            self._emit(Event._REDUCER_RESET, state)
            self._emit(Event._PBAR_CREATED, state)

            state["_batches_iter"] = iter(state["batches"])
            self._run_epoch()

            self._emit(Event._PBAR_CLOSED, state)
            self._emit(Event._REDUCER_COMPUTED, state)
            self._emit(Event.EPOCH_FINISHED, state)
            self._emit(Event._ETIMER_FINISHED, state)

        self._emit(Event.FINISHED, state)
        state["running"] = False

    def resume(self, repeat_last_batch: bool = False) -> None:
        """Resume runner starting from the current state.

        Args:
            repeat_last_batch: Whether to repeat processing the last batch. Ignored if the
                last epoch is finished (i.e. the batches have been exhausted).
        """
        state = self.state
        state["running"] = True

        finished_last_epoch = state["n_iters"] % len(state["batches"]) == 0

        if not finished_last_epoch:
            self._emit(Event._ETIMER_STARTED, state)
            self._emit(Event._PBAR_CREATED, state)

            if repeat_last_batch:
                self._emit(Event.BATCH, state)
                self._emit(Event._REDUCER_UPDATED, state)
                self._emit(Event._PBAR_UPDATED, state)

            self._run_epoch()

            self._emit(Event._PBAR_CLOSED, state)
            self._emit(Event._REDUCER_COMPUTED, state)
            self._emit(Event.EPOCH_FINISHED, state)
            self._emit(Event._ETIMER_FINISHED, state)

        while state["running"] and state["epoch"] < state["max_epoch"]:
            state["epoch"] += 1
            self._emit(Event._ETIMER_STARTED, state)
            self._emit(Event.EPOCH_STARTED, state)
            self._emit(Event._REDUCER_RESET, state)
            self._emit(Event._PBAR_CREATED, state)

            state["_batches_iter"] = iter(state["batches"])
            self._run_epoch()

            self._emit(Event._PBAR_CLOSED, state)
            self._emit(Event._REDUCER_COMPUTED, state)
            self._emit(Event.EPOCH_FINISHED, state)
            self._emit(Event._ETIMER_FINISHED, state)

        self._emit(Event.FINISHED, state)
        state["running"] = False

    def _run_epoch(self) -> None:
        state = self.state
        while state["running"]:
            try:
                state["batch"] = next(state["_batches_iter"])
            except StopIteration:
                break
            state["n_iters"] += 1
            self._emit(Event.BATCH, state)
            self._emit(Event._REDUCER_UPDATED, state)
            self._emit(Event._PBAR_UPDATED, state)

    def _emit(self, event: Event, state: dict) -> None:
        for callback in self._callbacks[event]:
            if not state["running"]:
                break
            callback(state)

# Copyright 2019 <NAME>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections import defaultdict
from typing import Any, Callable, Dict, Iterable, List

from .event import Event

Callback = Callable[[dict], None]


class Runner:
    """A neural network runner.

    A runner provides a thin abstraction of iterating over batches for several epochs,
    which is typically done in neural network training. To customize the behavior during
    a run, a runner provides a way to listen to events emitted during such run.
    To listen to an event, call `Runner.on` and provide a callback which will be called
    when the event is emitted. An event callback is a callable that accepts a `dict`
    and returns nothing. The `dict` is the state of the run. By default, the state contains:

    * ``batches`` - Iterable of batches which constitutes an epoch.
    * ``max_epoch`` - Maximum number of epochs to run.
    * ``n_iters`` - Current number of batch iterations.
    * ``running`` - A boolean which equals ``True`` if the runner is still running. Can
      be set to ``False`` to stop the runner earlier.
    * ``epoch`` - Current number of epoch.
    * ``batch`` - Current batch retrieved from ``state['batches']``.

    Attributes:
        state (dict): Runner's state that is passed to event callbacks.

    Note:
        Callbacks for an event are called in the order they are passed to `~Runner.on`.

    Caution:
        All the state contents above are required for a runner to function properly.
        You are free to change their values to suit your use cases better, but be careful.
    """

    def __init__(self) -> None:
        self.state: dict = {}
        self._callbacks: Dict[Event, List[Callback]] = defaultdict(list)

    def on(self, event: Event, callbacks=None):
        """Add single/multiple callback(s) to listen to an event.

        If ``callbacks`` is ``None``, this method returns a decorator which accepts
        a single callback for the event. If ``callbacks`` is a sequence of callbacks,
        they will all be added as listeners to the event *in order*.

        Args:
            event: Event to listen.
            callbacks: Callback(s) for the event.

        Returns:
            A decorator which accepts a callback, if ``callbacks`` is ``None``.
        """
        if callbacks is not None:
            cblist = self._callbacks[event]
            try:
                cblist.extend(callbacks)
            except TypeError:  # must be a single callback
                cblist.append(callbacks)
            return

        def decorator(cb: Callback) -> Callback:
            self._callbacks[event].append(cb)
            return cb

        return decorator

    def run(self, batches: Iterable[Any], max_epoch: int = 1) -> None:
        """Run on batches for a number of epochs.

        Args:
            batches: Batches to iterate over in an epoch.
            max_epoch: Maximum number of epochs to run.
        """
        state = self.state
        state.update(
            {
                "max_epoch": max_epoch,
                "batches": batches,
                "n_iters": 0,
                "running": True,
                "epoch": 0,
            }
        )

        self._emit(Event.STARTED, state)

        while state["running"] and state["epoch"] < state["max_epoch"]:
            state["epoch"] += 1
            self._emit(Event._ETIMER_STARTED, state)
            self._emit(Event.EPOCH_STARTED, state)
            self._emit(Event._REDUCER_RESET, state)
            self._emit(Event._PBAR_CREATED, state)

            state["_batches_iter"] = iter(state["batches"])
            self._run_epoch()

            self._emit(Event._PBAR_CLOSED, state)
            self._emit(Event._REDUCER_COMPUTED, state)
            self._emit(Event.EPOCH_FINISHED, state)
            self._emit(Event._ETIMER_FINISHED, state)

        self._emit(Event.FINISHED, state)
        state["running"] = False

    def resume(self, repeat_last_batch: bool = False) -> None:
        """Resume runner starting from the current state.

        Args:
            repeat_last_batch: Whether to repeat processing the last batch. Ignored if the
                last epoch is finished (i.e. the batches have been exhausted).
        """
        state = self.state
        state["running"] = True

        finished_last_epoch = state["n_iters"] % len(state["batches"]) == 0

        if not finished_last_epoch:
            self._emit(Event._ETIMER_STARTED, state)
            self._emit(Event._PBAR_CREATED, state)

            if repeat_last_batch:
                self._emit(Event.BATCH, state)
                self._emit(Event._REDUCER_UPDATED, state)
                self._emit(Event._PBAR_UPDATED, state)

            self._run_epoch()

            self._emit(Event._PBAR_CLOSED, state)
            self._emit(Event._REDUCER_COMPUTED, state)
            self._emit(Event.EPOCH_FINISHED, state)
            self._emit(Event._ETIMER_FINISHED, state)

        while state["running"] and state["epoch"] < state["max_epoch"]:
            state["epoch"] += 1
            self._emit(Event._ETIMER_STARTED, state)
            self._emit(Event.EPOCH_STARTED, state)
            self._emit(Event._REDUCER_RESET, state)
            self._emit(Event._PBAR_CREATED, state)

            state["_batches_iter"] = iter(state["batches"])
            self._run_epoch()

            self._emit(Event._PBAR_CLOSED, state)
            self._emit(Event._REDUCER_COMPUTED, state)
            self._emit(Event.EPOCH_FINISHED, state)
            self._emit(Event._ETIMER_FINISHED, state)

        self._emit(Event.FINISHED, state)
        state["running"] = False

    def _run_epoch(self) -> None:
        state = self.state
        while state["running"]:
            try:
                state["batch"] = next(state["_batches_iter"])
            except StopIteration:
                break
            state["n_iters"] += 1
            self._emit(Event.BATCH, state)
            self._emit(Event._REDUCER_UPDATED, state)
            self._emit(Event._PBAR_UPDATED, state)

    def _emit(self, event: Event, state: dict) -> None:
        for callback in self._callbacks[event]:
            if not state["running"]:
                break
            callback(state)

# Copyright 2019 <NAME> # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #     http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. A neural network runner.      A runner provides a thin abstraction of iterating over batches for several epochs,     which is typically done in neural network training. To customize the behavior during     a run, a runner provides a way to listen to events emitted during such run.     To listen to an event, call `Runner.on` and provide a callback which will be called     when the event is emitted. An event callback is a callable that accepts a `dict`     and returns nothing. The `dict` is the state of the run. By default, the state contains:      * ``batches`` - Iterable of batches which constitutes an epoch.     * ``max_epoch`` - Maximum number of epochs to run.     * ``n_iters`` - Current number of batch iterations.     * ``running`` - A boolean which equals ``True`` if the runner is still running. Can       be set to ``False`` to stop the runner earlier.     * ``epoch`` - Current number of epoch.     * ``batch`` - Current batch retrieved from ``state['batches']``.      Attributes:         state (dict): Runner's state that is passed to event callbacks.      Note:         Callbacks for an event are called in the order they are passed to `~Runner.on`.      Caution:         All the state contents above are required for a runner to function properly.         You are free to change their values to suit your use cases better, but be careful. Add single/multiple callback(s) to listen to an event.          If ``callbacks`` is ``None``, this method returns a decorator which accepts         a single callback for the event. If ``callbacks`` is a sequence of callbacks,         they will all be added as listeners to the event *in order*.          Args:             event: Event to listen.             callbacks: Callback(s) for the event.          Returns:             A decorator which accepts a callback, if ``callbacks`` is ``None``. # must be a single callback Run on batches for a number of epochs.          Args:             batches: Batches to iterate over in an epoch.             max_epoch: Maximum number of epochs to run. Resume runner starting from the current state.          Args:             repeat_last_batch: Whether to repeat processing the last batch. Ignored if the                 last epoch is finished (i.e. the batches have been exhausted).