# coding: utf-8
from __future__ import unicode_literals, division
from lab import *

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


##
## almost free oscillation
## θ(t) = Ae^(-γt)sin(ωt+φ)+θ0

G = sample(0.0421, 0.0514, 0.0421, 0.0392)
W = sample(2.6800, 2.6600, 2.6709, 2.6700)

g = ufloat(G.mean, G.std) # damping ratio γ
w = ufloat(W.mean, W.std) # pulsation ω

w0 = u.sqrt(w**2-g**2) # corrected value ω0
T  = 2*np.pi/w         # oscillation period

print('ω0: {}, γ: {}, T: {}'.format(w0, g, T))


##
## damped oscillation
##

G = sample(0.2280, 0.2130, 0.2180, 0.2150)
W = sample(2.6602, 2.6500, 2.6600, 2.6600)

g = ufloat(G.mean, G.std) # damping ratio γ
w = ufloat(W.mean, W.std) # pulsation ω

w0 = u.sqrt(w**2-g**2) # corrected value ω0
T  = 2*np.pi/w         # oscillation period

print('ω0: {}, γ: {}, T: {}'.format(w0, g, T))


##
## even more damped
##

G = sample(0.6060, 0.6150, 0.6057, 0.6060)
W = sample(2.6000, 2.6200, 2.6100, 2.5900)

g = ufloat(G.mean, G.std) # damping ratio γ
w = ufloat(W.mean, W.std) # pulsation ω

w0 = u.sqrt(w**2-g**2) # corrected value ω0
T  = 2*np.pi/w         # oscillation period

print('ω0: {}, γ: {}, T: {}'.format(w0, g, T))


##
## Driven oscillation
##

V  = np.array([1.8000, 1.9000, 2.0000, 2.1000, 2.2000,  # generator tension (V)
               2.3000, 2.3500, 3.0000, 3.2000, 3.4000])
A  = np.array([2.3900, 2.5400, 4.0000, 3.2900, 5.4500,  # oscillation amplitude (rad)
               6.2500, 4.5200, 3.7500, 2.4200, 1.8200])
T  = np.array([4.3100, 4.0600, 3.1000, 3.4000, 2.8300,  # oscillation period (s)
               2.7500, 2.9800, 1.9700, 1.8300, 1.7000])
Tf = np.array([4.3300, 4.1048, 3.0883, 3.4069, 2.8270,  # rotor period (s)
               2.7229, 2.9592, 1.9721, 1.8240, 1.7210])

w  = 2*np.pi/T  # disc angular velocity
wf = 2*np.pi/Tf # rotor angular velocity

# polynomial fit Y=aX^2+bX+c
Y = 1/A**2
X = w**2

y = 1/ufloat(A.mean(), 0.08)**2  # error estimation
a,b,c = polynomial(X, Y, 2, y.s) # polynomial coeffs

# plot Α - ω
x = np.linspace(1.4,3.8,250)
f = lambda w: 1/np.sqrt(a.n*w**4 + b.n*w**2 + c.n)

plt.xlabel('pulsazione (rad/s)')
plt.ylabel('ampiezza (rad)')
plt.scatter(w, A, color='#b99100')
plt.plot(x, f(x), '#c89d00')
plt.show()

# find curve parameters
M  = 1/u.sqrt(a)
w0 = (M**2 * c)**(1/4)
g  = 1/2*u.sqrt(b*M**2 + 2*w0**2)

print '''
  M: {}
  ω0: {}
  γ: {}
'''.format(M,w0,g)