from tensorflow.keras.layers import BatchNormalization
from keras.layers.convolutional import Conv2D
from keras.layers.convolutional import AveragePooling2D
from keras.layers.convolutional import MaxPooling2D
from keras.layers.core import Activation
from keras.layers.core import Dropout
from keras.layers.core import Dense
from keras.layers import Flatten, MaxPool2D
from keras.layers import Input
from keras.models import Model
from keras.layers import concatenate

from tensorflow.keras.optimizers import SGD

import tensorflow as tf
import tensorflow_probability as tfp

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten, Conv2D, MaxPooling2D
from tensorflow.keras.losses import SparseCategoricalCrossentropy
from tensorflow.keras.optimizers import RMSprop
from keras.datasets import mnist, fashion_mnist, cifar10
from tensorflow import keras
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from keras.utils import np_utils
from sklearn.preprocessing import LabelBinarizer

import copy
import pickle
import numpy as np
import os
import time
import random 
from pathlib import Path
import scipy.io
#################functions ###############
def accur(y_hat, obs):
  j = 0
  
  for i in range(y_hat.shape[0]):
 
    if np.argmax(y_hat, axis = 1)[i]== np.argmax(obs, axis=1)[i]:
      
      j = j+1
 
  return j/float(y_hat.shape[0])

def var_ratio(x):
  return  1 - max(x)
def max_ent(x):
  return -sum(x*np.log2(x))

def weight_un(x):
  return x/x.max()


def load_np(fil_name):
  with open(fil_name, "rb") as f:
    return np.load(f)


#########################################




batch_size = 100
batch_size_spn  = 32
nb_classes = 10


nb_epoch = 25

# input image dimensions
img_rows, img_cols = 32, 32

score=0
all_accuracy = 0
all_accuracy_spn = 0 
acquisition_iterations = 9

#use a large number of dropout iterations
dropout_iterations = 10
Queries = 1000


datapath =Path(__file__).resolve().parent.parent



# Load the data

train_raw = scipy.io.loadmat(os.path.join(datapath, "train_32x32.mat"))
test_raw = scipy.io.loadmat(os.path.join(datapath,"test_32x32.mat"))


# Load images and labels

train_images = np.array(train_raw['X'])
test_images = np.array(test_raw['X'])

train_labels = train_raw['y']
test_labels = test_raw['y']
train_images = np.moveaxis(train_images, -1, 0)
test_images = np.moveaxis(test_images, -1, 0)

train_images = train_images.astype('float64')
test_images = test_images.astype('float64')


# Convert train and test labels into 'int64' type

train_labels = train_labels.astype('int64')
test_labels = test_labels.astype('int64')




lb = LabelBinarizer()
# train_labels = lb.fit_transform(train_labels)
# test_labels = lb.fit_transform(test_labels)








########################## Setting##############################




batch_size = 100

nb_classes = 10







score=0
all_accuracy = 0 #The array to save accuracy of all acquisitions
acquisition_iterations = 9


dropout_iterations = 200
Queries = 1000
###########################
def load_np(fil_name):
  with open(fil_name, "rb") as f:
    return np.load(f)

###########################

Experiments_All_Accuracy = np.zeros(shape=(acquisition_iterations+1))
# (X_train_All, y_train_All), (X_test, y_test) = cifar10.load_data()

(X_train_All, y_train_All), (X_test, y_test) = (train_images, train_labels), (test_images, test_labels)

train_images /= 255.0
test_images /= 255.0








# X_train_All = X_train_All.reshape(X_train_All.shape[0],img_rows, img_cols, 3)
# X_test = X_test.reshape(X_test.shape[0],img_rows, img_cols,3)
random_split = np.asarray(random.sample(range(0,X_train_All.shape[0]), X_train_All.shape[0]))
# rand_path = Path(__file__).resolve().parent.parent.parent
# random_split = load_np(os.path.join(rand_path ,"rand_init_4.npy"))
X_train_All = X_train_All[random_split, :, :, :]
y_train_All = y_train_All[random_split]



lb = LabelBinarizer()
Y_test = lb.fit_transform(y_test)

Y_train_All = lb.fit_transform(y_train_All)

print(X_train_All.shape)









##########model and initial trining ################


# ######models

 

def define_model():
  model = Sequential()
  model.add(Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same', input_shape=(32, 32, 3)))
  model.add(BatchNormalization())
  model.add(Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same'))
  model.add(BatchNormalization())
  model.add(MaxPooling2D((2, 2)))
  model.add(Dropout(0.2))
  model.add(Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same'))
  model.add(BatchNormalization())
  model.add(Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same'))
  model.add(BatchNormalization())
  model.add(MaxPooling2D((2, 2)))
  model.add(Dropout(0.3))
  model.add(Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same'))
  model.add(BatchNormalization())
  model.add(Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same'))
  model.add(BatchNormalization())
  model.add(Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same'))
  model.add(BatchNormalization())
  model.add(MaxPooling2D((2, 2)))
  model.add(Dropout(0.4))
  model.add(Conv2D(512, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same'))
  model.add(BatchNormalization())
  model.add(Conv2D(512, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same'))
  model.add(BatchNormalization())
  model.add(Conv2D(512, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same'))
  model.add(BatchNormalization())
  model.add(MaxPooling2D((2, 2)))





  model.add(Dropout(0.5))
  model.add(Flatten())
  model.add(Dense(512, name="feature", activation='relu', kernel_initializer='he_uniform'))
  model.add(BatchNormalization(name="norm"))
  model.add(Dropout(0.5))
  model.add(Dense(10, activation='softmax'))
  # compile model
  opt = SGD(learning_rate=0.001, momentum=0.9)
  model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])
  return model


deterministic_model = define_model()

start = time.time()

hist = deterministic_model.fit(X_train_All, Y_train_All, batch_size=batch_size, epochs=nb_epoch, verbose=1)

score, acc = deterministic_model.evaluate(X_test, Y_test, verbose=1)
print("The accuracy for first acquisition:", acc)
predi = deterministic_model.predict(X_test, batch_size=batch_size)
all_accuracy = accur(predi, Y_test)
print("The accuracy for first acquisition:", all_accuracy)

# uncertainty test for deterministic

dropout_classes = deterministic_model.predict(X_test, batch_size=batch_size, verbose=1)

All_Dropout_Classes = dropout_classes
# lst_svhn = []
# for i in range(All_Dropout_Classes.shape[0]):
#     lst_svhn.append(max_ent(All_Dropout_Classes[i]))

##############################################################
score_All = np.zeros(shape=(X_test.shape[0], nb_classes))
print('Use trained model for test time dropout')
for d in range(dropout_iterations):
  print('Dropout Iteration', d)
  dropout_score = deterministic_model.predict(X_test, batch_size=batch_size, verbose=1)
  # dropout_classes = np.max(dropout_prob, axis=1).reshape(dropout_prob.shape[0],)
  # dropout_classes = np.argmax(dropout_prob, axis=1).reshape(dropout_prob.shape[0],)

  # dropout_classes = np.array([dropout_classes]).T
  # np.save('/Users/Riashat/Documents/Cambridge_THESIS/Code/Experiments/keras/active_learning/Acquisition_Functions/BCNN_Maximal_Uncertainty/Variation_Ratio/Dropout_Scores/'+'Dropout_Score_'+str(d)+'.npy',dropout_classes)

  score_All = score_All + dropout_score
Avg_Pi = np.divide(score_All, dropout_iterations)
Log_Avg_Pi = np.log2(Avg_Pi)
Entropy_Avg_Pi = - np.multiply(Avg_Pi, Log_Avg_Pi)
Entropy_Average_Pi = np.sum(Entropy_Avg_Pi, axis=1)
U_X = Entropy_Average_Pi

a_1d = U_X.flatten()
a_1d = a_1d.tolist()
#############################################################

print(len(a_1d))

file_name = "entropy_on_svhn-Mc_10"
open_file = open(file_name, "wb")

pickle.dump(a_1d, open_file)

open_file.close()

########## test on cifar10
(_, _), (X_test_CIFAR, y_test_CIFAR) = cifar10.load_data()
X_test_CIFAR = X_test_CIFAR.astype('float32')
X_test /= 255
y_test_CIFAR = np_utils.to_categorical(y_test_CIFAR, nb_classes)

all_accuracy = acc

dropout_classes = dropout_classes = deterministic_model.predict(X_test_CIFAR, batch_size=batch_size, verbose=1)

##############################################################
score_All = np.zeros(shape=(X_test_CIFAR.shape[0], nb_classes))
print('Use trained model for test time dropout')
for d in range(dropout_iterations):
  print('Dropout Iteration', d)
  dropout_score = deterministic_model.predict(X_test_CIFAR, batch_size=batch_size, verbose=1)
  # dropout_classes = np.max(dropout_prob, axis=1).reshape(dropout_prob.shape[0],)
  # dropout_classes = np.argmax(dropout_prob, axis=1).reshape(dropout_prob.shape[0],)

  # dropout_classes = np.array([dropout_classes]).T
  # np.save('/Users/Riashat/Documents/Cambridge_THESIS/Code/Experiments/keras/active_learning/Acquisition_Functions/BCNN_Maximal_Uncertainty/Variation_Ratio/Dropout_Scores/'+'Dropout_Score_'+str(d)+'.npy',dropout_classes)

  score_All = score_All + dropout_score
Avg_Pi = np.divide(score_All, dropout_iterations)
Log_Avg_Pi = np.log2(Avg_Pi)
Entropy_Avg_Pi = - np.multiply(Avg_Pi, Log_Avg_Pi)
Entropy_Average_Pi = np.sum(Entropy_Avg_Pi, axis=1)
U_X = Entropy_Average_Pi

a_1d = U_X.flatten()
a_1d = a_1d.tolist()

file_name = "entropy_on_cifar-Mc_10"
open_file = open(file_name, "wb")
print(len(a_1d))

pickle.dump(a_1d, open_file)

open_file.close()
