# -*- coding: utf-8 -*-

from __future__ import division
import matplotlib.pyplot as plt
import numpy as np
from math import pi
from RGB import *
from pylab import *
from matplotlib.widgets import Slider, Button, RadioButtons
import matplotlib.gridspec as gridspec

 #fonction tenant compte de la réponse non linéaire de l'oeil et variant entre 0 et 1
def ressenti(x):
    return np.sqrt(x)

couleur = True #respect de la couleur ou non (maximise le contraste si non)

lambda0 = 590
delta_lambda = 0.1

if couleur:
    WL = WavelengthToRGB(lambda0)
else:
    WL = [1,1,1]
    
champ = 0.1 #champ d'interférences en m

resol = 500 #résolution de l'image calculée
p = 500 #ordre d'interférence
delta0 = p*lambda0 #différence de marche au centre en nm

D = 1 #distance de projection en m
    

def E(delta): #formule de Fresnel
    k1 = 2.0*pi/((lambda0 - delta_lambda / 2)*1e-9)
    k2 = 2.0*pi/((lambda0 + delta_lambda / 2)*1e-9)
    return .25*(2.0+np.cos(k1*delta*1e-9) + np.cos(k2*delta*1e-9))

#calcule chaque composante de l'éclairement d'un point (y,z) de l'écran
def ecl(y,z):
    delta = delta0*D/np.sqrt(D**2+y**2+z**2)
    Ec = E(delta)
    R, G, B = Ec*WL[0],Ec*WL[1],Ec*WL[2]
    return np.array([ressenti(R),ressenti(G),ressenti(B)])

y = np.linspace(-champ,champ,resol+1) #tableau des coordonnées réelles des points
z = np.linspace(-champ,champ,resol+1) 
Y,Z = np.meshgrid(y,z) #donne 2 tableau contenant les coordonnées de chaque
#point en fonction de son index (deux dimensions)
 # on transpose ecl(Y,Z) pour qu'il soit [resol_y,resol_z,3]
# la dernière dimension donnant les 3 couches RGB
image1 = np.transpose(ecl(Y,Z), (1, 2, 0))


#création du graphe donnant les franges d'interférence
fig = plt.figure(1,figsize=(20,10))

gs = gridspec.GridSpec(2, 2,
                       width_ratios=[1,2],
                       height_ratios=[4,1]
                       )

fig.suptitle("Michelson éclairé par un doublet", fontsize=20)

ax1 = plt.subplot(gs[0])
ax1.set_title("Différence de marche au centre : %i $\mu m$" %(delta0/1000))
plt.gca().get_yaxis().set_visible(False)
xx = np.linspace(0,resol,5)
yy = np.linspace(-champ,champ,5)
plt.xticks(xx,yy)
plt.xlabel("distance en $m$")
im = plt.imshow(image1)


ax2 = plt.subplot(gs[1])

gs.tight_layout(fig, rect=(0,0,1,0.95))

delta = delta0*D/np.sqrt(D**2+y**2)
l, = plt.plot(y,E(delta), linewidth = 3, color = 'black')
plt.plot([-champ, champ],[0, 0], linewidth = 2, color = 'black')
plt.xlim(-champ, champ)
plt.ylim(-.05, 1.05)

axcolor = 'lightgoldenrodyellow'
axp = axes([0.25, 0.05, 0.65, 0.02], facecolor=axcolor)
axcoul = axes([0.25, 0.1, 0.65, 0.02], facecolor=axcolor)
axdeltalambda = axes([0.25, 0.15, 0.35, 0.02], facecolor=axcolor)

ordre = Slider(axp, "Ordre d'interférence au centre :", 0, 3000, valinit=p,valfmt='%i')
couleur = Slider(axcoul, "Longueur d'onde : ", 400, 780, valinit=lambda0,valfmt='%i $nm$')
deltalambda = Slider(axdeltalambda, "$\delta \lambda$ du doublet : ", 0, 10, valinit=delta_lambda,valfmt='%.1f $nm$')

def update_all():
    image1 = np.transpose(ecl(Y,Z), (1, 2, 0))
    im.set_data(image1)
    l.set_ydata(E(delta))
    draw()    

def update_ordre(val):
    global p,delta,delta0
    p = int(ordre.val)
    delta0 = p*lambda0
    delta = delta0*D/np.sqrt(D**2+y**2)
    ax1.set_title("Différence de marche au centre : %i $\mu m$" %(delta0/1000))
    update_all()
    
def update_couleur(val):
    global p,delta,lambda0,WL
    lambda0 = int(couleur.val)
    WL = WavelengthToRGB(lambda0)
    p = 0.1* int(10*delta0/lambda0)
    delta = delta0*D/np.sqrt(D**2+y**2)
    ordre.set_val(p)

def update_delta(val):
    global delta_lambda
    delta_lambda = deltalambda.val
    
    update_all()
   
ordre.on_changed(update_ordre)
couleur.on_changed(update_couleur)
deltalambda.on_changed(update_delta)

plt.show()