#!/usr/bin/env python

"""
Make plots of the calibration constants. For giggles.
"""

# Import the code needed to manage files.
import os, inspect, glob, argparse

# Because maths.
import numpy

from scipy.optimize import curve_fit

import matplotlib.pyplot as plt

def readfile(filename):

    print("* Analysing file '%s'" % (filename))

    vals = []

    f = open(filename, "r")
    ls = f.readlines()
    print("*--> Number of lines = %d" % len(ls))

    f.close()

    for i, l in enumerate(ls):
        arow = l.strip().split("\t")
        #print("*----> Row %3d: %10d pixels." % (i, len(arow)))
        for p in arow:
            vals.append(float(p))

    return vals

def gauss(x, *p):
    A, mu, sigma = p
    return A*numpy.exp(-(x-mu)**2/(2.*sigma**2))

def fittogauss(vals, mylabel, muguess, sdguess):

    hist, bin_edges = numpy.histogram(vals, density=True, bins=500)

    bin_centres = (bin_edges[:-1] + bin_edges[1:])/2

    # Define model function to be used to fit to the data above:

    # p0 is the initial guess for the fitting coefficients (A, mu and sigma above)
    p0 = [1., muguess, sdguess]

    coeff, var_matrix = curve_fit(gauss, bin_centres, hist, p0=p0)

    # Get the fitted curve
    hist_fit = gauss(bin_centres, *coeff)

    #plt.hist(avals, bins=100)
    plt.plot(bin_centres, hist, label=mylabel)
    plt.plot(bin_centres, hist_fit, label='Fit')

    # Finally, lets get the fitting parameters, i.e. the mean and standard deviation:
    print("* For %s:" % (mylabel))
    print("*---> Fitted mean = %f" % (coeff[1]))
    print("*---> Fitted standard deviation = %f" % (coeff[2]))

    plt.show()

    return coeff[0], coeff[1], coeff[2]

#
# The main program.
#
if __name__=="__main__":
    print("============================================")
    print("  CERN@school Calibration Constant Plotter  ")
    print("============================================")

    # Get the datafile path from the command line
    parser = argparse.ArgumentParser()
    parser.add_argument("datapath", \
        help="Path to the folder containing your data."  )
    args = parser.parse_args()

    # Set the data file path.
    datapath = args.datapath
    # a
    avals = readfile(datapath + "/a.txt")
    ascale, a, a_sd = fittogauss(avals, "a", 2.0, 0.2)

    # b
    bvals = readfile(datapath + "/b.txt")
    bscale, b, b_sd = fittogauss(bvals, "b", 80.0, 5.0)

    # c
    cvals = readfile(datapath + "/c.txt")
    cscale, c, c_sd = fittogauss(cvals, "c", 270.0, 10.0)

    # d
    tvals = readfile(datapath + "/t.txt")
    tscale, t, t_sd = fittogauss(tvals, "t", 0.0, 1.0)