# -*- coding: utf-8 -*-
"""
The ML (Neural Network) part of the solution to Huawei Honor Cup, where the task was to put a shuffled jigsaw puzzle together

Note: the actual problem was about grids with size up to 32x32, not 8x8 like the one in this folder

Initially was a .ipynb file. 

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1KAns0OO8ZGrLb_H3J3QyYp9NuDGLNOzP

"""

import numpy as np
import pandas as pd
import tensorflow as tf
import keras.preprocessing.image
from IPython.display import HTML
from keras.preprocessing.image import load_img
from keras.preprocessing.image import img_to_array
import sklearn.preprocessing
import sklearn.model_selection
import sklearn.metrics
import sklearn.linear_model
import sklearn.naive_bayes
import sklearn.tree
import sklearn.ensemble
import os
import datetime
import cv2
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.cm as cm

# %matplotlib inline

from google.colab import drive

drive.mount("/content/drive")

SZ = 512
szpix = 16
M = SZ // szpix
tf_size = (1 * szpix // 16) * 2
epochs = 35.0
train_width = 2
dirs = ["test2"]
grayscale = False
prefix = "/content/drive/My Drive/Honor"


def to_gray(img):
    gr = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    img = cv2.cvtColor(gr, cv2.COLOR_GRAY2RGB)
    return img


sizes = ["16", "32", "64"]
traindirs = []

for s in sizes:
    traindirs.append("{}/data_train/{}".format(prefix, s))


filelist = []
for dirname, _, filenames in os.walk("{}/data_train/{}-sources".format(prefix, 16)):
    for fn in filenames:
        filelist.append(os.path.join(dirname, fn))
filelist.sort()
images = []
for fn in filelist:
    print(fn)
    img = cv2.imread(fn)[..., ::-1] * 1.0
    if grayscale:
        img = to_gray(img)
    images.append(img)

images
print(np.shape(images))

cnt_train = M * (M - 1) * 2 * np.shape(images)[0]
print(cnt_train)
x_images = np.empty((cnt_train, szpix, 2 * train_width, 3))
y_images = np.empty((cnt_train), dtype=int)
curcnt = 0
for im in images:
    for i in range(M):
        for j in range(M - 1):
            my_end = im[
                i * szpix : (i + 1) * szpix,
                (j + 1) * szpix - train_width : (j + 1) * szpix,
            ]
            other_beg = im[
                i * szpix : (i + 1) * szpix,
                (j + 1) * szpix : (j + 1) * szpix + train_width,
            ]
            x_images[curcnt] = np.hstack((my_end, other_beg))
            y_images[curcnt] = 1
            curcnt += 1
            my_beg = im[
                i * szpix : (i + 1) * szpix, j * szpix : j * szpix + train_width
            ]
            i1, j1 = i, j + 1
            if False and j + 2 < M:
                j1 = j + 2
                other_beg = im[
                    i1 * szpix : (i + 1) * szpix, j1 * szpix : j1 * szpix + train_width
                ]
                x_images[curcnt] = np.hstack((my_end, other_beg))
            else:
                other_end = im[
                    i1 * szpix : (i1 + 1) * szpix,
                    (j1 + 1) * szpix - train_width : (j1 + 1) * szpix,
                ]
                x_images[curcnt] = np.hstack((other_end, my_beg))
            y_images[curcnt] = 0
            curcnt += 1
            # plt.imshow(x_images[-1]/255.0)

## normalize data

# convert one-hot encodings into labels
def one_hot_to_dense(labels_one_hot):
    return np.argmax(labels_one_hot, 1)


# convert class labels from scalars to one-hot vectors e.g. 1 => [0 1], 0 => [1 0]
def dense_to_one_hot(labels_dense, num_classes):
    num_labels = labels_dense.shape[0]
    index_offset = np.arange(num_labels) * num_classes
    labels_one_hot = np.zeros((num_labels, num_classes))
    labels_one_hot.flat[index_offset + labels_dense] = 1
    return labels_one_hot


# function to normalize data
def normalize_data(data):
    # scale features using statistics that are robust to outliers
    # rs = sklearn.preprocessing.RobustScaler()
    # rs.fit(data)
    # data = rs.transform(data)
    data = (data - data.mean()) / (data.std())  # standardisation
    # data=data/data.max()
    data = (data - data.min()) / (
        data.max() - data.min()
    )  # convert from [0:255] to [0.:1.]
    # data = ((data / 255.)-0.5)*2. # convert from [0:255] to [-1.:+1.]
    return data


# training and validation data
x_train_valid = normalize_data(x_images)

# use one-hot encoding for labels 0,1
y_train_valid = dense_to_one_hot(y_images, 2).astype(np.uint8)

# dictionaries for saving results
y_valid_pred = {}
y_train_pred = {}
y_test_pred = {}
train_loss, valid_loss = {}, {}
train_acc, valid_acc = {}, {}
cnf_valid_matrix = {}

print("x_train_valid.shape =", x_train_valid.shape)
print(
    "x_train_valid.min/mean/std/max = %.2f/%.2f/%.2f/%.2f"
    % (
        x_train_valid.min(),
        x_train_valid.mean(),
        x_train_valid.std(),
        x_train_valid.max(),
    )
)
print("")
print("y_train_valid.shape =", y_train_valid.shape)
print(
    "y_train_valid.min/mean/std/max = %.2f/%.2f/%.2f/%.2f"
    % (
        y_train_valid.min(),
        y_train_valid.mean(),
        y_train_valid.std(),
        y_train_valid.max(),
    )
)

## First try out some basic sklearn models
"""
# compute the accuracy of label predictions
def accuracy_from_dense_labels(y_target, y_pred):
    y_target = y_target.reshape(-1,)
    y_pred = y_pred.reshape(-1,)
    return np.mean(y_target == y_pred)

def predict_im(im):
    im1=im.reshape(szpix,2*train_width,3)
    a=im1[0:szpix,train_width-1:train_width,:]
    b=im1[0:szpix,train_width:train_width+1,:]
    diff=abs(a-b)
    return np.max(diff*diff)<0.06

def my_predict(x):
    return np.apply_along_axis(predict_im,1,x)

logreg = sklearn.linear_model.LogisticRegression(verbose=0, solver='lbfgs',
                                                 multi_class = 'multinomial')
decision_tree = sklearn.tree.DecisionTreeClassifier()
extra_trees = sklearn.ensemble.ExtraTreesClassifier(verbose=0, max_depth=4)
gradient_boost = sklearn.ensemble.GradientBoostingClassifier(verbose=0)
random_forest = sklearn.ensemble.RandomForestClassifier(verbose=0, max_depth=4, n_estimators=10)
gaussianNB = sklearn.naive_bayes.GaussianNB()


# store models in dictionary
base_models = {'logreg': logreg, 'extra_trees': extra_trees,
               'gradient_boost': gradient_boost, 'random_forest': random_forest, 
               'decision_tree': decision_tree, 'gaussianNB': gaussianNB}

# choose models for out-of-folds predictions
#take_models = ['logreg','random_forest','extra_trees','gaussianNB']
take_models = ['my','random_forest', 'extra_trees']

for mn in take_models:
    train_acc[mn] = []
    valid_acc[mn] = []
    cnf_valid_matrix[mn] = []

# start timer
start = datetime.datetime.now();
print(datetime.datetime.now().strftime('%d-%m-%Y %H:%M:%S'),': start training')
       
# cross validations
cv_num = 10 # cross validations default = 20 => 5% validation set
kfold = sklearn.model_selection.KFold(cv_num, shuffle=True) #, random_state=123)

for i,(train_index, valid_index) in enumerate(kfold.split(x_train_valid)):

    # start timer
    start = datetime.datetime.now();

    # train and validation data of original images
    x_train = x_train_valid[train_index].reshape(-1,szpix*2*train_width*3)
    y_train = y_train_valid[train_index]
    x_valid = x_train_valid[valid_index].reshape(-1,szpix*2*train_width*3)
    y_valid = y_train_valid[valid_index]

    for mn in take_models:

        if mn=='my':
            y_train_pred[mn]=my_predict(x_train)
            y_valid_pred[mn]=my_predict(x_valid)
            train_acc[mn].append(accuracy_from_dense_labels(y_train_pred[mn],
                                                        one_hot_to_dense(y_train)))
            valid_acc[mn].append(accuracy_from_dense_labels(y_valid_pred[mn],
                                                        one_hot_to_dense(y_valid)))
            cnf_valid_matrix_tmp = sklearn.metrics.confusion_matrix(
                y_pred = y_valid_pred[mn], 
            y_true = one_hot_to_dense(y_valid)).astype(np.float32)
        
        else:
        # create cloned model from base models
            model = sklearn.base.clone(base_models[mn])
            model.fit(x_train, one_hot_to_dense(y_train))

            # predictions
            y_train_pred[mn] = model.predict_proba(x_train)
            y_valid_pred[mn] = model.predict_proba(x_valid)
        
        # accuracies
            train_acc[mn].append(accuracy_from_dense_labels(one_hot_to_dense(y_train_pred[mn]),
                                                        one_hot_to_dense(y_train)))
            valid_acc[mn].append(accuracy_from_dense_labels(one_hot_to_dense(y_valid_pred[mn]),
                                                        one_hot_to_dense(y_valid)))
        
        # normalized confusion matrix
            cnf_valid_matrix_tmp = sklearn.metrics.confusion_matrix(
            y_pred = one_hot_to_dense(y_valid_pred[mn]), 
            y_true = one_hot_to_dense(y_valid)).astype(np.float32)
        cnf_valid_matrix_tmp[0,:] = cnf_valid_matrix_tmp[0,:]/cnf_valid_matrix_tmp[0,:].sum()
        cnf_valid_matrix_tmp[1,:] = cnf_valid_matrix_tmp[1,:]/cnf_valid_matrix_tmp[1,:].sum()
        cnf_valid_matrix[mn].append(cnf_valid_matrix_tmp)

        print(i,': '+mn+' train/valid accuracy = %.3f/%.3f'%(train_acc[mn][-1], 
                                                             valid_acc[mn][-1]))

print('running time for training: ', datetime.datetime.now() - start)
print('')
for mn in train_acc.keys():
    print(mn + ' : averaged train/valid accuracy = %.3f/%.3f'%(np.mean(train_acc[mn]),
                                                              np.mean(valid_acc[mn])))
"""


class nn_class:
    # class that implements the neural network

    # constructor
    def __init__(
        self,
        nn_name="tmp",
        log_step=0.1,
        keep_prob=0.33,
        mb_size=50,
        width=50,
        height=50,
        n_channel=3,
        n_output=2,
    ):

        # tunable hyperparameters for nn architecture
        self.s_f_conv1 = 3  # filter size of first convolution layer (default = 3)
        self.n_f_conv1 = (
            36
        )  # number of features of first convolution layer (default = 36)
        self.s_f_conv2 = 3  # filter size of second convolution layer (default = 3)
        self.n_f_conv2 = (
            36
        )  # number of features of second convolution layer (default = 36)
        self.s_f_conv3 = 3  # filter size of third convolution layer (default = 3)
        self.n_f_conv3 = (
            36
        )  # number of features of third convolution layer (default = 36)
        self.n_n_fc1 = (
            576
        )  # number of neurons of first fully connected layer (default = 576)
        self.n_channel = n_channel
        self.width = width
        self.height = height
        self.n_output = n_output

        # tunable hyperparameters for training
        self.mb_size = mb_size  # mini batch size
        self.keep_prob = keep_prob  # keeping probability with dropout regularization
        self.learn_rate_array = [
            10 * 1e-4,
            7.5 * 1e-4,
            5 * 1e-4,
            2.5 * 1e-4,
            1 * 1e-4,
            1 * 1e-4,
            1 * 1e-4,
            0.75 * 1e-4,
            0.5 * 1e-4,
            0.25 * 1e-4,
            0.1 * 1e-4,
            0.1 * 1e-4,
            0.075 * 1e-4,
            0.050 * 1e-4,
            0.025 * 1e-4,
            0.01 * 1e-4,
            0.0075 * 1e-4,
            0.0050 * 1e-4,
            0.0025 * 1e-4,
            0.001 * 1e-4,
        ]
        self.learn_rate_step_size = 3  # in terms of epochs

        # helper variables
        self.learn_rate = self.learn_rate_array[0]
        self.learn_rate_pos = 0  # current position pointing to current learning rate
        self.index_in_epoch = 0
        self.current_epoch = 0
        self.log_step = log_step  # log results in terms of epochs
        self.n_log_step = 0  # counting current number of mini batches trained on
        self.use_tb_summary = False  # True = use tensorboard visualization
        self.use_tf_saver = False  # True = use saver to save the model
        self.nn_name = nn_name  # name of the neural network
        self.perm_array = np.array([])  # permutation array

    # function to get the next mini batch
    def next_mini_batch(self):

        start = self.index_in_epoch
        self.index_in_epoch += self.mb_size
        self.current_epoch += self.mb_size / len(self.x_train)

        # adapt length of permutation array
        if not len(self.perm_array) == len(self.x_train):
            self.perm_array = np.arange(len(self.x_train))

        # shuffle once at the start of epoch
        if start == 0:
            np.random.shuffle(self.perm_array)

        # at the end of the epoch
        if self.index_in_epoch > self.x_train.shape[0]:
            np.random.shuffle(self.perm_array)  # shuffle data
            start = 0  # start next epoch
            self.index_in_epoch = self.mb_size  # set index to mini batch size

        end = self.index_in_epoch

        # use original data
        x_tr = self.x_train[self.perm_array[start:end]]
        y_tr = self.y_train[self.perm_array[start:end]]

        return x_tr, y_tr

    # weight initialization
    def weight_variable(self, shape, name=None):
        initial = tf.truncated_normal(shape, stddev=0.1)
        return tf.Variable(initial, name=name)

    # bias initialization
    def bias_variable(self, shape, name=None):
        initial = tf.constant(0.1, shape=shape)  #  positive bias
        return tf.Variable(initial, name=name)

    # 2D convolution
    def conv2d(self, x, W, name=None):
        return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding="SAME", name=name)

    # max pooling
    def max_pool_2x2(self, x, name=None):
        return tf.nn.max_pool(
            x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME", name=name
        )

    # attach summaries to a tensor for TensorBoard visualization
    def summary_variable(self, var, var_name):
        with tf.name_scope(var_name):
            mean = tf.reduce_mean(var)
            stddev = tf.sqrt(tf.reduce_mean(tf.square(var - mean)))
            tf.summary.scalar("mean", mean)
            tf.summary.scalar("stddev", stddev)
            tf.summary.scalar("max", tf.reduce_max(var))
            tf.summary.scalar("min", tf.reduce_min(var))
            tf.summary.histogram("histogram", var)

    # function to create the graph
    def create_graph(self):

        # reset default graph
        tf.reset_default_graph()

        # variables for input and output
        self.x_data_tf = tf.placeholder(
            dtype=tf.float32,
            shape=[None, self.height, self.width, self.n_channel],
            name="x_data_tf",
        )
        self.y_data_tf = tf.placeholder(
            dtype=tf.float32, shape=[None, self.n_output], name="y_data_tf"
        )

        # 1.layer: convolution + max pooling
        self.W_conv1_tf = self.weight_variable(
            [self.s_f_conv1, self.s_f_conv1, self.n_channel, self.n_f_conv1],
            name="W_conv1_tf",
        )  # (3,3,3,36)
        self.b_conv1_tf = self.bias_variable(
            [self.n_f_conv1], name="b_conv1_tf"
        )  # (36)
        self.h_conv1_tf = tf.nn.relu(
            self.conv2d(self.x_data_tf, self.W_conv1_tf) + self.b_conv1_tf,
            name="h_conv1_tf",
        )  # (.,50,50,36)
        self.h_pool1_tf = self.max_pool_2x2(
            self.h_conv1_tf, name="h_pool1_tf"
        )  # (.,25,25,36)

        # 2.layer: convolution + max pooling
        self.W_conv2_tf = self.weight_variable(
            [self.s_f_conv2, self.s_f_conv2, self.n_f_conv1, self.n_f_conv2],
            name="W_conv2_tf",
        )
        self.b_conv2_tf = self.bias_variable([self.n_f_conv2], name="b_conv2_tf")
        self.h_conv2_tf = tf.nn.relu(
            self.conv2d(self.h_pool1_tf, self.W_conv2_tf) + self.b_conv2_tf,
            name="h_conv2_tf",
        )  # (.,25,25,36)
        self.h_pool2_tf = self.max_pool_2x2(
            self.h_conv2_tf, name="h_pool2_tf"
        )  # (.,13,13,36)

        # 3.layer: convolution + max pooling
        self.W_conv3_tf = self.weight_variable(
            [self.s_f_conv3, self.s_f_conv3, self.n_f_conv2, self.n_f_conv3],
            name="W_conv3_tf",
        )
        self.b_conv3_tf = self.bias_variable([self.n_f_conv3], name="b_conv3_tf")
        self.h_conv3_tf = tf.nn.relu(
            self.conv2d(self.h_pool2_tf, self.W_conv3_tf) + self.b_conv3_tf,
            name="h_conv3_tf",
        )  # (.,13,13,36)
        self.h_pool3_tf = self.max_pool_2x2(
            self.h_conv3_tf, name="h_pool3_tf"
        )  # (.,7,7,36)

        # 4.layer: fully connected
        self.W_fc1_tf = self.weight_variable(
            [tf_size * self.n_f_conv3, self.n_n_fc1], name="W_fc1_tf"
        )  # (7*7*36, 1024)
        self.b_fc1_tf = self.bias_variable([self.n_n_fc1], name="b_fc1_tf")  # (1024)
        self.h_pool3_flat_tf = tf.reshape(
            self.h_pool3_tf, [-1, tf_size * self.n_f_conv3], name="h_pool3_flat_tf"
        )  # (.,1024)
        self.h_fc1_tf = tf.nn.relu(
            tf.matmul(self.h_pool3_flat_tf, self.W_fc1_tf) + self.b_fc1_tf,
            name="h_fc1_tf",
        )  # (.,1024)

        # add dropout
        self.keep_prob_tf = tf.placeholder(dtype=tf.float32, name="keep_prob_tf")
        self.h_fc1_drop_tf = tf.nn.dropout(
            self.h_fc1_tf, self.keep_prob_tf, name="h_fc1_drop_tf"
        )

        # 5.layer: fully connected
        self.W_fc2_tf = self.weight_variable(
            [self.n_n_fc1, self.n_output], name="W_fc2_tf"
        )  # (1024,1)
        self.b_fc2_tf = self.bias_variable([self.n_output], name="b_fc2_tf")  # (1024)
        self.z_pred_tf = tf.add(
            tf.matmul(self.h_fc1_drop_tf, self.W_fc2_tf),
            self.b_fc2_tf,
            name="z_pred_tf",
        )  # => (.,1)

        # cost function
        self.cross_entropy_tf = tf.reduce_mean(
            tf.nn.softmax_cross_entropy_with_logits(
                labels=self.y_data_tf, logits=self.z_pred_tf
            ),
            name="cross_entropy_tf",
        )

        # optimisation function
        self.learn_rate_tf = tf.placeholder(dtype=tf.float32, name="learn_rate_tf")
        self.train_step_tf = tf.train.AdamOptimizer(self.learn_rate_tf).minimize(
            self.cross_entropy_tf, name="train_step_tf"
        )

        # predicted probabilities in one-hot encoding
        self.y_pred_proba_tf = tf.nn.softmax(self.z_pred_tf, name="y_pred_proba_tf")

        # tensor of correct predictions
        self.y_pred_correct_tf = tf.equal(
            tf.argmax(self.y_pred_proba_tf, 1),
            tf.argmax(self.y_data_tf, 1),
            name="y_pred_correct_tf",
        )

        # accuracy
        self.accuracy_tf = tf.reduce_mean(
            tf.cast(self.y_pred_correct_tf, dtype=tf.float32), name="accuracy_tf"
        )

        # tensors to save intermediate accuracies and losses during training
        self.train_loss_tf = tf.Variable(
            np.array([]), dtype=tf.float32, name="train_loss_tf", validate_shape=False
        )
        self.valid_loss_tf = tf.Variable(
            np.array([]), dtype=tf.float32, name="valid_loss_tf", validate_shape=False
        )
        self.train_acc_tf = tf.Variable(
            np.array([]), dtype=tf.float32, name="train_acc_tf", validate_shape=False
        )
        self.valid_acc_tf = tf.Variable(
            np.array([]), dtype=tf.float32, name="valid_acc_tf", validate_shape=False
        )

        # number of weights and biases
        num_weights = (
            self.s_f_conv1 ** 2 * self.n_f_conv1 * self.n_channel
            + self.s_f_conv2 ** 2 * self.n_f_conv1 * self.n_f_conv2
            + self.s_f_conv3 ** 2 * self.n_f_conv2 * self.n_f_conv3
            + 4 * 4 * self.n_f_conv3 * self.n_n_fc1
            + self.n_n_fc1 * self.n_output
        )
        num_biases = self.n_f_conv1 + self.n_f_conv2 + self.n_f_conv3 + self.n_n_fc1
        print("num_weights =", num_weights)
        print("num_biases =", num_biases)

        return None

    def attach_summary(self, sess):

        # create summary tensors for tensorboard
        self.use_tb_summary = True
        self.summary_variable(self.W_conv1_tf, "W_conv1_tf")
        self.summary_variable(self.b_conv1_tf, "b_conv1_tf")
        self.summary_variable(self.W_conv2_tf, "W_conv2_tf")
        self.summary_variable(self.b_conv2_tf, "b_conv2_tf")
        self.summary_variable(self.W_conv3_tf, "W_conv3_tf")
        self.summary_variable(self.b_conv3_tf, "b_conv3_tf")
        self.summary_variable(self.W_fc1_tf, "W_fc1_tf")
        self.summary_variable(self.b_fc1_tf, "b_fc1_tf")
        self.summary_variable(self.W_fc2_tf, "W_fc2_tf")
        self.summary_variable(self.b_fc2_tf, "b_fc2_tf")
        tf.summary.scalar("cross_entropy_tf", self.cross_entropy_tf)
        tf.summary.scalar("accuracy_tf", self.accuracy_tf)

        # merge all summaries for tensorboard
        self.merged = tf.summary.merge_all()

        # initialize summary writer
        timestamp = datetime.datetime.now().strftime("%d-%m-%Y_%H-%M-%S")
        filepath = os.path.join(os.getcwd(), "logs", (self.nn_name + "_" + timestamp))
        self.train_writer = tf.summary.FileWriter(
            os.path.join(filepath, "train"), sess.graph
        )
        self.valid_writer = tf.summary.FileWriter(
            os.path.join(filepath, "valid"), sess.graph
        )

    def attach_saver(self):
        # initialize tensorflow saver
        self.use_tf_saver = True
        self.saver_tf = tf.train.Saver()

    # function to train the graph
    def train_graph(self, sess, x_train, y_train, x_valid, y_valid, n_epoch=1):

        # training and validation data
        self.x_train = x_train
        self.y_train = y_train
        self.x_valid = x_valid
        self.y_valid = y_valid

        # parameters
        mb_per_epoch = self.x_train.shape[0] / self.mb_size
        train_loss, train_acc, valid_loss, valid_acc = [], [], [], []

        # start timer
        start = datetime.datetime.now()
        print(datetime.datetime.now().strftime("%d-%m-%Y %H:%M:%S"), ": start training")
        print(
            "learnrate =",
            self.learn_rate,
            ", n_epoch =",
            n_epoch,
            ", mb_size =",
            self.mb_size,
            ", nn_name =",
            self.nn_name,
        )
        # looping over mini batches
        for i in range(int(n_epoch * mb_per_epoch) + 1):

            # adapt learn_rate
            if not self.learn_rate_pos == int(
                self.current_epoch // self.learn_rate_step_size
            ):
                self.learn_rate_pos = int(
                    self.current_epoch // self.learn_rate_step_size
                )
                self.learn_rate = self.learn_rate_array[self.learn_rate_pos]
                print(
                    datetime.datetime.now() - start,
                    ": set learn rate to %.6f" % self.learn_rate,
                )

            # get new batch
            x_batch, y_batch = self.next_mini_batch()

            # run the graph
            sess.run(
                self.train_step_tf,
                feed_dict={
                    self.x_data_tf: x_batch,
                    self.y_data_tf: y_batch,
                    self.keep_prob_tf: self.keep_prob,
                    self.learn_rate_tf: self.learn_rate,
                },
            )
            # store losses and accuracies
            if i % int(self.log_step * mb_per_epoch) == 0 or i == int(
                n_epoch * mb_per_epoch
            ):

                self.n_log_step += 1  # for logging the results

                feed_dict_train = {
                    self.x_data_tf: self.x_train[self.perm_array[: len(self.x_valid)]],
                    self.y_data_tf: self.y_train[self.perm_array[: len(self.y_valid)]],
                    self.keep_prob_tf: 1.0,
                }

                feed_dict_valid = {
                    self.x_data_tf: self.x_valid,
                    self.y_data_tf: self.y_valid,
                    self.keep_prob_tf: 1.0,
                }

                # summary for tensorboard
                if self.use_tb_summary:
                    train_summary = sess.run(self.merged, feed_dict=feed_dict_train)
                    valid_summary = sess.run(self.merged, feed_dict=feed_dict_valid)
                    self.train_writer.add_summary(train_summary, self.n_log_step)
                    self.valid_writer.add_summary(valid_summary, self.n_log_step)

                train_loss.append(
                    sess.run(self.cross_entropy_tf, feed_dict=feed_dict_train)
                )

                train_acc.append(
                    self.accuracy_tf.eval(session=sess, feed_dict=feed_dict_train)
                )

                valid_loss.append(
                    sess.run(self.cross_entropy_tf, feed_dict=feed_dict_valid)
                )

                valid_acc.append(
                    self.accuracy_tf.eval(session=sess, feed_dict=feed_dict_valid)
                )

                print(
                    "%.2f epoch: train/val loss = %.4f/%.4f, train/val acc = %.4f/%.4f"
                    % (
                        self.current_epoch,
                        train_loss[-1],
                        valid_loss[-1],
                        train_acc[-1],
                        valid_acc[-1],
                    )
                )

        # concatenate losses and accuracies and assign to tensor variables
        tl_c = np.concatenate(
            [self.train_loss_tf.eval(session=sess), train_loss], axis=0
        )
        vl_c = np.concatenate(
            [self.valid_loss_tf.eval(session=sess), valid_loss], axis=0
        )
        ta_c = np.concatenate([self.train_acc_tf.eval(session=sess), train_acc], axis=0)
        va_c = np.concatenate([self.valid_acc_tf.eval(session=sess), valid_acc], axis=0)

        sess.run(tf.assign(self.train_loss_tf, tl_c, validate_shape=False))
        sess.run(tf.assign(self.valid_loss_tf, vl_c, validate_shape=False))
        sess.run(tf.assign(self.train_acc_tf, ta_c, validate_shape=False))
        sess.run(tf.assign(self.valid_acc_tf, va_c, validate_shape=False))

        print("running time for training: ", datetime.datetime.now() - start)
        return None

    # save tensors/summaries
    def save_model(self, sess):

        # tf saver
        if self.use_tf_saver:
            # filepath = os.path.join(os.getcwd(), 'logs' , self.nn_name)
            filepath = os.path.join(os.getcwd(), self.nn_name)
            self.saver_tf.save(sess, filepath)

        # tb summary
        if self.use_tb_summary:
            self.train_writer.close()
            self.valid_writer.close()

        return None

    # forward prediction of current graph
    def forward(self, sess, x_data):
        y_pred_proba = self.y_pred_proba_tf.eval(
            session=sess, feed_dict={self.x_data_tf: x_data, self.keep_prob_tf: 1.0}
        )
        return y_pred_proba

    # function to load tensors from a saved graph
    def load_tensors(self, graph):

        # input tensors
        self.x_data_tf = graph.get_tensor_by_name("x_data_tf:0")
        self.y_data_tf = graph.get_tensor_by_name("y_data_tf:0")

        # weights and bias tensors
        self.W_conv1_tf = graph.get_tensor_by_name("W_conv1_tf:0")
        self.W_conv2_tf = graph.get_tensor_by_name("W_conv2_tf:0")
        self.W_conv3_tf = graph.get_tensor_by_name("W_conv3_tf:0")
        self.W_fc1_tf = graph.get_tensor_by_name("W_fc1_tf:0")
        self.W_fc2_tf = graph.get_tensor_by_name("W_fc2_tf:0")
        self.b_conv1_tf = graph.get_tensor_by_name("b_conv1_tf:0")
        self.b_conv2_tf = graph.get_tensor_by_name("b_conv2_tf:0")
        self.b_conv3_tf = graph.get_tensor_by_name("b_conv3_tf:0")
        self.b_fc1_tf = graph.get_tensor_by_name("b_fc1_tf:0")
        self.b_fc2_tf = graph.get_tensor_by_name("b_fc2_tf:0")

        # activation tensors
        self.h_conv1_tf = graph.get_tensor_by_name("h_conv1_tf:0")
        self.h_pool1_tf = graph.get_tensor_by_name("h_pool1_tf:0")
        self.h_conv2_tf = graph.get_tensor_by_name("h_conv2_tf:0")
        self.h_pool2_tf = graph.get_tensor_by_name("h_pool2_tf:0")
        self.h_conv3_tf = graph.get_tensor_by_name("h_conv3_tf:0")
        self.h_pool3_tf = graph.get_tensor_by_name("h_pool3_tf:0")
        self.h_fc1_tf = graph.get_tensor_by_name("h_fc1_tf:0")
        self.z_pred_tf = graph.get_tensor_by_name("z_pred_tf:0")

        # training and prediction tensors
        self.learn_rate_tf = graph.get_tensor_by_name("learn_rate_tf:0")
        self.keep_prob_tf = graph.get_tensor_by_name("keep_prob_tf:0")
        self.cross_entropy_tf = graph.get_tensor_by_name("cross_entropy_tf:0")
        self.train_step_tf = graph.get_operation_by_name("train_step_tf")
        self.z_pred_tf = graph.get_tensor_by_name("z_pred_tf:0")
        self.y_pred_proba_tf = graph.get_tensor_by_name("y_pred_proba_tf:0")
        self.y_pred_correct_tf = graph.get_tensor_by_name("y_pred_correct_tf:0")
        self.accuracy_tf = graph.get_tensor_by_name("accuracy_tf:0")

        # tensor of stored losses and accuricies during training
        self.train_loss_tf = graph.get_tensor_by_name("train_loss_tf:0")
        self.train_acc_tf = graph.get_tensor_by_name("train_acc_tf:0")
        self.valid_loss_tf = graph.get_tensor_by_name("valid_loss_tf:0")
        self.valid_acc_tf = graph.get_tensor_by_name("valid_acc_tf:0")

        return None

    # get losses of training and validation sets
    def get_loss(self, sess):
        train_loss = self.train_loss_tf.eval(session=sess)
        valid_loss = self.valid_loss_tf.eval(session=sess)
        return train_loss, valid_loss

    # get accuracies of training and validation sets
    def get_accuracy(self, sess):
        train_acc = self.train_acc_tf.eval(session=sess)
        valid_acc = self.valid_acc_tf.eval(session=sess)
        return train_acc, valid_acc

    # get weights
    def get_weights(self, sess):
        W_conv1 = self.W_conv1_tf.eval(session=sess)
        W_conv2 = self.W_conv2_tf.eval(session=sess)
        W_conv3 = self.W_conv3_tf.eval(session=sess)
        W_fc1_tf = self.W_fc1_tf.eval(session=sess)
        W_fc2_tf = self.W_fc2_tf.eval(session=sess)
        return W_conv1, W_conv2, W_conv3, W_fc1_tf, W_fc2_tf

    # get biases
    def get_biases(self, sess):
        b_conv1 = self.b_conv1_tf.eval(session=sess)
        b_conv2 = self.b_conv2_tf.eval(session=sess)
        b_conv3 = self.b_conv3_tf.eval(session=sess)
        b_fc1_tf = self.b_fc1_tf.eval(session=sess)
        b_fc2_tf = self.b_fc2_tf.eval(session=sess)
        return b_conv1, b_conv2, b_conv3, b_fc1_tf, b_fc2_tf

    # load session from file, restore graph, and load tensors
    def load_session_from_file(self, filename):
        tf.reset_default_graph()
        filepath = os.path.join(os.getcwd(), filename + ".meta")
        # filepath = os.path.join(os.getcwd(),'logs', filename + '.meta')
        saver = tf.train.import_meta_graph(filepath)
        print(filepath)
        sess = tf.Session()
        saver.restore(sess, mn)
        graph = tf.get_default_graph()
        self.load_tensors(graph)
        return sess

    # receive activations given the input
    def get_activations(self, sess, x_data):
        feed_dict = {self.x_data_tf: x_data, self.keep_prob_tf: 1.0}
        h_conv1 = self.h_conv1_tf.eval(session=sess, feed_dict=feed_dict)
        h_pool1 = self.h_pool1_tf.eval(session=sess, feed_dict=feed_dict)
        h_conv2 = self.h_conv2_tf.eval(session=sess, feed_dict=feed_dict)
        h_pool2 = self.h_pool2_tf.eval(session=sess, feed_dict=feed_dict)
        h_conv3 = self.h_conv3_tf.eval(session=sess, feed_dict=feed_dict)
        h_pool3 = self.h_pool3_tf.eval(session=sess, feed_dict=feed_dict)
        h_fc1 = self.h_fc1_tf.eval(session=sess, feed_dict=feed_dict)
        h_fc2 = self.z_pred_tf.eval(session=sess, feed_dict=feed_dict)
        return h_conv1, h_pool1, h_conv2, h_pool2, h_conv3, h_pool3, h_fc1, h_fc2


nn_name = ["nn0", "nn1", "nn2", "nn3", "nn4", "nn5", "nn6", "nn7", "nn8", "nn9"]

# cross validations
cv_num = 10  # cross validations default = 20 => 5% validation set
kfold = sklearn.model_selection.KFold(cv_num, shuffle=True, random_state=123)

for i, (train_index, valid_index) in enumerate(kfold.split(x_train_valid)):

    # start timer
    start = datetime.datetime.now()

    # train and validation data of original images
    x_train = x_train_valid[train_index]
    y_train = y_train_valid[train_index]
    x_valid = x_train_valid[valid_index]
    y_valid = y_train_valid[valid_index]

    # create neural network graph
    nn_graph = nn_class(
        log_step=0.1,
        nn_name=nn_name[i],
        height=szpix,
        width=2 * train_width,
        mb_size=50,
        keep_prob=0.33,
    )  # instance of nn_class
    nn_graph.create_graph()  # create graph
    nn_graph.attach_saver()  # attach saver tensors

    # start tensorflow session
    with tf.Session() as sess:

        # attach summaries
        nn_graph.attach_summary(sess)

        # variable initialization of the default graph
        sess.run(tf.global_variables_initializer())

        # training on original data
        nn_graph.train_graph(sess, x_train, y_train, x_valid, y_valid, n_epoch=epochs)

        # save tensors and summaries of model
        nn_graph.save_model(sess)

    break

print("total running time for training: ", datetime.datetime.now() - start)

filelist = []
for dirname, _, filenames in os.walk(
    "{}/data_{}_blank/{}".format(prefix, dirs[0], szpix)
):
    for fn in filenames:
        filelist.append(os.path.join(dirname, fn))
filelist.sort()
mn = "nn0"
nn_graph = nn_class()
sess = nn_graph.load_session_from_file(mn)
fout = open(dirs[0] + "_" + str(szpix) + "_neural", "w")
for indf, fn in enumerate(filelist):
    print(fn)
    # if "3308" not in fn and "3309" not in fn and "3320" not in fn:
    #    continue
    if indf % 10 == 0:
        print(indf)
    img = cv2.imread(fn)[..., ::-1] * 1.0
    if grayscale:
        img = to_gray(img)
    im = img
    cnt_train = M * M * (M * M - 1) * 2
    x_images = np.empty((cnt_train, szpix, 2 * train_width, 3))
    curcnt = 0
    for i in range(M):
        for j in range(M):
            for i1 in range(M):
                for j1 in range(M):
                    if i == i1 and j == j1:
                        continue
                    my_end = im[
                        i * szpix : (i + 1) * szpix,
                        (j + 1) * szpix - train_width : (j + 1) * szpix,
                    ]
                    other_beg = im[
                        i1 * szpix : (i1 + 1) * szpix,
                        j1 * szpix : j1 * szpix + train_width,
                    ]
                    x_images[curcnt] = np.hstack((my_end, other_beg))
                    curcnt += 1
                    my_end = im[
                        (i + 1) * szpix - train_width : (i + 1) * szpix,
                        j * szpix : (j + 1) * szpix,
                    ]
                    my_end = np.transpose(my_end, (1, 0, 2))
                    other_beg = im[
                        i1 * szpix : i1 * szpix + train_width,
                        j1 * szpix : (j1 + 1) * szpix,
                    ]
                    other_beg = np.transpose(other_beg, (1, 0, 2))
                    x_images[curcnt] = np.hstack((my_end, other_beg))
                    curcnt += 1
    x_images = normalize_data(x_images)
    # print(x_images[0,0])
    meme = nn_graph.forward(sess, x_images)
    curcnt = 0
    lst = []
    for i in range(M):
        for j in range(M):
            for i1 in range(M):
                for j1 in range(M):
                    if i == i1 and j == j1:
                        continue
                    if meme[curcnt, 1] >= 0.05:
                        lst.append((i * M + j, i1 * M + j1, 0, meme[curcnt, 1]))
                    curcnt += 1
                    if meme[curcnt, 1] >= 0.05:
                        lst.append((i * M + j, i1 * M + j1, 1, meme[curcnt, 1]))
                    curcnt += 1
    print(len(lst), file=fout)
    for x in lst:
        print(*x, file=fout)

# Commented out IPython magic to ensure Python compatibility.
import numpy as np
import pandas as pd
import tensorflow as tf
import keras.preprocessing.image
from IPython.display import HTML
from keras.preprocessing.image import load_img
from keras.preprocessing.image import img_to_array
import sklearn.preprocessing
import sklearn.model_selection
import sklearn.metrics
import sklearn.linear_model
import sklearn.naive_bayes
import sklearn.tree
import sklearn.ensemble
import os
import datetime
import cv2
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.cm as cm

# %matplotlib inline
SZ = 512
szpix = 16
M = SZ // szpix
tf_size = (1 * szpix // 16) * 2
epochs = 35.0
train_width = 2
dirs = ["test2"]
grayscale = False
prefix = "/content/drive/My Drive/Honor"


def to_gray(img):
    gr = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    img = cv2.cvtColor(gr, cv2.COLOR_GRAY2RGB)
    return img


def normalize_data(data):
    # scale features using statistics that are robust to outliers
    # rs = sklearn.preprocessing.RobustScaler()
    # rs.fit(data)
    # data = rs.transform(data)
    data = (data - data.mean()) / (data.std())  # standardisation
    # data=data/data.max()
    data = (data - data.min()) / (
        data.max() - data.min()
    )  # convert from [0:255] to [0.:1.]
    # data = ((data / 255.)-0.5)*2. # convert from [0:255] to [-1.:+1.]
    return data


def create_download_link(title="Result", filename="data.csv"):
    html = "<a href={filename}>{title}</a>"
    html = html.format(title=title, filename=filename)
    return HTML(html)


# create a link to download the dataframe which was saved with .to_csv method
create_download_link(filename=dirs[0] + "_" + str(szpix) + "_neural")
