#!/usr/bin/env python3
from pylab import *

sd = lambda x, y: np.sqrt(x**2+y**2)

def plot_comparison(dir1, dir2, outdir):

    f1d = np.loadtxt(f"{dir1}/out", unpack=True)
    f2d = np.loadtxt(f"{dir2}/out", unpack=True)
    s1d = np.loadtxt(f"{dir1}/mc.flx.std.spc", unpack=True, usecols=(4,5,6))
    s2d = np.loadtxt(f"{dir2}/mc.flx.std.spc", unpack=True, usecols=(4,5,6))

    z   = lambda f: f[0]
    edn = lambda f: f[1]
    eup = lambda f: f[2]

    diff = lambda f0, f1, f: f(f0)-f(f1)

    fsize=(12,8)
    figure(1,figsize=fsize).clf()
    fig, subs = subplots(2,2,num=1,sharey=True, sharex=False, figsize=fsize)

    subs[0][0].set_xlabel(r"$Edn$"+"[W/m2]")
    subs[0][0].errorbar(edn(f2d), z(f2d), xerr=edn(s2d)*3, linestyle='-', label=f'Edn {dir2}')
    subs[0][0].errorbar(edn(f1d), z(f1d), xerr=edn(s1d)*3, linestyle='-.', label=f'Edn {dir1}')

    subs[0][1].set_xlabel(r"$\Delta Edn$"+"[W/m2]")
    subs[0][1].errorbar(diff(f2d,f1d,edn), z(f2d), xerr=sd(edn(s1d),edn(s2d)), linestyle='-', label=f'Edn ({dir2} - {dir1})')

    subs[1][0].set_xlabel(r"$Eup$"+"[W/m2]")
    subs[1][0].errorbar(eup(f2d), z(f2d), xerr=eup(s2d)*3, linestyle='-', label=f'Eup {dir2}')
    subs[1][0].errorbar(eup(f1d), z(f1d), xerr=eup(s1d)*3, linestyle='-.', label=f'Eup {dir1}')

    subs[1][1].set_xlabel(r"$\Delta Eup$"+"[W/m2]")
    subs[1][1].errorbar(diff(f2d,f1d,eup), z(f2d), xerr=sd(eup(s1d),eup(s2d)), linestyle='-', label=f'Eup ({dir2} - {dir1})')

    [ s.set_ylabel("z [km]") for s in subs[:,0] ]
    [ s.legend(loc='best') for s in subs.flatten() ]
    tight_layout()
    plt.savefig(f'{outdir}/flx_comparison_{dir1}_{dir2}.pdf', bbox_inches='tight')

    [ s.set_ylim(0,13) for s in subs.flatten() ]
    plt.savefig(f'{outdir}/flx_comparison_lower13km_{dir1}_{dir2}.pdf', bbox_inches='tight')


def _main():
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('dir1')
    parser.add_argument('dir2')
    parser.add_argument('outdir')

    args = parser.parse_args()

    plot_comparison(args.dir1, args.dir2, args.outdir)


if __name__ == '__main__':
    _main()