#!/usr/bin/env python
import numpy as np

__doc__="""

%prog -m 1DME --ordering 0 --loop -o evol.pdf  PARAMETERFILE

Example parameter file:

m     -1      # logarithm, in [ev]
M1    12      # logarithm, in [GeV]
M2    12.6    # logarithm, in [GeV]
M3    13      # logarithm, in [GeV]
delta 213     # [deg]
a21    81     # [deg]
a31   476     # [deg]
x1     90     # [deg]
x2     87     # [deg]
x3    180     # [deg]
y1   -120     # [deg]
y2      0     # [deg]
y3   -120     # [deg]
t12    33.63  # [deg]
t13     8.52  # [deg]
t23    49.58  # [deg]
"""

def plotEvolution(data, f_out):
    import pylab
    pylab.plot(data[:,0], abs(data[:,1]), label="$N_{\\tau\\tau}$")
    pylab.plot(data[:,0], abs(data[:,2]), label="$N_{\mu\mu}$")
    pylab.plot(data[:,0], abs(data[:,3]), label="$N_{ee}$")
    pylab.plot(data[:,0], abs(data[:,4]), label="$\eta_B$")
    pylab.legend()
    pylab.xlabel("$z$")
    pylab.yscale("log")
    pylab.xscale("log")
    pylab.ylabel("$\left|N_{xx} \\right|, \eta_B$")
    pylab.tight_layout()
    pylab.savefig(f_out)




if __name__=="__main__":

    import optparse, os, sys
    op = optparse.OptionParser(usage=__doc__)
    op.add_option("-o", "--output",    dest="OUTPUT",      default=None, type=str, help="Output file name for evolution plots/data --- if not provided, only calculate etaB (default: %default)")
    op.add_option("-v", "--debug",     dest="DEBUG",       default=False, action="store_true", help="Turn on some debug messages")
    op.add_option("-m", "--model",     dest="MODEL",       default="1DME", help="Selection of of model (default: %default)")
    op.add_option("--zrange",          dest="ZRANGE",      default="0.1,100,1000", help="Ranges and steps of the evolution variable (default: %default)")
    op.add_option("--inv",             dest="INVORDERING", default=False, action='store_true', help="Use inverted mass ordering (default: %default)")
    op.add_option("--loop",            dest="LOOP",        default=False, action='store_true', help="Use loop-corrected Yukawa (default: %default)")
    opts, args = op.parse_args()


    if len(args)==0:
        print("No parameter space configuration given, exiting.")
        sys.exit(1)

    # Read parameter card and very explicit checks on parameter names
    # TODO also checks on values?
    import ulysses
    _, FIX = ulysses.readConfig(args[0])
    if len(FIX)!=16:
        print("Error, the number of parameters needs to be 16, user supplied {}, exiting".format(len(FIX)))
        sys.exit(1)

    pnames = ['m',  'M1', 'M2', 'M3', 'delta', 'a21', 'a31', 'x1', 'x2', 'x3', 'y1', 'y2', 'y3', 't12', 't13', 't23']


    for p in FIX.keys():
        if not p in pnames:
            print("Parameter {} in input file {} not recognised, exiting".format(p, args[0]))
            sys.exit(1)

    for p in pnames:
        if not p in FIX.keys():
            print("Parameter {} not provided in input file {}, exiting".format(p, args[0]))
            sys.exit(1)


    # Disect the zrange string
    zmin, zmax, zsteps = opts.ZRANGE.split(",")
    zmin=float(zmin)
    zmax=float(zmax)
    zsteps=int(zsteps)

    assert(zmin<zmax)
    assert(zsteps>0)


    LEPTO = ulysses.selectModel(opts.MODEL,
            zmin=zmin, zmax=zmax, zsteps=zsteps,
            ordering=int(opts.INVORDERING),
            loop=opts.LOOP,
            debug=opts.DEBUG
            )

    if opts.DEBUG:
        print(LEPTO)

    etaB = LEPTO(FIX)
    if opts.DEBUG:
        LEPTO.printParams()
        print(LEPTO.U)
    print("{}{}\n{}{}\n{}{}".format(
        "eta_b".ljust(      14), etaB,
        "Y_b".ljust(        14), etaB/7.03969,
        "Omega_b h^2".ljust(14), etaB/274.*1e10)
        )
    if opts.OUTPUT is not None:
        D=LEPTO.evolData
        E=np.sum(D[:,(1,2,3)], axis=1) * LEPTO.sphalfact
        DATA = np.column_stack((D,E))
        if opts.OUTPUT.endswith(".txt"):
            np.savetxt(opts.OUTPUT, DATA)
        elif opts.OUTPUT.endswith(".csv"):
            np.savetxt(opts.OUTPUT, DATA, delimiter=",")
        else:
            plotEvolution(DATA, opts.OUTPUT)
        if opts.DEBUG:
            print("Output written to {}".format(opts.OUTPUT))

