#!/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,
         'text.fontsize': 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']
    population = d[0]['# H psips']
    shift = d[0]['Shift']
    e_num = d[0]['\sum H_0j N_j']
    e_denom = d[0]['N_0']
    hf_energy = e_num/e_denom

    return iters, population, shift, e_num, e_denom, hf_energy

def plot_data(iters, population, shift, e_num, e_denom, hf_energy):

    fig = pylab.figure()

    gs1 = gridspec.GridSpec(4, 1)
    pylab.subplot(gs1[0,0])
    gs1.update(hspace=0.15)
    pylab.xlim(0,15000)
    pylab.ylim(8000,14000)
    ax = pylab.gca()
    ax.yaxis.set_major_locator(pylab.MaxNLocator(4))
    pylab.plot(iters, population, 'b-', linewidth=0.8)
    ax.xaxis.set_major_formatter(pylab.NullFormatter())
    ax.annotate(r'$\textrm{(a)}$', xy=(0.05,0.80), xycoords='axes fraction')
    fig.text(0.02, 0.96, r'$\textrm{No. of particles}$', rotation='vertical', size=9)

    gs2 = gridspec.GridSpec(4, 1)
    gs2.update(hspace=0.15)
    pylab.subplot(gs2[1,0])
    pylab.xlim(0,15000)
    pylab.ylim(-115,-95)
    pylab.gca().axes.get_yaxis().set_ticks([-115, -110, -105, -100])
    ax = pylab.gca()
    pylab.plot(iters, e_num, 'r-', linewidth=0.8)
    ax.xaxis.set_major_formatter(pylab.NullFormatter())
    ax.annotate(r'$\textrm{(b)}$', xy=(0.05,0.80), xycoords='axes fraction')
    fig.text(0.03, 0.69, r'$E \; \textrm{num.}/E_\textrm{h}$', rotation='vertical', size=9)

    gs3 = gridspec.GridSpec(4, 1)
    gs3.update(hspace=0.15)
    pylab.subplot(gs3[2,0])
    pylab.xlim(0,15000)
    pylab.ylim(450,520)
    pylab.gca().axes.get_yaxis().set_ticks([460, 480, 500])
    ax = pylab.gca()
    pylab.plot(iters, e_denom, 'b-', linewidth=0.8)
    ax.xaxis.set_major_formatter(pylab.NullFormatter())
    ax.annotate(r'$\textrm{(c)}$', xy=(0.05,0.80), xycoords='axes fraction')
    fig.text(0.03, 0.48, r'$E \; \textrm{denom.}$', rotation='vertical', size=9)

    gs4 = gridspec.GridSpec(4, 1)
    gs4.update(hspace=0.15)
    pylab.subplot(gs4[3,0])
    pylab.xlim(0,15000)
    pylab.ylim(-0.24,-0.18)
    pylab.gca().axes.get_yaxis().set_ticks([-0.24, -0.21, -0.18])
    ax = pylab.gca()
    pylab.plot(iters, hf_energy, 'r-', linewidth=1.0, label=r'$\textrm{Projected energy}$')
    #pylab.legend(loc=4, handlelength=3)
    pylab.xlabel(r'$\textrm{Iteration}$')
    fig.text(0.02, 0.25, r'$E_{\textrm{corr}}/E_\textrm{h}$', rotation='vertical', size=9)
    ax.annotate(r'$\textrm{(d)}$', xy=(0.05,0.80), xycoords='axes fraction')

    f = pylab.gcf()
    f.subplots_adjust(left=0.15, bottom=0.13, right=0.97, top=0.97, hspace=0)

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

if __name__ == '__main__':
    data_file = 'hande.out'
    iters, population, shift, e_num, e_denom, hf_energy = extract_data(data_file)
    plot_data(iters, population, shift, e_num, e_denom, hf_energy)
