# coding: utf-8
from __future__ import division, print_function, unicode_literals
from lab import *
import uncertainties.umath  as um
import uncertainties.unumpy as unp
import numpy                as np
import matplotlib.pyplot    as plt

from scipy import signal
from scipy import interpolate


##
## Reflection
##

## reflection
i = array(20, 30, 40, 45, 50, 60, 70, 80)                      # incidency angle (deg)
r = array(321, 300, 282, 270, 261, 241, 223, 203)*(-1) + 360-i # reflection angle (deg)

# fit and plot θi - θr
x = np.linspace(10, 80, 10)
k = simple_linear(i, r, 1)
f = lambda x: k.n*x
alpha = chi_squared_fit(i, r, f, 1)

print(mformat('''
# reflection
k: {}
χ² test:
  α={:.3f}
''', k, alpha))

plt.figure(1)
plt.clf()
plt.title('reflection')
plt.xlabel('incidency angle (deg)')
plt.ylabel('reflection angle (deg)')
plt.scatter(i, r, color='#003cad')
plt.plot(x, f(x), color='#0069ad')
plt.show()


##
## Refraction
##

# incidency angle
y  = ufloat(5.3, 0.1)
x  = ufloat(13.4, 0.1)
ti = 180/np.pi * um.atan(y/x)

# refraction angle
to = ufloat(180,1) - array(170, 168) + ti
n = to/ti

print(mformat('''
# refractive index
n₁: {} (styrene?)
n₂: {} (styrene)

t-test: α={:.2f}
''', n[0], n[1], check_measures(*n)))


##
## Brewster angle
##

# current, amplification factor, angle
data = array(
      (0.08,  3,  5),
      (0.10, 10, 10),
      (0.12,  3, 15),
      (0.38, 10, 20),
      (0.64, 10, 25),
      (0.62, 10, 30),
      (0.64, 10, 35),
      (0.44, 10, 40),
      (0.34, 10, 45),
      (0.38,  3, 50),
      (0.36,  3, 55),
      (0.32,  1, 60),
      (0.08,  1, 65))
  
t  = data[:,2]*np.pi/180 # angle
I  = data[:,0]*data[:,1] # normalized current
sI = 0.02*data[:,1]      # normalized error

# plot Ι - θ
plt.title('Brewster\'s angle')
plt.xlabel('incidency angle (rad)')
plt.ylabel('vertical polarization intensity (mA)')
plt.ylim(0, 8)
plt.scatter(t, I, color='#109d1d')
plt.errorbar(t, I, yerr=sI, linestyle='', color='#0c7616')
plt.show()


##
## Polarization - Malus Law
##

## distance 1

d = (104.8 - 20 + 13.5)*1e-2 # transmitter - receiver distance (m)
y = array(0,     15,    30,    45,    60,    90,    105,   120,   135,   150,   180)   # receiver angle (deg)
I = array(0.120, 0.100, 0.080, 0.050, 0.020, 0.000, 0.010, 0.020, 0.045, 0.100, 0.120) # current (mA)

# fit, plot γ - I, with I = I₀cos(γ²)
x = np.arange(0,180,1)
k = simple_linear(np.cos(np.pi/180*y)**2, I, 0.01)
f = lambda x: k.n*np.cos(np.pi/180*x)**2
alpha = chi_squared_fit(y, I, f, 0.01, s=1)

print(mformat('''
# malus law
## at d={}m
I₀: {}
χ² test:
  α={:.3f}''', d, k, alpha))

plt.figure(2)
plt.title('malus law')
plt.xlabel('horizontal angle (deg)')
plt.ylabel('intensity')
plt.xlim(0,180)
plt.ylim(0,0.2)
plt.scatter(y, I, color='#d6a800')
plt.plot(x, f(x), color='#e9b700')


## distance 2

d = (78.3 - 34.7 + 13.5)*1e-2 # transmitter - receiver angle (m)
y = array(0,     15,    30,    45,    60,    90,    105,   120,   135,   150,   180)   # receiver angle (deg)
I = array(0.175, 0.160, 0.120, 0.065, 0.035, 0.000, 0.010, 0.045, 0.085, 0.130, 0.175) # current (mA)

k = simple_linear(np.cos(np.pi/180*y)**2, I, 0.01)
f = lambda x: k.n*np.cos(np.pi/180*x)**2
alpha = chi_squared_fit(y, I, f, 0.01, s=1)

print(mformat('''
## at d={}m
I₀: {}
χ² test:
  α={:.3f}''', d, k, alpha))

plt.scatter(y, I, color='#178200')
plt.plot(x, f(x), color='#1ead00')


## distance 3

d = (69.2 - 34.7 + 13.5)*1e-2 # transmitter - receiver angle (m)
y = array(0,     15,    30,    45,    60,    90,    105,   120,   135,   150,   180)   # receiver angle (deg)
I = array(0.200, 0.180, 0.160, 0.085, 0.045, 0.000, 0.010, 0.050, 0.090, 0.150, 0.200) # current

k = simple_linear(np.cos(np.pi/180*y)**2, I, 0.01)
f = lambda x: k.n*np.cos(np.pi/180*x)**2
alpha = chi_squared_fit(y, I, f, 0.01, s=1)

print(mformat('''
## at d={}m
I₀: {}
χ² test:
  α={:.3f}
''', d, k, alpha))
plt.scatter(y, I, color='#003cad')
plt.plot(x, f(x), color='#0069ad')
plt.show()


## polarizer

t = array(0, 45, 90)
V = array(0.185, -0.080, -0.04)


##
## Standing waves
##

# receiver position from transmitter (m)
d = np.arange(30, 86.5, 0.5)*1e-2

# tension (V)
V = array(0.851, 1.019, 1.971, 0.810, 1.116, 1.591, 0.820, 1.064, 1.448, 0.804,
          1.196, 0.953, 0.772, 1.319, 0.846, 0.773, 1.211, 0.761, 0.806, 1.143,
          0.810, 1.049, 0.629, 0.785, 0.985, 0.571, 0.837, 0.816, 0.563, 0.867,
          0.739, 0.551, 0.839, 0.638, 0.542, 0.796, 0.586, 0.541, 0.828, 0.535, 
          0.544, 0.678, 0.556, 0.833, 0.546, 0.537, 0.790, 0.480, 0.511, 0.705, 
          0.451, 0.504, 0.637, 0.437, 0.524, 0.574, 0.422, 0.519, 0.525, 0.413, 
          0.517, 0.476, 0.396, 0.499, 0.416, 0.379, 0.474, 0.369, 0.371, 0.433, 
          0.319, 0.358, 0.383, 0.275, 0.345, 0.335, 0.245, 0.327, 0.297, 0.227,
          0.326, 0.264, 0.224, 0.322, 0.258, 0.235, 0.328, 0.242, 0.253, 0.341, 
          0.232, 0.249, 0.315, 0.199, 0.239, 0.264, 0.173, 0.236, 0.227, 0.143, 
          0.221, 0.188, 0.125, 0.206, 0.162, 0.112, 0.193, 0.130, 0.109, 0.193, 
          0.119, 0.111, 0.172)

# plot
plt.figure(5)
plt.viridis()
plt.title('standing waves')
plt.xlabel('receiver position (m)')
plt.ylabel('tension (V)')
plt.xlim(0.28,0.87)

# mark peaks
mins = signal.argrelextrema(V, np.less)[0]
maxs = signal.argrelextrema(V, np.greater)[0]
plt.scatter(d[mins], V[mins], marker='X', color='#d8a300')
plt.scatter(d[maxs], V[maxs], marker='X', color='#1ead00')

# spline interpolation
x = np.arange(d.min(), d.max(), 5e-4)
plt.plot(x, interpolate.spline(d, V, x), color='#0069ad')
plt.show()

# calculate wavelength
l1 = 2*np.diff(d[maxs]).mean()
l2 = 2*np.diff(d[mins]).mean()
l  = ufloat((l1 + l2)/2, 0.005)

print(mformat('''
# standing waves
λ: {} m
''', l))


##
## Double slit
##


t = 102+0.05*7 # total width (mm)
s =  20+0.05*9 # slits width (mm)
d = (t-s)*1e-3 # slits spacing (m)

# maxima order
n = array(1, 2, 3)

# peaks (right side)
r = array(200, 223, 262)*np.pi/180 - np.pi # 194, 204, 222

# peaks (left side)
l = np.pi - array(160, 136, 98)*np.pi/180 #166, 155, 137

# calculate wavelength
l = sample(*np.append(np.sin(r)/n, np.sin(l)/n)*d)

print(mformat('''
# double slit interference
λ: {} m
''', l.tval()))


##
## lloyd mirror
##

a = 37.5e-2 # transmitter position (from axis)
b = 32.5e-2 # receiver position (from axis)
m = 11.0e-2 # reflector position (first minimum)
h = array(14.4, 17.8, 20.8, 23.5)*1e-2 # reflector position (maxima)

# calculate wavelength (first distance)
n  = 1 + np.arange(1, len(h)+1)
l1 = sample(*(np.sqrt(a**2 + h**2) + np.sqrt(b**2 + h**2) - (a+b))/n).tval()

a = 27.5e-2 # transmitter position (from axis)
b = 27.5e-2 # receiver position (from axis)
m = 9.5e-2  # reflector position (first minimum)
h = array(12.6, 15.8, 18.4, 20.1)*1e-2 # reflector position (maxima)

# calculate wavelength (second distance)
n  = 1 + np.arange(1, len(h)+1)
l2 = sample(*(np.sqrt(a**2 + h**2) + np.sqrt(b**2 + h**2) - (a+b))/n).tval()


print(mformat('''
# lloyd mirror interference
λ₁: {} m
λ₂: {} m

compatibility test: α={:.3f}
λbest: {} m
''', l1, l2, check_measures(l1, l2), combine_measures(l1, l2)))


##
## Fabry-Perot
##

# reflection relative position (maxima)
m = array(19.4, 20.9, 22.3, 23.8, 25.3, 26.7, 28.0, 29.3, 30.7, 32.2)*1e-2

# calculate wavelength
l = sample(*np.diff(m)*2)

print(mformat('''
# fabry-perot interferometer
λ: {} m
''', l.tval()))


##
## Michelson
##

# interferometer arms length (m)
l1 = 0.307
l2 = 0.423
l3 = 0.246

# reflector position (m)
d = array(0.092, 0.106, 0.120, 0.133, 0.150, 0.165, 0.180, 0.194,
          0.220, 0.235, 0.245, 0.260, 0.276, 0.290, 0.305, 0.322,
          0.338, 0.347, 0.360, 0.376, 0.392, 0.410)

# fit and plot n - d
n   = np.arange(len(d))
a,b = linear(n, d, 2e-3)
f   = lambda x: a.n + b.n*x
alpha = chi_squared_fit(n, d, f, 4e-3)

print(mformat('''
# michelson interferometer
λ: {} m
χ² test: α={:.3f}
''', 2*b, alpha))

plt.figure(3)
plt.title('michelson interferometer')
plt.xlabel('n-th maximum')
plt.ylabel('reflector position (m)')
plt.scatter(n, d, color='#178200')
plt.plot(n, f(n), color='#1ead00')
plt.show()