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

#    BPM - Simulation 1+1d Faisceau Gaussien en espace libre avec une lentille
# --------()--------()--------()---.....---()------()------()------
#  dz/2 (n1)  dz  (n1)  ................................. (n1)  dz/2

# n1 ne doit pas depasser 1% en plus de n0

from IPython import get_ipython
ipython = get_ipython()
if ipython is not None:
    ipython.run_line_magic('matplotlib', 'qt5')     # pour affichage dans une fenêtre externe
    ipython.run_line_magic('matplotlib', 'inline')  # pour affichage dans une fenêtre interne


import numpy as np
from math import pi
import matplotlib.pyplot as plt
import sys

#### Constantes SI ####
largeur = 100.e-6           # Largeur de l'axe des abscisses
w0 = 1.e-6                  # Largeur du faisceau Gaussien
lambd = 633.e-9             # Longueur d'onde
Lr = (pi*w0**2) / lambd     # Longueur de Rayleigh
k = 2*pi/(lambd)            # Vecteur d'onde
N = 1001                   # Nombre de points de la fenêtre de calcul
i = complex (0,1)           # Nombre complexe
facteur = 0.001             # Pour le théorème de shannon
dz = facteur*Lr             # Distance de propagation
LargeurGuide =5.e-6           # Largeur guide (lentille) en m
LargeurGuidePt =int(N*LargeurGuide/largeur)         # Largeur guide (lentille) en nombre de points
milieu = int(((N-1)/2))

n0 = 1                      # Indice de réfraction milieu
n1 = 1.01    # Indice de réfraction du guide. S'il n'y a pas de guide mettre n1=n0

#### Paramètres BPM ####
a=-(N-1)/2                  # Centrage du nombre de points sur 0
b=(N-1)/2                   # idem
NbRayleigh = 5              # Distance de propagation totale en nombre de longueur de Rayleigh
x = np.linspace(a, b, N)*largeur/N #Creation de l'axe des abscisses
distance = NbRayleigh*Lr    # Distance de propagation mètres
Nb_cellules = int(NbRayleigh/facteur) # Nombre de cellules pour la BPM
z = np.linspace(0, distance, Nb_cellules)  # vecteur z de propagation mètres

#### Faisceau de départ ####
y = np.exp((-x**2)/(w0**2))         # Faisceau gaussien de départ
porte=np.zeros(N,dtype = int)
porte[milieu-LargeurGuidePt//2:milieu+LargeurGuidePt//2] = 1 


## Opérateur Lentille - (Il faut repasser en espace reel pour le multiplier le faisceau par la lentille)
Lens = np.ones(N,dtype = complex)
Lens[milieu-LargeurGuidePt//2:milieu+LargeurGuidePt//2] = np.exp(-i*k*dz*((n1/n0)-1)) # creation d'un guide de LargeurGuide points de large


## Opérateur de propagation dans un milieu de taille dz - A multiplier avec le faisceau dans l'espace des frequences
nf = np.linspace (a,b,N)
MFR_dz = np.exp((i/(2*k))*((2*pi/largeur)**2)*(nf**2)*dz) # matrice de frequence reduite


## Opérateur de propagation dans un milieu de taille dz/2 - A multiplier avec le faisceau dans l'espace des frequences
MFR_dz_2 = np.exp((i/(2*k))*((2*pi/largeur)**2)*(nf**2)*dz/2) # Matrice de frequence reduite


####Debut BPM ####
print(f'Distance simulée : {distance:.4e} mètres')
print("Lancement de la BPM ")
## Première FFT
Y        = np.fft.fftshift(np.fft.fft(y))   # Transformee de Fourier du faisceau + Decalage du spectre
Sortie1  = Y*MFR_dz_2                       #Traverse du premier espace libre (dz/2) avant une lentille


TABLEAU =np.zeros([N,Nb_cellules + NbRayleigh],dtype= complex)  # tableau vide dont le nb de colonnes correspond aux nombres de cellules de calcul BPM
progression = 0    # debut de la propagation
TABLEAU [:,0] = y  #Premiere colonne du tableau est le faisceau de depart


## Boucle de cellule en cellule

for j in range(1,Nb_cellules):
    if (int(j)%(Nb_cellules/5)==0): # compteur pour l'affichage
        print("Calcul à ",int(j/Nb_cellules *100)," %")

    Sortie2 = np.fft.ifft(Sortie1)*Lens     # Traversee de la lentille
    Sortie1 = np.fft.fft(Sortie2)*MFR_dz    # Traversee du milieu de propagation
    TABLEAU[:,j] = Sortie2                  # On inclut le faisceau apres la derniere cellule dans notre tableau
    progression = progression + dz          # On a donc avancé de dz sur notre chemin
    j+=1

print("Calcul à  100 %")
## Dernière cellule
Sortie2 = np.fft.ifft(Sortie1)*Lens         # Derniere lentillle
Sortie = np.fft.fft(Sortie2)*MFR_dz_2       # Dernier espace libre dz/2
TABLEAU [:,j+1] = np.fft.ifft(Sortie)       # Integration de notre faisceau a la sortie de la derniere cellule

##### Fin de la BPM



### Comparaison des faisceaux de depart et apres propagation
plt.plot(x,y,label='Faisceau initial')
plt.plot(x,np.abs(np.fft.ifft(Sortie)), label = 'Faisceau final')
plt.plot(x,porte, color = 'black',label = 'Guide onde (n=' + str(n1)+')')  
plt.title("Comparaison des faisceaux initiaux et finaux")
plt.xlabel('x (m)')
plt.ylabel('Intensity distribution (A.U.)')
plt.grid(True)
plt.legend()
plt.show()

### Affichage 1 + 1D
fig, ax = plt.subplots()
plt.imshow(np.abs(TABLEAU),aspect='auto',cmap='jet',extent=[z.min(),z.max(),x.min(),x.max()])
plt.ylabel('x (m)')
plt.xlabel('Distance de propagation z (m)')
plt.colorbar()
#plt.title(f'Propagation sur une distance de {distance:5.3e} m')
Line=plt.Line2D([(0, distance)], [(LargeurGuide/2,LargeurGuide/2)], linestyle='--',color='white', linewidth=2)  # Couleur: Rouge (255, 0, 0), Largeur: 2 pixels
ax.add_line(Line)
Line=plt.Line2D([(0, distance)], [(-LargeurGuide/2,-LargeurGuide/2)], linestyle='--',color='white', linewidth=2)  # Couleur: Rouge (255, 0, 0), Largeur: 2 pixels
ax.add_line(Line)
plt.show()