import os
from abc import ABC
from datetime import datetime

import numpy as np
import tensorflow as tf

from trainers.DLMODEL import DLMODEL
from utils.logger import Logger, Phase


class AEMODEL(DLMODEL, ABC):
    class Config(DLMODEL.Config):
        def __init__(self, modelname='AE'):
            super().__init__()
            self.modelname = modelname
            self.intermediateResolutions = [8, 8]
            self.outputWidth = 256
            self.outputHeight = 256
            self.numChannels = 3
            self.dropout = False
            self.dropout_rate = 0.2
            self.zDim = 128

    def __init__(self, sess, config=Config(), network=None):
        super().__init__(sess, config)
        self.losses = {}
        self.dropout = tf.placeholder(tf.bool, name='dropout')
        self.dropout_rate = tf.placeholder(tf.float32, name='dropout_rate')

        self.network = network
        self.checkpointDir = os.path.join(self.config.checkpointDir, self.network.__name__)
        self.logDir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'logs', self.network.__name__, self.model_dir,
                                   datetime.now().strftime('%Y%m%d_%H%M%S'))
        self.logger = Logger(self.sess, self.logDir)

    def log_to_tensorboard(self, epoch, scalars, visuals, phase: Phase, name='x'):
        for key in scalars.keys():
            scalars[key] = np.mean(scalars[key])

        if visuals:
            self.logger.summarize(epoch, phase=phase, summaries_dict={**scalars, **{name: np.vstack(visuals)[:50]}})

    def load_checkpoint(self):
        could_load, checkpoint_counter = self.load(self.checkpointDir)
        if could_load:
            last_epoch = checkpoint_counter
            print(" [*] Load SUCCESS")
        else:
            last_epoch = 0
            print(" [!] Load failed...")
        return last_epoch

    @property
    def model_dir(self):
        return "{}_d{}_s{}x{}_{}_b{}_z{}_{}".format(self.config.modelname, self.config.dataset,
                                                    self.config.outputWidth,
                                                    self.config.outputHeight,
                                                    self.network.__name__,
                                                    self.config.batchsize, self.config.zDim,
                                                    self.config.description)


def update_log_dicts(scalars, visuals, train_scalars, train_visuals):
    for k, v in list(scalars.items()):
        train_scalars[k].append(v)
    train_visuals.append(visuals)


def indicate_early_stopping(current_cost, best_cost, last_improvement):
    if current_cost < best_cost:
        best_cost = current_cost
        last_improvement = 0
        return best_cost, last_improvement, False
    else:
        last_improvement += 1
        if last_improvement >= 5:
            return best_cost, last_improvement, True
        return best_cost, last_improvement, False