import tensorflow as tf
import os
from datetime import datetime


class CNN:
    def __init__(self, pcg, nclasses=2, learning_rate=0.001,
                 epochs=5, batch_size=100, dropout=0.75, base_dir="/tmp",
                 model_name="cnn"):
        self.pcg = pcg
        self.nclasses = nclasses
        self.d_input = self.pcg.train.X.shape[1]
        self.learning_rate = learning_rate
        self.epochs = epochs
        self.batch_size = batch_size
        self.dropout = dropout
        self.nbatches = int(self.pcg.train.X.shape[0] / float(self.batch_size))
        self.model_name = model_name
        self.base_dir = base_dir

    def train(self):
        """
        Train a convolutional neural network over the input PCG dataset.
        This method is beefy: it is responsible for defining tensorflow
        variables, defining the training objective function, defining summary
        statistics creating the tensorflow session, running gradient
        descent and, ultimately, writing statistics

        In the future this will be refactored into more easily tested
        training segments.

        Parameters
        ----------
        None

        Returns
        -------
        None

        """
        print('begin train')
        print(self.__get_output_name())

        with tf.name_scope('input'):
            X = tf.placeholder(tf.float32, [None, self.d_input], name='X')
            y = tf.placeholder(tf.float32, [None, self.nclasses], name='y')
            do_drop = tf.placeholder(tf.float32, name='drop')

        with tf.name_scope('weights'):
            weights = {
                'wc1': tf.Variable(tf.random_normal([5, 1, 1, 32]), name='wc1'),
                'wc2': tf.Variable(tf.random_normal([5, 1, 32, 64]), name='wc2'),
                # 2 Max pools have taken original 10612 signal down to
                # 5306 --> 2653. Each max pool has a ksize=2.
                # 'wd1': tf.Variable(tf.random_normal([2653 * 1 * 64, 1024])),
                'wd1': tf.Variable(tf.random_normal([int(self.d_input / 4) * 1 * 64, 1024]), name='wd1'),
                'out': tf.Variable(tf.random_normal([1024, self.nclasses]), name='outW')
            }
        with tf.name_scope('biases'):
            biases = {
                'bc1': tf.Variable(tf.random_normal([32]), name='bc1'),
                'bc2': tf.Variable(tf.random_normal([64]), name='bc2'),
                'bd1': tf.Variable(tf.random_normal([1024]), name='bd1'),
                'out': tf.Variable(tf.random_normal([self.nclasses]), name='outB')
            }

        with tf.name_scope('pred'):
            pred = self.model1D(X, weights, biases, do_drop)

        with tf.name_scope('cost'):
            cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(pred, y, name='cost'))
            optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(cost)

        dim = tf.shape(y)[0]

        with tf.name_scope('sensitivity'):
            # sensitivity = correctly predicted abnormal / total number of actual abnormal
            abnormal_idxs = tf.cast(tf.equal(tf.argmax(pred, 1), 1), tf.float32)
            pred1d = tf.reshape(tf.slice(y, [0, 1], [dim, 1]), [-1])
            abn = tf.mul(pred1d, abnormal_idxs)
            sensitivity = tf.reduce_sum(abn) / tf.reduce_sum(pred1d)
            tf.scalar_summary('sensitivity', sensitivity)

        with tf.name_scope('specificity'):
            # specificity = correctly predicted normal / total number of actual normal
            normal_idxs = tf.cast(tf.equal(tf.argmax(pred, 1), 0), tf.float32)
            pred1d_n = tf.reshape(tf.slice(y, [0, 0], [dim, 1]), [-1])
            normal = tf.mul(pred1d_n, normal_idxs)
            specificity = tf.reduce_sum(normal) / tf.reduce_sum(pred1d_n)
            tf.scalar_summary('specificity', sensitivity)

        # Physionet score is the mean of sensitivity and specificity
        score = (sensitivity + specificity) / 2.0
        tf.scalar_summary('score', score)

        init = tf.initialize_all_variables()

        saver = tf.train.Saver()
        with tf.Session() as sess:
            sess.run(init)

            merged = tf.merge_all_summaries()
            train_writer = tf.train.SummaryWriter(os.path.join(self.base_dir, 'train'), sess.graph)

            for epoch in range(self.epochs):
                avg_cost = 0
                for batch in range(self.nbatches):
                    batch_x, batch_y = self.pcg.get_mini_batch(self.batch_size)
                    summary, _, c = sess.run([merged, optimizer, cost],
                                             feed_dict={X: batch_x,
                                                        y: batch_y,
                                                        do_drop: self.dropout})
                    train_writer.add_summary(summary, epoch*batch)
                    avg_cost += c
                avg_cost /= float(self.nbatches)
                print('Epoch %s\tcost %s' % (epoch, avg_cost))

                if epoch % 10 == 0:
                    acc, sens, spec = sess.run([score, sensitivity, specificity],
                                                feed_dict={X: self.pcg.test.X,
                                                           y: self.pcg.test.y,
                                                           do_drop: 1.})
                    print('Score %s\tSensitivity %s\tSpecificity %s' % (acc, sens, spec))

                    saver.save(sess, self.__get_output_name())
                    print('Epoch written')

    def __get_output_name(self):
        now = datetime.now()
        time_str = "-%s" % (now.date())  # now.hour, now.minute, now.second)
        model_path = os.path.join(self.base_dir, self.model_name + time_str + '.tnfl')
        return model_path

    def conv2d(self, x, w, b, strides=1):
        """
        A small helper function for calcualting a 1D convolution
        from tensorflow's conv2d method

        Parameters
        ----------
        x: tensorflow.placeholder
            The feature vector
        w: tensorflow.Variable
            The unknown weights to learn
        b: tensorflow.Variable
            The unknown biases to learn
        strides: int
            The length of the stride to use for convolution

        Returns
        -------
        tensorflow.Variable
            A convolution over the input feature vector

        """


        x = tf.nn.conv2d(x, w, strides=[1, strides, strides, 1], padding="SAME")
        x = tf.nn.bias_add(x, b)
        return tf.nn.relu(x)

    def model1D(self, x, weights, biases, dropout):
        """
        A Wrapper to chain several TensorFlow convolutional units together. This 1D model
        ultimately calls TensorFlow's conv2d, mapping a 1D feature vector to a collapsed
        2D convolution

        Parameters
        ----------
        x: tensorflow.placeholder
            A feature vector of size [None, no_features]

        weights: dict<str, tensorflow.Variable>
            Dictionary of Unknown weights to learn

        biases: dict<str, tensorflow.Variable>
            Dictionary of unknown biases to learn

        dropout: float
            the dropout fraction for convolutional units

        Returns
        -------
        out: tensorflow.Variable
            The result of applying multiple convolutional layers and
            a fully connected unit to the input feature vector

        """

        with tf.name_scope('reshape'):
            x = tf.reshape(x, shape=[-1, self.d_input, 1, 1])  # [n_images, width, height, n_channels]

        with tf.name_scope('conv1'):
            conv1 = self.conv2d(x, weights['wc1'], biases['bc1'])
            conv1 = tf.nn.max_pool(conv1, ksize=[1, 2, 1, 1], strides=[1, 2, 1, 1], padding='SAME')
            conv1 = tf.nn.relu(conv1)

        with tf.name_scope('conv2'):
            conv2 = self.conv2d(conv1, weights['wc2'], biases['bc2'])
            conv2 = tf.nn.max_pool(conv2, ksize=[1, 2, 1, 1], strides=[1, 2, 1, 1], padding="SAME")
            conv2 = tf.nn.relu(conv2)

        with tf.name_scope('fullyConnected'):
            d_layer1 = weights['wd1'].get_shape().as_list()[0]
            fc1 = tf.reshape(conv2, [-1, d_layer1])
            fc1 = tf.add(tf.matmul(fc1, weights['wd1']), biases['bd1'])
            fc1 = tf.nn.relu(fc1)
            fc1 = tf.nn.dropout(fc1, dropout)

            out = tf.add(tf.matmul(fc1, weights['out']), biases['out'])
        return out