#   Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
#
# 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 __future__ import print_function
from __future__ import absolute_import
from __future__ import unicode_literals

import os
import json
from functools import reduce
import six
from time import time
import shutil

import logging
import numpy as np
import paddle.fluid as F
import paddle.fluid.layers as L

from propeller import util
from propeller.types import StopException, ProgramPair
from propeller.paddle.train import hooks
from . import distribution

log = logging.getLogger(__name__)

__all__ = ['MonitoredExecutor', 'Saver']


class RunState(object):
    @classmethod
    def from_str(cls, s):
        j = json.loads(s)
        ret = RunState()
        ret._gstep = j['global_step']
        ret._time = j['time']
        ret._step = 0
        return ret

    def __init__(self):
        self._gstep = 0
        self._step = 0
        self._time = time()

    @property
    def gstep(self):
        return self._gstep

    @property
    def step(self):
        return self._step

    @property
    def time(self):
        return self._time

    def __repr__(self):
        return repr({'global_step': self._gstep, 'time': self._time})

    def serialize(self):
        return json.dumps({'global_step': self._gstep, 'time': self._time})

    def next(self):
        ret = RunState()
        ret._gstep = self._gstep + 1
        ret._step = self._step + 1
        ret._time = time()
        return ret


class Saver(object):
    def __init__(self,
                 save_dir,
                 exe,
                 program,
                 save_prefix='model',
                 max_ckpt_to_keep=None):
        if exe is not None:
            assert isinstance(
                exe, F.Executor
            ), 'expect normal executor to save, got executor of type %s' % repr(
                type(exe))
        self._exe = exe
        self._program = program
        self._save_dir = save_dir
        self._save_prefix = save_prefix
        self._max_ckpt_to_keep = 10 if max_ckpt_to_keep is None else max_ckpt_to_keep

        self.ckpt_info_path = os.path.join(save_dir, 'ckpt_info')

        if os.path.exists(self.ckpt_info_path):
            self.ckpt_list = [
                p.strip() for p in open(self.ckpt_info_path).readlines()
            ]
            log.debug('ckpt_list in this Saver: %s' % (self.ckpt_list))
        else:
            self.ckpt_list = []

    @property
    def last_ckpt(self):
        return self.ckpt_list[-1] if len(self.ckpt_list) else None

    def save(self, state):
        save_name = '%s_%d' % (self._save_prefix, state.gstep)
        save_dir = os.path.join(self._save_dir, save_name)
        tmp_dir = os.path.join(self._save_dir, 'tmp')
        try:
            shutil.rmtree(save_dir)
            shutil.rmtree(tmp_dir)
        except OSError:
            pass
        log.debug('saving step %d to %s' % (state.gstep, save_dir))
        F.io.save_persistables(self._exe, tmp_dir, self._program)
        shutil.move(tmp_dir, save_dir)
        meta = state.serialize()
        open(os.path.join(save_dir, 'meta'), 'w').write(meta)

        self.ckpt_list.append(save_name)
        if len(self.ckpt_list) > self._max_ckpt_to_keep:
            ckpt_to_keep = self.ckpt_list[-self._max_ckpt_to_keep:]
            ckpt_to_remove = set(self.ckpt_list) - set(ckpt_to_keep)
            self.ckpt_list = ckpt_to_keep
            for ckpt in ckpt_to_remove:
                ckpt_dir = os.path.join(self._save_dir, ckpt)
                if os.path.exists(ckpt_dir):
                    shutil.rmtree(ckpt_dir)
                    log.debug('No. of ckpt exceed %d, clean up: %s' %
                              (self._max_ckpt_to_keep, ckpt_dir))
        open(self.ckpt_info_path, 'w').write('\n'.join(self.ckpt_list))

    def restore(self, ckpt=-1):
        if not isinstance(ckpt, (int, ) + six.string_types):
            raise ValueError('ckpt type not understood %s' % repr(ckpt))
        if isinstance(ckpt, int):
            try:
                ckpt = self.ckpt_list[ckpt]
            except IndexError:
                raise ValueError('invalid restore ckpt number %d' % ckpt)
        if isinstance(ckpt, six.string_types):
            try:
                ckpt = self.ckpt_list.index(ckpt)
            except ValueError:
                raise ValueError('ckpt: %s not in ckpt list: %s' %
                                 (ckpt, self.ckpt_list))

        path = os.path.join(self._save_dir, self.ckpt_list[ckpt])
        meta_file = os.path.join(path, 'meta')
        if not os.path.exists(meta_file):
            raise RuntimeError('meta not found in restore dir: %s' % path)
        state = RunState.from_str(open(meta_file).read())
        log.info('restore from ckpt %s, ckpt-status: %s' % (path, repr(state)))

        def fn(v):
            vpath = os.path.join(path, v.name)
            if F.io.is_persistable(v):
                if os.path.exists(vpath):
                    return True
                else:
                    log.warning('var %s not found in checkpoint, ignored' %
                                v.name)
            return False

        F.io.load_vars(
            self._exe, path, main_program=self._program, predicate=fn)
        return state


class MonitoredExecutor(object):
    """A wrapper handling the train loop"""

    def __init__(
            self,
            executor,
            program,
            loss=None,  #must set in train
            state=None,
            run_config=None,  #none if not load
            run_hooks=[],
            warm_start_setting=None):
        if not isinstance(executor, F.Executor):
            raise ValueError('PE is no longer supported')
        if isinstance(executor, F.ParallelExecutor):
            raise ValueError('ParallelExecutor is deprecatd, use Executor')
        self._exe = executor
        self._hooks = run_hooks
        self._state = RunState()  # might be overwrite in freeze
        self._program = program
        self._loss = loss
        self._warm_start_setting = warm_start_setting
        self._saver = None  # will set in prepare
        self.result = None  # will set after train
        if run_config is not None:
            self._model_dir = run_config.model_dir
            self._save_dir = run_config.model_dir
            self._save_steps = run_config.save_steps
            self._skip_steps = run_config.skip_steps if run_config.skip_steps else 100
            self._save_prefix = 'model'
            self._max_ckpt = run_config.max_ckpt

    @property
    def state(self):
        return self._state

    def init_or_restore_variables(self):
        # The order of this 2 steps really matters
        # 1. init train

        F.Executor(F.cuda_places()[0]).run(self._program.startup_program)
        # 2. restore param
        if self._warm_start_setting is not None:
            if not os.path.exists(self._warm_start_setting.from_dir):
                raise ValueError('warm start dir not exists: %s' %
                                 self._warm_start_setting.from_dir)
            log.info("warm start from %s" % self._warm_start_setting.from_dir)
            if self._warm_start_setting.predicate_fn is not None:

                def fn(v):
                    ret = self._warm_start_setting.predicate_fn(v)
                    if ret:
                        log.info('warm start: %s' % v.name)
                    return ret

                F.io.load_vars(
                    F.Executor(F.cuda_places()[0]),
                    self._warm_start_setting.from_dir,
                    main_program=self._program.train_program,
                    predicate=fn)
            else:
                raise NotImplementedError()

        self._saver = Saver(
            self._model_dir,
            F.Executor(F.cuda_places()[0]),
            program=self._program.train_program,
            max_ckpt_to_keep=self._max_ckpt)
        if self._saver.last_ckpt is not None:
            self._state = self._saver.restore()

    def freeze(self):
        if self._loss is None:
            log.debug('will not freeze a program without loss')
            return
        if isinstance(self._program.train_program, F.compiler.CompiledProgram):
            log.debug('program has already been built')
            return
        exec_strategy = F.ExecutionStrategy()
        exec_strategy.num_threads = 4  #2 for fp32 4 for fp16
        exec_strategy.use_experimental_executor = True
        exec_strategy.num_iteration_per_drop_scope = 10  #important shit

        build_strategy = F.BuildStrategy()
        build_strategy.remove_unnecessary_lock = False
        #build_strategy.fuse_broadcast_ops = True
        build_strategy.num_trainers = distribution.status.num_replica
        build_strategy.trainer_id = distribution.status.replica_id
        build_strategy.memory_optimize = True

        log.info('replica id %d of %d' % (distribution.status.replica_id,
                                          distribution.status.num_replica))

        program = F.CompiledProgram(
            self._program.train_program).with_data_parallel(
                loss_name=self._loss.name,
                build_strategy=build_strategy,
                exec_strategy=exec_strategy)
        self._program = ProgramPair(
            train_program=program,
            startup_program=self._program.startup_program)

    def __enter__(self):
        log.debug('freezing program')
        self.freeze()
        log.debug('done freezing')
        log.info('********** Start Loop ************')
        # TODO init

        self.result = None
        for h in self._hooks:
            log.debug('train loop has hook %s' % h)
            h.before_train()
        return self

    def run(self, fetch_list=[], *args, **kwargs):
        #log.debug('Executor running step %d' % self._state.gstep)
        if self._hooks:
            fetch_list = [fetch_list]
            for h in self._hooks:
                #log.debug('calling hook.before_run %s' % h)
                fetch = h.before_run(self._state)
                fetch_list.append(fetch)
            fetch_list_len = map(len, fetch_list)
            fetch_list, schema = util.flatten(fetch_list)
            fetch_list = [
                f.name if not isinstance(f, six.string_types) else f
                for f in fetch_list
            ]
            #if len(set(fetch_list)) != len(fetch_list):
            #    log.error('strange shit happend when fetch list has idetity tensors %s' % fetch_list)
            res = self._exe.run(self._program.train_program,
                                fetch_list=fetch_list,
                                *args,
                                **kwargs)
            res = [self.merge_result(r) for r in res]
            #log.debug(res)

            res = util.unflatten(res, schema)
            ret, res = res[0], res[1:]
            for r, h in zip(res, self._hooks):
                #log.debug('calling hook.after_run')
                h.after_run(r, self._state)

            if any(map(lambda i: i.should_stop(self._state), self._hooks)):
                raise StopException('hook call stop')
        else:
            ret = self._exe.run(self._program.train_program,
                                fetch_list=fetch_list,
                                *args,
                                **kwargs)
        self._state = self._state.next()
        return ret

    def __exit__(self, err_type, err_value, trace):
        if (err_type is None) or isinstance(err_value, (
                F.core.EOFException, StopException, KeyboardInterrupt)):
            try:
                log.info('********** Stop Loop ************')
                self.result = []
                for h in self._hooks:
                    self.result.append(h.after_train())
            except Exception as e:
                log.exception('error occur after loop %s' % repr(e))
        else:
            log.info('********** Interupt Loop ************')
            log.exception('error occur during loop %s: %s' %
                          (err_type, err_value))

    def merge_result(self, ls):
        dev_count = len(self._program.train_program._places) if isinstance(
            self._program.train_program, F.compiler.CompiledProgram) else 1
        if dev_count == 1:
            return ls
        else:
            shape = (-1, ls.shape[0] // dev_count) + ls.shape[1:]
            ret = np.reshape(ls, shape).mean(axis=0)
            return ret