#!/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

    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()

    if args.save:
        # default \textwidth for a fullpage article in Latex is 16.50764 cm.
        # You can figure this out by compiling the following TeX document:
        #
        #    \documentclass{article}
        #    \usepackage{fullpage}
        #    \usepackage{layouts}
        #    \begin{document}
        #    textwidth in cm: \printinunitsof{cm}\prntlen{\textwidth}
        #    \end{document}

        width = 16.50764
        width /= 2.54 # cm -> inches
        # According to this page:
        # http://www-personal.umich.edu/~jpboyd/eng403_chap2_tuftegospel.pdf,
        # Tufte suggests an aspect ratio of 1.5 - 1.6.
        height = width/1.5
        FIGSIZE = (width,height)

        import matplotlib.pyplot as plt

        font = {'family':'serif', 'serif': ['computer modern roman']}
        plt.rc('font',**font)

        plt.rc('text', usetex=True)
    else:
        # on retina screens, the default plots are way too small
        # by using Qt5 and setting QT_AUTO_SCREEN_SCALE_FACTOR=1
        # Qt5 will scale everything using the dpi in ~/.Xresources
        import matplotlib
        matplotlib.use("Qt5Agg")

        import matplotlib.pyplot as plt

        # Default figure size. Currently set to my monitor width and height so that
        # things are properly formatted
        FIGSIZE = (13.78,7.48)

        # Make the defalt font bigger
        plt.rc('font', size=22)

    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(figsize=FIGSIZE)
        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(figsize=FIGSIZE)
        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(figsize=FIGSIZE)
        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()