import os
import sys
from datetime import datetime

import torch
import torch.nn as nn
import numpy as np

from bnn import BNN

import argparse

from utils.logging import get_logger

from data import uci_woval
import torch.nn.functional as F
import torchvision.transforms as transforms
from torch.autograd import Variable
from torch.utils.data.dataloader import DataLoader
from torch.utils.data.sampler import SubsetRandomSampler
from torchvision import datasets
from sklearn.metrics import roc_auc_score
from sklearn import metrics
import pandas as pd

import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt


################################### Hyper-parameters ###################################

lr_bnn = 1e-2
prior_coeff = 10
bnn_name_string = 'image_ifbnn'
uci_dataset_name_string = 'protein' #'yacht'
random_seed = 123
max_epoch_num = 100
save_model_freq = 5000
log_model_freq = 3000
test_interval = 2000
num_sample = 1000
epochs = 601

torch.manual_seed(random_seed)
np.random.seed(random_seed)


################################### Network Architecture ###################################

n_units = 800
n_hidden = 2
hidden_dims = [n_units] * n_hidden
activation_fn = 'tanh'


print("============================================================================================")
################################## set device ##################################

# set device to cpu or cuda
device = torch.device('cpu')
if (torch.cuda.is_available()):
    device = torch.device('cuda:0')
    torch.cuda.empty_cache()
    print("Device set to : " + str(torch.cuda.get_device_name(device)))
else:
    print("Device set to : cpu")

print("============================================================================================")
###################### logging ######################

#### log files for multiple runs are NOT overwritten

log_dir = "./logs"
if not os.path.exists(log_dir):
    os.makedirs(log_dir)

log_dir = log_dir + '/' + bnn_name_string + '/'
if not os.path.exists(log_dir):
    os.makedirs(log_dir)

#### get number of log files in log directory
current_num_files = next(os.walk(log_dir))[2]
run_num = len(current_num_files)

#### create new log file for each run
log_f_name = log_dir + '/' + bnn_name_string + "_" + str(run_num) + ".csv"

print("current logging run number for " + bnn_name_string + " : ", run_num)
print("logging at : " + log_f_name)


print("============================================================================================")
################### checkpointing ###################

run_num_pretrained = 0  #### change this to prevent overwriting weights in same env_name folder

directory = "./pretrained"
if not os.path.exists(directory):
    os.makedirs(directory)

directory = directory + '/' + bnn_name_string + '/'
if not os.path.exists(directory):
    os.makedirs(directory)

checkpoint_path = directory + "{}_{}_{}.pth".format(bnn_name_string, random_seed, run_num_pretrained)
print("save checkpoint path : " + checkpoint_path)

################### savefigures ###################

results_folder = "./results"
if not os.path.exists(results_folder):
    os.makedirs(results_folder)

figures_folder = results_folder + '/' + bnn_name_string + '/' + 'FWBI' + '/'
if not os.path.exists(figures_folder):
    os.makedirs(figures_folder)


print("============================================================================================")
############# print all hyperparameters #############

print("learning rate: ", lr_bnn)
print("coefficient of prior regularization: ", prior_coeff)
print("random seed: ", random_seed)
print("max number of epoches: ", max_epoch_num)


print("============================================================================================")
############################## load and normalize data ##############################

transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Lambda(lambda x: x * 255. / 126.),  # divide as in paper
        ])

transform2 = transforms.Compose([
            transforms.RandomHorizontalFlip(),
            transforms.RandomRotation(10),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])

train_data = datasets.MNIST(
            root='data',
            train=True,
            download=False,
            transform=transform)

test_data = datasets.MNIST(
            root='data',
            train=False,
            download=False,
            transform=transform)

train_data_fm = datasets.FashionMNIST(
            root='data2',
            train=True,
            download=False,
            transform=transform)

test_data_fm = datasets.FashionMNIST(
            root='data',
            train=False,
            download=False,
            transform=transform)

train_data_cf = datasets.CIFAR10(
            root='data',
            train=True,
            download=False,
            transform=transform2)

test_data_cf = datasets.CIFAR10(
            root='data',
            train=False,
            download=False,
            transform=transform2)
train_data_sv = datasets.SVHN(
            root='data',
            split='train',
            download=True,
            transform=transform2)

# test_data_sv = datasets.SVHN(
#             root='data',
#             train=False,
#             download=True,
#             transform=transform2)

training_num = len(train_data_fm)
test_num = len(test_data_fm)
input_dim = 28 * 28
# input_dim = 32 * 32 * 3
output_dim = 10
is_continuous = True
batch_size = 125
eval_batch_size = 1000
# print('train_data.shape: ', train_data.shape)
# obtain training indices that will be used for validation
valid_size = 1 / 6
indices = list(range(training_num))
split = int(valid_size * training_num)
train_idx, valid_idx = indices[split:], indices[:split]
oob_in_idx = train_idx[:10000]
train_idx = train_idx[40000:50000]

# define samplers for obtaining training and validation batches
train_sampler = SubsetRandomSampler(train_idx)
valid_sampler = SubsetRandomSampler(valid_idx)
oob_in_sampler = SubsetRandomSampler(oob_in_idx)

# miniset of training data and test data
mini_train_size = 1 / 12
mini_train_index = int(mini_train_size * training_num)
mini_train_idx = indices[:mini_train_index]
mini_train_sampler = SubsetRandomSampler(mini_train_idx)

indices_test = list(range(test_num))
mini_test_size = 1 / 2
mini_test_index = int(mini_test_size * test_num)
mini_test_idx = indices[:mini_test_index]
mini_test_sampler = SubsetRandomSampler(mini_test_idx)

############## define data loader #################

train_loader = torch.utils.data.DataLoader(
        train_data_fm,
        batch_size=batch_size,
        sampler=train_sampler,
        num_workers=1)
valid_loader = torch.utils.data.DataLoader(
    train_data_fm,
    batch_size=batch_size,
    sampler=valid_sampler,
    num_workers=1)
test_loader = torch.utils.data.DataLoader(
        test_data_fm,
        batch_size=eval_batch_size,
        num_workers=1)
oob_out_loader = torch.utils.data.DataLoader(
    train_data,
    batch_size=batch_size,
    sampler=valid_sampler, # valid_sampler
    num_workers=1)
oob_in_loader = torch.utils.data.DataLoader(
    train_data_fm,
    batch_size=batch_size,
    sampler=oob_in_sampler,
    num_workers=1)

mini_train_loader = torch.utils.data.DataLoader(
    train_data,
    batch_size=batch_size,
    sampler=mini_train_sampler,
    num_workers=1)

mini_test_loader = torch.utils.data.DataLoader(
    test_data,
    batch_size=eval_batch_size,
    sampler=mini_test_sampler,
    num_workers=1)

print("training_num = ", training_num, " input_dim = ", input_dim, " output_dim = ", output_dim)



loss = nn.MSELoss(reduction='mean')


print("============================================================================================")
################################## define model ##################################

bnn = BNN(input_dim, output_dim, hidden_dims, activation_fn, is_continuous, scaled_variance=True).to(device)

bnn_optimizer = torch.optim.Adam([{'params': bnn.parameters(), 'lr': lr_bnn}])

print("============================================================================================")
######### def train and evalution #############################
def train(loader):
    loss_sum = 0

    for batch_id, (data, target) in enumerate(loader):
        data, target = data.to(device), target.to(device)
        data = torch.reshape(data, (125, 784))
        # data = torch.reshape(data, (125, 3072))
        # print('data.shape: ', data.shape)
        # print('target.shape: ', target.shape)
        bnn_optimizer.zero_grad()
        pred_y, distance_prior = bnn.forward_kl(data)
        # pred_y = bnn.forward(data)
        # print('prior_distance: ', distance_prior)
        pred_y = F.softmax(pred_y, dim=1)
        # print(pred_y.shape)
        # print(target.shape)
        output = torch.log(pred_y)
        train_loss = F.nll_loss(output, target, reduction='sum') + prior_coeff * distance_prior
        # train_loss = F.nll_loss(output, target, reduction='sum')
        # print('train_loss: ', train_loss)
        loss_sum += train_loss / len(loader)
        train_loss.backward()
        bnn_optimizer.step()
    return loss_sum

def evaluate(model, loader, samples=1):
    acc_sum = 0
    model.eval()
    for idx, (data, target) in enumerate(loader):
        data, target = data.to(device), target.to(device)

        if samples == 1:
            output = model.forward(data)
            output = F.softmax(output, dim=1)
        else:
            output = model.forward(data)
            output = F.softmax(output, dim=1)
            for i in range(samples - 1):
                output += F.softmax(model.forward(data))

        predict = output.data.max(1)[1]
        acc = predict.eq(target.data).cpu().sum().item()
        acc_sum += acc
    return acc_sum / len(loader)


################################## start training ##################################

# track total training time
start_time = datetime.now().replace(microsecond=0)
print("Started training at (GMT) : ", start_time)

# logging file
log_f = open(log_f_name, "w+")
log_f.write('epoch,training_loss\n')

for epoch in range(epochs):
    train_loss = train(train_loader)
    # print('train_loss: ', train_loss)
    test_acc = evaluate(bnn, test_loader, samples=10)
    # print('test_acc: ', round(100 * (test_acc / eval_batch_size), 3))
    test_err = round((1 - test_acc / eval_batch_size), 4)
    # print('test_err: ', test_err)



    # print
    print("epoch : {} \t\t training loss \t\t : {} \t\t test_error \t\t : {}".format(epoch, train_loss,
                                                                                         test_err),
          datetime.now().replace(microsecond=0) - start_time)


    # log in logging file
    if epoch % log_model_freq == 0:
        print("--------------------------------------------------------------------------------------------")
        log_f.write('{},{}\n'.format(epoch, train_loss))
        log_f.flush()
        print("log saved")
        print("--------------------------------------------------------------------------------------------")

    # save model weights
    if epoch % save_model_freq == 0:
        print("--------------------------------------------------------------------------------------------")
        print("saving model at : " + checkpoint_path)
        bnn.save(checkpoint_path)
        print("model saved")
        print("--------------------------------------------------------------------------------------------")

    # if epoch % test_interval == 0:
    #
    #     bnn.eval()
    #
    #     samples_pred_test_y = bnn.forward_eval(original_x_test, num_sample).squeeze().detach()
    #
    #     mse_test_loss = loss(samples_pred_test_y, original_y_test)
    #
    #     bnn.train()

####### OOD_test ###############
log_upper_bound = np.log(output_dim)

entropies_in_sample=[]

for batch_id, (data, target) in enumerate(oob_in_loader):
    data, target = data.to(device), target.to(device)
    data = torch.reshape(data, (125, 784))
    # data = torch.reshape(data, (125, 3072))
    with torch.no_grad():
        p=bnn.forward(data)
        p=F.softmax(p, dim=1)
        e = torch.sum(- p * torch.log(p), dim=1)
        e = torch.nan_to_num(e, nan=0.0)
        # print(e)
        entropies_in_sample.append(e)

entropies_in_sample = torch.cat(entropies_in_sample, dim=0)
true_lables_in_sample = torch.zeros_like(entropies_in_sample)

entropies_out_sample = []
for batch_id, (data, target) in enumerate(oob_out_loader):
    data, target = data.to(device), target.to(device)
    data = torch.reshape(data, (125, 784))
    # data = torch.reshape(data, (125, 3072))
    with torch.no_grad():
        p=bnn.forward(data)
        p = F.softmax(p, dim=1)
        e = torch.sum(- p * torch.log(p), dim=1)
        # print(e)
        e = torch.nan_to_num(e, nan=0.0)
        entropies_out_sample.append(e)
entropies_out_sample = torch.cat(entropies_out_sample, dim=0)
true_lables_out_sample = torch.ones_like(entropies_out_sample)

in_sample_preds = entropies_in_sample/log_upper_bound
out_sample_preds = entropies_out_sample/log_upper_bound
preds = torch.cat([in_sample_preds, out_sample_preds], dim=0).cpu().numpy()
truths = torch.cat([true_lables_in_sample, true_lables_out_sample], dim=0).cpu().numpy()
AUC = roc_auc_score(truths, preds)
print('auc: ', AUC)

###############
# fpr, tpr, thresholds = metrics.roc_curve(truths, preds)
# fpr_kl_fm = pd.DataFrame.from_dict(fpr)
# fpr_kl_fm.to_csv(results_folder + '/' + bnn_name_string + '/' + 'fpr_kl2_fm.csv')
# tpr_kl_fm = pd.DataFrame.from_dict(tpr)
# tpr_kl_fm.to_csv(results_folder + '/' + bnn_name_string + '/' + 'tpr_kl2_fm.csv')



##############
# kl_fpr = pd.read_csv(results_folder + '/' + bnn_name_string + '/' + 'fpr_kl_fm.csv')
# kl_fpr = kl_fpr.to_numpy()
# kl_tpr = pd.read_csv(results_folder + '/' + bnn_name_string + '/' + 'tpr_kl_fm.csv')
# kl_tpr = kl_tpr.to_numpy()
#
# w_fpr = pd.read_csv(results_folder + '/' + bnn_name_string + '/' + 'fpr_w_fm.csv')
# w_fpr = w_fpr.to_numpy()
# w_tpr = pd.read_csv(results_folder + '/' + bnn_name_string + '/' + 'tpr_w_fm.csv')
# w_tpr = w_tpr.to_numpy()
#
# ifbnn_fpr = pd.read_csv(results_folder + '/' + bnn_name_string + '/' + 'fpr_ifbnn_fm.csv')
# ifbnn_fpr = ifbnn_fpr.to_numpy()
# ifbnn_tpr = pd.read_csv(results_folder + '/' + bnn_name_string + '/' + 'tpr_ifbnn_fm.csv')
# ifbnn_tpr = ifbnn_tpr.to_numpy()
#
# gwi_fpr = pd.read_csv(results_folder + '/' + bnn_name_string + '/' + 'fpr_gwi_fm.csv')
# gwi_fpr = gwi_fpr.to_numpy()
# gwi_tpr = pd.read_csv(results_folder + '/' + bnn_name_string + '/' + 'tpr_gwi_fm.csv')
# gwi_tpr = gwi_tpr.to_numpy()
#
# fbnn_fpr = pd.read_csv(results_folder + '/' + bnn_name_string + '/' + 'fpr_fbnn_fm.csv')
# fbnn_fpr = fbnn_fpr.to_numpy()
# fbnn_tpr = pd.read_csv(results_folder + '/' + bnn_name_string + '/' + 'tpr_fbnn_fm.csv')
# fbnn_tpr = fbnn_tpr.to_numpy()
#
# plt.clf()
# figure = plt.figure(figsize=(8, 5.5), facecolor='white')
# # init_plotting()
# kl_fpr = kl_fpr[:, 1]
# kl_tpr = kl_tpr[:, 1]
# w_fpr = w_fpr[:, 1]
# w_tpr = w_tpr[:, 1]
# ifbnn_fpr = ifbnn_fpr[:, 1]
# ifbnn_tpr = ifbnn_tpr[:, 1]
# gwi_fpr = gwi_fpr[:, 1]
# gwi_tpr = gwi_tpr[:, 1]
# fbnn_fpr = fbnn_fpr[:, 1]
# fbnn_tpr = fbnn_tpr[:, 1]
# plt.plot(kl_fpr, kl_tpr, label='KLBBB', color='darkorange')
# plt.plot(w_fpr, w_tpr, label='WBBB', color='navy')
# plt.plot(ifbnn_fpr, ifbnn_tpr, label='FWBI', color='deeppink')
# plt.plot(gwi_fpr, gwi_tpr, label='GWI', color='aqua')
# plt.plot(fbnn_fpr, fbnn_tpr, label='FBNN', color='green')
# plt.ylabel(r'True Positive Rate')
# plt.tight_layout()
# plt.xlabel('False Positive Rate')
# plt.tight_layout()
# plt.grid()
# plt.legend()
# plt.savefig(figures_folder + '/rocfm.pdf')




# roc_auc = metrics.auc(fpr, tpr)
# display = metrics.RocCurveDisplay(fpr=fpr, tpr=tpr, roc_auc=roc_auc, estimator_name='klbnn')
# display.plot()
# plt.savefig(figures_folder + '/oob_kl2bnn_fm.pdf')



log_f.close()

# print total training time
print("============================================================================================")
end_time = datetime.now().replace(microsecond=0)
print("Started training at (GMT) : ", start_time)
print("Finished training at (GMT) : ", end_time)
print("Total training time  : ", end_time - start_time)
print("Test err  : ", test_err)
print("============================================================================================")


