# -*- coding: utf-8 -*-
"""
Created on Sat Mar 10 15:14:47 2018

@author: tom.s (at) halina9000.com
"""
import numpy as np
import tensorflow as tf
from keras.models import Sequential
from keras.layers import Dense
from keras.initializers import RandomUniform
from keras.callbacks import Callback
from keras.optimizers import TFOptimizer, Adam

from keras import regularizers

from time import time
import os
import h5py


class BestAccs(Callback):
    """Saves weights after epoch if accuracies are better than previously.

    Criterion `better` is defined as follows:
    - both accuracies are >= min_accs,
    - difference between accuracies is <= diff,
    - current minimum accuracy (training or test) is better than previous one.

    # Arguments
        filepath (str): filepath (file name and directory) for saving
            weights.
        min_accs (float): minimum value of both acc and val_acc.
            Default: 0.0.
        diff (float): maximum difference between acc and val_acc.
            Default: 1.0
    """

    def __init__(self, filepath, min_accs=0.0, diff=1.0):
        super(BestAccs, self).__init__()
        self.filepath = filepath
        self.min_accs = min_accs
        self.diff = diff


    def on_train_begin(self, logs={}):
        self.best_acc = 0.0


    def on_epoch_end(self, epoch, logs={}):
        acc = logs.get('acc')
        val_acc = logs.get('val_acc')
        epoch_acc = np.minimum(acc, val_acc)

        if epoch_acc > np.max(self.best_acc):
            if np.absolute(acc - val_acc) <= self.diff:
                self.best_acc = epoch_acc
                if epoch_acc >= self.min_accs:
                    filepath = self.filepath.format(epoch_acc, **logs)
                    self.model.save_weights(filepath, overwrite=True)


def best_batch_size(train_x, train_y, epochs=200):
    """
    Determine optimal batch size.

    Parameters
    ----------
        train_x : np.array(float)
            Training `x` set (training features).
        train_y : np.array(int)
            Training `y` set (training labels).
        epochs : int, optional
            Number of epochs.

    Returns
    -------
        batch_size : int
            Size of most efficient (fastest) batch size.
        batch_exe_time: list
            List of execution time for given range of batch size.

    """
    lr = 0.005
    optimizer = TFOptimizer(tf.train.GradientDescentOptimizer(lr))

    batch_size_limit = int(np.log2(train_x.shape[0])) + 1
    batch_size_set = [2**x for x in range(5, batch_size_limit + 1)]


    model = Sequential()

    model.add(Dense(1,
                    kernel_initializer='zeros',
                    bias_initializer='zeros',
                    input_dim=train_x.shape[1],
                    activation='sigmoid'))

    model.compile(loss='binary_crossentropy',
                  optimizer=optimizer,
                  metrics=['accuracy'])

    # Measeure of execution time for each batch_size
    batches_exe_time = []
    for batch_size in batch_size_set:
        time_start = time()
        model.fit(train_x, train_y,
                  epochs=epochs,
                  batch_size=batch_size,
                  verbose=0)
        time_end = time() - time_start
        batches_exe_time.append([time_end, batch_size])
    batches_exe_time.sort()
    batch_size = batches_exe_time[0][1]
    return batch_size, batches_exe_time


def course_assignment(train_x, train_y, test_x, test_y,
                      file_output_path,
                      initializer='zeros', batch_size=256):
    """
    First Coursera Deep Learning course programming assignment in Keras.

    Training of single neuron with sigmoid activation function and set
    of images. Our goal is to teach that neuron to recognize images with cat.

    Parameters
    ----------
        train_x : np.array(float)
            Training `x` set (training features).
        train_y : np.array(int)
            Training `y` set (training labels).
        test_x : np.array(float)
            Test `x` set (test features).
        test_y : np.array(int)
            Test `y` set (test labels).
        file_output_path : str
            Path to store h5 files generated by custom BestAccs callback.
        initializer : str, optional
            Type of kernel and bias initializer.
        batch_size : int, optional
            Size of batch (amount of samples) for model fitting.

    Returns
    -------
        history : keras.callbacks.History object
            History of loss, accuracy, validation loss and validation
            accuracy during model fitting.

    """
    if not os.path.exists(file_output_path):
        os.makedirs(file_output_path)

    # Define form of h5 file name prefix
    metrics = '{:.2f}-{acc:.2f}-{val_acc:.2f}'

    filename = metrics + '-' + initializer + '.h5'
    path_and_filename = os.path.join(file_output_path, '')
    path_and_filename += filename   # workaround of os.path.join issue
    best_accs = BestAccs(path_and_filename, min_accs=0.7, diff=0.02)

    lr = 0.005
    optimizer = TFOptimizer(tf.train.GradientDescentOptimizer(lr))

    model = Sequential()

    model.add(Dense(1,
                    kernel_initializer=initializer,
                    bias_initializer=initializer,
                    input_dim=train_x.shape[1],
                    activation='sigmoid'))

    model.compile(loss='binary_crossentropy',
                  optimizer=optimizer,
                  metrics=['accuracy'])

    history = model.fit(train_x, train_y,
                        epochs=2000,
                        callbacks=[best_accs],
                        validation_data=(test_x, test_y),
                        batch_size=batch_size,
                        verbose=0)
    return history


def compare_kernels(files, file_output_path):
    """
    Compare two kernels saved as h5 files.

    Load kernels from h5 files and then calculate norm of kernels, norm
    of vector difference between kernels and the angle between kernels.

    Parameters
    ----------
        kernel_files : list(str)
            List of two kernel h5 files to compare.
        file_output_path : str
            Path where kernel h5 files were saved.

    Returns
    -------
        norms : list(float)
            List with norms of both compared kernels and norm of vector
            difference between them.
        angle : float
            Angle (radians) between two compared kernels.

    """
    kernels = []

    # Load kernel and bias from both best files
    for file in files:
        path_and_filename = os.path.join(file_output_path, file)
        h5f = h5py.File(path_and_filename, 'r')
        list_of_names = []
        h5f.visit(list_of_names.append)
        kernel = h5f[list_of_names[3]].value
        h5f.close()
        kernels.append(kernel)

    # Norms of kernel and norm of difference between kernels
    norms = []
    norms.append(np.linalg.norm(kernels[0]))
    norms.append(np.linalg.norm(kernels[1]))
    norms.append(np.linalg.norm(kernels[0] - kernels[1]))

    # Angle between kernels
    norms_product = np.linalg.norm(kernels[0]) * np.linalg.norm(kernels[1])
    angle_cos = np.dot(kernels[0].T, kernels[1])/ norms_product
    # Sometimes due to float type accuracy angle_cos becomes greater than 1
    if angle_cos > 1.0:
        angle_cos = 1.0
    angle = float(np.arccos(angle_cos))
    return norms, angle


def best_learning_rate(train_x, train_y, test_x, test_y, lr, batch_size=256):
    """
    Perform 10 random initializations of model, then compile and fit it.

    Model is initialized randomly (random_uniform), compiled and fitted.
    Operation is repeated 10 times. Then all histories are returned as a list.

    Parameters
    ----------
        train_x : np.array(float)
            Training `x` set (training features).
        train_y : np.array(int)
            Training `y` set (training labels).
        test_x : np.array(float)
            Test `x` set (test features).
        test_y : np.array(int)
            Test `y` set (test labels).
        lr : float
            Learning rate.
        batch_size : int, optional
            Size of batch (amount of samples) for model fitting.

    Returns
    -------
        history_set : list(keras.callbacks.History object)
            History of loss, accuracy, validation loss and validation
            accuracy during model fitting.

    """
    optimizer = TFOptimizer(tf.train.GradientDescentOptimizer(lr))
    history_set = []

    for i in range(10):
        model = Sequential()

        initializer = RandomUniform(minval=-1.0, maxval=1.0, seed=None)
        model.add(Dense(1,
                        kernel_initializer=initializer,
                        bias_initializer=initializer,
                        input_dim=train_x.shape[1],
                        activation='sigmoid'))

        model.compile(loss='binary_crossentropy',
                      optimizer=optimizer,
                      metrics=['accuracy'])

        history = model.fit(train_x, train_y,
                            epochs=1000,
                            validation_data=(test_x, test_y),
                            batch_size=batch_size,
                            verbose=0)

        history_set.append(history)
    return history_set


def sampling_hypersurface(train_x, train_y, test_x, test_y,
                          file_output_path,
                          suffix = '',
                          batch_size=256,
                          iterations=1000,
                          verbose=1):
    """
    Sampling hypersurface by different random initialization.

    Parameters
    ----------
        train_x : np.array(float)
            Training `x` set (training features).
        train_y : np.array(int)
            Training `y` set (training labels).
        test_x : np.array(float)
            Test `x` set (test features).
        test_y : np.array(int)
            Test `y` set (test labels).
        file_output_path : str
            Path to store h5 files generated by custom BestAccs callback.
        suffix: str, optional
            File name suffix.
        batch_size : int, optional
            Size of batch (amount of samples) for model fitting.
        iterations : int, optional
            Defines how many times defining, compilation and fitting
            procedure has to be executed.
        verbose : int, optional
            Verbosity level: 0 or 1. 0 means quite run, 1 means progress
            will be shown.


    """
    if not os.path.exists(file_output_path):
        os.makedirs(file_output_path)

    # Define form of h5 file name prefix
    metrics = '{:.3f}-{acc:.3f}-{val_acc:.3f}'

    initializer = RandomUniform(minval=-1.0, maxval=1.0, seed=None)
    lr = 0.1
    optimizer = TFOptimizer(tf.train.GradientDescentOptimizer(lr))

    for iteration in range(iterations):
        filename = metrics + '-iteration-' + str(iteration) + suffix + '.h5'
        path_and_filename = os.path.join(file_output_path, '')
        path_and_filename += filename   # workaround of os.path.join issue
        best_accs = BestAccs(path_and_filename, min_accs=0.88, diff=0.04)

        model = Sequential()

        model.add(Dense(1,
                        kernel_initializer=initializer,
                        bias_initializer=initializer,
                        input_dim=train_x.shape[1],
                        activation='sigmoid'))

        model.compile(loss='binary_crossentropy',
                      optimizer=optimizer,
                      metrics=['accuracy'])

        model.fit(train_x, train_y,
                  epochs=1000,
                  callbacks=[best_accs],
                  validation_data=(test_x, test_y),
                  batch_size=batch_size,
                  verbose=0)

        if verbose == 1:
            if (iteration + 1) % 100 == 0:
                end = '\n'
            else:
                end = ''
            if (iteration + 1)  % 10 == 0:
                print('|', end = end)
            else:
                print('.', end = end)

#%%
def best_model_evaluate(train_x, train_y, test_x, test_y,
                          file_output_path, file_name,
                          batch_size=256):
    """
    Evaluate accuracy on training and test set.

    Parameters
    ----------
        train_x : np.array(float)
            Training `x` set (training features).
        train_y : np.array(int)
            Training `y` set (training labels).
        test_x : np.array(float)
            Test `x` set (test features).
        test_y : np.array(int)
            Test `y` set (test labels).
        file_output_path : str
            Path to store h5 files generated by custom BestAccs callback.
        file_name : str
            Name of file that contains weights
        batch_size : int, optional
            Size of batch (amount of samples) for model fitting.

    Returns
    -------
        acc_train : float
            Accuracy on training set.
        acc_test : float
            Accuracy on test set.

    """

    initializer = 'zeros'
    lr = 0.01
    optimizer = TFOptimizer(tf.train.GradientDescentOptimizer(lr))
    model = Sequential()
    model.add(Dense(1,
                    kernel_initializer=initializer,
                    bias_initializer=initializer,
                    input_dim=train_x.shape[1],
                    activation='sigmoid'))

    model.compile(loss='binary_crossentropy',
                  optimizer=optimizer,
                  metrics=['accuracy'])

    file_output_path = 'samplingHypersurface'
    model.load_weights(os.path.join(file_output_path, file_name))

    metrics_train = model.evaluate(train_x, train_y, batch_size=256, verbose=0)
    acc_train = metrics_train[1]
    metrics_test = model.evaluate(test_x, test_y, batch_size=256, verbose=0)
    acc_test = metrics_test[1]
    return acc_train, acc_test
