#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt
# let's try some fitting
import pylab
from scipy import optimize
# to fit a normalized gaussian
from scipy.stats import norm
# scipy curve_fit:
# https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.curve_fit.html
# symbols in matplotlib:
# http://matplotlib.org/users/mathtext.html

# I haven't worked out yet, if there is in inbuild function
# to fit with scipy
def gauss(x, mn, sgm, c):
    return c * pylab.exp(-(x - mn)**2.0 / (2 * sgm**2))

def load_file(filename):
    """reads iputfile and puts data into an numpy array"""
    data = []
    inputfile = open(filename, "r")
    for line in inputfile.readlines():
        rawdata = line.strip()
        data.append(int(rawdata))
    return np.array(data)


def main():
    data = load_file("radio-counter-2017.txt")
    print "There are %u entries in the file." %len(data)
    fig, ax = plt.subplots(1,2,figsize=(12,6)) # width, height
    fig.canvas.set_window_title("Task 7")
    
    # raw data
    ncounts_raw, hbins, _ =  ax[0].hist(data, bins = 10, range = (80, 140), alpha = 0.3)
    # use middle of the bins to draw errors
    mid = 0.5*(hbins[1:] + hbins[:-1])
    yerrors = np.sqrt(ncounts_raw)
    ax[0].errorbar(mid, ncounts_raw, yerr=yerrors, fmt='None', ecolor='k')
    ax[0].set_ylim(ymin=0.0)
    ax[0].set_title("raw data")  
    ax[0].set_xlabel("number of recorded decays")
    ax[0].set_ylabel("frequency")
    # let's try and fit using scipy
    # x is the middle of the bins
    fit_input_x = mid
    # array with number of entries per bin  
    fit_input_y = ncounts_raw
    initial_mean_estimate = np.mean(mid)
    popt, pcov = optimize.curve_fit(gauss,
                                    fit_input_x,
                                    fit_input_y,
                                    p0=(initial_mean_estimate, 1, 1)) # 1 is default
       
    x_fit = pylab.linspace(fit_input_x[0], fit_input_x[-1], 50)
    y_fit = gauss(x_fit, *popt)

    ax[0].plot(x_fit, y_fit, lw=2, color="r")
    mean,sigma,const = popt
    print pcov
    perr = np.sqrt(np.diag(pcov))
    mean_err, sigma_err, const_err = perr
    print perr
           
    ax[0].text(120, 300, '$\\mu$=%.2f$\\pm$%.2f\n$\\sigma$=%.2f$\\pm$%.2f\nc=%.2f$\\pm$%.2f'
                %(mean, mean_err, sigma, sigma_err, const, const_err))
  
    
    
    # normalized histogram 
    ncounts_norm, hbins, _ = ax[1].hist(data, bins = 10, range = (80, 140), normed = 1, alpha = 0.3)
    mid = 0.5*(hbins[1:] + hbins[:-1])
    # len(hbins) -1 -> there seems to be an overflow bin
    bin_size = (ax[1].get_xlim()[1] - ax[1].get_xlim()[0])/(len(hbins) -1)
    # make sure both histograms are plotted with equivalent y-ranges
    y_max = ax[0].get_ylim()[1]/len(data)/bin_size
    yerrors = np.sqrt(ncounts_raw)/len(data)/bin_size
    ax[1].errorbar(mid, ncounts_norm, yerr=yerrors, fmt='None', ecolor='k')
    ax[1].set_ylim(ymin=0.0)
    ax[1].set_ylim(ymax=y_max)
    ax[1].set_title("normalised")
    ax[1].set_xlabel("number of recorded decays")
    ax[1].set_ylabel("frequency")
    # fitting
    fit_input_y = ncounts_norm
    initial_width_estimate = (mid[-1] - mid[0])/2.0
    popt, pcov = optimize.curve_fit(norm.pdf,
                                    fit_input_x,
                                    fit_input_y,
                                    p0=(initial_mean_estimate, initial_width_estimate))
 
    # plot fit  
    mean,sigma = popt
    # error from pcov
    perr = np.sqrt(np.diag(pcov))
    mean_err, sigma_err = perr
    x_fit = pylab.linspace(fit_input_x[0], fit_input_x[-1], 50)
    y_fit = norm.pdf(x_fit, *popt)
    ax[1].plot(x_fit, y_fit, lw=2, color="r")
    ax[1].text(120, 0.035, '$\\mu$=%.2f$\\pm$%.2f\n$\\sigma$=%.2f$\\pm$%.2f'
               %(mean, mean_err, sigma, sigma_err))
    # end of fitting

    plt.tight_layout()
    plt.show()    
    fig.savefig("task7.png")    


if __name__ == '__main__':
    main()