#!/usr/bin/env python
"""
Script to plot the probabilities for atmospheric neutrino oscillations. To run
it:

    $ ./plot-atmospheric-oscillations nue_osc_prob.txt
"""
from __future__ import print_function, division
import numpy as np

if __name__ == '__main__':
    import argparse
    from os.path import split, splitext
    from sddm import setup_matplotlib

    parser = argparse.ArgumentParser("script to plot atmospheric oscillations")
    parser.add_argument("filenames", nargs='+', help="oscillation probability filenames")
    parser.add_argument("--save", action="store_true", default=False, help="save plots")
    args = parser.parse_args()

    setup_matplotlib(args.save)

    import matplotlib.pyplot as plt

    for filename in args.filenames:
        head, tail = split(filename)
        root, ext = splitext(tail)

        e, z, pnue, pnum, pnut = np.genfromtxt(filename).T

        shape0 = len(np.unique(e))

        ee = e.reshape((shape0,-1))
        zz = z.reshape((shape0,-1))
        pnue = pnue.reshape((shape0,-1))
        pnum = pnum.reshape((shape0,-1))
        pnut = pnut.reshape((shape0,-1))

        levels = np.linspace(0,1,101)

        plt.figure()
        plt.contourf(ee,zz,pnue,levels=levels)
        plt.gca().set_xscale('log')
        plt.xlabel("Energy (GeV)")
        plt.ylabel("Cos(Zenith)")
        plt.colorbar()
        plt.tight_layout()
        if args.save:
            plt.savefig("%s_nue.pdf" % root)
            plt.savefig("%s_nue.eps" % root)
        else:
            plt.title(r"Probability to oscillate to $\nu_e$")
        plt.figure()
        plt.contourf(ee,zz,pnum,levels=levels)
        plt.gca().set_xscale('log')
        plt.xlabel("Energy (GeV)")
        plt.ylabel("Cos(Zenith)")
        plt.colorbar()
        plt.tight_layout()
        if args.save:
            plt.savefig("%s_num.pdf" % root)
            plt.savefig("%s_num.eps" % root)
        else:
            plt.title(r"Probability to oscillate to $\nu_\mu$")
        plt.figure()
        plt.contourf(ee,zz,pnut,levels=levels)
        plt.gca().set_xscale('log')
        plt.xlabel("Energy (GeV)")
        plt.ylabel("Cos(Zenith)")
        plt.colorbar()
        plt.tight_layout()
        if args.save:
            plt.savefig("%s_nut.pdf" % root)
            plt.savefig("%s_nut.eps" % root)
        else:
            plt.title(r"Probability to oscillate to $\nu_\tau$")

    plt.show()