import tensorflow as tf
import numpy as np
import os,sys
sys.path.insert(0, os.getcwd())

from dataload.LSA16h import LSA16h
from dataload.Dataset import Dataset
from dataload.BatchedDataset import BatchedDataset

import models

def load_data(input_folderpath,batch_size):


    lsa=LSA16h(input_folderpath)
    d=lsa.as_dataset()
    classes=max(d.y)+1
    image_size=d.x[0].shape
    train_dataset,test_dataset=d.split_stratified(0.8)
    bd=BatchedDataset(train_dataset.x,train_dataset.y,batch_size)
    return bd,train_dataset,test_dataset,image_size,classes

print("Loading data..")
batch_size=20
input_folderpath='../datasets/lsa32x32/lsa16h_nr/rgb_black_background'
bd,train_dataset,test_dataset,image_size,classes=load_data(input_folderpath,batch_size)
print('Image size %s' % str(image_size))

input_size=np.prod(image_size)
print("Preparing model..")
with tf.Graph().as_default():
  s =  tf.InteractiveSession()
  x = tf.placeholder(tf.float32, shape=[ None,image_size[0],image_size[1],image_size[2]])
  y = tf.placeholder(tf.int32, shape=[None])
  y_estimated,y_true= models.lsa_conv(x,y,image_size,classes)

  error = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=y_true, logits=y_estimated))

  train_step = tf.train.AdamOptimizer(1e-4).minimize(error)
  correct_prediction = tf.equal(tf.argmax(y_estimated, 1), tf.argmax(y_true, 1))
  accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

  init_op = tf.group(tf.global_variables_initializer())
  s.run(init_op)
  c = tf.train.Coordinator()
  threads = tf.train.start_queue_runners(sess=s, coord=c)

  train_feed_dict={
      x:train_dataset.x,
      y:train_dataset.y,
  }
  test_feed_dict={
      x:test_dataset.x,
      y:test_dataset.y,
  }
  keep_prob=0.3

  print("Training:\n")
  iterations=50000
  for i in range(iterations):

    batch_x,batch_y=bd.next_batch()
    batch_feed_dict={
        x:batch_x,
        y:batch_y,
    }
    if i% (iterations/5) == 0:
      test_accuracy  = accuracy.eval(feed_dict=test_feed_dict)
      train_accuracy = accuracy.eval(feed_dict=train_feed_dict)
      train_error    = error.eval(feed_dict=train_feed_dict)
      print("step %d, test accuracy %g, training accuracy %g, error %g"%(i,test_accuracy, train_accuracy,train_error))
    train_step.run(feed_dict=batch_feed_dict)
  print("Training finished")
  print("Dataset "+input_folderpath)
  print("Final test accuracy: %f" % accuracy.eval(feed_dict=test_feed_dict))
  #print("Final train accuracy: %f" % accuracy.eval(feed_dict=train_feed_dict))
