from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf

import time
import random
import os
import numpy as np
import sys
import math

from CreateDataset import tfTest, readSubject, readUnclassifiedImages
from sklearn.metrics import confusion_matrix

def deepnn(x):
  """
  Args:
    x: dimensions (N_examples, 384,384),

  Returns:
    A tuple (y, keep_prob). y is a tensor of shape (N_examples, 3), with values
    equal to the logits of classifying
    keep_prob is a scalar placeholder for the probability of dropout.
  """
  x_image = tf.reshape(x, [-1, 32, 32, 32, 1])

  W_conv1 = weight_variable([5, 5, 5, 1, 32])
  b_conv1 = bias_variable([32])
  h_conv1 = tf.nn.relu(conv3d(x_image, W_conv1) + b_conv1)
  h_pool1 = max_pool_2x2x2(h_conv1)

  W_conv2 = weight_variable([5, 5, 5, 32, 64])
  b_conv2 = bias_variable([64])

  h_conv2 = tf.nn.relu(conv3d(h_pool1, W_conv2) + b_conv2)
  h_pool2 = max_pool_2x2x2(h_conv2)

  W_fc1 = weight_variable([8 * 8 * 8 * 64, 1024])
  b_fc1 = bias_variable([1024])

  h_pool2_flat = tf.reshape(h_pool2, [-1, 8 * 8 * 8 * 64])
  h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, W_fc1) + b_fc1)

  keep_prob = tf.placeholder(tf.float32)
  h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)

  W_fc2 = weight_variable([1024, 3])
  b_fc2 = bias_variable([3])

  y_conv = tf.matmul(h_fc1_drop, W_fc2) + b_fc2
  return y_conv, keep_prob

def conv3d(x, W):
  """conv2d returns a 3d convolution layer with full stride."""
  return tf.nn.conv3d(x, W, strides=[1, 1, 1, 1, 1], padding='SAME')

def conv2d(x, W):
  """conv2d returns a 2d convolution layer with full stride."""
  return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')

def max_pool_2x2(x):
  """max_pool_2x2 downsamples a feature map by 2X."""
  return tf.nn.max_pool(x, ksize=[1, 2, 2, 1],
                        strides=[1, 2, 2, 1], padding='SAME')

def max_pool_2x2x2(x):
  """max_pool_2x2x2 downsamples a feature map by 2X."""
  return tf.nn.max_pool3d(x, ksize=[1, 2, 2, 2, 1],
                           strides=[1, 2, 2, 2, 1], padding='SAME')

def max_pool_4x4(x):
  """max_pool_2x2 downsamples a feature map by 2X."""
  return tf.nn.max_pool(x, ksize=[1, 4, 4, 1],
                        strides=[1, 4, 4, 1], padding='SAME')


def weight_variable(shape):
  """weight_variable generates a weight variable of a given shape."""
  initial = tf.truncated_normal(shape, stddev=0.1)
  return tf.Variable(initial)


def bias_variable(shape):
  """bias_variable generates a bias variable of a given shape."""
  initial = tf.constant(0.1, shape=shape)
  return tf.Variable(initial)


def softmax(z):
    assert len(z.shape) == 2
    s = np.max(z, axis=1)
    s = s[:, np.newaxis] # necessary step to do broadcasting
    e_x = np.exp(z - s)
    div = np.sum(e_x, axis=1)
    div = div[:, np.newaxis] # dito
    return e_x / div


def run():

  imagesDir = sys.argv[1].replace("\\", "/").replace("//", "/")
  modelDir = sys.argv[2].replace("\\", "/").replace("//", "/")
  outputDir = sys.argv[3].replace("\\", "/").replace("//", "/")
  
  print(imagesDir)
  print(modelDir)
  print(outputDir)

  # Create the model
  x = tf.placeholder(tf.float32, [None, 32, 32, 32])

  # Define loss and optimizer
  y_ = tf.placeholder(tf.float32, [None, 3])

  # Build the graph for the deep net
  y_conv, keep_prob = deepnn(x)

  cross_entropy = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=y_, logits=y_conv))
  train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)
  correct_prediction = tf.equal(tf.argmax(y_conv, 1), tf.argmax(y_, 1))
  accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
  probabilities = y_conv
  saver = tf.train.Saver()

  # Import data
  #rootDir = '/home/daten_stegmaier/Projects/2017/CNNTraining/Data/Simulated/'
  #rootDir = '/home/daten_stegmaier/Projects/2017/CNNTraining/Data/Meristem/'
  #rootDir = '/home/daten_stegmaier/Projects/2017/CNNTraining/Data/Simulated3D'
  #rootDir = 'C:/Temp/Images/'
  #outputDir = rootDir + '../Probabilities/'


  images, label, files = readUnclassifiedImages(imagesDir, -1)
  numTestImages = images.shape[0]
  miniBatchSize = 100
  numIterationsTest = max(1, int(math.ceil(numTestImages / miniBatchSize)))

  print(images.shape)
  print(label.shape)

  with tf.Session() as sess:
    #sess.run(tf.global_variables_initializer())

    saver.restore(sess, modelDir + "RACE3DModel.ckpt")
    print("Model loaded from file: %s" % modelDir + "RACE3DModel.ckpt")

    t0 = time.time()

    testBatch = np.array(range(0, images.shape[0]))

    maxTestImages = testBatch.shape[0]

    t0 = time.time()

    combinedAccuracyTest = 0.0
    labelsTest = []
    probabilitiesTest = []
    predictionsTest = []
    for j in range(0, numIterationsTest):
      currentIndices = range(j * miniBatchSize, min(maxTestImages, (j * miniBatchSize) + miniBatchSize))
      currentProbabilities = probabilities.eval(feed_dict={x: images[testBatch[currentIndices], :, :], y_: label[testBatch[currentIndices], :], keep_prob: 1.0})
      currentPredictions = np.argmax(currentProbabilities, 1)
      currentLabels = np.argmax(label[testBatch[currentIndices], :], 1)
      if (j == 0):
        labelsTest = currentLabels
        predictionsTest = currentPredictions
        probabilitiesTest = currentProbabilities
      else:
        labelsTest = np.concatenate((labelsTest, currentLabels), 0)
        predictionsTest = np.concatenate((predictionsTest, currentPredictions), 0)
        probabilitiesTest = np.concatenate((probabilitiesTest, currentProbabilities), 0)
      combinedAccuracyTest = combinedAccuracyTest + accuracy.eval(feed_dict={x: images[testBatch[currentIndices], :, :], y_: label[testBatch[currentIndices], :], keep_prob: 1.0})

    probabilitiesTestSoftmax = softmax(probabilitiesTest)

    for j in range(0, probabilitiesTestSoftmax.shape[0]):
      for k in range(0,3):
        if np.isnan(probabilitiesTestSoftmax[j,k]):
          probabilitiesTestSoftmax[j, k] = 1.0

      currentFileName = os.path.basename(files[j])
      currentFileName = currentFileName.replace(".tif", ".csv")
      currentFileName = currentFileName.replace("trainingData", "probabilities")
      file = open(outputDir + currentFileName, "w")
      file.write(str(probabilitiesTestSoftmax[j,0]) + ',' + str(probabilitiesTestSoftmax[j,1]) + ',' + str(probabilitiesTestSoftmax[j,2]))
      file.close()
      print(probabilitiesTestSoftmax[j,:])


    print('testing  accuracy %g' % (combinedAccuracyTest / numIterationsTest))
    confusionMatrixTest = confusion_matrix(labelsTest, predictionsTest)
    print(confusionMatrixTest)

    t1 = time.time()
    print( "avg. classification time per image snippet: %7.3f ms" % ((t1 - t0) * 1000 / maxTestImages) )
    t0 = time.time()

run()
