import pylab as pl
from intensity import *

class propagate :
    def fraunhofer(self, z, i) :
        
        ndx = i.wl*z/(i.dx*i.M)
        ndy = i.wl*z/(i.dy*i.N)
        ip = intensity(ndx, -i.M*ndx/2.0 , i.M*ndx/2.0,
                       ndy, -i.N*ndy/2.0 , i.N*ndy/2.0,
                       i.wl)
   
        ft = pl.fftshift(pl.fft2(i.i))
        pIs = pl.exp(1j*z*i.k)/(1j*i.wl*z)
        pDs = pl.exp((1j*i.k*((ip.xgrid**2)+(ip.ygrid**2)))/(2*z))        

        ip.i = pIs*pDs*ft

        return ip
 
    def singlefresnel(self,z,i):

        ndx = i.wl*z/(i.dx*i.M)
        ndy = i.wl*z/(i.dy*i.N)
        ip = intensity(ndx, -i.M*ndx/2.0 , i.M*ndx/2.0,
                       ndy, -i.N*ndy/2.0 , i.N*ndy/2.0,
                       i.wl)

        U = i.i*pl.exp((1j*i.k*((i.xgrid**2)+(i.ygrid**2)))/(2*z))
        
        ft = pl.fftshift(pl.fft2(U))
        
        pIs = pl.exp(1j*z*i.k)/(1j*i.wl*z)
        pDs = pl.exp((1j*i.k*((ip.xgrid**2)+(ip.ygrid**2)))/(2*z))        

        ip.i = pIs*pDs*ft
        
        return ip
    
    def distance(self,i,z):

        id = i

        id.i = (1/(1j*i.wl*z))*pl.exp(((1j*i.k)/(2*z))*(((x-xi)**2)+((y-eta)**2)))

        return id

    def lens(self,i,f):

        il = i

        il.i = i.i*pl.exp(((-1j*i.k)/(2*f))*((il.xgrid**2)+(il.ygrid**2)))

        return il