import numpy as np
import torch
import torch.nn as nn
import sys
import os

net, eps, alpha = sys.argv[1], float(sys.argv[2]), float(sys.argv[3])

std_errs, pgd_errs, rpgd_errs, pgd2_errs, rpgd2_errs = [], [], [], [], []
for seed in range(10):
    eval_std = "./results/evals/model_standard_net=%s_eps=%s_alpha=%s_seed=%s.pt"%(net,eps,alpha,seed)
    eval_pgd = "./results/evals/model_atk=pgd_net=%s_eps=%s_alpha=%s_seed=%s.pt"%(net,eps,alpha,seed)
    eval_rpgd = "./results/evals/model_atk=rpgd_net=%s_eps=%s_alpha=%s_seed=%s.pt"%(net,eps,alpha,seed)
    eval_pgd2 = "./results/evals/model_atk=pgd2_net=%s_eps=%s_alpha=%s_seed=%s.pt"%(net,eps,alpha,seed)
    eval_rpgd2 = "./results/evals/model_atk=rpgd2_net=%s_eps=%s_alpha=%s_seed=%s.pt"%(net,eps,alpha,seed)

    std_errs.append(torch.load(eval_std))
    pgd_errs.append(torch.load(eval_pgd))
    rpgd_errs.append(torch.load(eval_rpgd))
    pgd2_errs.append(torch.load(eval_pgd2))
    rpgd2_errs.append(torch.load(eval_rpgd2))
    
    
std_mean, std_std = np.mean(std_errs,axis=0), np.std(std_errs,axis=0)
pgd_mean, pgd_std = np.mean(pgd_errs,axis=0), np.std(pgd_errs,axis=0)
rpgd_mean, rpgd_std = np.mean(rpgd_errs,axis=0), np.std(rpgd_errs,axis=0)
pgd2_mean, pgd2_std = np.mean(pgd2_errs,axis=0), np.std(pgd2_errs,axis=0)
rpgd2_mean, rpgd2_std = np.mean(rpgd2_errs,axis=0), np.std(rpgd2_errs,axis=0)

print('std = ', std_mean)
print('pgd = ', pgd_mean)
print('rpgd = ', rpgd_mean)
print('pgd2 = ', pgd2_mean)
print('rpgd2 = ', rpgd2_mean)
