import pyhande
import matplotlib.pyplot as plt
import pandas as pd
# Start averaging after the population has stabilised in all calculations.
for (out, title) in (('calcs/ifciqmc/hubbard_ifciqmc.out', 'iFCIQMC (integer)'),
                     ('calcs/ifciqmc/hubbard_ifciqmc_real.out', 'iFCIQMC (real)')):
    calcs = pyhande.lazy.std_analysis([out], 30000, extract_psips=True)
    opt_blocks = [calc.opt_block for calc in calcs]
    estimates = pd.concat([opt.stack() for opt in opt_blocks], keys=range(len(opt_blocks)), axis=1).T
    estimates = estimates[estimates['# H psips']['mean'] > 10**4]
    plt.errorbar(estimates[('# H psips', 'mean')].values, estimates[('Proj. Energy', 'mean')].values,
                 xerr=estimates[('# H psips', 'standard error')].values,
                 yerr=estimates[('Proj. Energy', 'standard error')].values, label=title)

fciqmc_calc = pyhande.lazy.std_analysis(['calcs/fciqmc/hubbard_fciqmc.out'], 30000)[0]
fciqmc_energy = fciqmc_calc.opt_block['mean']['Proj. Energy']
fciqmc_err = fciqmc_calc.opt_block['standard error']['Proj. Energy']
plt.axhline(fciqmc_energy, label='FCIQMC', color='b')
plt.axhspan(fciqmc_energy+fciqmc_err, fciqmc_energy-fciqmc_err, color='b', alpha=0.25)

plt.xscale('log')
plt.xlabel('# particles')
plt.ylabel('Projected energy / $t$')
plt.legend()
ax = plt.gca()
ax.get_yaxis().get_major_formatter().set_useOffset(False)