import numpy as np
import matplotlib.pyplot as plt

#Ratio (Figure 1)
counters = np.loadtxt("./results/synthetic_mean/counters-1.txt")
nonpr = np.loadtxt("./results/synthetic_mean/nonpr-1.txt")
t3 = np.loadtxt("./results/synthetic_mean/t3-1.txt")
ratio = [x/y for x,y in zip(t3, nonpr)]

fig, ax = plt.subplots()
ax.plot(counters, ratio, marker="x", color='#E377C2')
plt.ylim(bottom=1.0)
ax.set_xlabel('n')
ax.set_ylabel('Cost of Privacy')
ax.set_title("Multivariate Mean Estimation")
plt.savefig('./plots/ratiomean.png')

#Headline comparison plotting (Figure 3)
counters = np.loadtxt("./results/synthetic_mean/counters-1.txt")
nonpr = np.loadtxt("./results/synthetic_mean/nonpr-1.txt")
naive = np.loadtxt("./results/synthetic_mean/naive-1.txt")
kv = np.loadtxt("./results/synthetic_mean/kv-1.txt")
dfmbg = np.loadtxt("./results/synthetic_mean/dfmbg-1.txt")
t1 = np.loadtxt("./results/synthetic_mean/t1-1.txt")
t2 = np.loadtxt("./results/synthetic_mean/t2-1.txt")
t3 = np.loadtxt("./results/synthetic_mean/t3-1.txt")
t4 = np.loadtxt("./results/synthetic_mean/t4-1.txt")

fig, ax = plt.subplots()
ax.plot(counters, nonpr, marker="x", label='Non-private', color='#1f77b4')
ax.plot(counters, naive, marker="x", label='Naive coordinatewise', color='#ff7f0e')
ax.plot(counters, kv, marker="x", label='KV', color='#2ca02c')
ax.plot(counters, dfmbg, marker="x", label='DFMBG', color='#8c564b')
ax.plot(counters, t1, marker="x", label='t = 1', color='#9467bd')
ax.plot(counters, t2, marker="x", label='t = 2', color='#d62728')
plt.ylim(0.0,1.5)
ax.set_xlabel('n')
ax.set_ylabel('L2 error')
ax.set_title("Baseline Comparisons")
ax.legend()
plt.savefig('./plots/1a.png')

fig, ax = plt.subplots()
ax.plot(counters, nonpr, marker="x", label='Non-private', color='#1f77b4')
ax.plot(counters, dfmbg, marker="x", label='DFMBG', color='#8c564b')
ax.plot(counters, t2, marker="x", label='t = 2', color='#d62728')
plt.ylim(0.0,0.3)
ax.set_xlabel('n')
ax.set_ylabel('L2 error')
ax.set_title("Baseline Comparisons")
ax.legend()
plt.savefig('./plots/1b.png')

fig, ax = plt.subplots()
ax.plot(counters, dfmbg -nonpr, marker="x", label='DFMBG', color='#8c564b')
ax.plot(counters, t2 - nonpr, marker="x", label='t = 2', color='#d62728')
plt.ylim(0.0,0.1)
ax.set_xlabel('n')
ax.set_ylabel('Excess L2 error over Non-private')
ax.set_title("Baseline Comparisons")
ax.legend()
plt.savefig('./plots/1c.png')

#Increasing R plotting (Figure 4)
counters = np.loadtxt("./results/synthetic_mean/counters-2.txt")
nonpr = np.loadtxt("./results/synthetic_mean/nonpr-2.txt")
naive = np.loadtxt("./results/synthetic_mean/naive-2.txt")
kv = np.loadtxt("./results/synthetic_mean/kv-2.txt")
dfmbg = np.loadtxt("./results/synthetic_mean/dfmbg-2.txt")
dfmbgu = np.loadtxt("./results/synthetic_mean/dfmbg-2unfair.txt")
t1 = np.loadtxt("./results/synthetic_mean/t1-2.txt")
t2 = np.loadtxt("./results/synthetic_mean/t2-2.txt")
t3 = np.loadtxt("./results/synthetic_mean/t3-2.txt")
t4 = np.loadtxt("./results/synthetic_mean/t4-2.txt")
t10 = np.loadtxt("./results/synthetic_mean/t10-2.txt")

fig, ax = plt.subplots()
ax.plot(counters, nonpr, marker="x", label='Non-private', color='#1f77b4')
ax.plot(counters, naive, marker="x", label='Naive coordinatewise', color='#ff7f0e')
ax.plot(counters, kv, marker="x", label='KV', color='#2ca02c')
ax.plot(counters, dfmbg, marker="x", label='DFMBG', color='#8c564b')
ax.plot(counters, t1, marker="x", label='t = 1', color='#9467bd')
ax.plot(counters, t2, marker="x", label='t = 2', color='#d62728')
plt.ylim(0.0,10000)
ax.set_xlabel('R')
ax.set_ylabel('L2 error')
ax.set_title("Effect of Increasing R")
ax.legend()
plt.savefig('./plots/2a.png')

fig, ax = plt.subplots()
ax.plot(counters, nonpr, marker="x", label='Non-private', color='#1f77b4')
ax.plot(counters, kv, marker="x", label='KV', color='#2ca02c')
ax.plot(counters, dfmbgu, marker="x", label='DFMBG (Advantaged)', color='#8c564b')
ax.plot(counters, t2, marker="x", label='t = 2', color='#d62728')
ax.plot(counters, t3, marker="x", label='t = 3', color='#e377c2')
ax.plot(counters, t4, marker="x", label='t = 4', color='#7f7f7f')
ax.plot(counters, t10, marker="x", label='t = 10', color='#bcbd22')
plt.ylim(0.0,2)
ax.set_xlabel('R')
ax.set_ylabel('L2 error')
ax.set_title("Effect of Increasing R")
ax.legend()
plt.savefig('./plots/2b.png')

fig, ax = plt.subplots()
ax.plot(counters, nonpr, marker="x", label='Non-private', color='#1f77b4')
ax.plot(counters, dfmbgu, marker="x", label='DFMBG (Advantaged)', color='#8c564b')
ax.plot(counters, t4, marker="x", label='t = 4', color='#7f7f7f')
ax.plot(counters, t10, marker="x", label='t = 10', color='#bcbd22')
plt.ylim(0.0,0.5)
ax.set_xlabel('R')
ax.set_ylabel('L2 error')
ax.set_title("Effect of Increasing R")
ax.legend()
plt.savefig('./plots/2c.png')

#Low dimension plotting (Figure 5)
counters = np.loadtxt("./results/synthetic_mean/counters-3.txt")
nonpr = np.loadtxt("./results/synthetic_mean/nonpr-3.txt")
naive = np.loadtxt("./results/synthetic_mean/naive-3.txt")
kv = np.loadtxt("./results/synthetic_mean/kv-3.txt")
dfmbg = np.loadtxt("./results/synthetic_mean/dfmbg-3.txt")
t1 = np.loadtxt("./results/synthetic_mean/t1-3.txt")
t2 = np.loadtxt("./results/synthetic_mean/t2-3.txt")
t3 = np.loadtxt("./results/synthetic_mean/t3-3.txt")
t4 = np.loadtxt("./results/synthetic_mean/t4-3.txt")

fig, ax = plt.subplots()
ax.plot(counters, nonpr, marker="x", label='Non-private', color='#1f77b4')
ax.plot(counters, naive, marker="x", label='Naive coordinatewise', color='#ff7f0e')
ax.plot(counters, kv, marker="x", label='KV', color='#2ca02c')
ax.plot(counters, dfmbg, marker="x", label='DFMBG', color='#8c564b')
ax.plot(counters, t1, marker="x", label='t = 1', color='#9467bd')
ax.plot(counters, t2, marker="x", label='t = 2', color='#d62728')
ax.set_xlabel('n')
ax.set_ylabel('L2 error')
ax.set_title("Low Dimensional Comparison")
ax.legend()
plt.savefig('./plots/3a.png')

fig, ax = plt.subplots()
ax.plot(counters, kv-nonpr, marker="x", label='KV', color='#2ca02c')
ax.plot(counters, dfmbg-nonpr, marker="x", label='DFMBG', color='#8c564b')
ax.plot(counters, t1-nonpr, marker="x", label='t = 1', color='#9467bd')
ax.plot(counters, t2-nonpr, marker="x", label='t = 2', color='#d62728')
plt.ylim(0.0,0.005)
ax.set_xlabel('n')
ax.set_ylabel('Excess L2 error')
ax.set_title("Low Dimensional Comparison")
ax.legend()
plt.savefig('./plots/3b.png')

#High dimension plotting (Figure 6)
counters = np.loadtxt("./results/synthetic_mean/counters-3.txt")
nonpr = np.loadtxt("./results/synthetic_mean/nonpr-4.txt")
dfmbg = np.loadtxt("./results/synthetic_mean/dfmbg-4.txt")
t2 = np.loadtxt("./results/synthetic_mean/t2-4.txt")


fig, ax = plt.subplots()
ax.plot(counters, nonpr, marker="x", label='Non-private', color='#1f77b4')
ax.plot(counters, dfmbg, marker="x", label='DFMBG', color='#8c564b')
ax.plot(counters, t2, marker="x", label='t = 2', color='#d62728')
ax.set_xlabel('n')
ax.set_ylabel('L2 error')
ax.set_title("High Dimensional Comparison")
ax.legend()
plt.savefig('./plots/4a.png')

fig, ax = plt.subplots()
ax.plot(counters, nonpr, marker="x", label='Non-private', color='#1f77b4')
ax.plot(counters, t2, marker="x", label='t = 2', color='#d62728')
ax.set_xlabel('n')
ax.set_ylabel('L2 error')
ax.set_title("High Dimensional Comparison")
ax.legend()
plt.savefig('./plots/4c.png')

startindex = 4
counters = counters[startindex:]
nonpr = nonpr[startindex:]
dfmbg = dfmbg[startindex:]
t2 = t2[startindex:]

fig, ax = plt.subplots()
ax.plot(counters, nonpr, marker="x", label='Non-private', color='#1f77b4')
ax.plot(counters, dfmbg, marker="x", label='DFMBG', color='#8c564b')
ax.plot(counters, t2, marker="x", label='t = 2', color='#d62728')
ax.set_xlabel('n')
ax.set_ylabel('L2 error')
ax.set_title("High Dimensional Comparison")
ax.legend()
plt.savefig('./plots/4b.png')


#Edge of ball comparisons
counters = np.loadtxt("./results/synthetic_mean/counters-5.txt")
nonpr = np.loadtxt("./results/synthetic_mean/nonpr-5.txt")
naive = np.loadtxt("./results/synthetic_mean/naive-5.txt")
kv = np.loadtxt("./results/synthetic_mean/kv-5.txt")
dfmbg = np.loadtxt("./results/synthetic_mean/dfmbg-5.txt")
t1 = np.loadtxt("./results/synthetic_mean/t1-5.txt")
t2 = np.loadtxt("./results/synthetic_mean/t2-5.txt")
t3 = np.loadtxt("./results/synthetic_mean/t3-5.txt")
t4 = np.loadtxt("./results/synthetic_mean/t4-5.txt")

fig, ax = plt.subplots()
ax.plot(counters, nonpr, marker="x", label='Non-private', color='#1f77b4')
ax.plot(counters, naive, marker="x", label='Naive coordinatewise', color='#ff7f0e')
ax.plot(counters, kv, marker="x", label='KV', color='#2ca02c')
ax.plot(counters, dfmbg, marker="x", label='DFMBG', color='#8c564b')
ax.plot(counters, t1, marker="x", label='t = 1', color='#9467bd')
ax.plot(counters, t2, marker="x", label='t = 2', color='#d62728')
plt.ylim(0.0,1.5)
ax.set_xlabel('n')
ax.set_ylabel('L2 error')
ax.set_title("Mean at Edge of Ball")
ax.legend()
plt.savefig('./plots/5a.png')

fig, ax = plt.subplots()
ax.plot(counters, nonpr, marker="x", label='Non-private', color='#1f77b4')
ax.plot(counters, dfmbg, marker="x", label='DFMBG', color='#8c564b')
ax.plot(counters, t2, marker="x", label='t = 2', color='#d62728')
plt.ylim(0.0,0.3)
ax.set_xlabel('n')
ax.set_ylabel('L2 error')
ax.set_title("Mean at Edge of Ball")
ax.legend()
plt.savefig('./plots/5b.png')

fig, ax = plt.subplots()
ax.plot(counters, dfmbg -nonpr, marker="x", label='DFMBG', color='#8c564b')
ax.plot(counters, t2 - nonpr, marker="x", label='t = 2', color='#d62728')
plt.ylim(0.0,0.1)
ax.set_xlabel('n')
ax.set_ylabel('Excess L2 error over Non-private')
ax.set_title("Mean at Edge of Ball")
ax.legend()
plt.savefig('./plots/5c.png')

#Vary privacy (Figure 7)
counters = np.loadtxt("./results/synthetic_mean/counters-6.txt")
nonpr = np.loadtxt("./results/synthetic_mean/nonpr-6.txt")
naive = np.loadtxt("./results/synthetic_mean/naive-6.txt")
kv = np.loadtxt("./results/synthetic_mean/kv-6.txt")
dfmbg = np.loadtxt("./results/synthetic_mean/dfmbg-6.txt")
t1 = np.loadtxt("./results/synthetic_mean/t1-6.txt")
t2 = np.loadtxt("./results/synthetic_mean/t2-6.txt")
t3 = np.loadtxt("./results/synthetic_mean/t3-6.txt")
t4 = np.loadtxt("./results/synthetic_mean/t4-6.txt")

fig, ax = plt.subplots()
plt.ylim(0.0,3.0)
plt.xscale("log")
ax.plot(counters, nonpr, marker="x", label='Non-private', color='#1f77b4')
ax.plot(counters, naive, marker="x", label='Naive coordinatewise', color='#ff7f0e')
ax.plot(counters, kv, marker="x", label='KV', color='#2ca02c')
ax.plot(counters, dfmbg, marker="x", label='DFMBG', color='#8c564b')
ax.plot(counters, t1, marker="x", label='t = 1', color='#9467bd')
ax.plot(counters, t2, marker="x", label='t = 2', color='#d62728')
ax.set_xlabel('rho')
ax.set_ylabel('L2 error')
ax.set_title("Effect of Varying Privacy")
ax.legend()
plt.savefig('./plots/6a.png')

lowprivlimit = 5
lpcounters = counters[:lowprivlimit]
lpnonpr = nonpr[:lowprivlimit]
lpt2 = t2[:lowprivlimit]
lpt3 = t3[:lowprivlimit]
lpt4 = t4[:lowprivlimit]

fig, ax = plt.subplots()
ax.plot(lpcounters, lpnonpr, marker="x", label='Non-private', color='#1f77b4')
ax.plot(lpcounters, lpt2, marker="x", label='t = 2', color='#d62728')
ax.plot(lpcounters, lpt3, marker="x", label='t = 3', color='#e377c2')
ax.plot(lpcounters, lpt4, marker="x", label='t = 4', color='#7f7f7f')
ax.set_xlabel('rho')
ax.set_ylabel('L2 error')
ax.set_title("Effect of Varying Privacy")
plt.xscale("log")
ax.legend()
plt.savefig('./plots/6b.png')

highprivlimit = 5
hpcounters = counters[highprivlimit:]
hpnonpr = nonpr[highprivlimit:]
hpkv = kv[highprivlimit:]
hpdfmbg = dfmbg[highprivlimit:]
hpt1 = t1[highprivlimit:]
hpt2 = t2[highprivlimit:]
hpt3 = t3[highprivlimit:]
hpt4 = t4[highprivlimit:]


fig, ax = plt.subplots()
plt.ylim(0.0,1.0)
ax.plot(hpcounters, hpnonpr, marker="x", label='Non-private', color='#1f77b4')
ax.plot(hpcounters, hpkv, marker="x", label='KV', color='#2ca02c')
ax.plot(hpcounters, hpdfmbg, marker="x", label='DFMBG', color='#8c564b')
ax.plot(hpcounters, hpt1, marker="x", label='t = 1', color='#9467bd')
ax.plot(hpcounters, hpt2, marker="x", label='t = 2', color='#d62728')
ax.plot(hpcounters, hpt3, marker="x", label='t = 3', color='#e377c2')
ax.set_xlabel('rho')
ax.set_ylabel('L2 error')
ax.set_title("Effect of Varying Privacy")
plt.xscale("log")
ax.legend()
plt.savefig('./plots/6c.png')


#Tail effect (Figure 8)
counters = np.loadtxt("./results/synthetic_mean/counters-7.txt")
nonpr = np.loadtxt("./results/synthetic_mean/gauss_nonpr-7.txt")
t2 = np.loadtxt("./results/synthetic_mean/gauss_pr-7.txt")
lap = np.loadtxt("./results/synthetic_mean/lap_pr-7.txt")
student = np.loadtxt("./results/synthetic_mean/t_pr-7.txt")


fig, ax = plt.subplots()
ax.plot(counters, t2-nonpr, marker="x", label='Private Gaussian', color='#d62728')
ax.plot(counters, lap-nonpr, marker="x", label='Private Laplace', color='#17becf')
ax.plot(counters, student-nonpr, marker="x", label='Private t-distribution', color='black')
ax.set_xlabel('n')
ax.set_ylabel('Excess L2 error (over non-private Gaussian)')
ax.set_title("Effect of Model Misspecification")
ax.legend()
plt.savefig('./plots/7a.png')
