# coding: utf-8

from __future__ import division, print_function, unicode_literals
from lab import *

import numpy as np
import matplotlib.pyplot as plt
import uncertainties.umath as u

# known constants
D  = ufloat(0.994, 0.001)
Ma = ufloat(1.000, 0.001)
Mb = ufloat(1.400, 0.001)

##
## Part 1
##

X = sample(0.067,0.080,0.100,0.135,0.150,0.170,0.196,0.227) # distance from 1

T1 = [ sample(2.0899,2.0842,2.0828,2.0823) # period 1
     , sample(2.0604,2.0575,2.0574,2.0555)
     , sample(2.0203,2.0186,2.0162,2.0170)
     , sample(1.9509,1.9493,1.9485,1.9470)
     , sample(1.9283,1.9260,1.9253,1.9208)
     , sample(1.9038,1.9009,1.8994,1.9007)
     , sample(1.8748,1.8731,1.8722,1.8733)
     , sample(1.8463,1.8485,1.8458,1.8428)
     ]

T2 = [ sample(2.0196,2.0177,2.0174,2.0179) # period 2
     , sample(2.0165,2.0159,2.0155,2.0151)
     , sample(2.0073,2.0064,2.0066,2.0070)
     , sample(1.9892,1.9899,1.9901,1.9913)
     , sample(1.9880,1.9872,1.9865,1.9865)
     , sample(1.9821,1.9815,1.9808,1.9828)
     , sample(1.9747,1.9751,1.9740,1.9730)
     , sample(1.9693,1.9683,1.9678,1.9672)
     ]

t1 = [i.mean for i in T1]
t2 = [i.mean for i in T2]
s1 = [i.stdm for i in T1]
s2 = [i.stdm for i in T1]

# quadratic fit
a1,b1,c1 = polynomial(X, t1, 2, sigma_Y=s1)
a2,b2,c2 = polynomial(X, t2, 2, sigma_Y=s2)

f = lambda x: a1*x**2+b1*x+c1
g = lambda x: a2*x**2+b2*x+c2
x = np.linspace(X.min, X.max, 100)


# plot T1, T2 - X
plt.figure(1)
plt.xlabel('distanza da 1 (m)')
plt.ylabel('periodo (s)')
plt.xlim(0.05,0.24)

plt.scatter(X, t1, color='#21812b')
plt.scatter(X, t2, color='#7755ff')
plt.plot(x, [f(i).n for i in x], color='#21812b')
plt.plot(x, [g(i).n for i in x], color='#7755ff')
plt.show()


# find intersection and calculate g
x  = ((b2-b1)-(u.sqrt)((b1-b2)**2-4*(a1-a2)*(c1-c2)))/(2*(a1-a2)) # intersection
T  = f(x)              # isochronic period
g1 = D/T**2*4*np.pi**2 # acceleration due to mass

print('''
Kater pendulum:
T: ({})s
g₁: ({})m/s²
'''.format(T, g1))


##
## Part 2
## free fall - direct measure of g

H = sample(0.300,0.400,0.500,0.550,0.650,0.750) # drop height

T = [ sample(0.241,0.242,0.241,0.248,0.251,0.252,0.249,0.243,0.248,0.244) # time to fall
    , sample(0.285,0.285,0.284,0.284,0.284,0.282,0.284,0.282,0.281,0.281)
    , sample(0.318,0.319,0.318,0.316,0.318,0.316,0.316,0.318,0.317,0.317)
    , sample(0.336,0.333,0.333,0.332,0.340,0.330,0.330,0.332,0.335,0.330)
    , sample(0.362,0.361,0.365,0.367,0.359,0.361,0.357,0.366,0.361,0.359)
    , sample(0.385,0.387,0.390,0.385,0.393,0.386,0.386,0.392,0.394,0.387)
    ]


t  = sample(i.mean for i in T)
St = sample(i.std  for i in T)

# linear fit y=kx where y=h, k=g/2, x=t^2
x = np.linspace(0.2, 0.4, 100)
k = simple_linear(t**2, H, 0.02)
f = lambda x: x**2 * k.n

# plot H - T
plt.figure(2)
plt.ylabel('altezza (m)')
plt.xlabel('tempo di caduta (s)')
plt.scatter(t, H)
plt.plot(x, f(x))
plt.show()

# calculate g 
g2 = 2*k

# χ² test
alpha = chi_squared_fit(t, H, f, St)

print('''
Free fall:
g₂: ({})m/s²
χ²: α={:.3f}, α>ε: {}
'''.format(g2, alpha, alpha>epsilon))


##
## Part 3
## compare g1, g2

alpha = check_measures(g1, g2)
g = combine_measures(g1, g2)

print('''
compare g₁ g₂:
α: {:.3f}, α>ε: {}

combine:
g_best: ({})m/s² 
'''.format(alpha, alpha>epsilon, g))