#!/usr/bin/env python

# ----------------------------------------------------------------------------------- #
#
#  ROOT macro for reading text (ASCII) files and fitting data written.
#
#  The data is from a damped harmonic oscillator, in this case a spring hanging with a
#  weight but (in HarmOsc2) also a round piece of cardboard, which increases the drag.
#
#  The basic fit is to a simple damped harmonic oscillator, which has parameters:
#   1: Offset constant (it is not swinging around 0!).
#   2: Amplitude
#   3: Damping coefficient
#   4: Frequency
#   5: Phase
#
#  Author: Troels C. Petersen (NBI)
#  Email:  petersen@nbi.dk
#  Date:   1st of October 2013
#
# ----------------------------------------------------------------------------------- #

from ROOT import *
from array import array

# Setting what to be shown in statistics box:
gStyle.SetOptStat("emr")
gStyle.SetOptFit(1111)

verbose = True
Nverbose = 10




# ----------------------------------------------------------------------------------- #
# Define files, lists, counters and read data:
# ----------------------------------------------------------------------------------- #

file = "data_HarmOsc1.txt"
time0 = 0.0      # For ensuring that time starts at 0.0s!
time = []
dist = []
Nread = 0

# Loop over files and open them
with open( file, 'r' ) as f : 

    # Loop over lines and extract variables of interest
    for line in f:
        line = line.strip()
        columns = line.split()
        if (len(columns) == 2) :
            
            # Put the values read into lists:
            if (Nread == 0) : time0 = float(columns[0])
            time.append(float(columns[0]) - time0)
            dist.append(float(columns[1]))
            if (verbose and Nread < Nverbose) :
                print "  Time = %6.3f    Dist = %6.3f "%(time[-1], dist[-1])
            Nread += 1

f.close()
print "  Number of entries read: ", Nread




# ----------------------------------------------------------------------------------- #
# Plot the data:
# ----------------------------------------------------------------------------------- #

t  = array( 'f', [0.0]*Nread )
et = array( 'f', [0.0]*Nread )
d  = array( 'f', [0.0]*Nread )
ed = array( 'f', [0.0]*Nread )

# Fill the arrays from the lists:
# -------------------------------
for i in range( Nread ) : 
    t[i]  = time[i]
    et[i] = 0.0
    d[i]  = dist[i]
    ed[i] = 0.1              # The size of the error on each point - initially WRONG!
    if not ((0.0 < time[i] < 100.0) and (-5.0 < dist[i] < 5.0)) :
        print "Warning: Strange value for time and/or dist!"

# Make a graph of the data:
graph = TGraphErrors(Nread, t, d, et, ed)


# Make a white canvas and draw the data in:
c0 = TCanvas("c0","", 50, 50, 1200, 500)

graph.SetLineWidth(2)
graph.SetMarkerStyle(20)
graph.SetMarkerColor(kBlue)
graph.SetMarkerSize(0.5)


# You need to write this functional form of a damped harmonic oscillator,
# possibly with corrections:
def func_osc(x, par) :
    return par[0] + par[1]*x[0]

fit = TF1("fit", func_osc, 0.0, 38.06, 2)
fit.SetParameters(0.0, 0.0)
fit.SetLineColor(kRed)
fit.SetLineWidth(2)
fit.SetNpx(2000)     # This is just to make sure that the function is drawn smoothly enough!
graph.Fit("fit", "R")
graph.Draw("AP")





raw_input( ' ... Press enter to exit ... ' )

# ----------------------------------------------------------------------------------- #
# 
# Questions:
# ----------
#
# ---------------- First data file ---------------
#
#  1) Look at the data file and plot and see if you can by eye estimate the size of the
#     uncertainty of the points. It is not easy, but you should be able to get it to
#     within a factor of 2-3.
# 
#  2) Once you have tried, set it to 0.0025, and try to fit a damped harmonic oscillator.
#     You will have to write the function yourself, and in particularly, you will need
#     to set the initial parameters quite accurately, for the fit to work. Possibly add
#     more parameters to improve on the fit quality!
# 
#  3) Try to fit only the range 20s-22s, and see how significant the damping is. Would
#     you from only two seconds data be able to tell, that there is a damping, and if
#     so, by how many sigma?
#
# ---------------- Second data file ---------------
# 
#  4) Now consider the data file "data_HarmOsc2.txt", and set the uncertainty on the
#     distance to be 0.0025. Again, make a simple harmonic oscillator fit run.
#     Now your job is to expand on the fitting function and introduce terms to include
#     various effects and thus reduce the Chi2.
#     See if you in a fit of the range [0.005,36.005] can get a better Chi2 than 6600.
# 
#  5) If you were to measure the oscillation frequency for (infinitely) small oscillations,
#     what would your estimate be?
# 
#  6) Can you detect a change in the distance uncertainty as a function of time?
#     I.e. does the errors needed to get a reasonable/constant Chi2 get better or worse
#     if you fit small ranges of time early and late in the experiment.
# 
# ---------------- Third data file ---------------
# 
#  7) The data file "data_HarmOsc3.txt" contains a different type of oscillation. Can
#     you fit this, and determine to which degree it is linear. End the fit range when
#     the oscillations reach 1% of the initial amplitude.
# 
# ----------------------------------------------------------------------------------- #