# vim: foldmethod=marker
from pymule import *
import pymule.mpl_axes_aligner
from tabulate import tabulate

setup(folder='em2em_muone_paper/out.tar.bz2')

## Load w/o bandcut{{{
setup(obs='0')
lo = scaleset(mergefks(sigma('em2em0')), alpha**2*conv)


nloNoVP = scaleset(mergefks(
    sigma('em2emFEE'), sigma('em2emFEM'), sigma('em2emFMM'),
    sigma('em2emREE15'), sigma('em2emREE35'),
    sigma('em2emREM'), sigma('em2emRMM')
), alpha**3*conv)
nlo = scaleset(mergefks(
    sigma('em2emFEE'), sigma('em2emFEM'), sigma('em2emFMM'),
    sigma('em2emREE15'), sigma('em2emREE35'),
    sigma('em2emREM'), sigma('em2emRMM'),
    anyxi=sigma('em2emA', obs='2')
), alpha**3*conv)


nnloNoVP = scaleset(mergefks(
    sigma('em2emFFEEEE'),
    sigma('em2emRFEEEE15'), sigma('em2emRFEEEE35'),
    sigma('em2emRREEEE1516'), sigma('em2emRREEEE3536')
), alpha**4*conv)
nnlo = scaleset(mergefks(
    sigma('em2emFFEEEE'),
    sigma('em2emRFEEEE15'), sigma('em2emRFEEEE35'),
    sigma('em2emRREEEE1516'), sigma('em2emRREEEE3536'),
    sigma('em2emAFEE', obs='0'),
    sigma('em2emAREE15', obs='0'),
    sigma('em2emAREE35', obs='0'),
    anyxi1=sigma('em2emNFEE', obs='0'),
    anyxi2=sigma('em2emAA', obs='0')
), alpha**4*conv)
#####################################################################}}}
## Load w/ bandcut{{{
setup(obs='1')

nlocutNoVP = scaleset(mergefks(
    sigma('em2emFEE', obs='0'),
    sigma('em2emFEM', obs='0'), sigma('em2emFMM', obs='0'),
    sigma('em2emREE15'), sigma('em2emREE35'),
    sigma('em2emREM'), sigma('em2emRMM')
), alpha**3*conv)
nlocut = scaleset(mergefks(
    sigma('em2emFEE', obs='0'),
    sigma('em2emFEM', obs='0'), sigma('em2emFMM', obs='0'),
    sigma('em2emREE15'), sigma('em2emREE35'),
    sigma('em2emREM'), sigma('em2emRMM'),
    anyxi=sigma('em2emA', obs='2')
), alpha**3*conv)


nnlocutNoVP = scaleset(mergefks(
    sigma('em2emFFEEEE', obs='0'),
    sigma('em2emRFEEEE15'), sigma('em2emRFEEEE35'),
    sigma('em2emRREEEE1516'), sigma('em2emRREEEE3536')
), alpha**4*conv)
nnlocut = scaleset(mergefks(
    sigma('em2emFFEEEE', obs='0'),
    sigma('em2emRFEEEE15'), sigma('em2emRFEEEE35'),
    sigma('em2emRREEEE1516'), sigma('em2emRREEEE3536'),
    sigma('em2emAFEE', obs='0'), sigma('em2emAREE15'), sigma('em2emAREE35'),
    anyxi1=sigma('em2emAA', obs='0'), anyxi2=sigma('em2emNFEE', obs='0')
), alpha**4*conv)
#####################################################################}}}
## Print numbers{{{
print(
    "run time %f days" % (
        (
            lo['time'] + nlo['time'] + nnlo['time']
            + nlocut['time'] + nnlocut['time']
        ) / 3600. / 24.)
    )

print(tabulate([
    ['\\sigma^(0)', printnumber(lo['value']), "", "", ""],
    [
        '\\sigma^(1)',
        printnumber(nlo['value']),
        printnumber(nlocut['value']),

        printnumber(dividenumbers(nlo['value'], lo['value'])),
        printnumber(dividenumbers(nlocut['value'], lo['value']))
    ],
    [
        '\\sigma^(2)',
        "+"+printnumber(nnlo['value']),
        "+"+printnumber(nnlocut['value']),

        printnumber(dividenumbers(
            nnlo['value'], plusnumbers(lo['value'], nlo['value'])
        )),
        printnumber(dividenumbers(
            nnlocut['value'], plusnumbers(lo['value'], nlocut['value'])
        ))
    ],
    [
        '\\sigma_2',
        printnumber(plusnumbers(lo['value'], nlo['value'], nnlo['value'])),
        printnumber(plusnumbers(
            lo['value'], nlocut['value'], nnlocut['value']
        ))
    ]
], headers=['', 'S1', 'S2', 'S1k', 'S2k']))
#####################################################################}}}
## Make pictures{{{
### $\theta_e${{{
#### w/o cut{{{
fig, (ax1, ax2, ax3) = kplot(
    {
        'lo':    mergebins(scaleplot(lo['thetae'],   1e-3)[:221], 4),
        'nlo':   mergebins(scaleplot(nlo['thetae'],  1e-3)[:221], 4),
        'nnlo':  mergebins(scaleplot(nnlo['thetae'], 1e-3)[:221], 4),
        'nlo2':  mergebins(scaleplot(nloNoVP['thetae'],  1e-3)[:221], 4),
        'nnlo2': mergebins(scaleplot(nnloNoVP['thetae'], 1e-3)[:221], 4)
    },
    labelx="$\\theta_e\,/\,{\\rm mrad}$",
    labelsigma="$\\D\\sigma/\\D\\theta_e\ /\ {\\rm\\upmu b}$",
    legend={
        'lo': '$\\sigma^{(0)}$',
        'nlo': '$\\sigma^{(1)}$',
        'nnlo': '$\\sigma^{(2)}$',
        'nlo2': '$\\sigma^{(1)}_\\text{no VP}$',
        'nnlo2': '$\\sigma^{(2)}_{\\text{no VP}}$'
    },
    legendopts={'what': 'u', 'loc': 'lower right'}
)


fig.savefig('plots/theta-e-no-cut.pdf')
#################################################}}}
#### w/ cut{{{
fig, (ax1, ax2, ax3) = kplot(
    {
        'lo':    mergebins(scaleplot(lo['thetae'],      1e-3)[:221], 4),
        'nlo':   mergebins(scaleplot(nlocut['thetae'],  1e-3)[:221], 4),
        'nnlo':  mergebins(scaleplot(nnlocut['thetae'], 1e-3)[:221], 4),
        'nlo2':  mergebins(scaleplot(nlocutNoVP['thetae'],  1e-3)[:221], 4),
        'nnlo2': mergebins(scaleplot(nnlocutNoVP['thetae'], 1e-3)[:221], 4)
    },
    labelx="$\\theta_e\,/\,{\\rm mrad}$",
    labelsigma="$\\D\\sigma/\\D\\theta_e\ /\ {\\rm\\upmu b}$",
    legend={
        'lo': '$\\sigma^{(0)}$',
        'nlo': '$\\sigma^{(1)}$',
        'nnlo': '$\\sigma^{(2)}$',
        'nlo2': '$\\sigma^{(1)}_\\text{no VP}$',
        'nnlo2': '$\\sigma^{(2)}_{\\text{no VP}}$'
    },
    legendopts={'what': 'u', 'loc': 'lower right'}
)

fig.savefig('plots/theta-e-yes-cut.pdf')
#################################################}}}
###########################################################}}}
### $t_{\mu\mu}${{{
#### w/o cut{{{
fig, (ax1, ax2, ax3) = kplot(
    {
        'lo':   mergebins(scaleplot(lo['tmm'][15:296],   1e+6), 8),
        'nlo':  mergebins(scaleplot(nlo['tmm'][15:296],  1e+6), 8),
        'nnlo': mergebins(scaleplot(nnlo['tmm'][15:296], 1e+6), 8)
    },
    labelx="$t_{\\mu}\,/\,{\\rm GeV}^2$",
    labelsigma="$\\D\\sigma/\\D t_{\\mu}\ /\ {\\rm\\upmu b}$",
    legend={
        'lo': '$\\sigma^{(0)}$',
        'nlo': '$\\sigma^{(1)}$',
        'nnlo': '$\\sigma^{(2)}$'
    },
    legendopts={'what': 'u', 'loc': 'lower right'}
)
ax1.set_yscale('log')
ax2.set_ylim(0.95-1, 1.02-1)
ax3.set_ylim(0.999-1, 1.0005-1)
pymule.mpl_axes_aligner.yaxes(ax2, ax3, 0)
fig.savefig('plots/tmm-no-cut.pdf')
#################################################}}}
#### w/ cut{{{
fig, (ax1, ax2, ax3) = kplot(
    {
        'lo':   mergebins(scaleplot(lo['tmm'][16:242],      1e-3), 8),
        'nlo':  mergebins(scaleplot(nlocut['tmm'][16:242],  1e-3), 8),
        'nnlo': mergebins(scaleplot(nnlocut['tmm'][16:242], 1e-3), 8)
    },
    labelx="$t_{\\mu}\,/\,{\\rm GeV}^2$",
    labelsigma="$\\D\\sigma/\\D t_{\\mu}\ /\ {\\rm\\upmu b}$",
    legend={
        'lo': '$\\sigma^{(0)}$',
        'nlo': '$\\sigma^{(1)}$',
        'nnlo': '$\\sigma^{(2)}$'
    },
    legendopts={'what': 'u', 'loc': 'lower right'}
)
ax1.set_yscale('log')
ax3.set_ylim(0.9994-1, 1.0005-1)
pymule.mpl_axes_aligner.yaxes(ax2, ax3, 0)
fig.savefig('plots/tmm-yes-cut.pdf')
#################################################}}}
###########################################################}}}
#####################################################################}}}