#!/usr/bin/python

import itertools
import math
import numpy as np
import sys


def error_fmt(x, y, final_str):

    if abs(y) < 1e-8:
        sys.stdout.write('%12.5f' % x)
    else:
        ndigits = -int(math.floor(math.log10(y)))
        if ndigits >= 0:
            x_fmt = '%12.' + str(ndigits) + 'f'

            sys.stdout.write(x_fmt % x)
            sys.stdout.write('(%d)' % int(round(y*(10**ndigits))))
        if ndigits < 0:
            x_fmt = '%12d'

            sys.stdout.write(x_fmt % x)
            sys.stdout.write('(%d)' % int(round(y,0)))

    sys.stdout.write(final_str)


def calc_stats_ratio(num, denom, num_err, denom_err):
    mean = num/denom
    error = math.sqrt((num_err/num)**2 + (denom_err/denom)**2) * abs(mean)
    return error


def print_percent_err(x, x_err, y, y_err, final_str):

    try:
        p = 100 * x/y
        p_err = 100 * calc_stats_ratio(x, y, x_err, y_err)
        error_fmt(p, p_err, final_str)
    except ZeroDivisionError:
        sys.stdout.write(' - ' + final_str)


if __name__ == '__main__':

    UHF=np.loadtxt("UHF.dat", usecols=[0,1], unpack=True)
    d2=np.loadtxt("x2Results.dat", usecols=[0,12,13], unpack=True)
    d3=np.loadtxt("x3Results.dat", usecols=[0,12,13], unpack=True)
    d4=np.loadtxt("x4Results.dat", usecols=[0,20,21], unpack=True)
    fciqmc=np.loadtxt("fciqmcResults.dat", usecols=[0,2,3], unpack=True)
    UHFd=dict(zip(UHF[0],UHF[1]))
    d2U=[UHFd[x] for x in d2[0]]
    d3U=[UHFd[x] for x in d3[0]]
    d4U=[UHFd[x] for x in d4[0]]
    dFU=[UHFd[x] for x in fciqmc[0]]

    #---Package necessary data into a dictionary----

    r_values = d4[0]

    # D refers to CCSD, T to CCSDT, Q to CCSDTQ and F to FCIQMC
    energies = {}
    energies['D'] = {}
    energies['T'] = {}
    energies['Q'] = {}
    energies['F'] = {}

    for r in r_values:
        energies['D'][r] = {}
        energies['D'][r]['Tot'] = 0.0
        energies['D'][r]['Cor'] = 0.0
        energies['D'][r]['Err'] = 0.0
        energies['T'][r] = {}
        energies['T'][r]['Tot'] = 0.0
        energies['T'][r]['Cor'] = 0.0
        energies['T'][r]['Err'] = 0.0
        energies['Q'][r] = {}
        energies['Q'][r]['Tot'] = 0.0
        energies['Q'][r]['Cor'] = 0.0
        energies['Q'][r]['Err'] = 0.0
        energies['F'][r] = {}
        energies['F'][r]['Tot'] = 0.0
        energies['F'][r]['Cor'] = 0.0
        energies['F'][r]['Err'] = 0.0

    for a, b, c, d in itertools.izip(d2[0], d2[1], d2U, d2[2]):
        energies['D'][a]['Tot'] = b + c + 129.0
        energies['D'][a]['Cor'] = b
        energies['D'][a]['Err'] = d
    for a, b, c, d in itertools.izip(d3[0], d3[1], d3U, d3[2]):
        energies['T'][a]['Tot'] = b + c + 129.0
        energies['T'][a]['Cor'] = b
        energies['T'][a]['Err'] = d
    for a, b, c, d in itertools.izip(d3[0], d4[1], d4U, d4[2]):
        energies['Q'][a]['Tot'] = b + c + 129.0
        energies['Q'][a]['Cor'] = b
        energies['Q'][a]['Err'] = d
    for a, b, c, d in itertools.izip(fciqmc[0], fciqmc[1], dFU, fciqmc[2]):
        energies['F'][a]['Tot'] = b + c + 129.0
        energies['F'][a]['Cor'] = b
        energies['F'][a]['Err'] = d

    #---Print latex code for table----

    sys.stdout.write('\\begin{table*}[t]\n')
    sys.stdout.write('\\begin{center}\n')
    sys.stdout.write('{\\footnotesize\n')
    sys.stdout.write('\\begin{tabular}{@{\extracolsep{4pt}}lccccccc@{}}\n')
    sys.stdout.write('\\hline\n')
    sys.stdout.write('\\hline\n')
    sys.stdout.write('& \\multicolumn{4}{c}{(Total energy + $129E_{\\textrm{h}}$)/E$_{\\textrm{h}}$} & \\multicolumn{3}{c}{Correlation energy recovered ($\%$)} \\\ \n')
    sys.stdout.write('\\cline{2-5} \\cline{6-8}\n')
    sys.stdout.write('$R/\\textrm{\\AA}$  & CCSD-MC  & CCSDT-MC  & CCSDTQ-MC  & i-FCIQMC & CCSD-MC & CCSDT-MC & CCSDTQ-MC \\\ \n')
    sys.stdout.write('\\hline\n')

    # Data section
    for r in r_values:
        sys.stdout.write('%5.3f  &' % r)
        error_fmt(energies['D'][r]['Tot'], energies['D'][r]['Err'], '  &')
        error_fmt(energies['T'][r]['Tot'], energies['T'][r]['Err'], '  &')
        error_fmt(energies['Q'][r]['Tot'], energies['Q'][r]['Err'], '  &')
        error_fmt(energies['F'][r]['Tot'], energies['F'][r]['Err'], '  &')
        print_percent_err(energies['D'][r]['Cor'], energies['D'][r]['Err'], energies['F'][r]['Cor'], energies['F'][r]['Err'],'  &')
        print_percent_err(energies['T'][r]['Cor'], energies['T'][r]['Err'], energies['F'][r]['Cor'], energies['F'][r]['Err'],'  &')
        print_percent_err(energies['Q'][r]['Cor'], energies['Q'][r]['Err'], energies['F'][r]['Cor'], energies['F'][r]['Err'],'  \\\\')
        sys.stdout.write('\n')

    # End of data section
    sys.stdout.write('\\hline\n')
    sys.stdout.write('\\hline\n')
    sys.stdout.write('\\end{tabular}\n')
    sys.stdout.write('}\n')
    sys.stdout.write('\\end{center}\n')
    sys.stdout.write('\\end{table*}\n')
