#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserve.
#
# 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 ....core import CPUPlace, EOFException
from .... import compiler
from ....framework import Variable
from .... import io
from .... import profiler
from .... import scope_guard
from ....data_feeder import DataFeeder
from ....log_helper import get_logger
from ....reader import DataLoaderBase
from ..graph import *
from .config import ConfigFactory
import numpy as np
from collections import Iterable
import time
import os
import logging
import sys
import pickle
import functools
import traceback

__all__ = ['Context', 'Compressor']

_logger = get_logger(
    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s')


def cached_reader(reader, sampled_rate, cache_path, cached_id):
    """
    Sample partial data from reader and cache them into local file system.
    Args:
        reader: Iterative data source.
        sampled_rate(float): The sampled rate used to sample partial data for evaluation. None means using all data in eval_reader. default: None.
        cache_path(str): The path to cache the sampled data.
        cached_id(int): The id of dataset sampled. Evaluations with same cached_id use the same sampled dataset. default: 0.
    """
    np.random.seed(cached_id)
    cache_path = os.path.join(cache_path, str(cached_id))
    _logger.debug('read data from: {}'.format(cache_path))

    def s_reader():
        if os.path.isdir(cache_path):
            for file_name in open(os.path.join(cache_path, "list")):
                yield np.load(
                    os.path.join(cache_path, file_name.strip()),
                    allow_pickle=True)
        else:
            os.makedirs(cache_path)
            list_file = open(os.path.join(cache_path, "list"), 'w')
            batch = 0
            dtype = None
            for data in reader():
                if batch == 0 or (np.random.uniform() < sampled_rate):
                    np.save(
                        os.path.join(cache_path, 'batch' + str(batch)), data)
                    list_file.write('batch' + str(batch) + '.npy\n')
                    batch += 1
                    yield data

    return s_reader


class Context(object):
    """
    The context in the process of compression.
    """

    def __init__(self,
                 place,
                 scope,
                 train_graph=None,
                 train_reader=None,
                 eval_graph=None,
                 eval_reader=None,
                 teacher_graphs=None,
                 train_optimizer=None,
                 distiller_optimizer=None,
                 search_space=None):
        """
        Args:
            place: The device place where the compression job running.
            scope: The scope used in compression job.
            train_graph: The graph with loss as output node.
            eval_graph: The graph used for evaluation.
            eval_reader: The data reader used for evaluation.
            teacher_graphs: The teacher graphs used in distillation strategies.
            train_optimizer: The optimizer used to append backward ops and
                             optimization ops into train_graph.
            distiller_optimizer: The optimizer used by distillation strategies.
        """
        # The total number of epoches to be trained.
        self.epoch = 0
        # Current epoch
        self.epoch_id = 0
        # Current batch
        self.batch_id = 0

        self.k_v = {}

        self.place = place
        self.scope = scope
        self.train_graph = train_graph
        self.train_reader = train_reader
        self.eval_graph = eval_graph
        self.eval_reader = eval_reader
        self.executor = None
        self.teacher_graphs = teacher_graphs
        self.train_optimizer = train_optimizer
        self.distiller_optimizer = distiller_optimizer
        self.optimize_graph = None
        self.cache_path = './eval_cache'
        self.eval_results = {}

        self.skip_training = False
        self.search_space = search_space

    def to_file(self, file_name):
        """
        Save the context into file.
        """
        data = {}
        data['epoch_id'] = self.epoch_id
        data['eval_results'] = self.eval_results
        with open(file_name, 'wb') as context_file:
            pickle.dump(data, context_file)

    def from_file(self, file_name):
        """
        Load the context from file.
        """
        with open(file_name, 'rb') as context_file:
            if sys.version_info < (3, 0):
                data = pickle.load(context_file)
            else:
                data = pickle.load(context_file, encoding='bytes')
            self.epoch_id = data['epoch_id']
            self.eval_results = data['eval_results']

    def eval_converged(self, metric_name, delta=0.001):
        """
        Check whether the training has been converged.
        Args:
            metric_name(str): The metric used to check convergence.
            delta(float): '(metric[k] - metric[k-1] / metric[k-1]) < delta'
                          means that the training has been converged.
        Returns:
            bool: True means the training has been converged.
        """
        # TODO(wanghaoshuang@baidu.com): enhence this method.
        if (metric_name not in self.eval_results
            ) or len(self.eval_results[metric_name]) < 2:
            return False
        results = self.eval_results[metric_name][-2:]
        _logger.info('Latest evaluations: {}'.format(results))
        return abs(results[1] - results[0]) / results[0] < delta

    def run_eval_graph(self, sampled_rate=None, cached_id=0):
        """
        Evaluate the current mode in context.
        Args:
            sampled_rate(float): The sampled rate used to sample partial data
            for evaluation. None means using all data in eval_reader. default: None.
            cached_id(int): The id of dataset sampled. Evaluations with same
                            cached_id use the same sampled dataset. default: 0.
        """
        _logger.info('Running evaluation')
        assert self.eval_graph is not None
        assert self.eval_reader is not None
        eval_graph = self.eval_graph.clone(for_test=True)

        executor = SlimGraphExecutor(self.place)
        results = []
        batch_id = 0
        s_time = time.time()
        reader = self.eval_reader
        if sampled_rate:
            assert (not isinstance(reader, Variable))
            assert (sampled_rate > 0)
            assert (self.cache_path is not None)
            _logger.info('sampled_rate: {}; cached_id: {}'.format(sampled_rate,
                                                                  cached_id))
            reader = cached_reader(reader, sampled_rate, self.cache_path,
                                   cached_id)

        if isinstance(reader, Variable) or (
                isinstance(reader, DataLoaderBase) and (not reader.iterable)):
            reader.start()
            try:
                while True:
                    result = executor.run(eval_graph, self.scope)
                    result = [np.mean(r) for r in result]
                    results.append(result)
                    if batch_id % 20 == 0:
                        _logger.info("batch-{}; {}={}".format(
                            batch_id, eval_graph.out_nodes.keys(), result))
                    batch_id += 1
            except EOFException:
                reader.reset()
        else:
            for data in reader():
                result = executor.run(eval_graph, self.scope, data=data)
                result = [np.mean(r) for r in result]
                results.append(result)
                if batch_id % 20 == 0:
                    _logger.info("batch-{}; {}={}".format(
                        batch_id, eval_graph.out_nodes.keys(), result))
                batch_id += 1

        result = list(np.mean(np.array(results), axis=0))
        _logger.info("Final eval result: {}={}".format(
            eval_graph.out_nodes.keys(), result))
        if not isinstance(result, Iterable):
            result = [result]
        _logger.info('Finish evaluation')
        return result, eval_graph.out_nodes.keys()

    def put(self, key, value):
        self.k_v[key] = value

    def get(self, key):
        return self.k_v.get(key)


class Compressor(object):
    """
    The pass used to compress model.
    """

    def __init__(self,
                 place,
                 scope,
                 train_program,
                 train_reader=None,
                 train_feed_list=None,
                 train_fetch_list=None,
                 eval_program=None,
                 eval_reader=None,
                 eval_feed_list=None,
                 eval_fetch_list=None,
                 eval_func=None,
                 save_eval_model=True,
                 prune_infer_model=None,
                 teacher_programs=[],
                 checkpoint_path=None,
                 train_optimizer=None,
                 distiller_optimizer=None,
                 search_space=None,
                 log_period=20):
        """
        Args:
            place(fluid.Place): The device place where the compression job running.
            scope(fluid.core.Scope): The scope used to run graph.
            train_program(Program): The main program to be compressed. It must have loss op.
            train_reader: The data reader used for training.
            train_feed_list(dict): A dict to indicate the input variable of the training program.
                                   The key is user-defined and human-readable name.
                                   The value is the name of Variable.
            train_fetch_list(dict): A dict to indicate the output variable of the training program.
                                   The key is user-defined and human-readable name.
                                   The value is the name of Variable.
            eval_program(Program): The program used for evaluation.
            eval_reader: The data reader used for evaluation. It can be None if eval_func is not None.
            eval_feed_list(dict): A dict to indicate the input variable of the evaluation program.
                                   The key is user-defined and human-readable name.
                                   The value is the name of Variable.
                                   It can be None if eval_func is not None.
            eval_fetch_list(dict): A dict to indicate the output variable of the evaluation program.
                                   The key is user-defined and human-readable name.
                                   The value is the name of Variable.
            eval_func(dict|function): Callback functions used to evaluate the compressed model.
                                   The eval_func is a dict, the key is user-defined name and the value is 
                                   a callback function. And the score returned from callback functions 
                                   can be referenced in config file by the key of eval_func.
                                   The args of callback function are compressed eval_program and scope which
                                   store the compressed parameters.
                                   Default: None.
            save_eval_model(bool): Whether to save eval model when saving checkpoints. Default: True.
            prune_infer_model(tuple|list): If prune_infer_model is not None, compressor will prune
                                   eval program into inference program according to inputs and outputs
                                   defined in prune_infer_model. prune_infer_model[0] is a list of input
                                   variables' names and prune_infer_model[1] is a list of output variables'
                                   names. If prune_infer_model is None, it will not save inference model.
                                   Default: None.
            teacher_programs: The teacher graphs used in distillation strategies.
            train_optimizer: The optimizer used to append backward ops and
                             optimization ops into train_graph.
            distiller_optimizer: The optimizer used by distillation strategies. In distillation strategy,
                                 this optimizer is used to minimize the combined loss of student-net and
                                 teacher-net while train_optimizer is used to minimize loss of
                                 student-net in fine-tune stage. 
            search_space(slim.nas.SearchSpace): The instance that define the searching space. It must inherit
                              slim.nas.SearchSpace class and overwrite the abstract methods.
            log_period(int): The period of print log of training.

        """
        assert train_feed_list is None or isinstance(
            train_feed_list, list
        ), "train_feed_list should be a list of tuple, such as [('image', image.name), ('label', gt.name)]"
        assert eval_feed_list is None or isinstance(
            eval_feed_list, list
        ), "eval_feed_list should be a list of tuple, such as [('image', image.name), ('label', gt.name)]"
        self.strategies = []
        self.epoch = 0
        self.place = CPUPlace() if place is None else place
        self.scope = scope
        self.train_graph = GraphWrapper(
            train_program, in_nodes=train_feed_list, out_nodes=train_fetch_list)
        self.eval_graph = GraphWrapper(
            eval_program, in_nodes=eval_feed_list, out_nodes=eval_fetch_list)
        self.train_reader = train_reader
        self.eval_reader = eval_reader
        self.eval_func = eval_func
        self.save_eval_model = save_eval_model
        self.prune_infer_model = prune_infer_model

        self.teacher_graphs = []
        for teacher in teacher_programs:
            self.teacher_graphs.append(GraphWrapper(teacher))

        self.checkpoint = None
        self.checkpoint_path = checkpoint_path
        self.eval_epoch = 1

        self.train_optimizer = train_optimizer
        self.distiller_optimizer = distiller_optimizer
        self.init_model = None

        self.search_space = search_space
        self.log_period = log_period
        assert (log_period > 0)

    def _add_strategy(self, strategy):
        """
        Add a strategy to current compress pass.
        Args:
            strategy: The strategy to be added into current compress pass.
        """
        self.strategies.append(strategy)
        self.epoch = max(strategy.end_epoch, self.epoch)

    def config(self, config_file):
        """
        Configure the compress pass from file with yaml format.
        Args:
            config_file(str): The config file in local file system.
        """
        factory = ConfigFactory(config_file)
        self.epoch = factory.compressor['epoch']
        for strategy in factory.compressor['strategies']:
            self._add_strategy(strategy)
        if 'checkpoint_path' in factory.compressor:
            self.checkpoint_path = factory.compressor['checkpoint_path']

        if 'init_model' in factory.compressor:
            self.init_model = factory.compressor['init_model']

        if 'eval_epoch' in factory.compressor:
            self.eval_epoch = factory.compressor['eval_epoch']
        assert (self.eval_epoch > 0)

    def _init_model(self, context):
        """
        Load model that has been compressed. 
        """
        if self.init_model and os.path.exists(self.init_model):
            exe = SlimGraphExecutor(context.place)
            with scope_guard(context.scope):
                context.train_graph.load_persistables(self.init_model, exe)
            flops = context.eval_graph.flops()
            conv_flops = context.eval_graph.flops(only_conv=True)
            context.eval_graph.update_param_shape(context.scope)
            context.eval_graph.update_groups_of_conv()
            _logger.info("conv flops: -{}".format(1 - float(
                context.eval_graph.flops(only_conv=True)) / conv_flops))
            _logger.info("total flops: -{}".format(1 - float(
                context.eval_graph.flops()) / flops))
            context.train_graph.update_param_shape(context.scope)
            context.train_graph.update_groups_of_conv()
            context.train_graph.infer_shape()
            _logger.info("Init model from: {}".format(self.init_model))

    def _load_checkpoint(self, context):
        """
        Load checkpoints from file.
        """
        _logger.debug('_load_checkpoint')
        strategies = self.strategies
        if self.checkpoint_path:
            if not os.path.exists(self.checkpoint_path):
                _logger.warning("Checkpints path doesn't exist: [{}]".format(
                    self.checkpoint_path))
                return context, strategies
            checkpoints = [
                dir for dir in os.listdir(self.checkpoint_path)
                if os.path.isdir(os.path.join(self.checkpoint_path, dir))
            ]
            _logger.debug('self.checkpoint_path: {}'.format(
                self.checkpoint_path))
            _logger.info('checkpoints: {}'.format(checkpoints))
            if len(checkpoints) > 0:
                latest = max([int(ck) for ck in checkpoints])
                latest_ck_path = os.path.join(self.checkpoint_path, str(latest))

                model_path = os.path.join(latest_ck_path, 'model')
                context_path = os.path.join(latest_ck_path, 'context')
                strategy_path = os.path.join(latest_ck_path, 'strategies')
                if os.path.exists(context_path):
                    context.from_file(context_path)
                    context.epoch_id += 1
                if os.path.exists(strategy_path):
                    with open(strategy_path, 'rb') as strategy_file:
                        if sys.version_info < (3, 0):
                            strategies = pickle.load(strategy_file)
                        else:
                            strategies = pickle.load(
                                strategy_file, encoding='bytes')
                assert (len(self.strategies) == len(strategies))
                for s, s1 in zip(self.strategies, strategies):
                    s1.__dict__.update(s.__dict__)

                for strategy in strategies:
                    strategy.restore_from_checkpoint(context)

                if os.path.exists(model_path):
                    exe = SlimGraphExecutor(context.place)
                    with scope_guard(context.scope):
                        context.optimize_graph.load_persistables(model_path,
                                                                 exe)
                    _logger.info("Loaded params from: {}".format(model_path))
        return context, strategies

    def _save_checkpoint(self, context):
        """
        Save checkpoints to file.
        """
        if context.epoch_id % 1 == 0 and self.checkpoint_path:
            checkpoint_path = os.path.join(self.checkpoint_path,
                                           str(context.epoch_id))
            model_path = os.path.join(checkpoint_path, 'model')
            eval_model_path = os.path.join(checkpoint_path, 'eval_model')
            context_path = os.path.join(checkpoint_path, 'context')
            strategy_path = os.path.join(checkpoint_path, 'strategies')
            if not os.path.isdir(model_path):
                os.makedirs(model_path)
            exe = SlimGraphExecutor(context.place)
            with scope_guard(context.scope):
                context.optimize_graph.save_persistables(model_path, exe)
                if self.save_eval_model:
                    context.eval_graph.save_model(eval_model_path, exe)
                if self.prune_infer_model:
                    context.eval_graph.save_infer_model(
                        eval_model_path,
                        exe,
                        self.prune_infer_model,
                        program_only=self.save_eval_model)

            context.to_file(context_path)
            with open(strategy_path, 'wb') as strategy_file:
                pickle.dump(self.strategies, strategy_file)
            _logger.info('Saved checkpoint to: {}'.format(checkpoint_path))

    def _train_one_epoch(self, context):
        """
        Train one epoch.
        """
        if context.skip_training:
            return
        executor = SlimGraphExecutor(self.place)

        if context.optimize_graph.compiled_graph is None:
            build_strategy = compiler.BuildStrategy()
            build_strategy.fuse_all_reduce_ops = False
            context.optimize_graph.compiled_graph = compiler.CompiledProgram(
                context.optimize_graph.program).with_data_parallel(
                    loss_name=context.optimize_graph.out_nodes['loss'],
                    build_strategy=build_strategy)

        if isinstance(context.train_reader, Variable) or (
                isinstance(context.train_reader, DataLoaderBase) and
            (not context.train_reader.iterable)):
            context.train_reader.start()
            try:
                while True:

                    for strategy in self.strategies:
                        strategy.on_batch_begin(context)
                    results = executor.run(context.optimize_graph,
                                           context.scope)
                    results = [float(np.mean(result)) for result in results]
                    if context.batch_id % self.log_period == 0:
                        _logger.info("epoch:{}; batch_id:{}; {} = {}".format(
                            context.epoch_id, context.batch_id,
                            context.optimize_graph.out_nodes.keys(
                            ), [round(r, 6) for r in results]))
                    for strategy in self.strategies:
                        strategy.on_batch_end(context)
                    context.batch_id += 1

            except EOFException:
                context.train_reader.reset()

        else:
            for data in context.train_reader():
                for strategy in self.strategies:
                    strategy.on_batch_begin(context)
                results = executor.run(context.optimize_graph,
                                       context.scope,
                                       data=data)
                results = [float(np.mean(result)) for result in results]
                if context.batch_id % self.log_period == 0:
                    _logger.info("epoch:{}; batch_id:{}; {} = {}".format(
                        context.epoch_id, context.batch_id,
                        context.optimize_graph.out_nodes.keys(
                        ), [round(r, 6) for r in results]))
                for strategy in self.strategies:
                    strategy.on_batch_end(context)
                context.batch_id += 1
        context.batch_id = 0

    def _eval(self, context):
        """
        Runing evaluation.
        """
        if self.eval_func is not None:
            for key in self.eval_func:
                func = self.eval_func[key]
                if key not in context.eval_results:
                    context.eval_results[key] = []
                context.eval_results[key].append(
                    func(self.eval_graph.program, self.scope))
        else:
            results, names = context.run_eval_graph()
            for name, result in zip(names, results):
                if name not in context.eval_results:
                    context.eval_results[name] = []
                context.eval_results[name].append(result)

    def run(self):
        """
        Execute compressing pass.
        """
        context = Context(
            place=self.place,
            scope=self.scope,
            train_graph=self.train_graph,
            train_reader=self.train_reader,
            eval_graph=self.eval_graph,
            eval_reader=self.eval_reader,
            teacher_graphs=self.teacher_graphs,
            train_optimizer=self.train_optimizer,
            distiller_optimizer=self.distiller_optimizer,
            search_space=self.search_space)
        self.context = context
        if self.teacher_graphs:
            context.put('teachers', self.teacher_graphs)
        self._init_model(context)
        if not context.optimize_graph:
            if context.train_optimizer:
                context.train_optimizer._name = 'train_opt'
                context.optimize_graph = context.train_graph.get_optimize_graph(
                    context.train_optimizer, context.place, context.scope)
            else:
                context.optimize_graph = context.train_graph

        context, self.strategies = self._load_checkpoint(context)

        for strategy in self.strategies:
            strategy.on_compression_begin(context)
        if 'MKLDNNPostTrainingQuantStrategy' in [
                i.__class__.__name__ for i in self.strategies
        ]:
            return None
        start = context.epoch_id
        for epoch in range(start, self.epoch):
            context.epoch_id = epoch
            try:
                for strategy in self.strategies:
                    strategy.on_epoch_begin(context)
                self._train_one_epoch(context)
                if self.eval_epoch and epoch % self.eval_epoch == 0:
                    self._eval(context)
                self._save_checkpoint(context)
                for strategy in self.strategies:
                    strategy.on_epoch_end(context)
            except Exception:
                _logger.error(traceback.print_exc())
                continue
        for strategy in self.strategies:
            strategy.on_compression_end(context)
        return context.eval_graph
