#!/usr/bin/env python

# ----------------------------------------------------------------------------------- #
#
#  Python script for reading text (ASCII) files and fitting data written.
#
#  This exercise is about more advanced fitting of data. The outset is the simplest
#  possible function, and there is a step before having a function fitting, which
#  even remotely does the job, which is then expanded to accommodate all the features
#  present in the data fitted. 
#  The data is from a spring hanging with a weight, where the air drag gives a light
#  damping.
#
#  The formula for the position (d) as a function of time (t) of a simple (under)damped
#  harmonic oscillator can be found here: d(t) = A * sin(w*t+phase) * exp(-gamma*t)
#
#  The challenge is the make the best fit!
#
#  Author: Troels C. Petersen (NBI)
#  Email:  petersen@nbi.dk
#  Date:   4th of January 2018
#
# ----------------------------------------------------------------------------------- #

from __future__ import division, print_function
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from scipy import stats
from iminuit import Minuit
from probfit import Chi2Regression
plt.close('all')

# Function to create a nice string output
def nice_string_output(names, values, extra_spacing = 2):
    max_values = len(max(values, key=len))
    max_names = len(max(names, key=len))
    string = ""
    for name, value in zip(names, values):
        string += "{0:s} {1:>{spacing}} \n".format(name, value,
                   spacing = extra_spacing + max_values + max_names - len(name))
    return string[:-2]

# Function to draw parameters from a fit on the figure in one line
def draw_fit_parameters(minuit_obj, axis, x, y, title=None, color='k', chi2=None, ndof=None, prob=None) :
    names = []
    values = []
    if title is not None :
        names += [title]
        values += ['']
    if chi2 is not None and ndof is not None and prob is not None :
        names += ['Chi2 / ndof','Prob']
        values += ["{:.3f} / {:d}".format(chi2, ndof), "{:.3f}".format(prob)]
    names += minuit_obj.values.keys()
    values += ["{:.3f} +/- {:.3f}".format(minuit_obj.values[val], minuit_obj.errors[val]) for val in minuit_obj.values.keys()]
    axis.text(x, y, nice_string_output(names, values), family='monospace', transform=axis.transAxes, fontsize=10, verticalalignment='top', color=color)



# ----------------------------------------------------------------------------------- #
# Set variables:
# ----------------------------------------------------------------------------------- #
SavePlots = False
verbose = True
Nverbose = 10
tmax = 38.0       # Maximum of time range fitted
filename = "data_HarmOsc1.txt"

# Load time and distance, assign uncertainty:
# -------------------------------------- #
time, dist = np.loadtxt(filename, unpack=True)
time -= time[0]      # For ensuring that time starts at 0.0s!
edist = np.ones_like(dist)*0.0042              # The size of the error on each point - initially WRONG!

# Check loaded data
# -------------------------------------- #
if verbose :
    for i in range(Nverbose) :
        print("  Time = %6.3f    Dist = %6.3f "%(time[i], dist[i]))
print("  Number of entries read: %d    Time of last read: %6.3f"%(len(time), time[-1]))

# Sanity check
for i in np.where((time < -0.001) | (time > 100.0) | (dist < -5.0) | (dist > 5.0))[0] :
    print("Warning: Strange value for time and/or dist!", i, time[i], dist[i])

# Plot the data:
# -------------------------------------- #
fig1 = plt.figure(figsize=(16, 10))   # Make an empty figure
ax1 = fig1.add_axes((.1,.3,.8,.6))    # Add the top subfigure
ax1.set_ylabel("Position")

# Make a graph of the data:
ax1.errorbar(time, dist, edist, fmt='k_', label='data', ecolor='k', elinewidth=1, capsize=1, capthick=1)
ax1.set_xlim(time[0],time[-1])
ax1.set_ylim(top=ax1.get_ylim()[1]*2.5)

# Fit the data:
# -------------------------------------- #
mask = time < tmax

# Make a fitting function for the graph:
def fit1(x, p0, p1) : return p0 + p1 * x

FitObject1 = Chi2Regression(fit1, time[mask], dist[mask], edist[mask])
Minuit1 = Minuit(FitObject1, pedantic=False, p0=0, p1=1, print_level=0)
Minuit1.migrad()
chi2 = Minuit1.fval
ndof = FitObject1.ndof
prob = stats.chi2.sf(chi2, ndof)
x_fit = np.linspace(0, tmax, 1000)
y_fit = fit1(x_fit,*Minuit1.args)
ax1.plot(x_fit, y_fit, 'r-')
draw_fit_parameters(Minuit1, ax1, 0.05, 0.95, 'Simple harmonic oscillator:','k',chi2,ndof,prob)

# Calculate residuals:
# -------------------------------------- #
dd1 = dist-fit1(time,*Minuit1.args)

if sum(dd1 > 0.020) < Nverbose :
    for i in range( len(dist) ) : 
        if (abs(dd1[i]) > 0.020) :
            print("  Large residual:  t = {:6.3f}     d = {:6.3f}   dd1 = {:6.3f}".format(time[i], dist[i], dd1[i]))
else :
    print("\n  WARNING: {:d} ({:3.1f}%) large residuals!".format(sum(dd1 > 0.020), 100.0*sum(dd1 > 0.020)/len(dd1)))

# Draw the residuals:
# -----------------------------------------
# Draw residuals as function of time in bottom subfigure
ax1.get_yaxis().get_ticklabels()[0].set_visible(False)
ax1.get_yaxis().get_ticklabels()[1].set_visible(False) # Remove bottom y-tick on top subfigure to prevent overlapping ticks
ax2 = fig1.add_axes((.1,.1,.8,.2),sharex=ax1)   # Add bottom subfigure for residuals, and have its x-axis follow the top figure when zooming manually
ax2.set_xlim(time[0],time[-1])
ax2.set_xlabel("Time elapsed [s]")
ax2.set_ylabel("Position residual")

ax2.errorbar(time, dd1, edist, fmt='r_', ecolor='r', elinewidth=1, capsize=1, capthick=1)
ax2.plot((0,tmax),(0,0), 'w--', lw=1, zorder=10)

# Draw histograms of residuals in a new inset figure
axins = inset_axes(ax1, 3, 3 , loc=5, bbox_to_anchor=(0.9, 0.5), bbox_transform=ax1.figure.transFigure)
plt.yticks([],[])
axins.hist(dd1, bins=120, range=(-0.03, 0.03), histtype='step', linewidth=1, color='r')
axins.set_ylim(top=axins.get_ylim()[1]*1.5)

axins.text(0.40, 0.95, "Residuals", transform=axins.transAxes, fontsize=10, verticalalignment='top')
axins.text(0.01, 0.90, nice_string_output(['simple:','Mean','RMS'], ["","{:.3f}".format(dd1.mean()),"{:.3f}".format(dd1.std(ddof=1))]), family='monospace', transform=axins.transAxes, fontsize=10, verticalalignment='top', color='k')

# Finalize the figure
plt.show(block=False)
if (SavePlots) :
    fig1.savefig("Fit_HarmOsc1.pdf", dpi=600)





# Finally, ensure that the program does not termine (and the plot disappears), before you press enter:
try:
    __IPYTHON__
except:
    raw_input('Press Enter to exit')



# ----------------------------------------------------------------------------------- #
# 
# Looking at the plot of data from the file "data_HarmOsc1.txt", there is little doubt
# that it is oscillating, and looking closer you will also see the slight damping.
# 
# Questions:
# ----------
#  1) Look at the data file and plot and see if you can by eye (or simple fits) 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. Try for 5-10 minutes and discuss it with your
#     neighbor before reading on!
#
#            - - - - - -     5-10 minutes (success or failure) later - - - - - -
#
#     If you didn't know how to estimate this uncertainty, then try to zoom in on a very
#     small part of the curve, where it should be possible to fit it with a line, or
#     possibly a 2nd or 3rd degree polynomial. Do so, and since you know that for a small
#     enough range of the data, this will be a reasonable PDF to use, extract the error
#     from the RMS of the residuals, which is roughly equivalent to choosing the errors,
#     that gives you a reasonable Chi2, however, based only on a very short range of the
#     data. If you want to check this, then using this error for all points, see if it
#     still fits well in some other (small) range.
# 
#  2) Once you have tried, set the error to 0.0035, and try to fit a damped harmonic
#     oscillator. You will have to write the function yourself (you can expand on the
#     existing linear function). Also remember, that you have to give it some reasonable
#     starting values. Do so and run the fit before reading on.
#
#            - - - - - -     15-20 minutes (success or failure) later - - - - - -
#
#     Did you manage to write the fitting function? Also, did you remember to put in a
#     parameter taking care of the offset from zero? If not, then look below these
#     questions, and "borrow" the parts that you need. Now run this fit, and see what
#     happens.
#
#            - - - - - -     10-15 minutes (success or failure) later - - - - - -
#
#  3) Did the fit converge? I imagine that it didn't (thought it might have), and my first
#     guess on why (apart from copying the below wrongly) would be initial parameters. You
#     need to set the initial parameters quite accurately, for the fit to work. Think about
#     this, i.e. in a 5D parameter space, how hard it is to find just the correct frequency
#     when all you got is a Chi2 value, and being just 5% wrong or so will give you nothing
#     of value!
#     So now try to evaluate what some good initial values for the fit would be. You can
#     start by simply making educated guesses, but if that fails, you can instead draw the
#     function choosing some parameters, until it starts looking like the data you have.
#     Note that especially the frequency requires a good initial value.
#     Also, once you have the fit running, you might want to see what the residuals look
#     like.
#
#  4) 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?
# 
# 
# 
# ---------------- Now consider the data file "data_HarmOsc2.txt" ---------------
# 
#  5) Now consider the data file "data_HarmOsc2.txt", and set the uncertainty on the
#     distance to be 0.0029. 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.
#     Set your fit to the range [0.005,36.005], and see how low a Chi2 you can get.
# 
#  Advanced question:
#  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.
# 
# 
# 
# ---------------- Now consider the data file "data_HarmOsc3.txt" ---------------
# 
#  7) The data file "data_HarmOsc3.txt" contains a different type of damping in the
#     oscillation. Fit this, and determine at which point the oscillation stops.
# 
# 
# ----------------------------------------------------------------------------------- #
#


#  # You need to write this functional form of a damped harmonic oscillator:
#  def fit1(x, p0, p1, p2, p3, p4) :
#      return p0 + p1*np.sin(p2*x+p3)*np.exp(-p4*x)