#!/usr/bin/env python

# --------------------------------------------------------- #
#  Simple python script to calculate g from pendulum data.
#   
#  Author: Troels C. Petersen (Niels Bohr Institute)
#  Date:   06-09-2013
#
#  Run by: ./g_Pendulum.py
# --------------------------------------------------------- #

from ROOT import *        # Import ROOT libraries (to use ROOT)
import math               # Import MATH libraries (for extended MATH functions)
from array import array   # Import the concept of arrays (needed in TGraphErrors)

SavePlots = False         # Determining if plots are saved or not

gStyle.SetOptStat("emr")
gStyle.SetOptFit(1111)


# --------------------------------------------------------- #
# Read the data:
# --------------------------------------------------------- #

# Define four arrays containing count and time (with errors):
# They will be used to draw and fit a graph below!
countT1  = array("f")
ecountT1 = array("f")
timeT1   = array("f")
etimeT1  = array("f")

# Open the data file. The below lines may NOT read your data format,
# but should be modifiable to do so:
with open('data_Pendulum_take1.txt', 'r') as file:
    for line in file:
        line = line.strip()               # This removes all "codes" that may be put into the file.
        line = line.split(":")            # This splits the line at all ":", so that you now have
                                          # a list of numbers in "line", which you can choose from.

        # First entry is the counter, and third and fourth, minutes and seconds, so:
        countT1.append(float(line[0]))
        ecountT1.append(0.0)              # I set the error on the count to 0
        timeT1.append(60.0*float(line[2]) + float(line[3]))
        etimeT1.append(0.10)              # I guess the error on the time to be 0.1 second!

        # Always print to see, if this is reasonable:
        print "  Count number: %3.0f    time measurement (s): %5.2f"%(countT1[-1], timeT1[-1])


# --------------------------------------------------------- #
# Draw output:
# --------------------------------------------------------- #

canvasT1 = TCanvas( "canvasT1", "canvasT1", 50, 50, 1200, 600 )

graph_T1 = TGraphErrors(len(countT1), countT1, timeT1, ecountT1, etimeT1)

fit_T1 = TF1("fit_T1", "[0] + [1]*x", 0, len(countT1)+1)     # A line as fitting function.
fit_T1.SetParameters(0.0, 7.0)         # Set the starting values of [0] and [1] to 0 and 7.0
fit_T1.SetLineColor(kRed)              # Set the line color to red.
fit_T1.SetLineWidth(2)                 # Set tne line width to 4.

graph_T1.Fit("fit_T1", "R")            # Make the fit with the range as set.
graph_T1.Draw("AP")                    # Draw the axis and points of the graph.

if (SavePlots):
    canvas.SaveAs("fit_T1.png")        # Save plot as "fit_T1.png" (format follow extension name)

raw_input('Press Enter to exit')