#!/usr/bin/env python
  # -*- coding: utf-8 -*-
  # Author: Titouan Parcollet
  
  from keras.layers import Input, Dense, Conv2D, MaxPooling2D, UpSampling2D
  import time
  import sys
  import numpy as np
  import keras
  from keras.models import Sequential,Model
  from keras.callbacks import Callback
  from keras.layers.core import Dense, Activation
  from keras.layers import Dropout, Input
  import tensorflow as tf
  from keras.backend.tensorflow_backend import set_session
  from keras.layers import Conv2D, MaxPooling2D
  from keras.layers import Conv1D, MaxPooling1D
  import os
  from complexnn import *
  
  print("##############################################")
  print("## Quaternion Convolutional encoder-decoder ##")
  print("## Titouan Parcollet, LIA 2017              ##")
  print("##############################################")
  
  ########  Static parameters 
  # size        : Number of train sample to choose
  # 		        -> 100,1000,5000
  # batch_size  : batch size during training
  # train_epoch : number of epochs
  # CIFAR_*     : path to pre-processed data files
  # output_*    : path where the output layer will be saved
  # n_input     : size of the input vector CIFAR = 32 * 32 * 3
  # n_classes   : size of the output vector
  
  size = 100
  batch_size = 5
  train_epoch = 100
  n_input = 3072
  n_classes = 3072
  CIFAR_TRAINING = "Resources/Corpus/"+str(size)+"_train.dat"
  CIFAR_TEST = "Resources/Corpus/"+str(size)+"_test.dat"
  if not os.path.exists("Results"):
  	    os.makedirs("Results")
  output_test = open("Results/"+str(size)+"_test.dat", "w")
  output_train = open("Results/"+str(size)+"_train.dat", "w")
  
  ######## Loading datasets
  train_file = open(CIFAR_TRAINING, "r").readlines()
  test_file = open(CIFAR_TEST, "r").readlines()
  train_data = np.zeros( (len(train_file), n_input) )
  train_label = np.zeros((len(train_file), n_classes))
  test_data = np.zeros( (len(test_file), n_input) )
  test_label = np.zeros((len(test_file), n_classes))
  
  print("Loading dataset ...")
  line_cpt =0
  for line in train_file:
      data_cpt =0
      for splitted in line.split("\t")[0].split(" "):
          if(data_cpt < n_input):
              train_data[line_cpt][data_cpt] = float(splitted)
              data_cpt+=1
      data_cpt=0
     
      for splitted in line.split("\t")[1].split(" "):
          train_label[line_cpt][data_cpt] = int(splitted)
          data_cpt+=1
      line_cpt+=1
  
  line_cpt =0
  for line in test_file:
      data_cpt =0
      for splitted in line.split("\t")[0].split(" "):
          if(data_cpt<n_input):
              test_data[line_cpt][data_cpt] = float(splitted)
              data_cpt+=1
      data_cpt=0
      for splitted in line.split("\t")[1].split(" "):
          test_label[line_cpt][data_cpt] = int(splitted)
          data_cpt+=1
   
      data_cpt+=1
      line_cpt+=1
  print("Data loaded :D")
  
  
  ######## Normalizing pixels
  train_data = train_data.astype('float32')
  test_data = test_data.astype('float32')
  train_data /= 255
  test_data /= 255
  train_label /= 255
  test_label /= 255
  
  ######## Reformating to match Quaternion format
  # Q = r + x + y + z
  # Q = 0 + R + G + B
  
  test_data = test_data.reshape(test_data.shape[0],32,32,3)
  test_label = test_label.reshape(test_data.shape[0],32,32,3)
  train_data = train_data.reshape(train_data.shape[0],32,32,3)
  train_label = train_label.reshape(train_data.shape[0],32,32,3)
  
  x=train_data.shape[0]
  y=train_data.shape[1]
  z=train_data.shape[2]
  zeros = np.zeros((x,y,z,1))
  train_data = np.concatenate((zeros, train_data), axis = 3)
  
  x=test_data.shape[0]
  y=test_data.shape[1]
  z=test_data.shape[2]
  zeros = np.zeros((x,y,z,1))
  test_data = np.concatenate((zeros, test_data), axis = 3)
  
  x=train_label.shape[0]
  y=train_label.shape[1]
  z=train_label.shape[2]
  zeros = np.zeros((x,y,z,1))
  train_label = np.concatenate((zeros, train_label), axis = 3)
  
  x=test_label.shape[0]
  y=test_label.shape[1]
  z=test_label.shape[2]
  zeros = np.zeros((x,y,z,1))
  test_label = np.concatenate((zeros, test_label), axis = 3)
  
  
  ######## Building model
  # A QuaternionConv2D layer has the same parameters than a real valued
  # convolutional layer. Be certain that the input can be divided by 4
  
  input_img = Input(shape=(32, 32, 4))
  x = QuaternionConv2D(32, kernel_size=(3, 3), activation='relu', padding='same')(input_img)
  x = MaxPooling2D((2, 2), padding='same')(x)
  x = QuaternionConv2D(16, (3, 3), activation='relu', padding='same')(x)
  
  encoded = MaxPooling2D((2, 2), padding='same')(x)
  
  x = QuaternionConv2D(16, (3, 3), activation='relu', padding='same')(encoded)
  x = UpSampling2D((2, 2))(x)
  x = QuaternionConv2D(32, kernel_size=(3,3), strides=(1, 1), activation='relu', padding='same')(x)
  x = UpSampling2D((2, 2))(x)
  decoded = QuaternionConv2D(1, (3, 3),activation='sigmoid', name='conv_last', padding='same')(x)
  
  ######## Compiling model
  autoencoder = Model(input_img, decoded)
  autoencoder.compile(optimizer='adam', loss='binary_crossentropy')
  
  ######## Training
  loss_history_cb = autoencoder.fit(train_data, train_label,
                  epochs=train_epoch,
                  batch_size=batch_size)
  
  ######## Recording Loss History
  loss_history = loss_history_cb.history["loss"]
  loss = np.array(loss_history)
  np.savetxt("loss_history_Q.txt", loss, delimiter=",")
  
  
  ######## Write the outputs for each images
  # Accordingly to the need of the script create_rgb.py
  #
  
  modulo_cpt =0
  output = autoencoder.predict(test_data[0:99])
  for element in output:
      for value in element:
          for other in value:
              for rgb in other:
                  if modulo_cpt %4 != 0:
                      output_test.write(str(float(rgb)*255)+" ")
                  modulo_cpt += 1
      output_test.write("
  ")
  output_test.close()
  
  modulo_cpt = 0
  output = autoencoder.predict(train_data[0:4900])
  for element in output:
      for value in element:
          for other in value:
              for rgb in other:
                  if modulo_cpt %4 != 0:
                      output_train.write(str(float(rgb)*255)+" ")
                  modulo_cpt +=1
      output_train.write("
  ")
  output_train.close()
  
  ######## Finally generate the reconstructed RGB pictures
  print("Reconstructing images : create_rgb.py "+str(size))
  create_rgb.reconstruct(str(size))
  print("That's all Folks !")