import numpy as np
#import pandas as pd
import matplotlib.pyplot as plt
import argparse, sys
import joblib
import warnings
warnings.filterwarnings('ignore')

import torch
import torch.nn as nn
from torch.autograd import Variable
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
cudnn.benchmark = True

from sklearn.metrics import roc_curve, auc, f1_score, precision_recall_curve, average_precision_score, ConfusionMatrixDisplay

from medmnistutils.evaluationmetrics import accuracy, roc, presenf1cfsmtx
from medmnistutils.medmnistdataloader import PathMNIST, OrganMNIST3D, PneumoniaMNIST, VesselMNIST3D, OCTMNIST

#from medmnistutils.jiaodaresnet import ResNet18 as jiaodaresnet18
#from nets.unknownthreedresnet import resnet18
from medmnistutils.blingblingresnet import resnet18 as blingblingresnet18
from medmnistutils.O2Uzidairesnet import ResNet18 as O2Uresnet18
from medmnistutils.yixianresnet import resnet18 as yixian3dresnet18


parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str, default='OCTMNIST', help='PathMNIST, OCTMNIST, PneumoniaMNIST, OrganMNIST3D, VesselMNIST3D')
parser.add_argument('--noise_rate', type=float, default=0.4, help='noise rate')
parser.add_argument('--batchsize', type=int, default=128, help='128')
parser.add_argument('--num_epochs', type=int, default=200, help='number of epochs')

#args = parser.parse_args(args=[])
args = parser.parse_args()


if args.dataset =='PathMNIST':    #2D, 9 classes, 89,996 / 10,004 / 7,180
    newtransform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[.5], std=[.5])])
    train_dataset = PathMNIST(split = 'train', root = '../../medmnistdata',  transform=newtransform, noise_rate=args.noise_rate)
    val_dataset = PathMNIST(split = 'val', root = '../../medmnistdata',  transform=newtransform)
    test_dataset = PathMNIST(split = 'test', root = '../../medmnistdata',  transform=newtransform)
    train_loader = DataLoader(train_dataset, batch_size=args.batchsize, drop_last = False, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=args.batchsize,  drop_last = False, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=args.batchsize, drop_last = False, shuffle=True)
    model = O2Uresnet18(input_channel=train_dataset.in_channels, n_outputs=train_dataset.num_classes)
    #model = blingblingresnet18(num_classes=train_dataset.num_classes)

if args.dataset =='OCTMNIST':   #2D, 4 classes,
    newtransform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[.5], std=[.5])])
    train_dataset = OCTMNIST(split = 'train', root = '../../medmnistdata',  transform=newtransform, noise_rate=args.noise_rate)
    val_dataset = OCTMNIST(split = 'val', root = '../../medmnistdata',  transform=newtransform)
    test_dataset = OCTMNIST(split = 'test', root = '../../medmnistdata',  transform=newtransform)
    train_loader = DataLoader(train_dataset, batch_size=args.batchsize, drop_last = False, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=args.batchsize,  drop_last = False, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=args.batchsize, drop_last = False, shuffle=True)
    model = O2Uresnet18(input_channel=train_dataset.in_channels, n_outputs=train_dataset.num_classes) 
    #model = blingblingresnet18(num_classes=train_dataset.num_classes)

elif args.dataset =='PneumoniaMNIST':  #2D, 2 class, 4,708 / 524 / 624
    newtransform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[.5], std=[.5])])
    train_dataset = PneumoniaMNIST(split = 'train', root = '../../medmnistdata',  transform=newtransform, noise_rate=args.noise_rate)
    val_dataset = PneumoniaMNIST(split = 'val', root = '../../medmnistdata',  transform=newtransform)
    test_dataset = PneumoniaMNIST(split = 'test', root = '../../medmnistdata',  transform=newtransform)
    train_loader = DataLoader(train_dataset, batch_size=args.batchsize, drop_last = False, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=args.batchsize,  drop_last = False, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=args.batchsize, drop_last = False, shuffle=True)
    model = O2Uresnet18(input_channel=train_dataset.in_channels, n_outputs=train_dataset.num_classes)
    #model = blingblingresnet18(num_classes=train_dataset.num_classes)
    
elif args.dataset =='OrganMNIST3D':   #3D, 11 class, 972 / 161 / 610
    train_dataset = OrganMNIST3D(split = 'train', root = '../../medmnistdata',  transform=None, noise_rate=args.noise_rate)
    val_dataset = OrganMNIST3D(split = 'val', root = '../../medmnistdata',  transform=None)
    test_dataset = OrganMNIST3D(split = 'test', root = '../../medmnistdata',  transform=None)
    train_loader = DataLoader(train_dataset, batch_size=args.batchsize, drop_last = False, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=args.batchsize,  drop_last = False, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=args.batchsize, drop_last = False, shuffle=True)
    model = yixian3dresnet18(num_classes = train_dataset.num_classes) 

elif args.dataset =='VesselMNIST3D':    #3D, 2 class, 1,335 / 192 / 382
    train_dataset = VesselMNIST3D(split = 'train', root = '../../medmnistdata',  transform=None, noise_rate=args.noise_rate)
    val_dataset = VesselMNIST3D(split = 'val', root = '../../medmnistdata',  transform=None)
    test_dataset = VesselMNIST3D(split = 'test', root = '../../medmnistdata',  transform=None)
    train_loader = DataLoader(train_dataset, batch_size=args.batchsize, drop_last = False, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=args.batchsize,  drop_last = False, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=args.batchsize, drop_last = False, shuffle=True)
    model = yixian3dresnet18(num_classes = train_dataset.num_classes) 

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

error = nn.CrossEntropyLoss()
learning_rate = 0.001 
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)   

###############################################################################
验证准确率列表 = []
测试准确率列表= []
###############################################################################

#main loop
for epoch in range(args.num_epochs):

    #train
    model.train()
    for images, labels, _ in train_loader:
        images, labels = images.to(device), labels.to(device)  
        labels = labels.squeeze().long() 
        outputs = model(images)
        loss = error(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
    #evaluation
    valaccuracy = accuracy(model, val_loader)
    testaccuracy = accuracy(model, test_loader)
    print('epoch', epoch+1, 'val accuracy', valaccuracy, 'test accuracy', testaccuracy)

###############################################################################
#以下都是不需要的
###############################################################################
    验证准确率列表.append(valaccuracy)
    测试准确率列表.append(testaccuracy)
    
实验名 = '20230924baselineexp1'
resultdict = dict()

#模型
resultdict['model'] = model

#acc变化图
resultdict['valacclist'] = 验证准确率列表
resultdict['testacclist'] = 测试准确率列表
验证准确率列表 = [x*100 for x in 验证准确率列表]
测试准确率列表 = [x*100 for x in 测试准确率列表]
plt.plot(验证准确率列表, label = 'validation set')
plt.plot(测试准确率列表, label = 'test set')
plt.xlim((0,200))
plt.ylim((0,100))
#plt.title('origingal method on ' + args.dataset + ' under noise rate ' + str(args.noise_rate))
plt.xlabel('Epoch')
plt.ylabel('Accuracy (%)')
acc变化图文件名 = 实验名 + '_acccurve_' + args.dataset + '_' + str(args.noise_rate) + '.png'
plt.legend()
plt.savefig(acc变化图文件名)
plt.show()

#ROC曲线图
resultdict['valfprdict'], resultdict['valtprdict'], resultdict['valaucdict'] = roc(model, val_loader)
resultdict['testfprdict'], resultdict['testtprdict'], resultdict['testaucdict'] = roc(model, test_loader)

plt.plot(resultdict['valfprdict']["micro"], resultdict['valtprdict']["micro"],
         label='validation set, AUC ' + str(round(100*resultdict['valaucdict']["micro"],2)))
plt.plot(resultdict['testfprdict']["micro"], resultdict['testtprdict']["micro"],
         label='test set, AUC ' + str(round(100*resultdict['testaucdict']["micro"],2)))
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.legend(loc="lower right")
ROC文件名 = 实验名 + '_roccurve_' + args.dataset + '_' + str(args.noise_rate) + '.png'
plt.savefig(ROC文件名)
plt.show()

#confusion matrix图
resultdict['valprecision'], resultdict['valrecall'], resultdict['valf1'], resultdict['valtruelist'], resultdict['valpredlist'], resultdict['valcfsmtx'] = presenf1cfsmtx(model, val_loader)
ConfusionMatrixDisplay.from_predictions(resultdict['valtruelist'], resultdict['valpredlist'], cmap = plt.cm.Blues, colorbar = False)
cfsmtx文件名 = 实验名 + '_valconfusionmatrix_' + args.dataset + '_' + str(args.noise_rate) + '.png'
plt.savefig(cfsmtx文件名)
plt.show()  
resultdict['testprecision'], resultdict['testrecall'], resultdict['testf1'], resultdict['testtruelist'], resultdict['testpredlist'], resultdict['testcfsmtx'] = presenf1cfsmtx(model, test_loader) 
ConfusionMatrixDisplay.from_predictions(resultdict['testtruelist'], resultdict['testpredlist'], cmap = plt.cm.Blues, colorbar = False)
cfsmtx文件名 = 实验名 + '_testconfusionmatrix_' + args.dataset + '_' + str(args.noise_rate) + '.png'
plt.savefig(cfsmtx文件名)
plt.show()  

#txt
txt文件名 = 实验名 + '_txt_' + args.dataset + '_' + str(args.noise_rate) + '.txt'
with open (txt文件名, 'a', encoding='utf-8') as txt:
    txt.write('最后一轮acc' + "\n" )
    txt.write(str(round(验证准确率列表[-1],2)) + "\n" )
    txt.write(str(round(测试准确率列表[-1],2)) + "\n" )
    txt.write('最后十轮acc平均' + "\n" )
    txt.write(str(round(sum(验证准确率列表[-11:-1])/len(验证准确率列表[-11:-1]),2)) + "\n" )
    txt.write(str(round(sum(测试准确率列表[-11:-1])/len(验证准确率列表[-11:-1]),2)) + "\n" )
    txt.write('precision' + "\n" )
    txt.write(str(round(100*resultdict['valprecision'],2)) + "\n" )
    txt.write(str(round(100*resultdict['testprecision'],2)) + "\n" )
    txt.write('recall' + "\n" )
    txt.write(str(round(100*resultdict['valrecall'],2)) + "\n" )
    txt.write(str(round(100*resultdict['testrecall'],2)) + "\n" )
    txt.write('f1' + "\n" )
    txt.write(str(round(100*resultdict['valf1'],2)) + "\n" )
    txt.write(str(round(100*resultdict['testf1'],2)) + "\n" )

#保存整个文件
resultdict文件名 = 实验名 + '_resultdict_' + args.dataset + '_' + str(args.noise_rate)
joblib.dump(resultdict, resultdict文件名)
