#!/usr/bin/env python
"""
Script for plotting solar fluxes from
http://www-pnp.physics.ox.ac.uk/~barr/fluxfiles/0403i/index.html.
"""
from __future__ import print_function, division
import numpy as np
import os

def splitext(path):
    """
    Like os.path.splitext() except it returns the full extension if the
    filename has multiple extensions, for example:

        splitext('foo.tar.gz') -> 'foo', '.tar.gz'
    """
    full_root, full_ext = os.path.splitext(path)
    while True:
        root, ext = os.path.splitext(full_root)
        if ext:
            full_ext = ext + full_ext
            full_root = root
        else:
            break

    return full_root, full_ext

# Taken from https://raw.githubusercontent.com/mwaskom/seaborn/c73055b2a9d9830c6fbbace07127c370389d04dd/seaborn/utils.py
def despine(fig=None, ax=None, top=True, right=True, left=False,
            bottom=False, offset=None, trim=False):
    """Remove the top and right spines from plot(s).

    fig : matplotlib figure, optional
        Figure to despine all axes of, default uses current figure.
    ax : matplotlib axes, optional
        Specific axes object to despine.
    top, right, left, bottom : boolean, optional
        If True, remove that spine.
    offset : int or dict, optional
        Absolute distance, in points, spines should be moved away
        from the axes (negative values move spines inward). A single value
        applies to all spines; a dict can be used to set offset values per
        side.
    trim : bool, optional
        If True, limit spines to the smallest and largest major tick
        on each non-despined axis.

    Returns
    -------
    None

    """
    # Get references to the axes we want
    if fig is None and ax is None:
        axes = plt.gcf().axes
    elif fig is not None:
        axes = fig.axes
    elif ax is not None:
        axes = [ax]

    for ax_i in axes:
        for side in ["top", "right", "left", "bottom"]:
            # Toggle the spine objects
            is_visible = not locals()[side]
            ax_i.spines[side].set_visible(is_visible)
            if offset is not None and is_visible:
                try:
                    val = offset.get(side, 0)
                except AttributeError:
                    val = offset
                _set_spine_position(ax_i.spines[side], ('outward', val))

        # Potentially move the ticks
        if left and not right:
            maj_on = any(
                t.tick1line.get_visible()
                for t in ax_i.yaxis.majorTicks
            )
            min_on = any(
                t.tick1line.get_visible()
                for t in ax_i.yaxis.minorTicks
            )
            ax_i.yaxis.set_ticks_position("right")
            for t in ax_i.yaxis.majorTicks:
                t.tick2line.set_visible(maj_on)
            for t in ax_i.yaxis.minorTicks:
                t.tick2line.set_visible(min_on)

        if bottom and not top:
            maj_on = any(
                t.tick1line.get_visible()
                for t in ax_i.xaxis.majorTicks
            )
            min_on = any(
                t.tick1line.get_visible()
                for t in ax_i.xaxis.minorTicks
            )
            ax_i.xaxis.set_ticks_position("top")
            for t in ax_i.xaxis.majorTicks:
                t.tick2line.set_visible(maj_on)
            for t in ax_i.xaxis.minorTicks:
                t.tick2line.set_visible(min_on)

        if trim:
            # clip off the parts of the spines that extend past major ticks
            xticks = ax_i.get_xticks()
            if xticks.size:
                firsttick = np.compress(xticks >= min(ax_i.get_xlim()),
                                        xticks)[0]
                lasttick = np.compress(xticks <= max(ax_i.get_xlim()),
                                       xticks)[-1]
                ax_i.spines['bottom'].set_bounds(firsttick, lasttick)
                ax_i.spines['top'].set_bounds(firsttick, lasttick)
                newticks = xticks.compress(xticks <= lasttick)
                newticks = newticks.compress(newticks >= firsttick)
                ax_i.set_xticks(newticks)

                xticks_minor = [t for t in ax_i.get_xticks(minor=True) if firsttick < t < lasttick]
                ax_i.set_xticks(xticks_minor,minor=True)

            yticks = ax_i.get_yticks()
            if yticks.size:
                firsttick = np.compress(yticks >= min(ax_i.get_ylim()),
                                        yticks)[0]
                lasttick = np.compress(yticks <= max(ax_i.get_ylim()),
                                       yticks)[-1]
                ax_i.spines['left'].set_bounds(firsttick, lasttick)
                ax_i.spines['right'].set_bounds(firsttick, lasttick)
                newticks = yticks.compress(yticks <= lasttick)
                newticks = newticks.compress(newticks >= firsttick)
                ax_i.set_yticks(newticks)

                yticks_minor = [t for t in ax_i.get_yticks(minor=True) if firsttick < t < lasttick]
                ax_i.set_yticks(yticks_minor,minor=True)

# Data is in the form: Fluxes in bins of neutrino energy (equally spaced bins
# in logE with 10 bins per decade with the low edge of the first bin at 100
# MeV) and zenith angle (20 bins equally spaced in cos(zenith) with bin width
# 0.1), integrated over azimuth. Note logE means log to base e. Fluxes below 10
# GeV are from the 3D calculation.

# The files all have the same format. After the initial comment lines (starting
# with a # character), the files contain one line per bin. No smoothoing
# between bins has been done. The columns are:
# 
#     1. Energy at centre of bin in GeV
#     2. Zenith files: Cos(zenith) at centre of bin
#        Azimuth files: Azimuth at centre of bin (degrees)
#     3. Flux in dN/dlogE in /m**2/steradian/sec
#     4. MC Statistical error on the flux
#     5. Number of unweighted events entering the bin (not too useful)

if __name__ == '__main__':
    import argparse
    import matplotlib
    import glob

    parser = argparse.ArgumentParser("plot solar fluxes")
    parser.add_argument("filenames", nargs='+', help="filenames of flux files")
    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)

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

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

        plt.rc('text', usetex=True)

    fig = plt.figure(figsize=FIGSIZE)

    colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
    linestyles = ['-','--']

    def key(filename):
        head, tail = os.path.split(filename)

        k = 0
        if tail.startswith('fmax'):
            k += 1
        if 'nue' in tail:
            k += 10
        elif 'nbe' in tail:
            k += 20
        elif 'num' in tail:
            k += 30
        elif 'nbm' in tail:
            k += 40
        elif 'nut' in tail:
            k += 50
        elif 'nbt' in tail:
            k += 60

        return k

    for filename in sorted(args.filenames,key=key):
        head, tail = os.path.split(filename)

        print(filename)

        data = np.genfromtxt(filename)

        shape1 = len(np.unique(data[:,0]))

        x = data[:,0].reshape((-1,shape1))
        y = data[:,1].reshape((-1,shape1))
        z = data[:,2].reshape((-1,shape1))

        # Convert to MeV
        x *= 1000.0
        z /= 1000.0

        zbins = np.linspace(-1,1,21)
        dz = zbins[1] - zbins[0]

        x = x[0]
        # Integrate over cos(theta) and multiply by 2*pi to convert 3D flux to
        # a total flux
        y = np.sum(z*dz,axis=0)*2*np.pi

        if 'sno_nue' in tail:
            plt.plot(x,y,color=colors[0],linestyle=linestyles[0],label=r'$\nu_e$')
        elif 'sno_nbe' in tail:
            plt.plot(x,y,color=colors[0],linestyle=linestyles[1],label=r'$\overline{\nu}_e$')
        elif 'sno_num' in tail:
            plt.plot(x,y,color=colors[1],linestyle=linestyles[0],label=r'$\nu_\mu$')
        elif 'sno_nbm' in tail:
            plt.plot(x,y,color=colors[1],linestyle=linestyles[1],label=r'$\overline{\nu}_\mu$')
        elif 'sno_nut' in tail:
            plt.plot(x,y,color=colors[2],linestyle=linestyles[0],label=r'$\nu_\tau$')
        elif 'sno_nbt' in tail:
            plt.plot(x,y,color=colors[2],linestyle=linestyles[1],label=r'$\overline{\nu}_\tau$')

    plt.gca().set_xscale("log")
    plt.gca().set_yscale("log")
    despine(fig,trim=True)
    plt.xlabel("$E$ (MeV)")
    plt.ylabel(r"$\mathrm{d}\Phi/\mathrm{d}E$ (1/$\mathrm{m}^2$/sec/MeV)")
    plt.legend()
    plt.tight_layout()

    if args.save:
        plt.savefig("irc01_atmospheric_flux.pdf")
        plt.savefig("irc01_atmospheric_flux.eps")

    plt.show()