import csv
import numpy as np
dims = (5,1,5,5)
PGD_tr_loss = np.zeros(dims)
PGD_tr_sa = np.zeros(dims)
PGD_tr_ra = np.zeros(dims)
PGD_test_loss = np.zeros(dims)
PGD_test_sa = np.zeros(dims)
PGD_test_ra = np.zeros(dims)

NPGD_tr_loss = np.zeros(dims)
NPGD_tr_sa = np.zeros(dims)
NPGD_tr_ra = np.zeros(dims)
NPGD_test_loss = np.zeros(dims)
NPGD_test_sa = np.zeros(dims)
NPGD_test_ra = np.zeros(dims)

for ii, i in enumerate([1, 2, 4, 8, 16]): # epsilon, size of attack
    for jj, j in enumerate(['1']): # learning rate fixed
        for kk, k in enumerate([2, 4, 8, 16, 32]): # how many steps of pgd
            for ll, l in enumerate([1, 2, 3, 4, 5]): # seed number
                file_name = 'accuracy/CIFAR10_PGD_pgd%d_eps%d_lr%s_seed%d.csv'%(k,i,j,l)
                file = open(file_name)
                csvreader = csv.reader(file)
                rows = []
                for row in csvreader:
                    rows.append(row)
                PGD_tr_loss[ii,jj,kk,ll] = rows[1][-1]
                PGD_tr_sa[ii,jj,kk,ll] = rows[2][-1]
                PGD_tr_ra[ii,jj,kk,ll] = rows[3][-1]
                PGD_test_loss[ii,jj,kk,ll] = rows[4][-1]
                PGD_test_sa[ii,jj,kk,ll] = rows[5][-1]
                PGD_test_ra[ii,jj,kk,ll] = rows[6][-1]

                
                file_name = 'accuracy/CIFAR10_NPGD_npgd%d_eps%d_lr%s_seed%d.csv'%(k,i,j,l)
                file = open(file_name)
                csvreader = csv.reader(file)
                rows = []
                for row in csvreader:
                    rows.append(row)
                NPGD_tr_loss[ii,jj,kk,ll] = rows[1][-1]
                NPGD_tr_sa[ii,jj,kk,ll] = rows[2][-1]
                NPGD_tr_ra[ii,jj,kk,ll] = rows[3][-1]
                NPGD_test_loss[ii,jj,kk,ll] = rows[4][-1]
                NPGD_test_sa[ii,jj,kk,ll] = rows[5][-1]
                NPGD_test_ra[ii,jj,kk,ll] = rows[6][-1]

                
PGD_test_sa = np.mean(PGD_test_sa,axis=3)
NPGD_test_sa = np.mean(NPGD_test_sa,axis=3)

print('epsilon=2, different number of attack iterates [2 4 8 16 32]')
print('PGD', PGD_test_sa[1])
print('NPGD', NPGD_test_sa[1])

print('epsilon=4, different number of attack iterates [2 4 8 16 32]')
print('PGD', PGD_test_sa[2])
print('NPGD', NPGD_test_sa[2])

print('epsilon=8, different number of attack iterates [2 4 8 16 32]')
print('PGD', PGD_test_sa[3])
print('NPGD', NPGD_test_sa[3])

print('epsilon=16, different number of attack iterates [2 4 8 16 32]')
print('PGD', PGD_test_sa[4])
print('NPGD', NPGD_test_sa[4])
