"""

Stefania Fresca, MOX Laboratory, Politecnico di Milano

April 2019

"""

import tensorflow as tf

import numpy as np

import scipy.io as sio

import time

import os

import utils

seed = 374

np.random.seed(seed)

class Net:

    def __init__(self, config):

        self.lr = config['lr']

        self.batch_size = config['batch_size']

        self.g_step = tf.Variable(0, dtype = tf.int32, trainable = False, name = 'global_step')

        self.n_data = config['n_data']

        self.n_train = int(0.8 * self.n_data)

        self.N_h = config['N_h']

        self.N_t = config['N_t']

        self.train_mat = config['train_mat']

        self.test_mat = config['test_mat']

        self.train_params = config['train_params']

        self.test_params = config['test_params']

        self.omega_h = config['omega_h']

        self.omega_n = config['omega_n']

        self.checkpoints_folder = config['checkpoints_folder']

        self.graph_folder = config['graph_folder']

        self.large = config['large']

        self.zero_padding = config['zero_padding']

        self.p = config['p']

        self.restart = config['restart']

    def get_data(self):

        with tf.name_scope('data'):

            self.X = tf.placeholder(tf.float32, shape = [None, self.N_h])

            self.Y = tf.placeholder(tf.float32, shape = [None, self.n_params])

            dataset = tf.data.Dataset.from_tensor_slices((self.X, self.Y))

            dataset = dataset.shuffle(self.n_data)

            dataset = dataset.batch(self.batch_size)

            iterator = dataset.make_initializable_iterator()

            self.init = iterator.initializer

            input, self.params = iterator.get_next()

            self.input = tf.reshape(input, shape = [-1, int(np.sqrt(self.N_h)), int(np.sqrt(self.N_h)), 1])

    def inference(self):

        raise NotImplementedError("Must be overridden with proper definition of forward path")

    def loss(self, u_h, u_n):

        with tf.name_scope('loss'):

            output = tf.reshape(self.input, shape = [-1, self.N_h])

            self.loss_h = self.omega_h * tf.reduce_mean(tf.reduce_sum(tf.pow(output - u_h, 2), axis = 1))

            self.loss_n = self.omega_n * tf.reduce_mean(tf.reduce_sum(tf.pow(self.enc - u_n, 2), axis = 1))

            self.loss = self.loss_h + self.loss_n

    def optimize(self):

        self.opt = tf.train.AdamOptimizer(self.lr).minimize(self.loss, global_step = self.g_step)

    def summary(self):

        with tf.name_scope('summaries'):

            self.summary = tf.summary.scalar('loss', self.loss)

    def build(self):

        self.get_data()

        self.inference()

        self.loss(self.u_h, self.u_n)

        self.optimize()

        self.summary()

    def train_one_epoch(self, sess, init, writer, epoch, step):

        start_time = time.time()

        sess.run(init, feed_dict = {self.X : self.S_train, self.Y : self.params_train})

        total_loss_h = 0

        total_loss_n = 0

        total_loss = 0

        n_batches = 0

        print('------------ TRAINING -------------', flush = True)

        try:

            while True:

                _, l_h, l_n, l, summary = sess.run([self.opt, self.loss_h, self.loss_n, self.loss, self.summary])

                writer.add_summary(summary, global_step = step)

                step += 1

                total_loss_h += l_h

                total_loss_n += l_n

                total_loss += l

                n_batches += 1

        except tf.errors.OutOfRangeError:

            pass

        print('Average loss_h at epoch {0} on training set: {1}'.format(epoch, total_loss_h / n_batches))

        print('Average loss_n at epoch {0} on training set: {1}'.format(epoch, total_loss_n / n_batches))

        print('Average loss at epoch {0} on training set: {1}'.format(epoch, total_loss / n_batches))

        print('Took: {0} seconds'.format(time.time() - start_time))

        return step

    def eval_once(self, sess, saver, init, writer, epoch, step):

        start_time = time.time()

        sess.run(init, feed_dict = {self.X : self.S_val, self.Y : self.params_val})

        total_loss_h = 0

        total_loss_n = 0

        total_loss = 0

        n_batches = 0

        print('------------ VALIDATION ------------')

        try:

            while True:

                l_h, l_n, l, summary = sess.run([self.loss_h, self.loss_n, self.loss, self.summary])

                writer.add_summary(summary, global_step = step)

                total_loss_h += l_h

                total_loss_n += l_n

                total_loss += l

                n_batches += 1

        except tf.errors.OutOfRangeError:

            pass

        total_loss_mean = total_loss / n_batches

        if total_loss_mean < self.loss_best:

            saver.save(sess, self.checkpoints_folder + '/Net', step)

        print('Average loss_h at epoch {0} on validation set: {1}'.format(epoch, total_loss_h / n_batches))

        print('Average loss_n at epoch {0} on validation set: {1}'.format(epoch, total_loss_n / n_batches))

        print('Average loss at epoch {0} on validation set: {1}'.format(epoch, total_loss_mean))

        print('Took: {0} seconds'.format(time.time() - start_time))

        return total_loss_mean

    def test_once(self, sess, init):

        start_time = time.time()

        sess.run(init, feed_dict = {self.X : self.S_test, self.Y : self.params_test})

        total_loss_h = 0

        total_loss_n = 0

        total_loss = 0

        n_batches = 0

        self.U_h = np.zeros(self.S_test.shape)

        print('------------ TESTING ------------')

        try:

            while True:

                l_h, l_n, l, u_h = sess.run([self.loss_h, self.loss_n, self.loss, self.u_h])

                self.U_h[self.batch_size * n_batches : self.batch_size * (n_batches + 1)] = u_h

                total_loss_h += l_h

                total_loss_n += l_n

                total_loss += l

                n_batches += 1

        except tf.errors.OutOfRangeError:

            pass

        print('Average loss_h on testing set: {0}'.format(total_loss_h / n_batches))

        print('Average loss_N on testing set: {0}'.format(total_loss_n / n_batches))

        print('Average loss on testing set: {0}'.format(total_loss / n_batches))

        print('Took: {0} seconds'.format(time.time() - start_time))

    #@profile (if memory profiling must be used)

    def train_all(self, n_epochs):

        if (not self.restart):

            utils.safe_mkdir(self.checkpoints_folder)

        saver = tf.train.Saver()

        train_writer = tf.summary.FileWriter('./' + self.graph_folder + '/train', tf.get_default_graph())

        test_writer = tf.summary.FileWriter('./' + self.graph_folder + '/test', tf.get_default_graph())

        print('Loading snapshot matrix...')

        if (self.large):

            S = utils.read_large_data(self.train_mat)

        else:

            S = utils.read_data(self.train_mat)

        idxs = np.random.permutation(S.shape[0])

        S = S[idxs]

        S_max, S_min = utils.max_min(S, self.n_train)

        utils.scaling(S, S_max, S_min)

        if (self.zero_padding):

            S = utils.zero_pad(S, self.p)

        self.S_train, self.S_val = S[:self.n_train, :], S[self.n_train:, :]

        del S

        print('Loading parameters...')

        params = utils.read_params(self.train_params)

        params = params[idxs]

        self.params_train, self.params_val = params[:self.n_train], params[self.n_train:]

        del params

        self.loss_best = 1

        count = 0

        with tf.Session(config = tf.ConfigProto(gpu_options = tf.GPUOptions(allow_growth = True))) as sess:

            sess.run(tf.global_variables_initializer())

            if (self.restart):

                ckpt = tf.train.get_checkpoint_state(os.path.dirname(self.checkpoints_folder + '/checkpoint'))

                if ckpt and ckpt.model_checkpoint_path:

                    print(ckpt.model_checkpoint_path)

                    saver.restore(sess, ckpt.model_checkpoint_path)

            step = self.g_step.eval()

            for epoch in range(n_epochs):

                step = self.train_one_epoch(sess, self.init, train_writer, epoch, step)

                total_loss_mean = self.eval_once(sess, saver, self.init, test_writer, epoch, step)

                if total_loss_mean < self.loss_best:

                    self.loss_best = total_loss_mean

                    count = 0

                else:

                    count += 1

                # early - stopping

                if count == 500:

                    print('Stopped training due to early-stopping cross-validation')

                    break

            print('Best loss on validation set: {0}'.format(self.loss_best))

        train_writer.close()

        test_writer.close()

        with tf.Session() as sess:

            sess.run(tf.global_variables_initializer())

            ckpt = tf.train.get_checkpoint_state(os.path.dirname(self.checkpoints_folder + '/checkpoint'))

            if ckpt and ckpt.model_checkpoint_path:

                print(ckpt.model_checkpoint_path)

                saver.restore(sess, ckpt.model_checkpoint_path)

            print('Loading testing snapshot matrix...')

            if (self.large):

                self.S_test = utils.read_large_data(self.test_mat)

            else:

                self.S_test = utils.read_data(self.test_mat)

            utils.scaling(self.S_test, S_max, S_min)

            if (self.zero_padding):

                self.S_test = utils.zero_pad(self.S_test, self.n)

            print('Loading testing parameters...')

            self.params_test = utils.read_params(self.test_params)

            self.test_once(sess, self.init)

            utils.inverse_scaling(self.U_h, S_max, S_min)

            utils.inverse_scaling(self.S_test, S_max, S_min)

            n_test = self.S_test.shape[0] // self.N_t

            err = np.zeros((n_test, 1))

            for i in range(n_test):

                num = np.sqrt(np.mean(np.linalg.norm(self.S_test[i * self.N_t : (i + 1) * self.N_t] - self.U_h[i * self.N_t : (i + 1) * self.N_t], 2, axis = 1) ** 2))

                den = np.sqrt(np.mean(np.linalg.norm(self.S_test[i * self.N_t : (i + 1) * self.N_t], 2, axis = 1) ** 2))

                err[i] = num / den

            print('Error indicator epsilon_rel: {0}'.format(np.mean(err)))