import matplotlib.pyplot as plt
import numpy as np
from scipy import integrate

def Fresnel2dreal(yp, xp, y, x, k, z):
    kernel = np.cos((k/(2*z))*((x-xp)**2+(y-yp)**2))
    return kernel

def Fresnel2dimag(yp, xp, y, x, k, z):
    kernel = np.sin((k/(2*z))*((x-xp)**2+(y-yp)**2))
    return kernel

c = 3e8
e0 = 8.85e-12

def plot1D():
    def genData(aperture, z, k, screen_range):
        y = 0

        xp1=yp1=-aperture/2
        xp2=yp2=aperture/2
    
        xs = np.linspace(-screen_range/2, screen_range/2, num=1000)
        intensities = []

        for x in xs:
            realpart, realerror = integrate.dblquad(Fresnel2dreal, xp1, xp2, yp1, yp2, args=(y, x, k, z))
            imagpart, imagerror = integrate.dblquad(Fresnel2dimag, xp1, xp2, yp1, yp2, args=(y, x, k, z))

            I = c*e0*(realpart**2+imagpart**2)
            intensities.append(I)

        return xs, intensities

    ax = plt.axes()
    xs, intensities = genData(2e-4, 0.005, 8.377e6, 0.002)
    ax.plot(xs, intensities)
    # xs, intensities = genData(2e-5, 0.05, 8.377e6, 0.015)
    # ax.plot(xs, intensities)
    plt.show()

def plot2Drectangular():
    def genData(aperture, z, k, screen_range):
        xp1=yp1=-aperture/2
        xp2=yp2=aperture/2
        
        xs = np.linspace(-screen_range/2, screen_range/2, num=40)
        ys = np.linspace(-screen_range/2, screen_range/2, num=40)

        intensities = []

        for y in ys:
            xIntensities = []
            for x in xs:
                realpart, realerror = integrate.dblquad(Fresnel2dreal, xp1, xp2, yp1, yp2, args=(y, x, k, z))
                imagpart, imagerror = integrate.dblquad(Fresnel2dimag, xp1, xp2, yp1, yp2, args=(y, x, k, z))

                I = c*e0*(realpart**2+imagpart**2)
                xIntensities.append(I)
            intensities.append(xIntensities)
        intensities = np.array(intensities)
        return intensities

    intensity = genData(2e-4, 0.005, 8.377e6, 0.002)
    extents = (-0.01,0.01,-0.01,0.01)

    plt.imshow(intensity,vmin=0.0,vmax=1.0*intensity.max(),extent=extents,origin="lower",cmap="nipy_spectral_r")
    plt.colorbar()
    plt.show()

def plot2Dcircular():
    def genData(aperture, z, k, screen_range):
        xp1=-aperture/2
        xp2=aperture/2

        def yp1func(xp):
            return -np.sqrt((aperture/2)**2-(xp**2))

        def yp2func(xp):
            return np.sqrt((aperture/2)**2-(xp**2))

        xs = np.linspace(-screen_range/2, screen_range/2, num=50)
        ys = np.linspace(-screen_range/2, screen_range/2, num=50)

        intensities = []

        for y in ys:
            xIntensities = []
            for x in xs:
                realpart, realerror = integrate.dblquad(Fresnel2dreal, xp1, xp2, yp1func, yp2func, args=(y, x, k, z))
                imagpart, imagerror = integrate.dblquad(Fresnel2dimag, xp1, xp2, yp1func, yp2func, args=(y, x, k, z))

                I = c*e0*(realpart**2+imagpart**2)
                xIntensities.append(I)
            intensities.append(xIntensities)
        intensities = np.array(intensities)
        return intensities

    intensity = genData(2e-5, 0.05, 8.377e6, 0.015)
    extents = (-0.01,0.01,-0.01,0.01)

    plt.imshow(intensity,vmin=0.0,vmax=1.0*intensity.max(),extent=extents,origin="lower",cmap="nipy_spectral_r")
    plt.colorbar()
    plt.show()

def monte():
    N = 1000

    def doubleInteg(x, y, xp, yp, z, k, aperture):
        values = []
        for i in range(len(xp)):
            if (xp[i]**2+yp[i]**2) > aperture:
                values.append(0)
            else:
                value = np.exp(((1j*k)/(2*z))*((x-xp)**2+(y-yp)**2))
                values.append(value.real)
        return np.array(values)
    
    def monteCarlo(x, y, z, k, aperture):
       xp = np.random.uniform(low=(-aperture/2), high=aperture/2, size=N) 
       yp = np.random.uniform(low=(-aperture/2), high=aperture/2, size=N)
       values = doubleInteg(x, y, xp, yp, z, k , aperture)
       mean = values.sum()/N
       meansq = (values*values).sum()/N
       integral = aperture*mean
       error = aperture*np.sqrt((meansq-mean*mean)/N)
       return integral, error
         
    def genData(aperture, z, k, screen_range):

        xs = np.linspace(-screen_range/2, screen_range/2, num=50)
        ys = np.linspace(-screen_range/2, screen_range/2, num=50)

        constant = k/2*np.pi*z

        intensities = []
        
        for y in ys:
            xIntensities = []
            for x in xs:
                integral, error = monteCarlo(x, y, z, k, aperture)
                I = c*e0*constant*integral
                xIntensities.append(I)
            intensities.append(xIntensities)
        intensities = np.array(intensities)
        return intensities
    
    intensity = genData(2e-4, 0.005, 8.377e6, 0.002)
    extents = (-0.001,0.001,-0.001,0.001)

    plt.imshow(intensity,vmin=0.0,vmax=1.0*intensity.max(),extent=extents,origin="lower",cmap="nipy_spectral_r")
    plt.colorbar()
    plt.show()

monte()