#!/usr/bin/env python

# ----------------------------------------------------------------------------------- #
#  Python macro for showing how to numerically solve for Delta(theta) in the ball on
#  incline experiment:  a_normal / sin(theta + dtheta) = a_reverse / sin(theta - dtheta)
#
#  Author: Troels C. Petersen (petersen@nbi.dk),
#  Date:   29th of November 2017
# ----------------------------------------------------------------------------------- #

from __future__ import print_function, division        #
import numpy as np                                     # Matlab like syntax for linear algebra and functions
import matplotlib.pyplot as plt                        # Plots and figures like you know them from Matlab
from iminuit import Minuit, describe, Struct
from probfit import BinnedLH, Chi2Regression, Extended # Helper tool for fitting
from math import sqrt, cos, sin, tan, atan, pi

r = np.random
r.seed(42)

# Degrees to radians:
deg2rad   = 2.0 * pi / 360.0


# ----------------------------------------------------------------------------------- #
# The binary search:
# ----------------------------------------------------------------------------------- #
# Assume that you have done two "runs": One in the NORMAL direction, and one turning the
# setup around to the REVERSE direction. You have fitted each of these for the acceleration,
# and you have also measured the angle theta (but not dtheta) using trigonometry.
#
# Now you're faced with solving the below equation to obtain dtheta (with errors propagated):
#          a_normal / sin(theta + dtheta) = a_reverse / sin(theta - dtheta)

theta     = [6.705 * deg2rad, 0.10 * deg2rad]
a_normal  = [0.675, 0.008]
a_reverse = [0.776, 0.008]

# The difference between the two accelerations is due to the table, which is not level,
# but has an unknown angle, dtheta. In order to determine dtheta, a binary search is
# performed, as this the analytical solution is "hard".

# We assume that dtheta is small, i.e. in the range [0, 0.1] radians, and we denote the
# difference between the two sides of the equation by "dist". We want to get this distance
# to be smaller than 10^-6 (max_dist_allowed):
max_dist_allowed = 0.000001
dist = 99.9       # Initial distance just needs to be above the lower limit
dtheta = 0.02     # Initial value
ddtheta = 0.01    # Initial step
N = 0             # Number of iterations

print("\n  -------------- Results from each iteration of a single example -----------------")

while (abs(dist) > max_dist_allowed and N < 100) :
    # Calculate the new distance:
    dist = (a_normal[0] / sin(theta[0] - dtheta)) - (a_reverse[0] / sin(theta[0] + dtheta))

    # Ask which side of the equation is greater, and adjust dtheta accordingly (and half the step size):
    if (dist > 0.0) : 
        dtheta -= ddtheta 
    else :
        dtheta += ddtheta 
    ddtheta = ddtheta / 2.0
    N += 1         # Count number of iterations, and stop at 100, or the while loop could go infinite!!!

    # Print the status:
    print("  N_iter = {0:2d}   dist = {1:12.9f}   ddtheta = {2:11.9f}     dtheta = {3:9.7f}".format(N, dist, ddtheta, dtheta))


# ----------------------------------------------------------------------------------- #
# The error propagation:
# ----------------------------------------------------------------------------------- #

# We have obtained the central value for dtheta, but we would also like to know its
# uncertainty. This is simply done by repeating the above procedure many times with
# variations of the input within the input uncertainties, and then seeing the impact
# on the variation of the results in theta:

print("\n  -------------- Results from many examples -----------------")

list_dtheta = []
for i in xrange( 1000 ) :
# In order to get the uncertainty on dtheta, it is calculated many times for choices
# of the input parameters (varied without their uncertainties).
    an_used = r.normal( a_normal[0], a_normal[1] )
    ar_used = r.normal( a_reverse[0], a_reverse[1] )
    theta_used = r.normal( theta[0], theta[1] )

    dist = 99.9       # Initial distance just needs to be above the lower limit
    dtheta = 0.02
    ddtheta = 0.01
    N = 0
    while (abs(dist) > max_dist_allowed and N < 100) :
        dist = (an_used / sin(theta_used - dtheta)) - (ar_used / sin(theta_used + dtheta))
        if (dist > 0.0) : 
            dtheta -= ddtheta 
        else :
            dtheta += ddtheta 
        ddtheta = ddtheta / 2.0
        N += 1
    if (i < 25) :
        print("   {0:3d}:  Nite = {1:2d}   dist = {2:12.9f}   ddtheta = {3:11.9f}     dtheta = {4:7.5f}".format(i, N, dist, ddtheta, dtheta))
    list_dtheta.append(dtheta)


# Before assuming anything about the result, one should always plot the result:
Nbins = 200
dtheta_min = 0.00
dtheta_max = 0.02
binwidth = (dtheta_max - dtheta_min) / Nbins

fig, ax = plt.subplots(figsize=(12, 6))
hist_dtheta = ax.hist(list_dtheta, bins=Nbins, range=(dtheta_min, dtheta_max), histtype='step', linewidth=2, label=r"Values for $\Delta \theta$", color='blue')
ax.set_xlabel(r"$\Delta \theta$ [rad]")
ax.set_ylabel("Frequency / 0.0001 rad")
ax.set_title(r"Values for $\Delta \theta$ (obtained numerically with variations of input)")
ax.legend(loc='upper right')
plt.show(block=False)


# The distribution is Gaussian (as expected), and in order to extract a value and uncertainty for dtheta,
# one could just use the mean and RMS of this distribution:
dtheta_mean = np.array(list_dtheta).mean()
err_dtheta_mean = np.array(list_dtheta).std()
print("  The final result for dtheta = {0:.5f} +- {1:.5f} rad = {2:.3f} +- {3:.3f} deg   ({4:5.3f})".format(dtheta_mean, err_dtheta_mean, dtheta_mean/deg2rad, err_dtheta_mean/deg2rad, err_dtheta_mean/dtheta_mean))
"""
"""

# Hold script until you want to exit:
raw_input( ' ... Press enter to exit ... ' )