# %% Importing libraries

import torch
import torch.nn as nn
import torch.optim as optim

import matplotlib.pyplot as plt
import numpy as np
import scipy.io as sio

import time
import os
import random
import copy

import loss_landscapes
import loss_landscapes.metrics

from mpl_toolkits.mplot3d import axes3d, Axes3D # ignore warning, we need this
from sklearn.decomposition import PCA

from CNN_functions import files_tensor, getListOfFiles, batch_generator, string_tensor, \
        sec_max, weights_init, plot_confusion_matrix, weights_to_list, list_to_weights

import pandas as pd
import seaborn as sb

from scipy.special import softmax

from evt_fitting_V2 import weibull_tailfitting
from compute_openmax_V2 import recalibrate_scores, computeOpenMaxProbability

# %%  ----- Delete old net -----------------------------------------------------------------------------------------------

if 'net' in locals():
    del net

# %%  ----- Set Parameters and Folders -----------------------------------------------------------------------------------

CUDA_flag = False

# Place where output folders are created:
DT = r'C:\Users\ddavidse\Desktop'

# Location of original data:
dirName = r'C:\SharedData\Datasets for CNN\100x100 downsampled from 150x150'

# Location of fooling data:
dataset_fool_dir = r'C:\SharedData\Datasets for CNN\elephant set'
#dataset_fool_dir = r'C:\SharedData\Datasets for CNN\mirrored set'

BS = 50             # batch size
epochs = 12         # number of epochs for training
fig_dpi = 150       # dpi for saved figures
fig_font_size = 12  # font size for figures

# %% ----- Control randomness -----------------------------------------------------------------------------------

random_initialization = True        # set False to use saved_weights
saved_weights = r'D:\ddavidse\Desktop\5 class network runs\network runs - background - full data\Network_run_16\initial_weights.pth'

set_random_seed = False             # set True to use seed_value
seed_value = 784410

# %% ----- Configure dataset split ---------------------------------------------------------------------------------------

ValFrac = 1/6       # set fraction of total dataset used for validation
TestFrac = 1/6      # set fraction of total dataset used for testing

# %% ----- Select features -----------------------------------------------------------------------------------------------

batchnorm_flag = True                   # make True to use batch normalization
misclassified_images_flag = False       # make True to output misclassified images
misclassified_outputs_flag = False      # make True to output misclassified net output values to run info.txt
loss_landscape_flag = False             # make True to create a loss landscape based on random planes
PCA_flag = False                        # make True to create a loss landscape based on PCA

baseline_flag = False                   # make True to use probability thresholding
MAV_flag = True                         # make True to do all the MAV magic
openmax_flag = True                     # make True to use OpenMax (requires MAV_flag to be true as well)

# %% ----- Loss landscapes package settings ------------------------------------------------------------------------------

STEPS = 40      # resolution of landscape, higher = more detail

# %% ----- PCA settings --------------------------------------------------------------------------------------------------

FilterSteps = 40
filter_normalization_flag = True
distance_multiplier = 2
bowl = False
contour_center = False

# %%  ----- CUDA initialization ------------------------------------------------------------------------------------------

if CUDA_flag:
    torch.set_default_tensor_type('torch.cuda.FloatTensor')

# %% ----- Create output folder ------------------------------------------------------------------------------------------

os.chdir(DT)
mfname = 'Network_runs'

if not os.path.exists(mfname):
    os.mkdir(mfname)

fdir = '{}\\{}'.format(DT,mfname)

counter = 0
tv = 0

while tv == 0:

    counter += 1
    fname = '{}\\Network_run_{}'.format(fdir, counter)

    if not os.path.exists(fname):
        os.mkdir(fname)
        tv = 1

# %% ----- detecting number of images per class --------------------------------------------------------------------------

os.chdir(dirName)

DirListTop = os.listdir()
names = DirListTop

N = []

for x in names:

    os.chdir(x)
    DirListSub = os.listdir()
    N.append(len(DirListSub))
    os.chdir(dirName)

NC = np.cumsum(N)

# %% ----- detecting file size -------------------------------------------------------------------------------------------

os.chdir(names[0])
filelist = os.listdir()
testfile = filelist[0]

testimage = sio.loadmat(testfile)
testimage2 = list(testimage.values())
testimage3 = testimage2[3]

input_size = len(testimage3)

# %% ----- Obtaining and storing data ------------------------------------------------------------------------------------

# Data
listOfFiles = getListOfFiles(dirName)
listOfImages = files_tensor(listOfFiles)

# Labels
label = np.zeros(sum(N))
labels = [x for x in range(len(names))]

for i in range(0,NC[0]):
    label[i] = labels[0]

for k in range(1,len(N)):
    for i in range(NC[k-1],NC[k]):
        label[i] = labels[k]

# Data and labels
shuffle_list = list(zip(listOfImages, label))
dataset_orig = [listOfImages, label]

# %% ----- Creating sets with equal class distribution -------------------------------------------------------------------

if set_random_seed == True:
    SeedVal = seed_value
else:
    SeedVal = random.randrange(1,1000000)

random.seed(SeedVal)

zipZ = [[] for x in range(len(N))]
zip_train = [[] for x in range(len(N))]
zip_val = [[] for x in range(len(N))]
zip_test = [[] for x in range(len(N))]

zipZ[0] = shuffle_list[0:NC[0]]

for k in range(1,len(N)):
    zipZ[k] = shuffle_list[NC[k-1]:NC[k]]

for k in range(len(N)):
    random.shuffle(zipZ[k])

TrainFrac = 1 - ValFrac - TestFrac

for k in range(len(N)):
    zip_train[k] = zipZ[k][0:int(TrainFrac*N[k])]
    zip_val[k] = zipZ[k][int(TrainFrac*N[k]):int((TrainFrac+ValFrac)*N[k])]
    zip_test[k] = zipZ[k][int((TrainFrac+ValFrac)*N[k]):N[k]]

train = [x for s in zip_train for x in s]
val = [x for s in zip_val for x in s]
test = [x for s in zip_test for x in s]

files1, labels1 = zip(*test)
test_set = [files1, labels1]

files2, labels2 = zip(*train)
train_set = [files2, labels2]

files3, labels3 = zip(*val)
val_set = [files3, labels3]

# %% ----- Adjusting sets to fit in batches (optional code, not important) -----------------------------------------------

diff_val = BS - (len(val_set[0]) % BS);

if diff_val == 2:

    c_images = []
    c_labels = []

    c_images.append(train_set[0][0])
    c_labels.append(train_set[1][0])

    c_images.append(train_set[0][len(train_set[0]) - 1])
    c_labels.append(train_set[1][len(train_set[1]) - 1])

    c_images = tuple(c_images)
    c_labels = tuple(c_labels)

    val_set[0] = val_set[0] + c_images
    val_set[1] = val_set[1] + c_labels

    train_set[0] = train_set[0][1:len(train_set[0])-1]
    train_set[1] = train_set[1][1:len(train_set[1])-1]

elif diff_val == 1:

    c_image = train_set[0][len(train_set[0]) - 1]
    c_label = train_set[1][len(train_set[1]) - 1]

    t0list = list(val_set[0])
    t0list.append(c_image)
    val_set[0] = tuple(t0list)

    t1list = list(val_set[1])
    t1list.append(c_label)
    val_set[1] = tuple(t1list)

    train_set[0] = train_set[0][0:len(train_set[0])-2]
    train_set[1] = train_set[1][0:len(train_set[1])-2]

diff_test = BS - (len(test_set[0]) % BS);

if diff_test == 2:

    c_images = []
    c_labels = []

    c_images.append(train_set[0][0])
    c_labels.append(train_set[1][0])

    c_images.append(train_set[0][len(train_set[0]) - 1])
    c_labels.append(train_set[1][len(train_set[1]) - 1])

    c_images = tuple(c_images)
    c_labels = tuple(c_labels)

    test_set[0] = test_set[0] + c_images
    test_set[1] = test_set[1] + c_labels

    train_set[0] = train_set[0][1:len(train_set[0])-1]
    train_set[1] = train_set[1][1:len(train_set[1])-1]

elif diff_test == 1:

    c_image = train_set[0][len(train_set[0]) - 1]
    c_label = train_set[1][len(train_set[1]) - 1]

    t0list = list(test_set[0])
    t0list.append(c_image)
    test_set[0] = tuple(t0list)

    t1list = list(test_set[1])
    t1list.append(c_label)
    test_set[1] = tuple(t1list)

    train_set[0] = train_set[0][0:len(train_set[0])-1]
    train_set[1] = train_set[1][0:len(train_set[1])-1]

diff_train = len(test_set[0]) % BS

if diff_train < 5:

    c_image = train_set[0][len(test_set[0]) - 1]
    c_label = train_set[1][len(test_set[1]) - 1]

    t0list = list(train_set[0])
    t0list.append(c_image)
    train_set[0] = tuple(t0list)

    t1list = list(train_set[1])
    t1list.append(c_label)
    train_set[1] = tuple(t1list)

    test_set[0] = test_set[0][0:len(test_set[0])-1]
    test_set[1] = test_set[1][0:len(test_set[1])-1]

# %% ----- Batch generator and loaders ----------------------------------------------------------------------------------

train_loader = batch_generator (BS, train_set)
test_loader = batch_generator (BS, test_set)
val_loader = batch_generator (BS, val_set)

# %% ----- Soft coding input size ----------------------------------------------------------------------------------------

if input_size % 2 == 0:
    Size_1 = input_size / 2
else:
    Size_1 = (input_size - 1) / 2

if Size_1 % 2 == 0:
    Size_2 = Size_1 / 2
else:
    Size_2 = (Size_1 - 1) / 2

Size_2 = int(Size_2)

# %% ----- Defining the neural nets --------------------------------------------------------------------------------------

#note: nn.Conv2D(in, out, kernel, stride=1, padding)
#note: nn.MaxPool2d(kernel, stride, padding)

class BN_Net(nn.Module):

    def __init__(self):
        super(BN_Net, self).__init__()

        self.features = nn.Sequential(
                                      nn.Conv2d(1,5,5,1,2),
                                      nn.MaxPool2d(2,2),
                                      nn.ReLU(inplace=True),
                                      nn.BatchNorm2d(5),
                                      nn.Conv2d(5,8,5,1,2),
                                      nn.MaxPool2d(2,2),
                                      nn.ReLU(inplace=True),
                                      nn.BatchNorm2d(8)
                                         )

        self.classifier = nn.Sequential(
                                        nn.Linear(8*Size_2*Size_2, 120),
                                        nn.ReLU(inplace=True),
                                        nn.BatchNorm1d(120),
                                        nn.Linear(120,84),
                                        nn.ReLU(inplace=True),
                                        nn.BatchNorm1d(84),
                                        nn.Linear(84,5)
                                        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(-1, 8*Size_2*Size_2)
        x = self.classifier(x)

        return x

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()

        self.features = nn.Sequential(
                                      nn.Conv2d(1,5,5,1,2),
                                      nn.MaxPool2d(2,2),
                                      nn.ReLU(inplace=True),
                                      nn.Conv2d(5,8,5,1,2),
                                      nn.MaxPool2d(2,2),
                                      nn.ReLU(inplace=True)
                                         )

        self.classifier = nn.Sequential(
                                        nn.Linear(8*Size_2*Size_2, 120),
                                        nn.ReLU(inplace=True),
                                        nn.Linear(120,84),
                                        nn.ReLU(inplace=True),
                                        nn.Linear(84,5)
                                        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(-1, 8*Size_2*Size_2)
        x = self.classifier(x)

        return x

# %% ----- Initialize or load weights -----------------------------------------------------------------------------

if batchnorm_flag:
    net = BN_Net()
else:
    net = Net()

if CUDA_flag:
    net.cuda()

if random_initialization:
    net.apply(weights_init)
else:
    checkpoint = torch.load(saved_weights)
    net.load_state_dict(checkpoint['model_state_dict'])

# %% ----- Loss and optimizer ---------------------------------------------------------------------------------------

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(),lr=0.001, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False)
#optimizer = optim.RMSprop(net.parameters(), lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0, centered=False)
#optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0, dampening=0, weight_decay=0, nesterov=False)
model_initial = copy.deepcopy(net)

# %% ----- Training the net ---------------------------------------------------------------------------------------------

start = time.time()

running_loss_history = []
running_corrects_history = []
val_running_loss_history = []
val_running_corrects_history = []

StateVecList = []

for e in range(epochs):

    print('\nepoch :', (e+1))
    running_loss = 0.0
    running_corrects = 0.0
    val_running_loss = 0.0
    val_running_corrects = 0.0
    train_loader = batch_generator(BS, train_set)
    val_loader = batch_generator(BS, val_set)

    train_amount = 0.0
    val_amount = 0.0

    for inputs, labels in train_loader:

        labels = string_tensor(labels)
        inputs = torch.stack(inputs)
        inputs = torch.unsqueeze(inputs, 1)
        outputs = net(inputs)
        loss = criterion(outputs, labels.long())

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        _, preds = torch.max(outputs, 1)
        running_loss += loss.item()
        running_corrects += torch.sum(preds == labels.long().data)

        train_amount += BS

    if PCA_flag:
        #StateVecList.append(weights_to_list(net))
        print('nvm')

    with torch.no_grad():
        for val_inputs, val_labels in val_loader:

            val_inputs = torch.stack(val_inputs)
            val_inputs = torch.unsqueeze(val_inputs, 1)
            val_labels = string_tensor(val_labels)
            val_outputs = net(val_inputs)
            val_loss = criterion(val_outputs, val_labels.long())

            _, val_preds = torch.max(val_outputs, 1)
            val_running_loss += val_loss.item()
            val_running_corrects += torch.sum(val_preds == val_labels.long().data)

            val_amount += BS

    epoch_loss = running_loss
    epoch_acc = running_corrects
    running_loss_history.append(epoch_loss)
    running_corrects_history.append(epoch_acc)
    val_epoch_loss = val_running_loss
    val_epoch_acc = val_running_corrects
    val_running_loss_history.append(val_epoch_loss)
    val_running_corrects_history.append(val_epoch_acc)

    print('\ntraining loss: {:.4f} '.format(epoch_loss))
    print('validation loss: {:.4f}'.format(val_epoch_loss))

if CUDA_flag:
    torch.cuda.synchronize()

end = time.time()
print('\n-----------------------------------------------------------------------')
print('\nTime of training: {:d} s'.format(round((end - start))))

# %% ----- Prediction and training stats and graphs --------------------------------------------------------------------

plt.rcParams.update({'font.size': fig_font_size})

correct = [0,0,0]
total = [0,0,0]
testsize = [0,0,0]
AV_all = [[],[],[],[],[]]

with torch.no_grad():
    for data in test_loader:
        images_0_0, labels_0 = data
        images_0_1 = torch.stack(images_0_0)
        images_0_2 = torch.unsqueeze(images_0_1, 1)
        outputs_0 = net(images_0_2)
        _, predicted_0 = torch.max(outputs_0.data, 1)
        total[0] += string_tensor(labels_0).size(0)
        testsize[0] += 1

        correct_index_0 = predicted_0 == string_tensor(labels_0).long()
        correct[0] += correct_index_0.sum().item()

        if MAV_flag:
            for x in range(len(outputs_0)):
                if correct_index_0[x]:
                    AV_all[predicted_0[x]].append(outputs_0[x])

    train_loader = batch_generator(BS, train_set)
    val_loader = batch_generator(BS, val_set)

    for data in train_loader:
        images_1_0, labels_1 = data
        images_1_1 = torch.stack(images_1_0)
        images_1_2 = torch.unsqueeze(images_1_1, 1)
        outputs_1 = net(images_1_2)
        _, predicted_1 = torch.max(outputs_1.data, 1)
        total[1] += string_tensor(labels_1).size(0)
        testsize[1] += 1

        correct_index_1 = predicted_1 == string_tensor(labels_1).long()
        correct[1] += correct_index_1.sum().item()

        if MAV_flag:
            for x in range(len(outputs_1)):
                if correct_index_1[x]:
                    AV_all[predicted_1[x]].append(outputs_1[x])

    for data in val_loader:
        images_2_0, labels_2 = data
        images_2_1 = torch.stack(images_2_0)
        images_2_2 = torch.unsqueeze(images_2_1, 1)
        outputs_2 = net(images_2_2)
        _, predicted_2 = torch.max(outputs_2.data, 1)
        total[2] += string_tensor(labels_2).size(0)
        testsize[2] += 1

        correct_index_2 = predicted_2 == string_tensor(labels_2).long()
        correct[2] += correct_index_2.sum().item()

        if MAV_flag:
            for x in range(len(outputs_2)):
                if correct_index_2[x]:
                    AV_all[predicted_2[x]].append(outputs_2[x])

if MAV_flag:

    MAV = np.empty((5,5))

    for i in range(5):
        AV_array = np.array([list(x.cpu().numpy()) for x in AV_all[i]])
        MAV[i] = np.mean(AV_array, axis=0)

"""
================================================
================================================
===== SECONDARY DATASET LOOPS FOR MAV
================================================
================================================
"""

AV_dist_correct = [[],[],[],[],[]]
AV_dist_wrong = []

if MAV_flag:

    train_loader = batch_generator(BS, train_set)
    val_loader = batch_generator(BS, val_set)
    test_loader = batch_generator(BS, test_set)

    with torch.no_grad():
        for data in test_loader:
            images_0_0, labels_0 = data
            images_0_1 = torch.stack(images_0_0)
            images_0_2 = torch.unsqueeze(images_0_1, 1)
            outputs_0 = net(images_0_2)
            _, predicted_0 = torch.max(outputs_0.data, 1)

            correct_index_0 = predicted_0 == string_tensor(labels_0).long()

            for x in range(len(outputs_0)):
                avdiff = outputs_0[x].cpu().numpy() - MAV[predicted_0[x]]
                avdist = np.linalg.norm(avdiff)

                if correct_index_0[x]:
                    AV_dist_correct[predicted_0[x]].append(avdist)
                else:
                    AV_dist_wrong.append(avdist)

        for data in train_loader:
            images_1_0, labels_1 = data
            images_1_1 = torch.stack(images_1_0)
            images_1_2 = torch.unsqueeze(images_1_1, 1)
            outputs_1 = net(images_1_2)
            _, predicted_1 = torch.max(outputs_1.data, 1)

            correct_index_1 = predicted_1 == string_tensor(labels_1).long()

            for x in range(len(outputs_1)):
                avdiff = outputs_1[x].cpu().numpy() - MAV[predicted_1[x]]
                avdist = np.linalg.norm(avdiff)

                if correct_index_1[x]:
                    AV_dist_correct[predicted_1[x]].append(avdist)
                else:
                    AV_dist_wrong.append(avdist)

        for data in val_loader:
            images_2_0, labels_2 = data
            images_2_1 = torch.stack(images_2_0)
            images_2_2 = torch.unsqueeze(images_2_1, 1)
            outputs_2 = net(images_2_2)
            _, predicted_2 = torch.max(outputs_2.data, 1)

            correct_index_2 = predicted_2 == string_tensor(labels_2).long()

            for x in range(len(outputs_0)):
                avdiff = outputs_2[x].cpu().numpy() - MAV[predicted_2[x]]
                avdist = np.linalg.norm(avdiff)

                if correct_index_2[x]:
                    AV_dist_correct[predicted_2[x]].append(avdist)
                else:
                    AV_dist_wrong.append(avdist)

    AV_dist_tot = [x for i in AV_dist_correct for x in i]

    AVCA = np.array(AV_dist_tot)
    AVC_max = np.max(AVCA)
    AVC_min = np.min(AVCA)
    AVC_avg = np.mean(AVCA)
    AVC_std = np.std(AVCA)

    AVWA = np.array(AV_dist_wrong)
    AVW_max = np.max(AVWA)
    AVW_min = np.min(AVWA)
    AVW_avg = np.mean(AVWA)
    AVW_std = np.std(AVWA)

    test_loader = batch_generator(BS, test_set)

    with torch.no_grad():

        ImagesChecked = 0
        ImagesCheckedCorrect = 0

        for data in test_loader:
            images_3_0, labels_3 = data
            images_3_1 = torch.stack(images_3_0)
            images_3_2 = torch.unsqueeze(images_3_1, 1)
            outputs_3 = net(images_3_2)
            _, predicted_3 = torch.max(outputs_3.data, 1)

            correct_index_3 = predicted_3 == string_tensor(labels_3).long()

            for x in range(len(outputs_3)):
                avdiff = outputs_3[x].cpu().numpy() - MAV[predicted_3[x]]
                avdist = np.linalg.norm(avdiff)
                if avdist < AVW_min:
                #if avdist < 2.5:
                    ImagesChecked += 1
                    if correct_index_3[x]:
                        ImagesCheckedCorrect += 1

    df1 = pd.DataFrame({'1':AVCA})
    df2 = pd.DataFrame({'2':AVWA})
    df = pd.concat([df1, df2], ignore_index=True, axis=1)
    df.columns = ['correct','incorrect']

    MAV_dist_plot = plt.figure(figsize=(9,7))
    plt.grid(1, which='major')
    plt.grid(1, which='minor', color='k', linestyle='-', alpha=0.08)
    plt.minorticks_on()
    sbplot = sb.stripplot(data=df)
    plt.ylabel('distance to MAV')
    x = plt.gca().axes.get_xlim()
    plt.plot(x, len(x)*[AVW_min],'r')

# %%

"""
================================================
================================================
===== FOOLING DATASET
================================================
================================================

"""

if MAV_flag:

    thresh_flag = True
    thresh = 4.7

    os.chdir(dataset_fool_dir)

    DirListTop_fool = os.listdir()
    names_fool = DirListTop_fool

    N_fool = []

    for x in names_fool:

        os.chdir(x)
        DirListSub_fool = os.listdir()
        N_fool.append(len(DirListSub_fool))
        os.chdir(dataset_fool_dir)

    NC_fool = np.cumsum(N_fool)

    listOfFiles_fool = getListOfFiles(dataset_fool_dir)
    listOfImages_fool = files_tensor(listOfFiles_fool)

    label_fool = np.zeros(sum(N_fool))
    labels_fool = [x for x in range(len(names_fool))]

    for i in range(0,NC_fool[0]):
        label_fool[i] = labels_fool[0]

    for k in range(1,len(N_fool)):
        for i in range(NC_fool[k-1],NC_fool[k]):
            label_fool[i] = labels_fool[k]

    dataset_fool = [listOfImages_fool, label_fool]

    list_1 = AV_dist_tot
    list_2 = AV_dist_wrong

    data_loader = batch_generator(BS, dataset_fool)

    dist_test = []

    with torch.no_grad():
        for data in data_loader:
            images, labels = data
            images = torch.stack(images)
            images = torch.unsqueeze(images, 1)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)

            correct_index = predicted == string_tensor(labels).long()

            for x in range(len(outputs)):
                avdiff = outputs[x].cpu().numpy() - MAV[predicted[x]]
                avdist = np.linalg.norm(avdiff)

                dist_test.append(avdist)

    df0 = pd.DataFrame({'0':list_1 + list_2})
    df1 = pd.DataFrame({'1':list_1})
    df2 = pd.DataFrame({'2':list_2})
    df3 = pd.DataFrame({'3':dist_test})

    df = pd.concat([df0, df3], ignore_index=True, axis=1)
    df.columns = ['original set','fooling set']

    MAV_dist_plot = plt.figure(figsize=(9,7))
    plt.grid(1, which='major')
    plt.grid(1, which='minor', color='k', linestyle='-', alpha=0.08)
    plt.minorticks_on()
    sbplot = sb.stripplot(data=df, jitter=0.15)
    plt.ylabel('distance to MAV')

    if thresh_flag:

        x = plt.gca().axes.get_xlim()
        plt.plot(x, len(x)*[thresh],'r')

        list_1_thr = [x for x in list_1 if x < thresh]
        list_2_thr = [x for x in list_2 if x < thresh]
        dist_test_thr = [x for x in dist_test if x < thresh]

        L1C = len(list_1)
        L1W = len(list_2)
        L1T = len(dist_test)

        L2C = len(list_1_thr)
        L2W = len(list_2_thr)
        L2T = len(dist_test_thr)

        A1 = round(100 * L1C / (L1C + L1W), 1)
        A2 = round(100 * L2C / (L2C + L2W), 1)

        P1 = round(100 * (1 - (L2C + L2W) / (L1C + L1W)), 1)
        P2 = round(100 * (1 - L2T / L1T), 1)

        print('\n-------------------------------------------------------------------')

        os.chdir(fname)

        path = '{}\\MAV_distance_dist.png'.format(fname)
        MAV_dist_plot.savefig(path, dpi=fig_dpi)

        f = open('MAV results.txt','w+')
        f.write('--------------------------------------------------------------------------------------')
        f.write('\n\nAccuracy on the original dataset: \t\t{} %'.format(A1))
        f.write('\nAccuracy on thresholded original dataset: \t{} %'.format(A2))
        f.write('\n\nThreshold: {}'.format(thresh))
        f.write('\n\nPart of the original dataset thrown away: \t{} %'.format(P1))
        f.write('\nPart of the secondary dataset thrown away: \t{} %'.format(P2))
        f.write('\n\n--------------------------------------------------------------------------------------')
        f.close()

# %%

"""
================================================
================================================
===== OPENMAX
================================================
================================================
"""

if openmax_flag:

    tail = 2
    open_thresh = 0.2
    alpha = 4

    maxarray = np.array([1,2,3,4,5])

    weibull_models = weibull_tailfitting(AV_dist_correct, MAV, names, tailsize=tail, distance_type='euclidean')

    data_loader_fool = batch_generator(BS, dataset_fool)

    n_open_set = 0
    n_set = 0
    SMM_fool = []
    SMM_orig = []

    with torch.no_grad():
        for data in data_loader_fool:
            images, labels = data
            images = torch.stack(images)
            images = torch.unsqueeze(images, 1)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)

            correct_index = predicted == string_tensor(labels).long()

            for x in outputs:
                openmax_probs = recalibrate_scores(weibull_models, x.numpy(), alpharank=alpha)
                n_set += 1

                SMM = max(openmax_probs[:-1])
                SMM_fool.append(SMM)

                if SMM < open_thresh:
                    n_open_set += 1

    open_throw = round(100 * n_open_set / n_set, 1)

    data_loader_orig = batch_generator(BS, dataset_orig)

    with torch.no_grad():
        for data in data_loader_orig:
            images, labels = data
            images = torch.stack(images)
            images = torch.unsqueeze(images, 1)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)

            correct_index = predicted == string_tensor(labels).long()

            for x in outputs:
                openmax_probs = recalibrate_scores(weibull_models, x.numpy(), alpharank=alpha)
                n_set += 1

                SMM = max(openmax_probs[:-1])
                SMM_orig.append(SMM)

                if SMM < open_thresh:
                    n_open_set += 1

    open_throw_2 = round(100 * n_open_set / n_set, 1)

    df1 = pd.DataFrame({'1':SMM_orig})
    df2 = pd.DataFrame({'2':SMM_fool})
    df = pd.concat([df1, df2], ignore_index=True, axis=1)
    df.columns = ['original','fool']

    openmax_plot = plt.figure(figsize=(9,7))
    plt.grid(1, which='major')
    plt.grid(1, which='minor', color='k', linestyle='-', alpha=0.08)
    plt.minorticks_on()
    sbplot = sb.stripplot(data=df)
    plt.ylabel('max prob 1:5')
    x = plt.gca().axes.get_xlim()
    plt.plot(x, len(x)*[open_thresh],'r')

    print('\n-------------------------------------------------------------------------------------')
    print('\nOpenMax threw out {} % of the original dataset'.format(open_throw_2))
    print('OpenMax threw out {} % of the fooling dataset'.format(open_throw))
    print('\nDifference: {:.1f} %'.format(open_throw - open_throw_2))

    os.chdir(fname)

    path = '{}\openmax_plot.png'.format(fname)
    openmax_plot.savefig(path, dpi=fig_dpi)

    f = open('openmax info.txt','w+')
    f.write('-----------------------------------------------------------------------------------------')
    f.write('\n\nparameters used:')
    f.write('\n\ntail = {}'.format(tail))
    f.write('\nopen_thresh = {}'.format(open_thresh))
    f.write('\nalpha = {}'.format(alpha))
    f.write('\n\n-----------------------------------------------------------------------------------------')
    f.write('\n\nResults:')
    f.write('\n\nOpenMax threw out {} % of the fooling dataset'.format(open_throw))
    f.write('\nOpenMax threw out {} % of the original dataset'.format(open_throw_2))
    f.write('\n\nDifference: {:.1f} %'.format(open_throw - open_throw_2))
    f.write('\n\n-----------------------------------------------------------------------------------------')
    f.close()

    os.chdir(DT)

# %%

"""
================================================
================================================
===== SECONDARY DATASET LOOPS FOR BASELINE
================================================
================================================
"""

if baseline_flag:

    train_loader = batch_generator(BS, train_set)
    val_loader = batch_generator(BS, val_set)
    test_loader = batch_generator(BS, test_set)

    unc_correct = []
    unc_wrong = []

    with torch.no_grad():

        for data in test_loader:
            images, labels = data
            images = torch.stack(images)
            images = torch.unsqueeze(images, 1)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)

            correct_index = predicted == string_tensor(labels).long()

            for x in range(len(outputs)):

                probs = softmax(outputs[x].cpu().numpy())
                prob = np.max(probs)
                uncert = 100*(1 - prob)

                if correct_index[x]:
                    unc_correct.append(uncert)
                else:
                    unc_wrong.append(uncert)

        for data in train_loader:
            images, labels = data
            images = torch.stack(images)
            images = torch.unsqueeze(images, 1)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)

            correct_index = predicted == string_tensor(labels).long()

            for x in range(len(outputs)):

                probs = softmax(outputs[x].cpu().numpy())
                prob = np.max(probs)
                uncert = 100*(1 - prob)

                if correct_index[x]:
                    unc_correct.append(uncert)
                else:
                    unc_wrong.append(uncert)

        for data in val_loader:
            images, labels = data
            images = torch.stack(images)
            images = torch.unsqueeze(images, 1)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)

            correct_index = predicted == string_tensor(labels).long()

            for x in range(len(outputs)):

                probs = softmax(outputs[x].cpu().numpy())
                prob = np.max(probs)
                uncert = 100*(1 - prob)

                if correct_index[x]:
                    unc_correct.append(uncert)
                else:
                    unc_wrong.append(uncert)

    df1 = pd.DataFrame({'1':unc_correct})
    df2 = pd.DataFrame({'2':unc_wrong})
    df = pd.concat([df1, df2], ignore_index=True, axis=1)
    df.columns = ['correct','wrong']

    baseline_plot = plt.figure(figsize=(9,7))
    plt.grid(1, which='major')
    plt.grid(1, which='minor', color='k', linestyle='-', alpha=0.08)
    plt.minorticks_on()
    sbplot = sb.stripplot(data=df, jitter=.15)
    plt.ylabel('uncertainty in %')
    plt.title('baseline uncertainty plot')

total_tot = sum(total)
correct_tot = sum(correct)
testsize_tot = sum(testsize)

print('\nAccuracy of the network on the test images: {:.1f} %'.format((100 * correct[0] / total[0])))
print('Test result obtained from {} images coming from {} batches'.format(total[0], testsize[0]))
print('\nAccuracy of the network on the total set of all images: {:.1f} %'.format((100 * correct_tot / total_tot)))
print('Test result obtained from {} images coming from {} batches'.format(total_tot, testsize_tot))

if MAV_flag:
    print('\nAmount of images in high confidence set: {}'.format(ImagesChecked))
    print('This is {:.1f}% of the test set'.format(100*ImagesChecked / total[0]))

    print('\nAmount of correctly classified images in high confidence set: {}'.format(ImagesCheckedCorrect))
    print('Accuracy on high confidence set: {:.1f} %'.format(100*ImagesCheckedCorrect / ImagesChecked))

#print('\nAmount of correct predictions in test 1: {}'.format(correct[0]))

if MAV_flag:
    print('\n-------------------------------------------------------------------------------------')
    print('\nAccuracy on the original dataset: \t\t{} %'.format(A1))
    print('Accuracy on thresholded original dataset: \t{} %'.format(A2))
    print('\nThreshold: {}'.format(thresh))
    print('\nPart of the original dataset thrown away: \t{} %'.format(P1))
    print('Part of the secondary dataset thrown away: \t{} %'.format(P2))
    if openmax_flag:
        print('\n-------------------------------------------------------------------------------------')
        print('\nOpenMax threw out {} % of the original dataset'.format(open_throw_2))
        print('OpenMax threw out {} % of the fooling dataset'.format(open_throw))
        print('\nDifference: {:.1f} %'.format(open_throw - open_throw_2))

val_plot = plt.figure(figsize=(9,7))
plt.grid(1, which='major')
plt.grid(1, which='minor', color='k', linestyle='-', alpha=0.08)
plt.minorticks_on()
plt.plot(running_loss_history, 'r',label='training loss')
plt.plot(val_running_loss_history,'b', label='validation loss')
plt.autoscale(enable=True, axis='x', tight=True)
plt.legend()

v_len = len(val_set[0])
t_len = len(train_set[0])

acc_plot = plt.figure(figsize=(9,7))
plt.grid(1, which='major')
plt.grid(1, which='minor', color='k', linestyle='-', alpha=0.08)
plt.minorticks_on()
train_corrected = [100*float(x)/(train_amount) for x in running_corrects_history]
plt.plot(train_corrected,'r', label='training accuracy [%]')
val_corrected = [100*float(x)/(val_amount) for x in val_running_corrects_history]
plt.plot(val_corrected,'b', label='validation accuracy [%]')
plt.autoscale(enable=True, axis='x', tight=True)
plt.legend()

# %% ----- confusion matrix --------------------------------------------------------------------------------------------
# code inspiration from https://deeplizard.com/learn/video/0LhiS6yu2qQ

test_loader = batch_generator (BS, test_set)
all_preds = torch.tensor([])
all_labels = torch.tensor([])

misc_preds_t = []
misc_labels_t = []
misc_img_t = []
misc_outputs_t = []
maxdif_t = []

for inputs, labels in test_loader:

        inputs = torch.stack(inputs)
        inputs = torch.unsqueeze(inputs, 1)

        labels = string_tensor(labels)
        all_labels = torch.cat((all_labels, labels),dim=0)

        outputs = net(inputs)
        all_preds = torch.cat((all_preds, outputs),dim=0)

        imlist = list(inputs)
        outputs1 = outputs.argmax(dim=1)
        vec_equal = outputs1 == labels

        if misclassified_outputs_flag:

            misc_outputs_0 = outputs[~vec_equal]
            misc_outputs_1 = misc_outputs_0.cpu().detach().numpy().tolist()
            misc_outputs = [list(map(lambda x: round(x,1), i)) for i in misc_outputs_1]
            misc_outputs_t += misc_outputs

            maxdif = [max(x) - sec_max(x) for x in misc_outputs_1]
            maxdif_t += maxdif

        if misclassified_images_flag:

            misc_ind = [x for x in range(BS) if vec_equal[x] == False]

            misc_img = [imlist[x] for x in misc_ind]
            misc_img_t += misc_img

            misc_labels = [labels[x] for x in misc_ind]
            misc_labels_t += misc_labels

            misc_preds = [outputs1[x] for x in misc_ind]
            misc_preds_t += misc_preds

AL = all_labels
AL2 = np.array(list(AL), dtype=np.int)
AL3 = torch.tensor(AL2)

TP = all_preds.argmax(dim=1)
TP2 = np.array(list(TP), dtype=np.int)
TP3 = torch.tensor(TP2)

ConfAcc = sum(TP3 == AL3).item()
#print('Amount of correct predictions in test 2: {}'.format(ConfAcc))
print('\n-----------------------------------------------------------------------------------')

stacked = torch.stack((AL3, TP3), dim=1)

cmt = torch.zeros(len(N),len(N), dtype=torch.int64)

for p in stacked:
    tl, pl = p.tolist()
    cmt[tl, pl] = cmt[tl, pl] + 1

cmt2 = cmt.cpu()

if len(names) == 3:
    conf_fig_size = (8,8)
elif len(names) == 4:
    conf_fig_size = (9,9)
else:
    conf_fig_size = (10,10)

conf_fig = plt.figure(figsize = conf_fig_size)
plot_confusion_matrix(cmt2, names)

conf_fig_2 = plt.figure(figsize = conf_fig_size)
plot_confusion_matrix(cmt2, names, normalize='row', percentage=True)

conf_fig_3 = plt.figure(figsize = conf_fig_size)
plot_confusion_matrix(cmt2, names, normalize='full', percentage=True)

# %% ----- loss landscape ----------------------------------------------------------------------------------------------
# source: https://github.com/marcellodebernardi/loss-landscapes/blob/master/examples/core-features.ipynb

if loss_landscape_flag:
    start = time.time()

    train_loader = batch_generator (BS, train_set)
    x, y = iter(train_loader).__next__()
    x = torch.stack(x)
    x = torch.unsqueeze(x, 1)
    y = torch.tensor(y).long()
    metric = loss_landscapes.metrics.Loss(criterion, x, y)

    LCP = loss_landscapes.random_plane(net, metric, 100, STEPS, normalization='filter', deepcopy_model=True)

    if CUDA_flag:
        torch.cuda.synchronize()

    end = time.time()
    print('\nTime of calculating loss landscape: {:d} s'.format(round((end - start))))

    loss_con = plt.figure()
    plt.contour(LCP, levels=50)
    plt.title('Loss Contours around Trained Model')
    plt.show()

    loss_surf_1 = fig = plt.figure(figsize=(9,7))
    ax = plt.axes(projection='3d')

    X = np.array([[j for j in range(STEPS)] for i in range(STEPS)])
    Y = np.array([[i for _ in range(STEPS)] for i in range(STEPS)])
    ax.plot_surface(X, Y, LCP, rstride=1, cstride=1, cmap='coolwarm', edgecolor='none')
    ax.set_title('Surface Plot of Loss Landscape')
    ax.set_xlabel(r'$\theta$', fontsize=18, labelpad=10)
    ax.set_ylabel(r"$\theta '$", fontsize=18, labelpad=10)
    ax.set_zlabel('Loss', fontsize=18, labelpad=10, rotation=90)

    loss_surf_2 = fig = plt.figure(figsize=(9,7))
    ax = plt.axes(projection='3d')

    ax.plot_surface(X, Y, LCP, rstride=1, cstride=1, cmap='coolwarm', edgecolor='none')
    #ax.set_title('Surface Plot of Loss Landscape')
    ax.view_init(30, 45)
    ax.set_xlabel(r'$\theta$', fontsize=18, labelpad=10)
    ax.set_ylabel(r"$\theta '$", fontsize=18, labelpad=10)
    ax.set_zlabel('Loss', fontsize=18, labelpad=10, rotation=90)

# %% ----- PCA ----------------------------------------------------------------------------------------------------------

if PCA_flag:

    PCA_time_start = time.time()
    StateVecArray = np.array(StateVecList)
    TrainedNetVector = StateVecArray[epochs-1]

    pca = PCA(n_components=2)
    PC = pca.fit_transform(StateVecArray)
    PC_norm = []

    for i in range(len(PC[0])):
        col = PC[:,i]
        col_fixed = col / np.linalg.norm(col)
        PC_norm.append(col_fixed)

    SVA_1 = []
    for i in range(len(StateVecArray[0])):
        weightvec = StateVecArray[:,i]
        SVA_1.append(np.dot(weightvec, PC_norm[0]))

    SVA_2 = []
    for i in range(len(StateVecArray[0])):
        weightvec = StateVecArray[:,i]
        SVA_2.append(np.dot(weightvec, PC_norm[1]))

    SVA_1 = np.array(SVA_1)
    SVA_2 = np.array(SVA_2)

    if filter_normalization_flag:

        A = net.state_dict()
        S = [x.size() for x in A.values()]

        SL = [list(x) for x in S]

        L = len(SL)
        Sizes = []
        WeightTensors = []
        ite = 0

        for x in range(L):
            ite += 1
            xsize = 1
            for number in SL[x]:
                xsize = xsize * number
            Sizes.append(xsize)
            Sizes_C = np.cumsum(Sizes)

        Seg0 = TrainedNetVector[0:Sizes_C[0]]
        Seg0_Norm = np.linalg.norm(Seg0)

        Seg1 = SVA_1[0:Sizes_C[0]]
        Seg1_Norm = np.linalg.norm(Seg1)

        Seg2 = SVA_2[0:Sizes_C[0]]
        Seg2_Norm = np.linalg.norm(Seg2)

        SVA_1[0:Sizes_C[0]] = SVA_1[0:Sizes_C[0]] * Seg0_Norm / Seg1_Norm
        SVA_2[0:Sizes_C[0]] = SVA_2[0:Sizes_C[0]] * Seg0_Norm / Seg2_Norm

        TestNorm0 = np.linalg.norm(TrainedNetVector[0:Sizes_C[0]])
        TestNorm1 = np.linalg.norm(SVA_1[0:Sizes_C[0]])
        TestNorm2 = np.linalg.norm(SVA_2[0:Sizes_C[0]])

        print('\nTestNorm of first filter from core: \t\t{}'.format(TestNorm0))
        print('TestNorm of first filter from PCA vector 1: \t{}'.format(TestNorm1))
        print('TestNorm of first filter from PCA vector 2: \t{}'.format(TestNorm2))

        for x in range(len(Sizes_C) - 1):

            Seg0 = TrainedNetVector[Sizes_C[x]:Sizes_C[x+1]]
            Seg0_Norm = np.linalg.norm(Seg0)

            Seg1 = SVA_1[Sizes_C[x]:Sizes_C[x+1]]
            Seg1_Norm = np.linalg.norm(Seg1)

            Seg2 = SVA_2[Sizes_C[x]:Sizes_C[x+1]]
            Seg2_Norm = np.linalg.norm(Seg2)

            SVA_1[Sizes_C[x]:Sizes_C[x+1]] = SVA_1[Sizes_C[x]:Sizes_C[x+1]] * Seg0_Norm / Seg1_Norm
            SVA_2[Sizes_C[x]:Sizes_C[x+1]] = SVA_2[Sizes_C[x]:Sizes_C[x+1]] * Seg0_Norm / Seg2_Norm

    RIstep = int(round(FilterSteps/2))

    loss_array = np.zeros([FilterSteps+1,FilterSteps+1])

    for i in range(-RIstep, RIstep+1):
        for j in range(-RIstep, RIstep+1):
            NetAdd_1 = distance_multiplier*(i/RIstep)*SVA_1
            NetAdd_2 = distance_multiplier*(j/RIstep)*SVA_2
            NetWeights = TrainedNetVector + NetAdd_1 + NetAdd_2

            net_updated = list_to_weights(NetWeights, net, True)

            if CUDA_flag:
                net_updated.cuda()

            val_running_loss = 0.0
            val_loader = batch_generator (BS, val_set)

            with torch.no_grad():
                for val_inputs, val_labels in val_loader:

                    val_inputs = torch.stack(val_inputs)
                    val_inputs = torch.unsqueeze(val_inputs, 1)
                    val_labels = string_tensor(val_labels)
                    val_outputs = net_updated(val_inputs)
                    val_loss = criterion(val_outputs, val_labels.long())

                    val_running_loss += val_loss.item()

            loss_array[i+RIstep, j+RIstep] = val_running_loss

    x = np.arange(-RIstep+1, RIstep+1)
    y = np.arange(-RIstep, RIstep+1)
    X,Y = np.meshgrid(x,y)

    loss_array_cor = loss_array[:,1:]

    if bowl:

        tv0 = loss_array[round(FilterSteps/2) - 1, FilterSteps - 1]
        tv1 = loss_array[round(FilterSteps/2) - 1, 0]
        tv2 = loss_array[0, round(FilterSteps/2) - 1]
        tv3 = loss_array[FilterSteps - 1, round(FilterSteps/2) - 1]

        bowlmax = max([tv0, tv1, tv2, tv3])
        bowlmin = min([tv0, tv1, tv2, tv3])

        for vi in range(FilterSteps):
            for vj in range(FilterSteps):
                if loss_array[vi,vj] > bowlmin:
                    loss_array[vi,vj] = bowlmin

    PCA_fig_1 = plt.figure(figsize=(10,7))
    ax = PCA_fig_1.gca(projection='3d')
    PCA_surf = ax.plot_surface(X,Y, loss_array_cor, cmap='coolwarm')
    ax.set_xlabel(r'X', fontsize=20)
    ax.set_ylabel(r'Y', fontsize=20)
    ax.set_zlabel(r'Z', fontsize=20)

    PCA_fig_2 = plt.figure(figsize=(10,7))
    ax = PCA_fig_2.gca(projection='3d')
    PCA_surf_2 = ax.plot_surface(X,Y, loss_array_cor, cmap='coolwarm')
    ax.set_xlabel(r'X', fontsize=20)
    ax.set_ylabel(r'Y', fontsize=20)
    ax.set_zlabel(r'Z', fontsize=20)
    ax.view_init(30, 45)

    PCA_fig_3 = plt.figure(figsize=(10,7))
    PCA_cont = plt.contour(X,Y, loss_array_cor, 100)
    ax = PCA_fig_3.gca()
    ax.set_xlabel(r'X', fontsize=20)
    ax.set_ylabel(r'Y', fontsize=20)
    if contour_center:
        hlinex = [-19,20]
        hliney = [0,0]
        plt.plot(hlinex, hliney,'k')
        vlinex = [0,0]
        vliney = [-20,20]
        plt.plot(vlinex, vliney,'k')

        loss_array_cor[0][0] = loss_array_cor[0][1]
        LMIN = np.min(loss_array_cor)
        LMINLOC = np.argmin(loss_array_cor)
        RowLen = len(loss_array_cor[0])
        min_remain = LMINLOC % RowLen
        min_y = (LMINLOC - min_remain) / RowLen
        min_x = min_remain
        plt.plot(min_x - 20, min_y - 20, 'rx', markersize=15)

    PCA_time_end = time.time()
    PCA_time = PCA_time_end - PCA_time_start
    print('\nPCA time: {} s'.format(round(PCA_time,1)))

# %% ----- Saving results -----------------------------------------------------------------------------------------------

os.chdir(fname)

if MAV_flag:
    mavpath1 = '{}\\MAV.mat'.format(fname)
    mavpath2 = '{}\\MAV_dist_cor.mat'.format(fname)
    mavpath3 = '{}\\MAV_dist_inc.mat'.format(fname)
    sio.savemat(mavpath1, {'data':MAV})
    sio.savemat(mavpath2, {'data':AV_dist_correct})
    sio.savemat(mavpath3, {'data':AV_dist_wrong})

    f = open('AV stats.txt','w+')
    f.write('Correct predictions:')
    f.write('\n\nAV_max = {:.2f}'.format(AVC_max))
    f.write('\nAV_min = {:.2f}'.format(AVC_min))
    f.write('\nAV_avg = {:.2f}'.format(AVC_avg))
    f.write('\nAV_std = {:.2f}'.format(AVC_std))

    f.write('\n\nWrong predictions:')
    f.write('\n\nAV_max = {:.2f}'.format(AVW_max))
    f.write('\nAV_min = {:.2f}'.format(AVW_min))
    f.write('\nAV_avg = {:.2f}'.format(AVW_avg))
    f.write('\nAV_std = {:.2f}'.format(AVW_std))
    f.close()

f = open('run info.txt','w+')
A = str(net)
f.write('-----------------------------------------------------------------------------------------\n\n')
f.write(A)
f.write('\n\n-----------------------------------------------------------------------------------------')
f.write('\n\nData used:\n')

for i in range(len(N)):
    f.write('\nClass {}: {} images'.format(names[i], N[i]))

f.write('\n\nTrain set size: {} images'.format(len(train_set[0])))
f.write('\nVal set length: {} images'.format(len(val_set[0])))
f.write('\nTest set size: \t{} images'.format(len(test_set[0])))
f.write('\n\nTotal size of dataset: {} images'.format(sum(N)))

f.write('\n\n-----------------------------------------------------------------------------------------')

f.write('\n\ninput_size = {}'.format(input_size))
f.write('\nBS = {}'.format(BS))
f.write('\nepochs = {}'.format(epochs))
f.write('\nSTEPS = {}'.format(STEPS))
f.write('\nSeedVal = {}'.format(SeedVal))

f.write('\n\nLoss function: {}'.format(criterion))
f.write('\n\nOptimizer: \n{}'.format(optimizer))

f.write('\n\n-----------------------------------------------------------------------------------------')

f.write('\n\nAccuracy of the network on the test images: {:.1f} %'.format((100 * correct[0] / total[0])))
f.write('\nTest result obtained from {} images coming from {} batches'.format(total[0], testsize[0]))
f.write('\n\nAccuracy of the network on the total set of all images: {:.1f} %'.format((100 * correct_tot / total_tot)))
f.write('\nTest result obtained from {} images coming from {} batches'.format(total_tot, testsize_tot))

if MAV_flag:
    f.write('\n\nAmount of images in high confidence set: {}'.format(ImagesChecked))
    f.write('\nThis is {:.1f}% of the test set'.format(100*ImagesChecked / total[0]))

    f.write('\n\nAmount of correctly classified images in high confidence set: {}'.format(ImagesCheckedCorrect))
    f.write('\nAccuracy on high confidence set: {:.1f} %'.format(100*ImagesCheckedCorrect / ImagesChecked))

if PCA_flag:
    f.write('\n\n-----------------------------------------------------------------------------------------')
    f.write('\n\nPCA - filter normalization: {}'.format(filter_normalization_flag))
    f.write('\nPCA - distance multiplier: {}'.format(distance_multiplier))
    f.write('\nPCA - FilterSteps: {}'.format(FilterSteps))

if misclassified_outputs_flag:
    f.write('\n\n-----------------------------------------------------------------------------------------')
    f.write('\n\nMisclassified net outputs: ')
    f.write('\n')
    maxdif_t2 = [round(x,1) for x in maxdif_t]
    for x in range(len(misc_outputs_t)):
        f.write('\n{} \t- maxdif = {}'.format(misc_outputs_t[x], maxdif_t2[x]))
f.write('\n\n-----------------------------------------------------------------------------------------')

f.close()

cwd = os.getcwd()

torch.save({'model_state_dict': net.state_dict()
            }, r'{}\\final_weights.pth'.format(cwd))

torch.save({'model_state_dict': model_initial.state_dict()
            }, r'{}\\initial_weights.pth'.format(cwd))

path  = '{}\\loss_curve.png'.format(cwd)
val_plot.savefig(path, dpi=fig_dpi)
path  = '{}\\accuracy_curve.png'.format(cwd)
acc_plot.savefig(path, dpi=fig_dpi)
path  = '{}\\confusion_matrix.png'.format(cwd)
conf_fig.savefig(path, dpi=fig_dpi)
path  = '{}\\confusion_matrix_normalize_row.png'.format(cwd)
conf_fig_2.savefig(path, dpi=fig_dpi)
path  = '{}\\confusion_matrix_normalize_full.png'.format(cwd)
conf_fig_3.savefig(path, dpi=fig_dpi)

if PCA_flag:
    path  = '{}\\PCA_loss_landscape_1.png'.format(cwd)
    PCA_fig_1.savefig(path, dpi=fig_dpi)
    path  = '{}\\PCA_loss_landscape_2.png'.format(cwd)
    PCA_fig_2.savefig(path, dpi=fig_dpi)
    path  = '{}\\PCA_loss_contour.png'.format(cwd)
    PCA_fig_3.savefig(path, dpi=fig_dpi)

if loss_landscape_flag:
    path  = '{}\\loss_surface_contour.png'.format(cwd)
    loss_con.savefig(path, dpi=fig_dpi)
    path  = '{}\\loss_surface_angle_1.png'.format(cwd)
    loss_surf_1.savefig(path, dpi=fig_dpi)
    path  = '{}\\loss_surface_angle_2.png'.format(cwd)
    loss_surf_2.savefig(path, dpi=fig_dpi)

if MAV_flag:
    path = '{}\\MAV_distance_dist.png'.format(cwd)
    MAV_dist_plot.savefig(path, dpi=fig_dpi)

if misclassified_images_flag:
    mcfname = 'misclassified images'
    os.mkdir(mcfname)
    mcfdir = '{}\\{}'.format(fname, mcfname)
    os.chdir(mcfdir)

    for i in range(len(misc_preds_t)):

        imagets0 = misc_img_t[i][0]
        imagets = imagets0.cpu()
        predts = int(misc_preds_t[i])
        predname = names[predts]
        labts = int(misc_labels_t[i])
        labname = names[labts]
        fstag = '{} - label {}, classified as {}'.format(i, labname, predname)

        sf = plt.figure(figsize=(6,6))
        plt.imshow(imagets, cmap='jet')
        path = '{}\\{}.png'.format(mcfdir, fstag)
        sf.savefig(path)

if baseline_flag:
    path = '{}\\baseline_uncertainty_plot.png'.format(cwd)
    baseline_plot.savefig(path, dpi=fig_dpi)

    sio.savemat('unc_correct.mat', {'data':unc_correct})
    sio.savemat('unc_wrong.mat', {'data':unc_wrong})

os.chdir(DT)
#os.startfile(fname)


# coding:utf-8

from scrapy.selector import HtmlXPathSelector
from scrapy.contrib.spiders import CrawlSpider
from scrapy.http.request import Request
from youtube.items import youtubeItem

class youtubeK18Spider(CrawlSpider):
    name = 'youtube_k18'
    allowed_domains = ['www.youtube.com']
    start_urls = ['http://www.youtube.com/education?hl=en']

    def parse(self, response):
        filter_cate = ['Primary & Secondary Education']
        hxs = HtmlXPathSelector(response)
        sites = hxs.select('//div[@class=\'edu-mixed-collection\']')
        for site in sites:
            item = youtubeItem()
            category = site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/text()').extract()[0]
            category = category.replace(u'»',' ')
            category = category.strip();
            is_target = category in filter_cate
            if is_target == True:
                item['category1'] = category
                url = 'http://www.youtube.com/education' + site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/@href').extract()[0]
                request = Request(url, meta={'item': item}, callback=self.parse_category2)
                yield request

    def parse_category2(self, response):
        hxs = HtmlXPathSelector(response)
        sites = hxs.select('//div[@class=\'yt-uix-slider browse-collection yt-rounded\']')
        for site in sites:
            item = response.meta['item']
            has_data = len(site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/text()'))
            if has_data == 0:
                continue
            category = site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/text()').extract()[0]
            category = category.replace(u'»',' ')
            category = category.strip();
            url = site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/@href').extract()[0]
            if url != '':
                item['category2'] = category
                url = 'http://www.youtube.com/education' + url
                request = Request(url, meta={'item': item}, callback=self.parse_category3)
                yield request

    def parse_category3(self, response):
        hxs = HtmlXPathSelector(response)
        sites = hxs.select('//div[@class=\'yt-uix-slider browse-collection yt-rounded\']')
        for site in sites:
            item = response.meta['item']
            has_data = len(site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/text()'))
            if has_data == 0:
                continue
            category = site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/text()').extract()[0]
            category = category.replace(u'»',' ')
            category = category.strip();
            url = site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/@href').extract()[0]
            if url != '':
                item['category3'] = category
                url = 'http://www.youtube.com/education' + url
                request = Request(url, meta={'item': item}, callback=self.parse_category4)
                yield request

    def parse_category4(self, response):
        hxs = HtmlXPathSelector(response)
        site_first =  hxs.select('//button[@class=\'yt-uix-button-toggled yt-uix-button\']')
        sites = hxs.select('//button[@class=\' yt-uix-button\']')
        sites.append(site_first)
        filter_urls = []
        for site in sites:
            item = response.meta['item']
            has_data = len(site.select('./@href').extract())
            if has_data == 0:
                continue
            url = site.select('./@href').extract()[0]
            url = 'http://www.youtube.com/education' + url
            is_filter_target = url in filter_urls
            if is_filter_target == True:
                continue
            filter_urls.append(url)
            request = Request(url, meta={'item': item}, callback=self.parse_detail1)
            yield request

    def parse_detail1(self, response):
        hxs = HtmlXPathSelector(response)
        sites = hxs.select('/html/body/div/div[2]/div/div[2]/div[3]/div/div/ul/li')
        for site in sites:
            item = response.meta['item']
            url = site.select('div/div/div/h3/a/@href').extract()[0]
            url = 'http://www.youtube.com' + url
            request = Request(url, meta={'item': item}, callback=self.parse_detail2)
            yield request

    def parse_detail2(self, response):
        hxs = HtmlXPathSelector(response)
        sites = hxs.select('/html/body/div/div[2]/div/div/div[3]/div/ol/li')
        for site in sites:
            item = response.meta['item']
            item['img'] = site.select('div/a/span/span/img/@src').extract()
            item['detail_title1'] = site.select('h3/a/text()').extract()
            item['detail_link'] = 'http://www.youtube.com' + site.select('h3/a/@href').extract()[0]
            yield item

        sites = hxs.select('/html/body/div/div[2]/div/div/div[3]/div/div[2]/div[2]/div/ul/li')
        for site in sites:
            item = response.meta['item']
            item['img'] = site.select('a/span/span/span/img/@src').extract()
            item['detail_title1'] = site.select('a/span[2]/text()').extract()
            has_data = len(site.select('a/@href').extract())
            #TODO load more
            if has_data == 0:
                continue
            item['detail_link'] = 'http://www.youtube.com' + site.select('a/@href').extract()[0]
            yield item


# -*- coding: utf-8 -*-
"""
Created on Mon Apr  2 00:58:47 2018

@author: IBM
"""

#import bokeh
from bokeh.plotting import figure, output_file, show
from bokeh.layouts import gridplot
from bokeh.models import Range1d, LinearAxis, PrintfTickFormatter, ColumnDataSource
#from bokeh.models.widgets import Slider

import pandas as pd
import numpy as np
from datetime import datetime, timedelta

from ExecQuery import ExecQuery
from config import RunGUI_config
data_config = RunGUI_config['data']
plot_config = RunGUI_config['plot']
#ticker = 'JNJ'

def main():
    PlotAll(plot_config, data_config, 'JNJ')

#########################################
### The Function
#########################################
def PlotAll(plot_config, data_config, ticker):
    output_dir = plot_config['output_dir']

    px_series_df = QueryGUIData(data_config, ticker, 'px_series')
    #px_series_df_curr = px_series_df[px_series_df['date']==px_series_df['date'].max()]
    px_dist_crude_df = QueryGUIData(data_config, ticker, 'px_dist_crude')
    px_dist_manual_df = QueryGUIData(data_config, ticker, 'px_dist_manual')

    #1. the first plot - price and volume time series
    px_series_fig = PlotPriceVolume(plot_config, px_series_df)
    #show(px_series_fig)

    #2. price and prob based on the crude model
    px_dist_crude_fig = PlotPdf(plot_config, px_dist_crude_df, 'crude')
    #show(px_dist_crude_fig)

    #3. price and prob based on the manual model
    px_dist_manual_fig = PlotPdf(plot_config, px_dist_manual_df, 'manual')
    #show(px_dist_manual_fig)

    p = gridplot([[px_series_fig, None]
                , [px_dist_crude_fig, px_dist_manual_fig]])
    output_file(output_dir+"px_dist.html", title="px_dist")
    show(p)

#########################################
### The Secondary Plot Functions
#########################################
def PlotPriceVolume(plot_config, px_series_results):
    px_series_df = px_series_results['df']
    volume_unit_str = px_series_results['volume_unit_str']
    px_series_source = ColumnDataSource(data=dict(price=px_series_df['price'], volume=px_series_df['volume'], date=px_series_df['date']))
    #type(px_series_source.data['volume'])

    plot_options = dict(width=int(2.5*plot_config['width']), plot_height=plot_config['height'], tools=plot_config['tools'])
    px_series_fig = figure(**plot_options, x_axis_type="datetime", title = "Price Volume Time Series")
    px_series_fig.y_range = Range1d(start=0.95*px_series_df['price'].min(), end=1.05*px_series_df['price'].max())
    px_series_fig.extra_y_ranges = {"volume": Range1d(start=0, end=1.2*px_series_df['volume'].max())}
    print(px_series_df['volume'], px_series_df['volume'].max())
    px_series_fig.add_layout(LinearAxis(y_range_name='volume'), 'right')
    px_series_fig.yaxis[1].formatter = PrintfTickFormatter(format="%5.1f " + volume_unit_str)

    px_series_fig.line('date', 'price', source=px_series_source, line_width=plot_config['line_width'], color='olivedrab', alpha=0.9)
    px_series_fig.vbar(bottom=0, top='volume', x='date', source=px_series_source, width=1.5, color='grey', alpha=0.5, y_range_name='volume')

    return px_series_fig

def PlotPdf(plot_config, px_dist_df, dist_type='Crude'):
    plot_options = dict(width=plot_config['width'], plot_height=plot_config['height'], tools=plot_config['tools'])

    px_dist_fig = figure(**plot_options, title = "Px Dist Plot - " + dist_type)
    px_dist_fig.line(px_dist_df['price'], px_dist_df['prob'], line_width=plot_config['line_width'], color='steelblue', alpha=0.9)

    #show(px_dist_fig)
    return px_dist_fig

#########################################
### The Secondary Query Functions
#########################################
def QueryGUIData(data_config, ticker, data_type, normalize=False):
    db = data_config['database']

    if data_type == 'px_series':
        px_series_table = data_config['px_series_table']
        px_series_query = """select date as date_raw, adj_close as price, volume as volume_raw
                             from """+px_series_table+"""
                             where ticker = '"""+ticker+"""'
                             order by ticker, date;"""
        print(px_series_query)
        px_series_df = ExecQuery(db, px_series_query, True, True)

        px_series_df['date'] = pd.to_datetime(px_series_df['date_raw'], format='%Y-%m-%d') + timedelta(hours=16)

        volume_unit = 1e6 if px_series_df['volume_raw'].max()>1e6 else 1e3
        px_series_df['volume'] = px_series_df['volume_raw']/volume_unit
        volume_unit_str = 'M' if volume_unit==1e6 else 'K'

        return {'df':px_series_df[['date','price','volume']], 'volume_unit_str':volume_unit_str}

    elif data_type == 'px_dist_crude':
        px_dist_table = data_config['px_dist_table']

        max_dt_subquery = """(
                                select ticker, price, max(asof_datetime) as max_dt
                                from """+px_dist_table+"""
                                where ticker='"""+ticker+"""'
                                and model='crude'
                                and not deprecated
                                group by 1,2
                             )"""
        px_dist_query = """ select p.price, p.prob
                            from """+px_dist_table+""" p
                            join """+max_dt_subquery+""" s
                            on p.price=s.price and p.asof_datetime=s.max_dt #and p.ticker=s.ticker
                            where p.ticker='"""+ticker+"""'
                            and p.model='crude'
                            and not deprecated
                            order by price
                         ;"""
        print(px_dist_query)
        px_dist_crude_df = ExecQuery(db, px_dist_query, True, True)
        if normalize:
            px_dist_crude_df['prob'] = px_dist_crude_df['prob']/px_dist_crude_df['prob'].sum()
        return px_dist_crude_df

    elif data_type == 'px_dist_manual':
        px_dist_table = data_config['px_dist_table']

        max_dt_subquery = """(
                                select ticker, price, max(asof_datetime) as max_dt
                                from """+px_dist_table+"""
                                where ticker='"""+ticker+"""'
                                and model='manual'
                                and not deprecated
                                group by 1,2
                             )"""
        px_dist_query = """ select p.price, p.prob
                            from """+px_dist_table+""" p
                            join """+max_dt_subquery+""" s
                            on p.price=s.price and p.asof_datetime=s.max_dt #and p.ticker=s.ticker
                            where p.ticker='"""+ticker+"""'
                            and p.model='manual'
                            and not deprecated
                            order by price
                         ;"""
        print(px_dist_query)
        px_dist_manual_df = ExecQuery(db, px_dist_query, True, True)
        if normalize:
            px_dist_manual_df['prob'] = px_dist_manual_df['prob']/px_dist_manual_df['prob'].sum()
        return px_dist_manual_df

def InsertPXDistManual(data_config, ticker, raw_str):
    """
        the input should look like "123.12 0.1;123.23 0.15"
        - if duplicated, will take the latter one
        - if not able to convert to float, will skip
    """
    #raw_str = "123.12 0.1;123.23 0.15;123.12 0.3;123df 0.2" ###DEBUG
    #ticker = 'JNJ' ###DEBUG

    # 1. process the string
    # 1.1 if the string is "clear all", mark everything for this ticker to be deprecated
    if raw_str == 'clear all':
        deprecate_query = 'update '+data_config['px_dist_table']+" set deprecated=True where ticker = '"+ticker+"' and model = 'manual'; commit;"
        print(deprecate_query)
        ExecQuery(data_config['database'], deprecate_query, False, False)

    # 1.2 subtract numbers, and proceed if there is any ligit numbers
    px_dist_dict = {x.split(' ')[0]:x.split(' ')[1] for x in raw_str.split(';') if CanToFloat(x.split(' ')[0]) and CanToFloat(x.split(' ')[1])}
    if len(px_dist_dict)==0: return None

    # 2. construct the query
    query_list = []
    query_prepend = ', '.join(['current_timestamp() as asof_datetime',"'manual' as model", "'"+ticker+"' as ticker"])

    for i in px_dist_dict.keys():
        query_list.append('(select ' +str(i)+' as price, '+str(px_dist_dict[i])+' as prob)')

    select_query = 'select '+query_prepend+', temp.*, false as deprecated from (\n'+ '\nUNION ALL\n'.join(query_list) +'\n) as temp'
    insert_query = 'insert into '+data_config['px_dist_table']+' \n'+ select_query + '; commit;'

    print(insert_query)
    ExecQuery(data_config['database'], insert_query, False, False)

def InsertPxDistCrude(data_config, ticker):

    deprecate_query = 'update '+data_config['px_dist_table']+" set deprecated=True where ticker = '"+ticker+"' and model = 'crude'; commit;"
    print(deprecate_query)
    ExecQuery(data_config['database'], deprecate_query, False, False)

    query_list = []
    px_dist_dict = {'0':'0.2',
                    '0.1':'0.1','-0.1':'0.1',
                    '0.2':'0.08','-0.2':'0.08',
                    '0.3':'0.06','-0.3':'0.06',
                    '0.5':'0.05','-0.5':'0.05',
                    '1':'0.04','-1':'0.04',
                    '2':'0.03','-2':'0.03'
                   }
    for i in px_dist_dict.keys():
        query_list.append('(select avg(adj_close)+('+i+'*stddev(adj_close)) as price, '+px_dist_dict[i]+' as prob from '+data_config['px_series_table']+" where ticker='"+ticker+"')")

    insert_query = """insert into """ +data_config['px_dist_table']+ """
                    select current_timestamp() as asof_datetime, 'crude' as model, '"""+ticker+"""' as ticker, temp.*, false as deprecated
                    from
                    (\n""" + '\nUNION ALL\n'.join(query_list) + """\n) as temp
                    ; commit;
                   """
    print(insert_query)
    ExecQuery(data_config['database'], insert_query, False, False)

def InsertPxDistConviction(data_config, ticker, model, conviction):
    # conviction should be a string that can be translated into float
    insert_query = """insert into """ +data_config['px_dist_conviction_table']+ """
                      select current_timestamp() as asof_datetime,
                           '"""+model+"""' as model,
                           '"""+ticker+"""' as ticker,
                           """+conviction+"""as conviction,
                           false as deprecated; commit;
                   """
    print(insert_query)
    ExecQuery(data_config['database'], insert_query, False, False)

#########################################
### The Tertiary Helper Functions
#########################################
def StrToFloat(str):
    try:
        f = float(str)
    except ValueError:
        f = None
    return f

def CanToFloat(str):
    try:
        float(str)
        return True
    except ValueError:
        return False

#########################################
### Execute
#########################################
if __name__ == '__main__':
    main()


import json
from channels import Group
from channels.sessions import channel_session
from channels import Channel
from channels.auth import channel_session_user_from_http, channel_session_user

from .models import Room , Status
@channel_session
@channel_session_user_from_http
def ws_connect(message):
    message.reply_channel.send({"accept": True})
    prefix, label = message['path'].strip('/').split('/')
    room = Room.objects.get(label=label)
    #'get' will create a Entry in Status for a particular user in
    #a particular table
    get=Status.objects.create(room=label,uname=message.user)
    Group(label).add(message.reply_channel)
    message.channel_session['room'] = room.label
    #Stat will store the user's status id
    message.channel_session['stat']= get.id
    ws_status(message, get.id,"online")

@channel_session
def ws_receive(message):
    label = message.channel_session['room']
    room = Room.objects.get(label=label)
    data = json.loads(message['text'])
    #'msg' will store the data in to the db and send it
    #back to socket to display
    msg = room.messages.create(**data)
    Group(label).send({'text': json.dumps(msg.as_dict())})
@channel_session
def ws_status(message, id, stat):
    label = message.channel_session['room']
    u_id=Status.objects.get(id=id)
    Group(label).send({"text": json.dumps({
            "state": stat,
            "u_name": u_id.uname,
        }),  })
#ws_status will send online/Offline Status to the channel !!

@channel_session
def ws_disconnect(message):
    label = message.channel_session['room']
    id=message.channel_session['stat']
    ws_status(message,id,"offline")
    u_id=Status.objects.get(id=id)
    Status.objects.filter(uname=u_id.uname).delete()
    #It will clear the online status of the user
    Group(label).discard(message.reply_channel)


import math
import cv2
import numpy as np

#Computer the final output values,
#angle 1 is the Yaw to the target
#distance is the distance to the target
#angle 2 is the Yaw of the Robot to the target

def compute_output_values(rvec, tvec):
    '''Compute the necessary output distance and angles'''

    # The tilt angle only affects the distance and angle1 calcs
    # This is a major impact on calculations
    tilt_angle = math.radians(0)

    x = tvec[0][0]
    z = math.sin(tilt_angle) * tvec[1][0] + math.cos(tilt_angle) * tvec[2][0]

    # distance in the horizontal plane between camera and target
    distance = math.sqrt(x**2 + z**2)

    # horizontal angle between camera center line and target
    angleInRad = math.atan2(x, z)
    angle1 = math.degrees(angleInRad)

    rot, _ = cv2.Rodrigues(rvec)
    rot_inv = rot.transpose()
    pzero_world = np.matmul(rot_inv, -tvec)
    angle2InRad = math.atan2(pzero_world[0][0], pzero_world[2][0])
    angle2 = math.degrees(angle2InRad)

    return distance, angle1, angle2


import random, board, itertools, math
from functools import lru_cache

"""
Custom exceptions
"""
class FullBoard(Exception):
  pass

class InvalidMoveDirection(Exception):
  pass

class InvalidSpawnLocation(Exception):
  pass

"""
Table initialization
"""

move_left_table = {}
move_right_table = {}
score_table = {}
heuristic_table = {}
loss_penalty = 0

heuristic_weights = (3.5, 11, 4, 47, 700, 270, 200000)

# special-case exponentiation where 2^0 = 0
def exp(n):
  return 0 if n == 0 else 2 ** n

# converts a vector representing the numbers in a row to the corresponding
# 2 bytes representing that row
def row_to_bits(row):
  res = 0
  for n in row:
    # Shift over by a tile
    res <<= 4
    # Or in the value for the current tile
    res |= int(math.log(n, 2) if n != 0 else 0)
  return res

"""
Create move tables
"""
def init_moves():
  print("Generating bitboard move tables...")
  for vector in itertools.product(range(16), repeat = 4):
    # convert the vector into the corresponding 2 bytes representing that row
    row = 0
    for n in vector:
      row <<= 4
      row |= n

    """ Move table creation """
    # define the tables of the row – merging operates on non-logs, so have
    # to exponentiate
    new_vec_left, score_table[row] = board.merge([exp(n) for n in vector])
    new_vec_right = list(reversed(board.merge([exp(n) for n in reversed(vector)])[0]))

    move_left_table[row] = row_to_bits(new_vec_left)
    move_right_table[row] = row_to_bits(new_vec_right)

# initiate heuristic table
# takes a list of heuristic weights
# weights has definitions of:
#   idx 0: power scale for sum penalty
#   idx 1: weight for sum penalty
#   idx 2: power scale for monotonicty
#   idx 3: weight for monotonicty
#   idx 4: weight for smoothness
#   idx 5: weight for empty tiles
#   idx 6: loss penalty
def init_heuristics(weights = heuristic_weights):
  print("Generating bitboard heuristics...")
  supower, suweight, mpower, mweight, sweight, eweight, loss_penalty = weights
  for vector in itertools.product(range(16), repeat = 4):

    """
    Heuristic table creation
    Four heuristics: want high-value tiles, smoothness, empty tiles, and monotonicity
    """

    row = 0
    for n in vector:
      row <<= 4
      row |= n

    sum = smoothness = mono_left = mono_right = empty = 0

    # score tile values and emptiness
    for n in vector:
      sum += n ** supower
      if n == 0: empty += 1

    # score smoothness
    counter = 0
    for i in range(3):
      # might want to do difference of exponents, but probably not
      if vector[i] == vector[i+1]:
        counter += 1
      elif counter > 0:
        sum += 1 + counter
        counter = 0

    if counter > 0: sum += 1 + counter

    # score monotonicity
    for i in range(3):
      if vector[i] > vector[i+1]:
        mono_left += vector[i] ** mpower - vector[i+1] ** mpower
      else:
        mono_right += vector[i+1] ** mpower - vector[i] ** mpower

    heuristic_table[row] = loss_penalty + \
                           empty * eweight - \
                           smoothness * sweight - \
                           min(mono_left, mono_right) * mweight - \
                           sum * suweight

def init(weights=heuristic_weights):
  init_moves()
  init_heuristics(weights)

@lru_cache(maxsize = 1000000)
def heuristic_value(board):
  total_value = 0
  # do rows
  for offset in (0, 16, 32, 48):
    total_value += heuristic_table[(board >> offset) & 0xFFFF]
  # do columns
  for n in range(4):
    cur_col = 0
    for offset in (0, 16, 32, 48):
      # shift to the current tile
      cur_col <<= 4
      # or in the value of the tile at col index n
      cur_col |= (board >> (offset + 4 * n)) & 0xF
    total_value += heuristic_table[cur_col]
  return total_value

# boards are represented as an integer of up to 64 bits
# every four bits is the log of the value of a cell - first four bits are the
# first cell, first 16 bits are the first row, etc. a 0 represents an emtpy tile

def new_board(): return 0

"""
Check if two boards are equal
"""
def equal(b1, b2):
  return b1 == b2

"""
Return a list of ints corresponding to the points where empty tiles begin.
Works right to left, i.e. 0 means the bottom right tile is empty, 1 means
the second-from right tile in the bottom row is empty, etc.
"""
def empty_tiles(board):
  empty_pos = []
  for cur_pos in range(16):
    if not ((board >> (cur_pos * 4)) & 0xF): empty_pos.append(cur_pos)
  return empty_pos

"""
Returns board with a tile randomly spawned.
"""
def spawn_tile(board):
  try:
    spawn_pos = random.choice(empty_tiles(board))
    # log values, so 2 is 4 and 1 is 2
    tile = 2 if random.random() < 0.1 else 1
    return board | (tile << 4 * spawn_pos)
  except IndexError:
    raise FullBoard("Can't spawn tiles on a full board")

def spawn_manual(board, value, loc):
  real_value = int(math.log(value, 2))
  return board | (value << (4 * loc))

"""
Returns board moved in the given direction - L, R, U, or D, along with the
increase in score for moving in that direction
"""
def move(board, dir):
  res = 0
  score_inc = 0
  if dir == "L":
    for offset in (0, 16, 32, 48):
      cur_row = (board >> offset) & 0xFFFF
      res |= move_left_table[cur_row] << offset
      score_inc += score_table[cur_row]

  elif dir == "R":
    for offset in (0, 16, 32, 48):
      cur_row = (board >> offset) & 0xFFFF
      res |= move_right_table[cur_row] << offset
      score_inc += score_table[cur_row]

  elif dir == "U":

    for n in range(4):
      cur_col = 0
      for offset in (0, 16, 32, 48):
        # shift to the current tile
        cur_col <<= 4
        # or in the value of the tile at col index n
        cur_col |= (board >> (offset + 4 * n)) & 0xF
      # look up the replacement for the current column
      replacement_col = move_right_table[cur_col]
      # increment the score
      score_inc += score_table[cur_col]
      # or the replacement column into res
      for offset in (0, 16, 32, 48):
        # 48 - offset + adjust because the col has to put back in reverse
        res |= ((replacement_col >> (int(offset / 4))) & 0xF) << (48 - offset + 4 * n)

  elif dir == "D":
    for n in range(4):
      cur_col = 0
      for offset in (0, 16, 32, 48):
        # shift to the current tile
        cur_col <<= 4
        # or in the value of the tile at col index n
        cur_col |= (board >> (offset + 4 * n)) & 0xF
      # look up the replacement for the current column
      replacement_col = move_left_table[cur_col]
      # increment the score
      score_inc += score_table[cur_col]
      # or the replacement column into res
      for offset in (0, 16, 32, 48):
        # 48 - offset + adjust because the col has to put back in reverse
        res |= ((replacement_col >> (int(offset / 4))) & 0xF) << (48 - offset + 4 * n)

  else:
    # invalid direction given, raise an exception with the given direction
    raise InvalidMoveDirection(dir)

  return res, score_inc

"""
Returns True if no further moves are possible and False otherwise
"""
def game_over(board):
  # easiest but probably not most efficient solution is to just
  # move in every direction and check for changes

  # for speed, check if the board is empty first: quicker than moving
  # in all four directions
  if len(empty_tiles(board)) > 0: return False
  for dir in ("L", "R", "U", "D"):
    if (move(board, dir))[0] != board: return False
  # board is full and no moves are possible – game is over
  return True

"""
Returns the maximum tile present on a board
"""
def max_tile(board):
  max = 0
  while board:
    if board & 0xF > max: max = board & 0xF
    board >>= 4
  return (0 if max == 0 else 2**max)

"""
Converts the board state to string format
"""
def string_of_board(board):
  # for now, give five digits for each tile
  row_div = "-" + "-" * 6 * (4) + "\n"
  state_string = row_div
  for offset in (48, 32, 16, 0):
    state_string += "|"
    for i in (12, 8, 4, 0):
      cur = (board >> offset + i) & (0xF)
      rep = "" if cur == 0 else str(2**cur)
      state_string += (rep.rjust(5) + "|")
    state_string += "\n" + row_div
  # trim trailing newline
  return state_string[:-1]


import pytest
import pdb

from TradeJournal.trade import Trade

@pytest.fixture
def t_object():
    '''Returns a Trade object'''

    td=Trade(
        start="2017-04-20T14:00:00",
        end="2017-04-26T14:00:00",
        pair="AUD/USD",
        type="long",
        timeframe="H8",
        strat="counter_b2",
        id= "AUD_USD 20APR2017H8"
        )
    return td

@pytest.fixture
def unfisished_t_object():
    ''' Returns a Trade object without the end defined'''

    td = Trade(
        start="2018-12-12 10:00:00",
        entry=1.57488,
        SL=1.56685,
        TP=1.58724,
        pair="EUR/AUD",
        type="long",
        timeframe="H12",
        strat="counter_b2",
        id="EUR_AUD 12DEC2018H12"
    )
    return td

def test_fetch_candlelist(t_object):
    '''
    This test checks the function to return a CandleList object
    corresponding to this trade
    '''

    cl=t_object.fetch_candlelist()
    assert cl.clist[0].openBid==0.7521
    assert cl.clist[0].highBid==0.75464

def test_run_trade(unfisished_t_object):
    '''
    This test checks the progression of the Trade
    and checks if the outcome attribute is correctly
    defined.
    '''

    unfisished_t_object.run_trade()
    assert unfisished_t_object.outcome=='success'


import os, sys, time
import numpy as np
import magent
from magent.model import BaseModel
from magent.builtin.tf_model import DeepQNetwork
from keras import backend as K
from keras.optimizers import Adam
import tensorflow as tf
import random
from ReplayBuffer_v2 import ReplayBuffer
from keras.layers import Dense, Dropout, Conv2D, Input, Lambda, Flatten, TimeDistributed, merge
from keras.layers import Add, Reshape, MaxPooling2D, Concatenate, Embedding, RepeatVector
from keras.models import Model
from keras.layers.core import Activation
from keras.utils import np_utils,to_categorical
from keras.engine.topology import Layer
from keras.callbacks import TensorBoard
from magent.builtin.tf_model import DeepQNetwork

np.random.seed(16)

def Adjacency(state):
    adj = []
    dis = []
    for j in range(20):
        dis.append([state[j][-2],state[j][-1],j])
    for j in range(20):
        f = []
        for r in range(len(dis)):
            f.append([(dis[r][0]-dis[j][0])**2+(dis[r][1]-dis[j][1])**2,r])
        f.sort(key=lambda x:x[0])
        y = []
        for r in range(4):
            y.append(f[r][1])
        y = to_categorical(y,num_classes=20)
        adj.append(y)
    return adj

def observation(state1,state2):
    state = []
    for j in range(20):
        state.append(np.hstack(((state1[j][0:11,0:11,1]-state1[j][0:11,0:11,5]).flatten(),state2[j][-1:-3:-1])))
    return state

def MLP():
    In_0 = Input(shape=[123])

    h = Dense(512, activation='relu',kernel_initializer='random_normal')(In_0)
    h = Dense(128, activation='relu',kernel_initializer='random_normal')(h)

    h = Reshape((1,128))(h)

    model = Model(input=In_0,output=h)
    return model

def MultiHeadsAttModel(l=2, d=128, dv=16, dout=128, nv = 8 ):

    v1 = Input(shape = (l, d))
    q1 = Input(shape = (l, d))
    k1 = Input(shape = (l, d))
    ve = Input(shape = (1, l))

    v2 = Dense(dv*nv, activation = "relu",kernel_initializer='random_normal')(v1)
    q2 = Dense(dv*nv, activation = "relu",kernel_initializer='random_normal')(q1)
    k2 = Dense(dv*nv, activation = "relu",kernel_initializer='random_normal')(k1)

    v = Reshape((l, nv, dv))(v2)
    q = Reshape((l, nv, dv))(q2)
    k = Reshape((l, nv, dv))(k2)
    v = Lambda(lambda x: K.permute_dimensions(x, (0,2,1,3)))(v)
    k = Lambda(lambda x: K.permute_dimensions(x, (0,2,1,3)))(k)
    q = Lambda(lambda x: K.permute_dimensions(x, (0,2,1,3)))(q)

    att = Lambda(lambda x: K.batch_dot(x[0],x[1] ,axes=[3,3]) / np.sqrt(dv))([q,k])# l, nv, nv
    att = Lambda(lambda x: K.softmax(x))(att)
    out = Lambda(lambda x: K.batch_dot(x[0], x[1],axes=[3,2]))([att, v])
    out = Lambda(lambda x: K.permute_dimensions(x, (0,2,1,3)))(out)

    out = Reshape((l, dv*nv))(out)

    T = Lambda(lambda x: K.batch_dot(x[0],x[1]))([ve,out])

    out = Dense(dout, activation = "relu",kernel_initializer='random_normal')(T)
    model = Model(inputs=[q1,k1,v1,ve], outputs=out)
    return model

def Q_Net(action_dim):
    I1 = Input(shape = (1, 128))
    I2 = Input(shape = (1, 128))
    I3 = Input(shape = (1, 128))

    h1 = Flatten()(I1)
    h2 = Flatten()(I2)
    h3 = Flatten()(I3)

    h = merge([h1,h2,h3],mode='concat')
    V = Dense(action_dim,kernel_initializer='random_normal')(h)

    model = Model(input=[I1,I2,I3],output=V)
    return model

path="data/battle_model"
map_size=100
capacity = 200000
batch_size = 256
totalTime = 0
TAU = 0.01
LRA = 0.0001
param = None
alpha = 0.6
GAMMA = 0.96
n_episode = 100000
max_steps = 300
episode_before_train = 200
n_agent=20
magent.utility.init_logger("battle")
env = magent.GridWorld("battle", map_size=30)
env.set_render_dir("build/render")
handles = env.get_handles()
sess = tf.Session()
K.set_session(sess)
n = len(handles)
n_actions=env.get_action_space(handles[0])[0]
i_episode=0
buff=ReplayBuffer(capacity)
l=40

print(env.get_action_space(handles[0])[0])
print(env.get_action_space(handles[1])[0])
#f = open('log.txt','w')

######build the model#########
cnn = MLP()
m1 = MultiHeadsAttModel(l=4)
m2 = MultiHeadsAttModel(l=4)
q_net = Q_Net(action_dim = 9)
vec = np.zeros((1,4))
vec[0][0] = 1

In= []
for j in range(n_agent):
    In.append(Input(shape=[123]))
    In.append(Input(shape=(4,20)))
In.append(Input(shape=(1,4)))
feature = []
for j in range(n_agent):
    feature.append(cnn(In[j*2]))

feature_ = merge(feature,mode='concat',concat_axis=1)

relation1 = []
for j in range(n_agent):
    T = Lambda(lambda x: K.batch_dot(x[0],x[1]))([In[j*2+1],feature_])
    relation1.append(m1([T,T,T,In[40]]))

relation1_ = merge(relation1,mode='concat',concat_axis=1)

relation2 = []
for j in range(n_agent):
    T = Lambda(lambda x: K.batch_dot(x[0],x[1]))([In[j*2+1],relation1_])
    relation2.append(m2([T,T,T,In[40]]))

V = []
for j in range(n_agent):
    V.append(q_net([feature[j],relation1[j],relation2[j]]))

model = Model(input=In,output=V)
model.compile(optimizer=Adam(lr = 0.00003), loss='mse')
model.summary()

######build the target model#########
cnn_t = MLP()
m1_t = MultiHeadsAttModel(l=4)
m2_t = MultiHeadsAttModel(l=4)
q_net_t = Q_Net(action_dim = 9)
In_t= []
for j in range(n_agent):
    In_t.append(Input(shape=[123]))
    In_t.append(Input(shape=(4,20)))
In_t.append(Input(shape=(1,4)))

feature_t = []
for j in range(n_agent):
    feature_t.append(cnn_t(In_t[j*2]))

feature_t_ = merge(feature_t,mode='concat',concat_axis=1)

relation1_t = []
for j in range(n_agent):
    T = Lambda(lambda x: K.batch_dot(x[0],x[1]))([In_t[j*2+1],feature_t_])
    relation1_t.append(m1_t([T,T,T,In_t[40]]))

relation1_t_ = merge(relation1_t,mode='concat',concat_axis=1)

relation2_t = []
for j in range(n_agent):
    T = Lambda(lambda x: K.batch_dot(x[0],x[1]))([In_t[j*2+1],relation1_t_])
    relation2_t.append(m2_t([T,T,T,In_t[40]]))

V_t = []
for j in range(n_agent):
    V_t.append(q_net_t([feature_t[j],relation1_t[j],relation2_t[j]]))

model_t = Model(input=In_t,output=V_t)

tf_model = DeepQNetwork(env, handles[1], 'trusty-battle-game-l', use_conv=True)
tf_model.load("data/battle_model", 0, 'trusty-battle-game-l')
###########playing#############
while i_episode<n_episode:
    alpha*=0.996
    if alpha<0.01:
        alpha=0.01
    print(i_episode)
    i_episode=i_episode+1
    env.reset()
    #env.add_walls(method="random", n=map_size * map_size * 0.03)
    env.add_agents(handles[0], method="random", n=20)
    env.add_agents(handles[1], method="random", n=12)
    step_ct = 0
    done = False
    n = len(handles)
    obs  = [[] for _ in range(n)]
    ids  = [[] for _ in range(n)]
    action = [[] for _ in range(n)]
    nums = [env.get_num(handle) for handle in handles]
    steps = 0
    score = 0
    loss = 0
    dead = [0,0]

    while steps<max_steps:
        steps+=1
        i=0
        obs[i] = env.get_observation(handles[i])
        adj = Adjacency(obs[i][1])
        flat_ob = observation(obs[i][0],obs[i][1])
        ob=[]
        for j in range(n_agent):
            ob.append(np.asarray([flat_ob[j]]))
            ob.append(np.asarray([adj[j]]))
        ob.append(np.asarray([vec]))
        acts = model.predict(ob)
        action[i]=np.zeros(n_agent,dtype = np.int32)
        for j in range(n_agent):
            if np.random.rand()<alpha:
                action[i][j]=random.randrange(n_actions)
            else:
                action[i][j]=np.argmax(acts[j])
        env.set_action(handles[i], action[i])

        obs[1] = env.get_observation(handles[1])
        ids[1] = env.get_agent_id(handles[1])
        acts = tf_model.infer_action(obs[1], ids[1], 'e_greedy')
        env.set_action(handles[1], acts)
        done = env.step()

        next_obs = env.get_observation(handles[0])
        flat_next_obs = observation(next_obs[0],next_obs[1])
        rewards = env.get_reward(handles[0])
        score += sum(rewards)
        if steps%3 ==0:
            buff.add(flat_ob, action[0], flat_next_obs, rewards, done, adj)

        if (i_episode-1) % 10 ==0:
            env.render()
        if max_steps == steps:
            print(dead[0],end='\t')
            print(dead[1],end='\t')
            print(score/300,end='\t')
            #f.write(str(dead[i])+'\t'+str(score[i]/300)+'\t')
            #f.write(str(loss/100)+'\n')
            print(loss/100,end='\n')
        env.clear_dead()

        ############add to n_agent##############
        idd = n_agent - len(env.get_agent_id(handles[0]))
        if idd>0:
            env.add_agents(handles[0], method="random", n=idd)
            dead[0]+=idd
        idd = 12 - len(env.get_agent_id(handles[1]))
        if idd>0:
            env.add_agents(handles[1], method="random", n=idd)
            dead[1]+=idd

        if i_episode < episode_before_train:
            continue
        if steps%3 != 0:
            continue
        #############training###########
        batch = buff.getBatch(128)
        states,actions,rewards,new_states,dones,adj=[],[],[],[],[],[]
        for i_ in  range(n_agent*2+1):
            states.append([])
            new_states.append([])
        for e in batch:
            for j in range(n_agent):
                states[j*2].append(e[0][j])
                states[j*2+1].append(e[5][j])
                new_states[j*2].append(e[2][j])
                new_states[j*2+1].append(e[5][j])
            states[40].append(vec)
            new_states[40].append(vec)
            actions.append(e[1])
            rewards.append(e[3])
            dones.append(e[4])

        actions = np.asarray(actions)
        rewards = np.asarray(rewards)
        dones = np.asarray(dones)

        for i_ in  range(n_agent*2+1):
            states[i_]=np.asarray(states[i_])
            new_states[i_]=np.asarray(new_states[i_])

        q_values = model.predict(states)
        target_q_values = model_t.predict(new_states)

        for k in range(len(batch)):
            if dones[k]:
                for j in range(n_agent):
                    q_values[j][k][actions[k][j]] = rewards[k][j]
            else:
                for j in range(n_agent):
                    q_values[j][k][actions[k][j]] =rewards[k][j] + GAMMA*np.max(target_q_values[j][k])

        history=model.fit(states, q_values,epochs=1,batch_size=128,verbose=0)
        his=0
        for (k,v) in history.history.items():
            his+=v[0]
        loss+=(his/20)
        #########train target model#########
        weights = cnn.get_weights()
        target_weights = cnn_t.get_weights()
        for w in range(len(weights)):
            target_weights[w] = TAU * weights[w] + (1 - TAU)* target_weights[w]
        cnn_t.set_weights(target_weights)

        weights = q_net.get_weights()
        target_weights = q_net_t.get_weights()
        for w in range(len(weights)):
            target_weights[w] = TAU * weights[w] + (1 - TAU)* target_weights[w]
        q_net_t.set_weights(target_weights)

        weights = m1.get_weights()
        target_weights = m1_t.get_weights()
        for w in range(len(weights)):
            target_weights[w] = TAU * weights[w] + (1 - TAU)* target_weights[w]
        m1_t.set_weights(target_weights)

        weights = m2.get_weights()
        target_weights = m2_t.get_weights()
        for w in range(len(weights)):
            target_weights[w] = TAU * weights[w] + (1 - TAU)* target_weights[w]
        m2_t.set_weights(target_weights)

        #######save model###############
    model.save('gdn.h5')