#!/usr/bin/python

import matplotlib
import numpy
import matplotlib.gridspec as gridspec
import pyhande

fig_size = (5,4)
matplotlib.rcParams.update(
        {'legend.borderpad':0.3,
         'legend.labelspacing':0.3,
         'legend.fontsize':'medium',
         'text.usetex':True,
         'axes.labelsize': 11,
         'font.size': 11,
         'legend.fontsize': 11,
         'xtick.labelsize': 11,
         'ytick.labelsize': 11,
         'figure.figsize': fig_size}
)

import pylab

def extract_data(data_file):
   
    (d, md) = pyhande.lazy.zeroT_qmc([data_file])

    # warning: this assumes that output file only contains one calculation. 
    iters = d[0]['iterations']
    e_num = d[0]['\sum H_0j N_j']
    e_denom = d[0]['N_0']
    energy = e_num/e_denom

    return iters, energy

def plot_data(iters_no, energy_no, iters_hub, energy_hub):

    fig = pylab.figure()

    gs1 = gridspec.GridSpec(2, 1)
    pylab.subplot(gs1[0,0])
    gs1.update(hspace=0.18)
    pylab.xlim(0,187100)
    pylab.ylim(-0.346,-0.341)
    ax = pylab.gca()
    ax.yaxis.set_major_locator(pylab.MaxNLocator(4))
    ax.axvline(x=93550,color='black',linestyle='--')
    pylab.plot(iters_no, energy_no, 'b-', linewidth=0.5)
    ax.annotate(r'$\textrm{(a)}$', xy=(0.85,0.80), xycoords='axes fraction')
    fig.text(0.02, 0.9, r'$\textrm{Correlation energy}/E_\textrm{h}$', rotation='vertical', size=9)

    gs2 = gridspec.GridSpec(2, 1)
    gs2.update(hspace=0.18)
    pylab.subplot(gs2[1,0])
    pylab.xlim(0,100000)
    pylab.ylim(-0.3302,-0.3292)
    #pylab.gca().axes.get_yaxis().set_ticks([-0.24, -0.21, -0.18])
    ax = pylab.gca()
    ax.yaxis.set_major_locator(pylab.MaxNLocator(4))
    ax.axvline(x=50000,color='black',linestyle='--')
    pylab.plot(iters_hub, energy_hub, 'r-', linewidth=0.5)
    pylab.xlabel(r'$\textrm{Iteration}$')
    fig.text(0.02, 0.45, r'$\textrm{Correlation energy}/t$', rotation='vertical', size=9)
    ax.annotate(r'$\textrm{(b)}$', xy=(0.85,0.80), xycoords='axes fraction')

    f = pylab.gcf()
    f.subplots_adjust(left=0.175, bottom=0.11, right=0.95, top=0.97, hspace=0)

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

if __name__ == '__main__':
    data_file_1 = 'no.out'
    data_file_2 = 'hubbard.out'
    iters_no, energy_no = extract_data(data_file_1)
    iters_hub, energy_hub = extract_data(data_file_2)
    plot_data(iters_no, energy_no, iters_hub, energy_hub)
