import json

import os

from abc import abstractmethod

import matplotlib.pyplot

import numpy as np

import tensorflow as tf

# Baseline class for all your Deep Learning needs with TensorFlow #

class DLMODEL(object):

    class Config(object):

        def __init__(self):

            self.modelname = ''

            self.model_config = {}

            self.checkpointDir = None

            self.description = ''

            self.batchsize = 6

            self.useTensorboard = True

            self.tensorboardPort = 8008

            self.useMatplotlib = False

            self.debugGradients = False

            self.tfSummaryAfter = 100

            self.dataset = ''

            self.beta1 = 0.5

    def __init__(self, sess, config=Config()):

        """

        Args:

          sess: TensorFlow session

          config: (optional) a DLMODEL.Config object with your options

        """

        self.sess = sess

        self.config = config

        self.variables = {}

        self.curves = {}  # For plotting via matplotlib

        self.phase = tf.placeholder(tf.bool, name='phase')

        self.handles = {}

        if self.config.useMatplotlib:

            self.handles['curves'] = matplotlib.pyplot.figure()

            self.handles['samples'] = matplotlib.pyplot.figure()

        self.saver = None

        self.losses = None

    @abstractmethod

    def train(self, dataset):

        """Train a Deep Neural Network"""

    def initialize_variables(self):

        uninitialized_var_names_raw = set(self.sess.run(tf.report_uninitialized_variables()))

        uninitialized_var_names = [v.decode() for v in uninitialized_var_names_raw]

        variables_to_initialize = [v for v in tf.global_variables() if v.name.split(':')[0] in uninitialized_var_names]

        self.sess.run(tf.initialize_variables(variables_to_initialize))

        print("Initialized all unitialized variables.")

    @property

    def model_dir(self):

        return "{}_d{}_b{}_{}".format(self.config.modelname, self.config.dataset, self.config.batchsize, self.config.description)

    def save(self, checkpoint_dir, step):

        model_name = self.config.modelname + ".model"

        checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)

        # Create checkpoint directory, if it does not exist

        if not os.path.exists(checkpoint_dir):

            os.makedirs(checkpoint_dir)

        # Save the current model state to the checkpoint directory

        self.saver.save(self.sess,

                        os.path.join(checkpoint_dir, model_name),

                        global_step=step)

        # Save the config to a json file, such that you can inspect it later

        with open(os.path.join(checkpoint_dir, 'Config-{}.json'.format(step)), 'w') as outfile:

            try:

                json.dump(self.config.__dict__, outfile)

            except:

                print("Failed to save config json")

        # Save the curves to a np file such that we can recover and monitor the entire training process

        np.save(os.path.join(checkpoint_dir, 'Curves.npy'), self.curves)

    def load(self, checkpoint_dir, iteration=None):

        import re

        print(" [*] Reading checkpoints...")

        checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)

        # Load training curves, if any

        curves_file = os.path.join(checkpoint_dir, 'Curves.npy')

        if os.path.isfile(curves_file):

            self.curves = np.load(curves_file, allow_pickle=True).item()

        if iteration is not None:

            self.saver.restore(self.sess, os.path.join(checkpoint_dir, self.config.modelname + '.model-' + str(iteration)))

            counter = iteration

            return True, counter

        ckpt = tf.train.get_checkpoint_state(checkpoint_dir)

        if ckpt and ckpt.model_checkpoint_path:

            ckpt_name = os.path.basename(ckpt.model_checkpoint_path)

            self.saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))

            counter = int(next(re.finditer('(\d+)(?!.*\d)', ckpt_name)).group(0))

            print(" [*] Success to read {}".format(ckpt_name))

            return True, counter

        else:

            print(" [*] Failed to find a checkpoint")

            return False, 0

    @staticmethod

    def create_optimizer(loss, var_list=(), learningrate=0.001, type='ADAM', beta1=0.05, momentum=0.9, name='optimizer', minimize=True, scope=None):

        if type == 'ADAM':

            optim = tf.train.AdamOptimizer(learningrate, beta1=beta1, name=name)

        elif type == 'SGD':

            optim = tf.train.GradientDescentOptimizer(learningrate)

        elif type == 'MOMENTUM':

            optim = tf.train.MomentumOptimizer(learning_rate=learningrate, momentum=momentum)

        elif type == 'RMS':

            optim = tf.train.RMSPropOptimizer(learning_rate=learningrate, momentum=momentum)

        else:

            raise ValueError('Invalid optimizer type')

        if minimize:

            update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

            with tf.control_dependencies(update_ops):

                train_op = optim.minimize(loss, var_list=var_list)

            return train_op

        else:

            return optim

    @staticmethod

    def get_number_of_trainable_params():

        def inner_get_number_of_trainable_params(_scope):

            total_parameters = 0

            _variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, _scope)

            for variable in _variables:

                # shape is an array of tf.Dimension

                shape = variable.get_shape()

                variable_parametes = 1

                for dim in shape:

                    variable_parametes *= dim.value

                total_parameters += variable_parametes

            return total_parameters

        variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, "")

        scopes = list(set(map(lambda variable: os.path.split(os.path.split(variable.name)[0])[0], variables)))

        for scope in scopes:

            if scope != '':

                print(f'#Params in {scope}: {inner_get_number_of_trainable_params(scope)}')

        print(f'#Params in total: {inner_get_number_of_trainable_params("")}')