#!/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
from sddm.plot import despine
from sddm import splitext

# 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
    from sddm import setup_matplotlib

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

    setup_matplotlib(args.save)

    import matplotlib.pyplot as plt

    fig = plt.figure()

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