# importing all needed packages
from scipy.optimize import leastsq
from numpy.lib import loadtxt
from math import pi, e
from pylab import *

data_filename_1 = "Total_data.txt" # file structured with two columns, separated by a tab
data_filename_2 = "Background_data.txt"

def load_data_1():
    """Return a two-dimensional array containing the data from the first column
       and the data from the second column. Preserve the structure when
       converting to an array.
    """

    return loadtxt(data_filename_1)

def load_data_2():
    """Return a two-dimensional array containing the data from the first column
       and the data from the second column. Preserve the structure when
       converting to an array.
    """

    return loadtxt(data_filename_2)

def peval(p, x):
    """Calculate the number of decays in each time interval in array times,
       based on parameters in list p.
    """

    return p[0] * e**(-p[1] * x) + p[2]

def load_data():
    load_data1 = load_data_1()
    load_data2 = load_data_2()

    ave2 = sum(data_2[:,1])/len(data_2[:,1])

    load_data_3 = zeros(shape(load_data1))
    count =0    #integer counter for current index in array
    for i in load_data1[:,1]:   #subtracts the total activity by the average background activity
        load_data_3[count][0]= count + 1
        load_data_3[count][1] = round(load_data1[count][1]-ave2, 2)
        count += 1
    #print load_data1
    #print load_data2
    #print load_data_3
    return load_data_3

# function to calculate the residuals at a given set of parameters p
# the second and third parameters are actual data
def residuals(p, x, y):
    return y - peval(p, x)


if __name__ == '__main__':
    #Total Activity Data
    data_1 = load_data_1() # data will be an array, with columns and rows
    times_1 = data_1[:, 0] # gets data satisfying requirements of: any row, column 0
    counts_1 = data_1[:, 1] # gets data satisfying requirements of: any row, column 1
    delta_counts_1 = (sum(counts_1)*20)**0.5
    print sum(counts_1)
    print delta_counts_1
    #Background Activity Data
    data_2 = load_data_2() # data will be an array, with columns and rows
    times_2 = data_2[:, 0] # gets data satisfying requirements of: any row, column 0
    counts_2 = data_2[:, 1] # gets data satisfying requirements of: any row, column 1
    delta_counts_2 = sum(counts_2)**0.5
    #Total - Background Activity Data
    data_3 = load_data()
    times_3 = data_3[:,0]
    counts_3 = data_3[:,1]
    delta_counts_3 = (sum(counts_1)+sum(counts_2))**0.5


    #Create estimates for initial parameters
    i = 0
    while counts_3[i] > max(counts_3)/e:
        i += 1
    tau = times_3[i]
    gamma = 1 / tau
    p0 = (max(counts_3), gamma, 0) # initial guesses for p[0], p[1] in function 'calculation'
    success = False
    while not success:
        p_final, cov_x, info, mesg, success = leastsq(residuals, p0, args=(times_3, counts_3), full_output = True)


   # Calculation of goodness of fit
    dof = len(times_3) - 3 # Degrees of freedom, we used 3 parameters
    chi_squared = sum( (residuals(p_final, times_3, counts_3)/delta_counts_3)**2 )


    # Use estimate of covariance to get stdev for each parameter
    stdev = [((cov_x[i, i])**0.5) for i in xrange(3)] # create list of sqrts of variance of each parameter
    # Print parameter results
#   #print "Parameters are:\nA = %0.5f +/- %0.5f\ngamma = %0.5f +/- %0.5f\nomega = %0.5f +/- %0.5f\nalpha = %0.5f +/- %0.5f\npsi = %0.5f +/- %0.5f\n" % (p_final[0], stdev[0], p_final[1], stdev[1], p_final[2], stdev[2], p_final[3], stdev[3], p_final[4], stdev[4])
    print "Chi squared over v is %0.3f" % (chi_squared / dof)
    print stdev
    print p_final



    #Data plotting
    #Ba-137m Decay Plot

xlabel("Time Intervals")
    ylabel("Activity (Bq)")
    title("Ba-137m Decay")

    plot(times_3, counts_3) # Plotting the actual data
#    print peval(p_final, times_3)
    plot(times_3, peval(p_final, times_3), 'ro-') # Plotting the fit
#    errorbar(times_3, counts_3, delta_counts_3) # Plotting the error bars
    #show () # Displays the graph

    #Total Activity Plot


    xlabel("Time Intervals")
    ylabel("Activity (Bq)")
    title("Total Activity")

    plot(times_1, counts_1)
    #show()


    #Background Activity Plot
    xlabel("Time Intervals")
    ylabel("Activity (Bq)")
    title("Background Activity")

    plot(times_2, counts_2)
    show()












    #Plotting of histogram of residuals
    #hist(residuals(p_final, times_3, counts_3))
    #show()

    # Residual plotting
    #xlabel("Time Intervals")
    #ylabel("Residual (Bq)")
    #plot(times_3, residuals(p_final, times_3, counts_3)) # Plotting residuals
    #show() # Displays the graph