
import scipy.special
from scipy.stats import norm, binom_test,sem
from statsmodels.stats.proportion import proportion_confint
import math
import numpy as np
import torch
import matplotlib.pyplot as plt

def get_cvar_cert_time_t(estimate, t, eps,sigma):
	erf = scipy.special.erf(math.sqrt(t+1) * eps/(2*math.sqrt(2)*sigma))
	cvar = 1. if estimate > erf else estimate/erf
	return cvar * erf

def get_exact_time_t(estimate, t, eps,sigma):
	return norm.cdf(norm.ppf(estimate) - math.sqrt(t+1) * eps/(sigma))

def _lower_confidence_bound(NA: int, N: int, alpha: float) -> float:
    return proportion_confint(NA, N, alpha=2 * alpha, method="beta")[0]
def get_exact_total(estimate, eps,sigma):
	return norm.cdf(norm.ppf(estimate) -  eps/(sigma))
import argparse



plt.figure(figsize=(8.4,4.8))

# for sigma in [0.2,0.4]:
# 	data = torch.load('mountain_car_sigma_'+str(sigma)+'/best_model.zip_evals_10000.pth')
# 	prob = torch.tensor(data).float().mean()
# 	lcb = _lower_confidence_bound(int(prob*10000), 10000, 0.05)
# 	vals = [lcb]
# 	for eps in np.arange(0.05,0.85,0.05):
# 		vals.append(get_exact_total(lcb,eps,sigma))
# 	plt.plot(np.arange(0.0,0.85,0.05),vals, color=('blue' if sigma ==0.2 else 'cornflowerblue'), linestyle=('-' if sigma ==0.2 else '--') ,label="σ = " + str(sigma))#"Policy Smoothing: Certified\nLower Bound (σ = " + str(sigma) +')')

attack_mags_nonzero = [1.,2.,3.,4.]

attack_vals =  [torch.tensor(torch.load('mountain_car_sigma_0.0/best_model.zip_evals_10000.pth')).float().mean().item()]
attack_sems = [sem(torch.tensor(torch.load('mountain_car_sigma_0.0/best_model.zip_evals_10000.pth')))]
for attack_mag in attack_mags_nonzero:
	attack_val = None
	attack_sem = None
	for i,thresh in enumerate([0.001,0.0001,0.00001]):
		cur_val = (torch.tensor( torch.load('mountain_car_sigma_0.0/best_model.zip_evals_250_attack_eps_'+str(attack_mag)+'_steps_100_attack_step_0.01_threshold_'+str(thresh)+'.pth')).float().mean().item())
		if (attack_val is None or cur_val < attack_val):
			attack_val = cur_val
			attack_sem = (sem(torch.tensor( torch.load('mountain_car_sigma_0.0/best_model.zip_evals_250_attack_eps_'+str(attack_mag)+'_steps_100_attack_step_0.01_threshold_'+str(thresh)+'.pth'))))
	attack_vals.append(attack_val)
	attack_sems.append(attack_sem)
attack_mags = [0] + attack_mags_nonzero
plt.errorbar([x for x in attack_mags],attack_vals,  yerr= attack_sems, color='red',  linestyle ="--",label="Undefended"  )



styles = ['-.','--',":"]
attack_mags_nonzero = [1.,2.,3.,4.]
for j,sigma in enumerate([0.2]):#,0.4]):
	attack_vals =  [torch.tensor(torch.load('mountain_car_sigma_'+str(sigma)+'/best_model.zip_evals_10000.pth')).float().mean().item()]
	attack_sems = [sem(torch.tensor(torch.load('mountain_car_sigma_'+str(sigma)+'/best_model.zip_evals_10000.pth')))]
	for attack_mag in attack_mags_nonzero:
		attack_val = None
		attack_sem = None
		for i,thresh in enumerate([0.001,0.0001,0.00001]):
			cur_val = (torch.tensor( torch.load('mountain_car_sigma_'+str(sigma)+'/best_model.zip_smooth_evals_250_attack_eps_'+str(attack_mag)+'_steps_100_attack_step_0.01_threshold_'+str(thresh)+'_num_smoothing_points_128.pth')).float().mean().item())
			if (attack_val is None or cur_val < attack_val):
				attack_val = cur_val
				attack_sem = (sem(torch.tensor( torch.load('mountain_car_sigma_'+str(sigma)+'/best_model.zip_smooth_evals_250_attack_eps_'+str(attack_mag)+'_steps_100_attack_step_0.01_threshold_'+str(thresh)+'_num_smoothing_points_128.pth'))))
		attack_vals.append(attack_val)
		attack_sems.append(attack_sem)
	attack_mags = [0] + attack_mags_nonzero
	plt.errorbar([x for x in attack_mags],attack_vals,  yerr= attack_sems, color='blue',  linestyle ="-",label="Policy Smoothing (σ = " + str(sigma) + ')')



plt.legend()
plt.title('(d) Mountain Car', fontsize=18)
plt.xlim(0,4)
plt.xlabel('Perturbation Budget', fontsize=14)
plt.ylim(0,1.1)
plt.ylabel('Success Rate', fontsize=14)
plt.savefig('mountain_car_sigma_attacks.png',dpi=400,bbox_inches='tight')
