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

@author: arash
"""
import matplotlib.pyplot as plt
import torch, pickle, glob
from data_loader import my_collate, CustomIterableDataset, CustomDataset, LineDatasetBalanced
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, f1_score, roc_curve, auc
import numpy as np
import torch.nn as nn

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


def test(model):

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(device)
    #model=torch.load('modelLine3.pt')
    model.eval()
    #model=nn.DataParallel(model, device_ids=[0,1,2,3], output_device=[0])
    #model.to(device)
    
    #dataset = CustomIterableDataset(data_link='test_processed_LLVM3.p')
    res=[]     
    
    files=glob.glob('target_lined*.p')
    for i in range(len(files)):
    #for i in range(10,len(files)):
        torch.cuda.empty_cache() 
        if i!=9:
            print('files   ',i)
            dataset=LineDatasetBalanced(files[i])
        
            X=DataLoader(dataset, batch_size = 1)
         
            for local_batch, context, local_labels in X: 
                
                
                tar = model(local_batch.float().cuda(), context.float().cuda())
                
                res.append([local_labels.cpu().numpy()[0][0], F.sigmoid(tar).detach().cpu().numpy()[0][0]])
            
            
        
        
    res=np.array(res)
    return res
    
    
# pickle.dump(res,open('accuracy_test_line.p','wb'))


# res[res<0.5]=0
# res[res>=0.5]=1
# print(accuracy_score(res[:,0],res[:,1]))   
# print(f1_score(res[:,0],res[:,1]))   

# res=pickle.load(open('accuracy_test_line.p','rb'))
# fpr, tpr, _ = roc_curve(res[:,0], res[:,1])
# plt.plot(fpr, tpr, lw=2)


from sklearn.metrics import confusion_matrix
from sklearn.metrics import accuracy_score, roc_curve, auc, f1_score
import pickle

X=pickle.load(open('res_line_5_runs.p','rb'))
res=X[4]
# 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)




