# ----------------------------------------------------------------------------
# Copyright 2014 Nervana Systems Inc.
# 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.
# ----------------------------------------------------------------------------
import os
import sys
import logging
import h5py
from collections import deque
from neon import NervanaObject
from neon.util.persist import save_obj
from timeit import default_timer
from neon.layers import Convolution
import time
import numpy as np
logger = logging.getLogger(_
_name__)


class Callbacks(NervanaObject):

    """
    Container class for storing and iterating over
 callbacks.

    Attributes:
        callbacks (list): Ordered set of Callback objects to be run.
    """

    def __init__(self, model, train_set,
                 output_file=None,
                 eval_freq=None,
                 progress_bar=True,
                 epochs=None,
                 save_path=None,
                 serialize=0,
                 history=1,
                 model_file=None,
                 eval_set=None,
                 metric=None):
        """
        Create a callbacks container with the default callbacks.

        Arguments:
            model (Model): the model object
            train_set (DataIterator): the training dataset
            output_file (string, optional): path to save callback data to
            eval_freq (int, optional): how often (in epochs) to run evaluation
            progress_bar (bool): control whether a progress bar
            callback is created.  Defaults to True.
            epochs (int): how many epochs the model will train for (default: None)
            save_path (string): file path to save model snapshots (default: None)
            serialize (int): serialize model every N epochs (default: 0)
            history (int): number of checkpoint files to retain (default: 1)
            model_file(string, optional): file to load weights (serialized model) from
            eval_set (DataIterator, optional): the dataset upon which to evaluate loss
                                               or metric
            metric (Metric, optional):  metric to evaluate
       """
        self.callbacks = list()
        self.epoch_marker = 0
        if output_file is None:
            if hasattr(self, 'callback_data'):
                del self.callback_data
            self.callback_data = h5py.File("no_file", driver='core', backing_store=False)
        else:
            if os.path.isfile(output_file):
                logger.warn("Overwriting output file %s", output_file)
                os.remove(output_file)
            self.callback_data = h5py.File(output_file, "w")

        if model_file:
            model.load_params(model_file)

        self.model = model
        self.train_set = train_set

        self.add_callback(TrainCostCallback(self.callback_data, self.model))

        if progress_bar:
            self.add_callback(ProgressBarCallback(self.callback_data, model, train_set))

        if eval_freq:
            if not eval_set:
                err_msg = 'Evaluation frequency specified but no eval set provided`!'
                logger.exception(err_msg)
                raise ValueError(err_msg)

            ecb = LossCallback(self.callback_data, model, eval_set, eval_freq)
            self.add_callback(ecb, insert_pos=0)
            if metric:
                ecb = MetricCallback(self.callback_data, model, eval_set, metric, eval_freq)
                self.add_callback(ecb, insert_pos=None)

        if save_path:
            serialize_interval = serialize if serialize > 1 else 1
            checkpoint_schedule = range(0, epochs, serialize_interval)
            if not checkpoint_schedule:
                err_msg = 'For %d epochs and schedule %d, model will not be serialized' % (
                          epochs, serialize)
                logger.exception(err_msg)
                raise ValueError(err_msg)
            scb = SerializeModelCallback(model, save_path, checkpoint_schedule, history)
            self.add_callback(scb)

        self.add_callback(TrainLoggerCallback(self.callback_data, model))
        self.add_callback(RunTimerCallback(self.callback_data, model))

    def add_deconv_callback(self, train_set, valid_set, max_fm=16, dataset_pct=25):
        """
        Convenience function to create and add a deconvolution callback. The data can be used for
        visualization.

        Arguments:
            train_set (DataIterator): the train dataset to use
            valid_set (DataIterator): the validation dataset to use
        """
        self.add_callback(DeconvCallback(self.callback_data, self.model,
                                         train_set, valid_set,
                                         max_fm=max_fm, dataset_pct=dataset_pct))

    def add_save_best_state_callback(self, path):
        """
        Convenience function to create and add a save best state callback.

        Arguments:
            path (string): where to save the best model state.
        """
        self.add_callback(SaveBestStateCallback(self.callback_data, self.model, path))

    def add_early_stop_callback(self, stop_func):
        """
        Convenience function to create and add an early stopping callback.

        Arguments:
            stop_func (function): function to determine when to stop.
        """
        self.add_callback(EarlyStopCallback(self.callback_data, self.model, stop_func))

    def add_hist_callback(self, plot_per_mini=False):
        self.callbacks.append(HistCallback(self.callback_data, self.model,
                              plot_per_mini=plot_per_mini))

    def add_callback(self, callback, insert_pos=None):
        """
        Add a user supplied callback. Since callbacks are run serially and share data,
        order can matter.  If the default behavior (to append the callback) is not
        sufficient, insert position can be controlled.

        Arguments:
            callback (Callback): callback object to be registered
            insert_pos (int, optional): position in the list to insert the callback.
                                        Defaults to None, meaning append
        """
        if insert_pos is None:
            self.callbacks.append(callback)
        else:
            self.callbacks.insert(insert_pos, callback)

    def on_train_begin(self, epochs):
        """
        Call all registered callbacks' on_train_begin functions
        """
        # data iterator wraps around to avoid partial minibatches
        # callbacks producing per-minibatch data need a way to preallocate buffers
        config = self.callback_data.create_group('config')
        total_minibatches = -((-self.train_set.ndata * epochs) // self.be.bsz)
        config.attrs['total_minibatches'] = total_minibatches
        config.attrs['total_epochs'] = epochs

        time_markers = self.callback_data.create_group("time_markers")
        time_markers.create_dataset("minibatch", (epochs,))

        for c in self.callbacks:
            c.on_train_begin(epochs)

    def on_train_end(self):
        """
        Call all registered callbacks' on_train_end functions
        """
        for c in self.callbacks:
            c.on_train_end()

        self.callback_data.close()

    def on_epoch_begin(self, epoch):
        """
        Call all registered callbacks' on_epoch_begin functions

        Arguments:
            epoch (int): index of epoch that is beginning
        """
        for c in self.callbacks:
            if c.should_fire(epoch, c.epoch_freq):
                c.on_epoch_begin(epoch)

    def on_epoch_end(self, epoch):
        """
        Call all register