#!/usr/bin/python

import matplotlib.pyplot as plt
import numpy as np
import matplotlib
import matplotlib.gridspec as gridspec
import pylab

fig_size = (3.4,2.9)

matplotlib.rcParams.update(
        {'legend.borderpad':0.3,
         'legend.labelspacing':0.3,
         'text.usetex':True,
         'axes.linewidth' : 0.8,
         'axes.labelsize': 8,
         'text.fontsize': 8,
         'font.size': 8,
         'legend.fontsize': 8,
         'xtick.labelsize': 8,
         'xtick.major.width' : 0.2,
         'ytick.labelsize': 8,
         'ytick.major.width' : 0.2,
         'figure.figsize': fig_size}
)

UHF=np.loadtxt("UHF.dat", usecols=[0,1], unpack=True)
d2=np.loadtxt("x2Results.dat", usecols=[0,12,13], unpack=True)
d3=np.loadtxt("x3Results.dat", usecols=[0,12,13], unpack=True)
d4=np.loadtxt("x4Results.dat", usecols=[0,20,21], unpack=True)
fciqmc=np.loadtxt("fciqmcResults.dat", usecols=[0,2,3], unpack=True)
UHFd=dict(zip(UHF[0],UHF[1]))
d2U=[UHFd[x] for x in d2[0]]
d4U=[UHFd[x] for x in d4[0]]
dFU=[UHFd[x] for x in fciqmc[0]]
print "max err",max(max(d2[2]),max(d3[2]),max(d4[2]))

#matplotlib.rcParams['text.usetex'] = True
#font = {'size' : 8}
#fig, ax = plt.subplots(1,1,figsize=[3.34646, 2.50984252])
#matplotlib.rc('font', **font)

gs1 = gridspec.GridSpec(2, 1)
pylab.subplot(gs1[0,0])
gs1.update(hspace=0.25)
pylab.xlim(0.8, 2.5)
pylab.ylim(-129.65,-129.3)
ax = pylab.gca()
ax.xaxis.set_major_formatter(pylab.NullFormatter())
ax.yaxis.set_major_locator(pylab.MaxNLocator(4))
ax.set_ylabel(r'$\textrm{Energy}/E_{\textrm{h}}$')
pylab.errorbar(d2[0],d2[1]+d2U,d2[2],linewidth=0.5,ls='solid',marker="+",color='g',label=r'$\textrm{CCSD-MC}$')
pylab.errorbar(d3[0],d3[1]+UHF[1],d3[2],linewidth=0.5,ls='dashed',marker="x",color='r',label=r'$\textrm{CCSDT-MC}$')
pylab.errorbar(d4[0],d4[1]+d4U,d4[2],linewidth=0.5,ls='solid',marker="x",color='b',label=r'$\textrm{CCSDTQ-MC}$')
pylab.legend(markerscale=0.5,loc='lower right',fontsize=8)

gs2 = gridspec.GridSpec(2, 1)
gs2.update(hspace=0.25)
pylab.subplot(gs2[1,0])
pylab.xlim(0.8, 2.5)
pylab.ylim(-129.65,-129.3)
ax = pylab.gca()
ax.yaxis.set_major_locator(pylab.MaxNLocator(4))
ax.set_xlabel(r'$R/\textrm{\AA}$')
ax.set_ylabel(r'$\textrm{Energy}/E_{\textrm{h}}$')
pylab.errorbar(d4[0],d4[1]+d4U,d4[2],linewidth=0.5,ls='solid',marker="x",color='b',label=r'$\textrm{CCSDTQ-MC}$')
pylab.errorbar(fciqmc[0],fciqmc[1]+dFU,fciqmc[2],linewidth=0.5,ls='dashed',color='r',marker="+",label=r'$\textrm{i-FCIQMC}$')
pylab.legend(markerscale=0.5,loc='lower right',fontsize=8)

f = pylab.gcf()
f.subplots_adjust(left=0.18, bottom=0.13, right=0.96, top=0.92, hspace=0)

pylab.savefig('no.eps')
