#!/usr/bin/env python
# coding: utf-8

# ## Fraunhoferova difrakcija na dvostrukoj pukotini
# #### Kod je pisan u Python programskom jeziku koristeći Jupyter notebook

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
import math

from scipy.fftpack import fft      #paket za FFT
from scipy.fftpack import ifftshift,  fftshift

from mpl_toolkits.mplot3d import Axes3D   #paket za 3d plot
plt.style.use('dark_background')          # izbor boja za plot
plt.rcParams['image.cmap'] = 'plasma'


# In[2]:


def Fraunhofer_input(lambd,w,z):
    N=1048   #rezolucija
    L=5.    # širina zaslona na kojoj se javlja difrakciona slika

    lamb=lambd*10**(-6) #(mm) talasna dužina svjetlosti
#    w širina otvora (mm)
#    z= udaljenost otvor - zaslon (mm)
    k=2*np.pi/lamb   #talasni broj 

    delta_src=L/N     #prostorna frekvencija
    
    #definisati ravan otvora
    xv=np.linspace(-N/2,N/2-1,N)
    yv=np.linspace(-N/2,N/2-1,N)
    [x1,y1]=np.meshgrid(xv,yv)
    x1=x1*delta_src  #(mm)
    y1=y1*delta_src  # (mm) 
 
    #definisati dvostruku pukotinu
    X=np.where(np.abs(x1)<w/4,0,0)
    Y=np.where(np.abs(y1)<10*w,1,0)
    X[np.logical_and(x1<-1*w, x1>-3*w)]=1.0
    X[np.logical_and(x1<3*w, x1>1*w)]=1.0

    u1=X #elekticno polje na pukotini
    ap1=u1

# Fourierova transformacija  izračunavanje integrala
    g=u1
    G=fftshift(np.fft.fft(g))
    h=np.exp(1j*k*z)/1j*lamb*z*np.exp(1j*np.pi/(lamb*z)*x1**2)
    H=fftshift(h)
    u2=H*G
    identity=np.ones((N,N))
    u2=u2*identity
    
    # Intenzitet zračenja (normiran na jedinicu)
    I=1./(lamb*z)**2*(np.abs(u2))**2
    I_1d=np.zeros(N, dtype='complex128')
    for a in range (1,N):
        I_1d[a]=I[0][a]
        
    
    I_1d=I_1d/I_1d.max()
        
    return(u2,I_1d,y1,ap1) #OUTPUT: električno polje na zaslonu, intenzitet, koordinate ravni otvora i polje na otvoru


# In[3]:


lamb=632.
w=0.05
z=20.

u1,I1,x1,ap1=Fraunhofer_input(lamb,2*w,1000.)
u2,I2,x2,ap2=Fraunhofer_input(lamb,w/5,1000.)
u3,I3,x3,ap3=Fraunhofer_input(lamb,w,1000.)

I1=I1/I1.max()
I2=I2/I2.max()
I3=I3/I3.max()


# In[7]:


# PLOTANJE relativne veličine otvora, difrakcione slike i intenziteta zračenja

fig2=plt.figure(figsize=(11,10)) 

ax1=fig2.add_subplot(4,3,1)
plt.imshow(np.abs(ap1))
plt.xlim(324,724)
plt.ylim(460,588)
plt.axis('off')

ax1=fig2.add_subplot(4,3,2)
plt.imshow(np.abs(ap2))
plt.xlim(324,724)
plt.ylim(460,588)
plt.axis('off')

ax1=fig2.add_subplot(4,3,3)
plt.imshow(np.abs(ap3))
plt.xlim(324,724)
plt.ylim(460,588)
plt.axis('off')


ax1=fig2.add_subplot(4,3,4)
plt.title('(A) z = 1 m, w = 0.1 mm')
plt.imshow(np.abs(u1))
plt.xlim(324,724)
plt.ylim(460,588)
plt.axis('off')
plt.xlabel('(A)')

ax2=fig2.add_subplot(4,3,5)
plt.title('(B) z = 1 m, w = 0.01 mm')
plt.imshow(np.abs(u2))
plt.xlim(324,724)
plt.ylim(460,588)
plt.xlabel('(B)')
plt.axis('off')

ax3=fig2.add_subplot(4,3,6)
plt.title('(C) z = 1 m, w = 0.05 mm')
plt.imshow(np.abs(u3))
plt.xlim(324,724)
plt.ylim(460,588)
plt.xlabel('(C)')
plt.axis('off')

ax4=fig2.add_subplot(2,1,2)
plt.title('Fraunhoferova difrakcija - dvostruka pukotina')

ax4_1=ax4.plot(x1,np.abs(I1),'r-')
ax4_2=ax4.plot(x2,np.abs(I2),'b-')
ax4_3=ax4.plot(x2,np.abs(I3),'y:')

ax4.text(-0.6,0.6,'(A)', color='red',fontsize=16)
ax4.text(0.4,0.6,'(B)', color='blue',fontsize=16)
ax4.text(0.08,0.6,'(C)', color='yellow',fontsize=16)

plt.xlabel('X(mm)')
plt.ylabel('Relativni intenzitet')
fig2.savefig('Fraunhofer_2slit.png')
plt.show()


# In[ ]: