import matplotlib.pyplot as plt
import numpy as np
from scipy import stats
import pdb
from sklearn.linear_model import LinearRegression



opt_bal = np.array([1.89, 2.45, 3.89, 3.05, 1.45])
opt = np.array([1.11, 1.13, 3.11, 1.41, 1.87])

ci_bal = np.array([2.2, 1.93, 1.7, 2.24, 2.90])
ci = np.array([0.79, 1.11, 1.62, 1.32, 0.67])

fs = 15

corr_bal, _ = stats.pearsonr(opt_bal, ci_bal)
corr, _ = stats.pearsonr(opt, ci)
print(corr_bal, corr)

fig, ax = plt.subplots()
ax.scatter(opt_bal, ci_bal)
ax.scatter(opt, ci)
ax.legend(['BAL', 'UA-L2O'], fontsize=fs)

reg = LinearRegression().fit(opt_bal.reshape(-1, 1), ci_bal)
X = np.array([opt_bal.min()-0.5, opt_bal.max()+0.5])
ax.plot(X, X*reg.coef_+reg.intercept_)

reg = LinearRegression().fit(opt.reshape(-1, 1), ci)
X = np.array([opt.min()-0.5, opt.max()+0.5])
ax.plot(X, X*reg.coef_+reg.intercept_)


ax.set_xlabel(r'$E \Vert \mathrm{x}^* - \mathrm{x}_\mathrm{true} \Vert_2$', fontsize=fs)
ax.set_ylabel(r'$\mathrm{ub}_{0.9} - \mathrm{lb}_{0.9}$', fontsize=fs)


plt.tight_layout()
plt.savefig('./calibration_protein.pdf')

