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

__doc__="""

Scan of EtaB in a single variable

%prog -m 1DME --ordering 0 --loop -o scan.pdf  PARAMETERFILE -n 100

Example paramter file:

m     -3   -1 # logarithm, in [ev]
K1    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 plotEtaB(data, f_out, pname):
    import pylab

    pylab.plot(data[:,0], 1e10*data[:,1])
    pylab.xlabel(pname)
    pylab.ylabel("$\eta_B\cdot 10^{10}$")
    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="scan.pdf", type=str, help="Output file name for evolution plots/data (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("-n", "--n-scan",    dest="NSCAN",       default=30, type=int, help="Number of point to scan (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)

    # Make sure specified file actually exists
    if not os.path.exists(args[0]):
        print("Specified input file {} does not exist, exiting.".format(args[0]))

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

    pfile, gdict = ulysses.tools.parseArgs(args)

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

    # Read parameter card and very explicit checks on parameter names
    RNG, FIX, isCasIb = ulysses.readConfig(pfile)

    assert(len(RNG)==1)
    pscan = list(RNG.keys())[0]
    pmin, pmax = RNG[pscan]

    if isCasIb:
        if len(FIX) != len(LEPTO.pnames)-1:
            print("Error, the number of fixed parameters needs to be {}, user supplied {}, exiting".format(len(LEPTO.pnames)-1, len(FIX)))
            sys.exit(1)

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

        if pscan not in LEPTO.pnames:
            print("Scan-parameter {} in input file {} not recognised, exiting".format(pscan, args[0]))

    else:
        print("beware!")
        LEPTO.isCasasIbarrra = False

    if opts.DEBUG:
        print(LEPTO)

    PP = np.linspace(pmin, pmax, opts.NSCAN)
    EE = []
    print("Scanning {} in [{},{}] for {} values".format(pscan, pmin, pmax, opts.NSCAN))

    from progressbar import ProgressBar
    pbar = ProgressBar()

    for p in pbar(PP):
        FIX[pscan] = p
        etaB = LEPTO(FIX)
        EE.append(etaB)

    DATA = np.empty((opts.NSCAN,2))
    DATA[:,0]=PP
    DATA[:,1]=EE

    if opts.OUTPUT is not None:
        if opts.OUTPUT.endswith(".txt"):
            np.savetxt(opts.OUTPUT, DATA)
        elif opts.OUTPUT.endswith(".csv"):
            np.savetxt(opts.OUTPUT, DATA, delimiter=",")
        else:
            plotEtaB(DATA, opts.OUTPUT, pscan)
        if opts.DEBUG:
            print("Output written to {}".format(opts.OUTPUT))
