#!/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("\n")
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("\n")
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 !")