from matplotlib.offsetbox import TextArea, DrawingArea, OffsetImage, AnnotationBbox

import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import pandas as pd
import numpy as np
from jolymer import pymol_utility as pu
from jolymer import jocolors as joc

plt.style.use('../jowebstyle.mplstyle')

# Protein charge of folded and unfolded state as a function of pH
# The pI is  9.66 (folded) and  9.54 (unfolded)
phdf = pd.read_csv("protein-charge.csv", sep="\s+")

# Uncomment to recreate all the electrostatic maps:
# for pH in range(15):
#     pu.surfacecharge('1s81', pH=pH)


fig, ax = plt.subplots(figsize=(4, 4))
ax.set_prop_cycle(color=[joc.Magenta, joc.Cyan, joc.GreenYellow],
                  marker=['o', '+', 'x'])

ax.plot(phdf.pH, phdf.unfolded, label='unfolded')
ax.plot(phdf.pH, phdf.folded, label='folded')
# for pH in [2,8,10,12]:
showph = [1, 4, 6, 8, 10, 12]
for incex, row  in phdf.iterrows():
    ph = row.pH
    if ph in showph:
        # imsize = 0.6
        im = plt.imread(f'1s81_ph{int(ph)}.png')
        im_arr = np.array(im)
        imagebox = OffsetImage(im_arr, zoom=0.15)
        ab = AnnotationBbox(imagebox, (row.pH, row.folded),
                            bboxprops={'edgecolor':'none','alpha':0.1})
        ax.add_artist(ab)

plt.grid()
ax.legend()
ax.set_xlabel('pH')
ax.set_ylabel('protein charge [$k_BT/e$]')
ax.set_ylim(-25, 25)
# plt.draw()
# plt.savefig('add_picture_matplotlib_figure.png',bbox_inches='tight')
plt.savefig('ph-charge.svg')
plt.show()