#!/usr/bin/env python

"""
CERN@school - spectra maker
"""

# Import the argument parser functionality.
import argparse

# ...for the colour scaling.
from matplotlib import colorbar, colors

# ...for the plotting.
import matplotlib.pyplot as plt

# ...for nice text.
from matplotlib import rc

# Uncomment to use LaTeX for the plot text.
rc('font',**{'family':'serif','serif':['Computer Modern']})
rc('text', usetex=True)

#... for the MATH.
import numpy as np

# Custom imports.

from spectra import FluxSpectrum

#
# The main event!
#
if __name__ == "__main__":

    print("============================================")
    print("     CERN@school Spectra Maker (Python)     ")
    print("============================================")

    dbg = True

    # Get the datafile paths from the command line.
    parser = argparse.ArgumentParser()
    parser.add_argument("inputFilePath", help="The path of the ROOT Timepix data file.")
    parser.add_argument("outputPath", help="The path for the output files."         )
    args = parser.parse_args()

    # Process the input arguments.

    ## The full path of the input Mafalda ROOT file to be analysed.
    datapath = args.inputFilePath

    outputpath = args.outputPath

    # Update the user.
    if dbg:
        print
        print("* Input path is: '%s'" % (datapath))
        print("* The output is being written to: '%s'" % (outputpath))

    sf = open(datapath, 'r')
    ls = sf.readlines()
    sf.close()

    ## The energy bin line read in from the file.
    lEs = ls[0].split(',')

    # Remove the "Energy" label.
    del lEs[0]

    # Get the number of energy bins and remove it.
    numEs = int(lEs[0])
    del lEs[0]

    # Get the energy units.
    eUnit = lEs[-1]
    del lEs[-1]

    Es = {}

    # Get the actual energy bin edges.
    for i, E in enumerate(lEs):
        Es[i] = float(E.strip())

    if dbg:
        print("*")
        print("* Energy bins:")
        print("*--> Number of bins = %d (%d)" % (numEs, len(lEs)))
        print Es

    ## The integrated flux spectra line read in from the file.
    lIFSs = ls[1].split(',')

    # Remove the B field and L values.
    del lIFSs[0]
    del lIFSs[0]

    ## The integrated flux values.
    IFSs = {}

    for i, IFS in enumerate(lIFSs):
        IFSs[i] = float(IFS.strip())

    print lIFSs
    if dbg:
        print("*")
        print("* Integrated flux values:")
        print("*--> Number of values = %d" % (len(lIFSs)))
        print IFSs

    #plt.plot(Es.values(), IFSs.values())
    #
    #plt.yscale('log')
    #
    #plt.show()

    lowerbound = 0.2 # MeV
    upperbound = 7.0 # MeV

    fs = FluxSpectrum("proton", Es, IFSs, lowerbound, upperbound, True)

    s = fs.MakeSource()


    r = 5.0 # cm

    ## Surface area exposed = 1/2 pi r2
    Afac = 0.5 * np.pi * r * r
    print("* Area factor is %f" % (Afac))

    e_total_flux_within_bounds, e_total_flux = fs.GetTotalFlux()

    print("* Total electron flux = %f cm-2 s-1" % (e_total_flux_within_bounds))

    e_persec = Afac * e_total_flux_within_bounds # s-1

    print("* Electrons per second = %f" % (e_persec))

    frame_rate = 0.25 # s 4 Hz, 1/5 of 20 Hz minimum frame rate (LUCID DSS)

    e_per_frame = e_persec * frame_rate # Number of electrons.

    print("* Particles per frame = %f" % (e_per_frame))

    #print s

    fo = open('source_out.in','w')
    fo.write(s)
    fo.close()