#!/usr/bin/python

import matplotlib
import matplotlib.gridspec as gridspec
import numpy

fig_size = (4.0,4.0*(numpy.sqrt(5)-1.0)/2.5)

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

import pylab

def extract_data(data_file):

    f = open(data_file)

    pops = []
    energies = []
    errors = []

    for line in f:
        if '#' not in line:
            d = line.split()
            pops.append(float(d[1]))
            energies.append(float(d[5]))
            errors.append(float(d[6]))

    return(pops, energies, errors)

def plot_data(pops, energies, errors):

    fig = pylab.figure()

    pylab.xlim(2000, 180000)
    pylab.ylim(-0.2185, -0.2124)
    pylab.xscale('log')
    pylab.axhline(-0.217925, color='black', linestyle='dashed', linewidth=0.5, label=r'$\textrm{FCI}$')

    (_, caps, _) = pylab.errorbar(pops, energies, yerr=errors, color='blue', marker='x', markerfacecolor='blue', markeredgecolor='blue', markersize=3.0, mew=0.5, linewidth=0.5, capsize=1.5, label=r'$\textrm{i-FCIQMC}$')
    for cap in caps:
        cap.set_color('blue')
        cap.set_markeredgewidth(0.6)

    fig.text(0.53, 0.05, r'$\textrm{Particle population}$', ha='center', va='center', fontsize='medium')
    fig.text(0.04, 0.57, r'$\textrm{Correlation energy}/E_{\textrm{h}}$', ha='center', va='center', rotation='vertical', fontsize='medium')

    pylab.legend(numpoints=1)

    f = pylab.gcf()
    f.subplots_adjust(left=0.17, bottom=0.17, right=0.98, top=0.96, hspace=0)

    pylab.savefig('converge.eps', dpi=(1200))

if __name__ == '__main__':
    data_file = 'results.dat'
    pops, energies, errors = extract_data(data_file)
    plot_data(pops, energies, errors)
