#!/usr/bin/env python

# --------------------------------------------------------- #
#  Python/ROOT script to fit and illustrate period measurement.
#   
#  Author: Troels C. Petersen (Niels Bohr Institute)
#  Date:   06-09-2014
# --------------------------------------------------------- #

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)

Blinding  = False         # Determining if random offset is added to result for blinding
SavePlots = True         # Determining if plots are saved or not

gStyle.SetOptStat("")
gStyle.SetOptFit(0)
# gStyle.SetStatFormat("8.6g");

gStyle.SetStatX(0.49)
gStyle.SetStatY(0.88)

r = TRandom3()
r.SetSeed(1)

# Text to write on the plot:
text = TLatex()
text.SetNDC()
text.SetTextFont(72)
text.SetTextColor(1)




# --------------------------------------------------------- #
# A lot of analysis...
# --------------------------------------------------------- #


# --------------------------------------------------------- #
# Draw output - Troels time measurement number 1 (T1):
# --------------------------------------------------------- #

countT1  = array("f")
ecountT1 = array("f")
timeT1   = array("f")
etimeT1  = array("f")
timeresT1  = array("f")
etimeresT1 = array("f")

with open('data_Pendulum_Troels1.txt', 'r') as file:
    for line in file:
        line = line.strip()
        line = line.split(":")
        countT1.append(float(line[0]))
        ecountT1.append(0.0)
        timeT1.append(60.0*float(line[2]) + float(line[3]))
        etimeT1.append(0.06)         # Guess at the error
        print "  Count number: %3.0f    time measurement (s): %5.2f"%(countT1[-1], timeT1[-1])

canvasT1 = TCanvas( "canvasT1", "canvasT1", 50, 50, 900, 600 )
canvasT1.DrawFrame(0.0, -40.0, 26.0, 200.0)
graph_timeT1 = TGraphErrors(len(countT1), countT1, timeT1, ecountT1, etimeT1)
graph_timeT1.GetXaxis().SetRangeUser(0.0, 26.0)
graph_timeT1.GetYaxis().SetRangeUser(-40.0, 200.0)
graph_timeT1.SetTitle("")
graph_timeT1.GetXaxis().SetTitle("Measurement number")
graph_timeT1.GetYaxis().SetTitle("Time elapsed (s)")
graph_timeT1.SetMarkerStyle(20)
graph_timeT1.SetMarkerSize(0.8)
fit_timeT1 = TF1("fit_timeT1", "[0] + [1]*x", 0.5, len(countT1)+0.5)
fit_timeT1.SetParameters(0.0, 7.0)
fit_timeT1.SetParNames("Offset (s)", "Period (s)")
fit_timeT1.SetLineColor(kRed)
fit_timeT1.SetLineWidth(2)

graph_timeT1.Fit("fit_timeT1", "R")


# Find the residuals, and display them on the same graph with scale on right side.
# Scale values and errors up by factor, to match new scale (i.e. be visible):
factor = 250.0

hist_timeresT1 = TH1F("hist_timeresT1", "", 14, -0.14, 0.14)
for i in range (0, len(countT1)) :
    timeresT1.append(factor * (timeT1[i] - fit_timeT1.Eval(countT1[i])))
    etimeresT1.append(factor * etimeT1[i])
    hist_timeresT1.Fill(timeresT1[-1]/factor)
    print i, countT1[i], timeT1[i],  fit_timeT1.Eval(countT1[i]), timeresT1[-1], etimeresT1[-1]
graph_timeresT1 = TGraphErrors(len(countT1), countT1, timeresT1, ecountT1, etimeresT1)
graph_timeresT1.SetLineColor(kBlue)
graph_timeresT1.SetLineWidth(2)
graph_timeresT1.SetMarkerStyle(20)
graph_timeresT1.SetMarkerSize(0.8)

graph_timeT1.Draw("AP same")
graph_timeresT1.Draw("P same")



# Draw an axis on the right side:
min = -0.12
max =  0.12
axis = TGaxis(26.0, min*factor, 26.0, max*factor, min, max, 510, "+L");
# axis.SetTitle("Time residual (s)")
axis.SetTitleOffset(1.3)
axis.SetTitleColor(kBlue)
axis.CenterTitle()
axis.SetLineColor(kBlue);
axis.SetLineWidth(2);
axis.SetLabelColor(kBlue);
axis.Draw();

line = TLine(0.0, 0.0, 26.0, 0.0)
line.Draw()
line1 = TLine(0.0, factor*0.032, 26.0, factor*0.032)
line1.SetLineStyle(3)
line1.Draw()
line2 = TLine(0.0, -factor*0.032, 26.0, -factor*0.032)
line2.SetLineStyle(3)
line2.Draw()


# Text:
text.SetTextSize(0.030)
text.DrawLatex(0.14, 0.85, "Result of the fit:")
text.SetTextSize(0.045)
text.DrawLatex(0.14, 0.80, "Offset = %6.4f #pm %6.4f s"%(fit_timeT1.GetParameter(0), fit_timeT1.GetParError(0)))
text.SetTextColor(kRed)
text.DrawLatex(0.14, 0.75, "Period = %6.4f #pm %6.4f s"%(fit_timeT1.GetParameter(1), fit_timeT1.GetParError(1)))

text.SetTextColor(kBlack)
text.SetTextSize(0.030)
text.DrawLatex(0.14, 0.69, "Uncertainty on time measurements")
text.DrawLatex(0.14, 0.66, "obtained from RMS of residuals")
# text.SetTextSize(0.045)
# text.DrawLatex(0.14, 0.60, "#sigma(t) = 0.06 s")


text.SetTextSize(0.04)
text.SetTextColor(kRed)
text.DrawLatex(0.54, 0.84, "Time measurements (s)")

text.SetTextSize(0.04)
text.SetTextColor(kBlue)
text.DrawLatex(0.67, 0.15, "Time residuals (s)")

text.SetTextSize(0.022)
text.SetTextColor(kBlack)
text.DrawLatex(0.68, 0.675, "Distribution of time residuals")

text.SetTextSize(0.018)
text.SetTextColor(kBlack)
text.DrawLatex(0.69, 0.13, "Dashed lines show #pm 1#sigma")



# Draw histogram in its on pad:
pad = TPad("pad", "", 0.65, 0.30, 0.90, 0.67, 0, 0, 0);
pad.SetTopMargin(0.0);
pad.SetLineColor(0);
pad.SetFillColor(0);
pad.GetFrame().SetFillColor(0);
pad.GetFrame().SetBorderSize(0);
pad.GetFrame().SetLineColor(kWhite);
pad.Draw();
pad.cd(); 

fit_timeresT1 = TF1("fit_timeresT1", "gaus", -0.12, 0.12)
hist_timeresT1.Fit("fit_timeresT1", "RL")
hist_timeresT1.GetXaxis().SetTitle("Time residual (s)")
hist_timeresT1.GetYaxis().SetTitle("Frequency")
hist_timeresT1.Draw("e")


if (SavePlots):
    canvasT1.SaveAs("fit_period.pdf")


# --------------------------------------------------------- #
raw_input('Press Enter to exit')