#!/usr/bin/python

import numpy
import matplotlib
import math

fig_size = (3.4,3.4*(numpy.sqrt(5)-1.0)/2.0)
matplotlib.rcParams.update(
        {'legend.borderpad':0.25,
         'legend.labelspacing':0.25,
         'legend.fontsize':'medium',
         'text.usetex':True,
         'axes.labelsize': 8,
         'font.size': 8,
         'legend.fontsize': 8,
         'xtick.labelsize': 8,
         'ytick.labelsize': 8,
         'figure.figsize': fig_size}
)
matplotlib.rcParams['text.usetex'] = True

import pylab

kpoints, H00_DFT, ccsd, ccsd_err, ccsdt, ccsdt_err, ccsdtq, ccsdtq_err, fciqmc, fciqmc_err, H00_HF, ccsd_Jirka_pyscf = numpy.genfromtxt("energies.dat", delimiter=',',unpack=True)
kpoints_pyscf, ccsd_pyscf = numpy.genfromtxt("pyscfpaper_figure1ccsd.dat", usecols=(0,2), unpack=True)

#Convert pyscf data from eV to Eh:
ccsd_pyscf = 0.0367493*ccsd_pyscf
#Subtract difference of H00 from HF orbs to H00 from DFT orbs
print ccsd, ccsdt, ccsdtq, fciqmc
ccsd = ccsd + H00_DFT - H00_HF
ccsdt = ccsdt + H00_DFT - H00_HF
ccsdtq = ccsdtq + H00_DFT - H00_HF
fciqmc = fciqmc + H00_DFT - H00_HF
print ccsd, ccsdt
print H00_DFT, H00_HF

f = pylab.gcf()
ax = pylab.gca()
ax.set_xlim([0.9,2.5])
ax.set_ylim([-0.26,-0.225])
(_, caps, _) = pylab.errorbar(kpoints**(1.0/3.0), ccsd, yerr=ccsd_err, color='blue', marker='x', markerfacecolor='blue', markeredgecolor='blue', markersize=3.0, mew=0.5, linestyle='--', linewidth=0.5, capsize=1.5, label=r'$\textrm{CCSD-MC}$')
for cap in caps:
    cap.set_color('blue')
    cap.set_markeredgewidth(0.6)

(_, caps, _) = pylab.errorbar(kpoints**(1.0/3.0), ccsdt, yerr=ccsdt_err, color='green', marker='.', markerfacecolor='green', markeredgecolor='green', markersize=3.0, mew=0.5, linestyle='-.', linewidth=0.5, capsize=1.5, label=r'$\textrm{CCSDT-MC}$')
for cap in caps:
    cap.set_color('green')
    cap.set_markeredgewidth(0.6)

(_, caps, _) = pylab.errorbar(kpoints**(1.0/3.0), ccsdtq, yerr=ccsdtq_err, color='cyan', marker='o', markerfacecolor='cyan', markeredgecolor='cyan', markersize=3.0, mew=0.5, linestyle='-', linewidth=0.5, capsize=1.5, label=r'$\textrm{CCSDTQ-MC}$')
for cap in caps:
    cap.set_color('cyan')
    cap.set_markeredgewidth(0.6)

(_, caps, _) = pylab.errorbar(kpoints**(1.0/3.0), fciqmc, yerr=fciqmc_err, color='C5', marker='8', markerfacecolor='C5', markeredgecolor='C5', markersize=3.0, mew=0.5, linestyle='-', linewidth=0.5, capsize=1.5, label=r'$\textrm{FCIQMC}$')
for cap in caps:
    cap.set_color('C5')
    cap.set_markeredgewidth(0.6)

pylab.plot(kpoints**(1.0/3.0), ccsd_Jirka_pyscf, color='C1', marker='v', markerfacecolor='C1', markeredgecolor='C1', markersize=3.0, mew=0.5, linestyle='-', linewidth=0.5, label=r'$\textrm{CCSD-PySCF}$')
pylab.plot(kpoints_pyscf, ccsd_pyscf, color='red', marker='1', markerfacecolor='red', markeredgecolor='red', markersize=3.0, mew=0.5, linestyle=':', linewidth=0.5, label=r'$\textrm{CCSD-McClain et al.}$')

ax.set_xlabel(r'$(\#\bf{k}\ \mathrm{points})^{1/3}$')
ax.set_ylabel(r'$E_\mathrm{tot.} - E_\mathrm{HF}/E_{\textrm{h}}$')
ax.legend(loc='best', numpoints=1, fontsize=7)
#https://matplotlib.org/users/usetex.html
pylab.rc('text', usetex=True)

f.subplots_adjust(left=0.18, bottom=0.18, right=0.96, top=0.97, hspace=0)

pylab.savefig('corr_diamond.eps', dpi=1200, bbox_inches='tight')
