import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from ase.db import connect

# connect to the database
db = connect("dmc.db")

# plot details
mpl.rcParams["figure.figsize"] = (4,4)
mpl.rcParams["text.usetex"] = True
mpl.rcParams["font.family"] = "serif"
mpl.rcParams["font.serif"] = "Computer Modern Roman"

# colors
colors = ['#800080',
          '#aa0000',
          '#000080',
          '#008080',
          '#ac9d93',
          '#ff6600',
          '#00ffff',
          '#00ff00']

# fetch data from database

# CO
CO = []
labels = []
COerr = []
for row in db.select(type="adsorbate", adsorbate="CO"):
    CO.append(row.adsorption_free_energy)
    labels.append(row.element)
    COerr.append(row.uncertainty)

# CH3O
CH3O = []
CH3Oerr = []
for row in db.select(type="adsorbate", adsorbate="CH3O"):
    CH3O.append(row.adsorption_free_energy)
    CH3Oerr.append(row.uncertainty)

# CO goes on the x-axis
X = CO
Xerr = COerr

# CH3O goes on the y-axis
Y = CH3O
Yerr = CH3Oerr

# DMC reaction free energy
g_dmc = db.get(type="molecule", molecule="DMC").reaction_free_energy

# create the figure
fig = plt.figure()
ax = fig.add_subplot(111)

# scatter plot
ax.errorbar(X, Y, ls='None', xerr=Xerr, yerr=Yerr, color="black", marker='None', zorder=3, label=None)

# 0 V line
X = np.linspace(0, 2, 32, endpoint=True)
Y = [0 for x in X]
ax.plot(X, Y, ls="none", marker="^", markersize=2, markeredgecolor=colors[0], color=colors[0], zorder=2, label=r'$\Delta G_{\rm{M}}=0$')

# Replacing CO with CH3O line
X = np.linspace(-2, 0, 32, endpoint=False)
Y = [x for x in X]
ax.plot(X, Y, ls="none", marker="v", markersize=2, markeredgecolor=colors[0], color=colors[0], zorder=2, label=r'$\Delta G_{\rm{M}}=\Delta G_{\rm{CO}}$')

# DMC line
X = np.linspace(-2, 2, 64, endpoint=True)
Y = [-0.5*x + 0.5*g_dmc for x in X]
ax.plot(X, Y, ls="--", color=colors[1], zorder=2, label=r'$\Delta G_{\rm{M}}=\Delta G_{\rm{DMC}}^{\rm{des}}$')

# CH3OCO scaling
Y = [0.64*x+0.57 for x in X]
ax.plot(X, Y, ls="-.", color=colors[2], zorder=2, label=r'$\Delta G_{\rm{M}}=\Delta G_{\rm{MF}}$')

# CH2OH scaling
X = np.linspace(-2, 2, 64, endpoint=True)
Y = [0.55*x+1.09 for x in X]
ax.plot(X, Y, ls=':', color=colors[3], zorder=2, label=r'$\Delta G_{\rm{M}}=\Delta G_{\rm{HM}}$')

# Label the points
ax.text(CO[0]+0.1, CH3O[0]+0.07, labels[0], fontsize=10) # Ag
ax.text(CO[1]+0.05, CH3O[1]-0.14, labels[1], fontsize=10) # Au
ax.text(CO[2]+0.05, CH3O[2]+0.1, labels[2], fontsize=10) # Cu
ax.text(CO[3]+0.07, CH3O[3]+0.1, labels[3], fontsize=10) # Ir
ax.text(CO[4]+0.05, CH3O[4]+0.1, labels[4], fontsize=10) # Ni
ax.text(CO[5]+0.1, CH3O[5]+0.15, labels[5], fontsize=10) # Pd
ax.text(CO[6]+0.07, CH3O[6]+0.1, labels[6], fontsize=10) # Pt
ax.text(CO[7]-0.25, CH3O[7]-0.20, labels[7], fontsize=10) # Rh
ax.text(CO[8]+0.05, CH3O[8]-0.2, labels[8], fontsize=10) # Ru
ax.text(CO[9]+0.05, CH3O[9]-0.2, labels[9], fontsize=10) # Sn

# Fill
# DMC stays on the surface
X = np.linspace(-1.5, 1.5, 64, endpoint=True)
ax.fill_between(X, -1.5, [-0.5*x + 0.5*g_dmc for x in X], color=colors[4], alpha=0.5, edgecolor='none', zorder=1)
# Too strong CO adsorption
X = np.linspace(-2, -0.3, 64, endpoint=True)
ax.fill_between(X, 2, [-0.5*x + 0.5*g_dmc for x in X], color=colors[4], alpha=0.5, edgecolor='none', zorder=1)
# Bad selectivity due to CH2OHCO
X = np.linspace(-0.3, 0.5, 64, endpoint=True)
ax.fill_between(X, 2, [0.55*x+1.09 for x in X], color=colors[5], alpha=0.5, edgecolor='none', zorder=1)
# Selectivity to DMO
X = np.linspace(-0.3, 0.5, 64, endpoint=True)
ax.fill_between(X, [0.55*x+1.09 for x in X], [0.64*x+0.57 for x in X], color=colors[6], alpha=0.5, edgecolor='none', zorder=1)
# DMC zone
X = np.linspace(-0.3, 0.5, 64, endpoint=True)
ax.fill_between(X, [-0.5*x + 0.5*g_dmc for x in X], [0.64*x+0.57 for x in X], color=colors[7], alpha=0.5, edgecolor='none', zorder=1)
# Too weak CO binding
X = np.linspace(0.5, 1.5, 64, endpoint=True)
ax.fill_between(X, [-0.5*x + 0.5*g_dmc for x in X], 1.5, color=colors[4], alpha=0.5, edgecolor='none', zorder=1)

# Legend
ax.legend(loc="lower right", fontsize=10)

# Axes
ax.set_xlim(-1.5, 1.5)
ax.set_ylim(-1.5, 1.5)
ax.set_ylabel(r'$\Delta G_{\rm{M}}$ (eV)')
ax.set_xlabel(r'$\Delta G_{\rm{CO}}$ (eV)')

fig.tight_layout()
fig.savefig('thermodynamic_analysis.png', dpi=300)