# -*- coding: utf-8 -*-
"""
Created on Mon Sep 20 14:12:14 2021

@author: arash
"""

import torch, pickle
from data_loader import my_collate, CustomIterableDataset, CustomDataset
import random
import numpy as np
from torch.utils.data import DataLoader, TensorDataset
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, roc_curve, auc, f1_score

import numpy as np
import torch.nn as nn

#model=torch.load('model2.pt',map_location=torch.device('cpu') )

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
model=torch.load('model2N12.pt')
model.eval()
#model=nn.DataParallel(model, device_ids=[0], output_device=[0])
model.to(device)

#dataset = CustomIterableDataset(data_link='test_processed_LLVM3.p')

dataset = CustomDataset(data_link='test_processed_LLVM3.p')

X=DataLoader(dataset, batch_size = 1, collate_fn=my_collate, num_workers=1)
 
      
res=[]        
for local_batch, local_labels in X: 
    #print(local_batch)
    local_batch=torch.permute(local_batch,(1,0,2))
    #print(local_labels.numpy())
    
    #model.window=local_batch.size()[0]
    #print(local_batch.size())
    #model.hn=model.initHidden().cuda()
    #model.hn=torch.zeros(([2,model.window,100]))
    tar = model(local_batch.float().cuda())
    #print('************************')
    
    #print([local_labels.numpy()[0], F.sigmoid(tar).detach().cpu().numpy()[0][0]])
    res.append([local_labels.numpy()[0], F.sigmoid(tar).detach().cpu().numpy()[0][0]])
    
    
    #L=[]
    #for l in local_labels:
    #    L.append(l.item())
    #res.append([F.sigmoid(tar), local_labels, accuracy_score(,F.sigmoid(tar).detach().cpu())])
    
res=np.array(res)
pickle.dump(res,open('accuracy_test.p','wb'))
   

from sklearn.metrics import confusion_matrix
res=pickle.load(open('accuracy_test.p','rb'))
X=pickle.load(open('5_runs_res_whole_random.p','rb'))
res=X[1]
# fpr, tpr, _ = roc_curve(res[:,0], res[:,1])
# plt.plot(fpr, tpr, lw=2)



pred=res[:,1]
pred[pred>=0.5]=1
pred[pred<0.5]=0
print('F1: ',f1_score(res[:,0], pred, average='binary'))
print('Accuracy: ', accuracy_score(res[:,0],res[:,1]))

cm=confusion_matrix(res[:,0], pred)
TP=cm[0,0]
FP=cm[0,1]
FN=cm[1,0]
TN=cm[1,1]
# FP = cm.sum(axis=0) - np.diag(cm)  
# FN = cm.sum(axis=1) - np.diag(cm)
# TP = np.diag(cm)
# TN = cm.sum() - (FP + FN + TP)

FPR=FP/(FN+TN)
FNR=FN/(FN+TP)
print('FPR ', FPR)
print('FNR ', FNR)

