import matplotlib.pyplot as plt
import numpy as np
from scipy import integrate
from tqdm import tqdm #Import all needed modules

def Fresnel2dreal(yp, xp, y, x, k, z): #Define functions for integral kernels
    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 #Universal constants
e0 = 8.85e-12

def plot1D(aperture, z, k, screen_range, resolution): #Function for part 1
    def genData(aperture, z, k, screen_range, resolution): #Function to generate data
        y = 0 #As in 1D

        xp1=yp1=-aperture/2 #Set range over which we are integrating
        xp2=yp2=aperture/2
    
        xs = np.linspace(-screen_range/2, screen_range/2, num=resolution) #Generate values to calculate for
        intensities = []

        constant = k/2*np.pi*z #Relative intensity constant
        
        for x in tqdm(xs): #Loop through all values of x
            realpart, realerror = integrate.dblquad(Fresnel2dreal, xp1, xp2, yp1, yp2, args=(y, x, k, z), epsabs=1e-10, epsrel=1e-10) #Calculate both real and imaginary parts
            imagpart, imagerror = integrate.dblquad(Fresnel2dimag, xp1, xp2, yp1, yp2, args=(y, x, k, z), epsabs=1e-10, epsrel=1e-10)

            I = c*e0*((realpart*constant)**2+(imagpart*constant)**2) #Combine parts and constants to get intensity
            intensities.append(I) #Add calculated intensity to list
        return xs, intensities

    ax = plt.axes()
    xs, intensities = genData(aperture, z, k, screen_range, resolution) #Plot intensity against distance
    ax.plot(xs, intensities)
    plt.xlabel("Position (m)")
    plt.ylabel("Relative Intensity")
    plt.title("1D diffraction")
    plt.show()

def plot2Drectangular(aperture, z, k, screen_range, resolution): #Function for part 2
    def genData(aperture, z, k, screen_range, resolution): #Function to generate data
        xp1=yp1=-aperture/2 #Set integration limits
        xp2=yp2=aperture/2
        
        xs = np.linspace(-screen_range/2, screen_range/2, num=resolution) #Generate values to integrate for
        ys = np.linspace(-screen_range/2, screen_range/2, num=resolution)

        intensities = []

        constant = k/2*np.pi*z #Relative intensity constant
        completion = 0

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

                I = c*e0*((realpart*constant)**2+(imagpart*constant)**2)#Combine both parts and constants
                xIntensities.append(I)
            intensities.append(xIntensities)
        intensities = np.array(intensities)
        return intensities

    intensity = genData(aperture, z, k, screen_range, resolution) #Generate the data
    extents = (-screen_range/2,screen_range/2,-screen_range/2,screen_range/2) #Set limits of the plot

    plt.imshow(intensity,vmin=0.0,vmax=1.0*intensity.max(),extent=extents,origin="lower",cmap="nipy_spectral_r") #Plot the graphs
    plt.colorbar()
    plt.xlabel("X Position (m)")
    plt.ylabel("Y Position (m)")
    plt.title("2D diffraction with a rectangular aperture")
    plt.show()

def plot2Dcircular(aperture, z, k, screen_range, resolution): #Function for part 3
    def genData(aperture, z, k, screen_range, resolution):
        xp1=-aperture/2 #Set integration limits
        xp2=aperture/2

        def yp1func(xp):
            return -np.sqrt((aperture/2)**2-(xp**2)) #Define y limits in terms of x

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

        xs = np.linspace(-screen_range/2, screen_range/2, num=resolution) #Generate values to integrate for
        ys = np.linspace(-screen_range/2, screen_range/2, num=resolution)

        intensities = []

        constant = k/2*np.pi*z

        for y in tqdm(ys):
            xIntensities = []
            for x in xs:
                realpart, realerror = integrate.dblquad(Fresnel2dreal, xp1, xp2, yp1func, yp2func, args=(y, x, k, z), epsabs=1e-10, epsrel=1e-10) #Calculate functions using circular limits
                imagpart, imagerror = integrate.dblquad(Fresnel2dimag, xp1, xp2, yp1func, yp2func, args=(y, x, k, z), epsabs=1e-10, epsrel=1e-10)

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

    intensity = genData(aperture, z, k, screen_range, resolution)
    extents = (-screen_range/2,screen_range/2,-screen_range/2,screen_range/2)
    
    plt.imshow(intensity,vmin=0.0,vmax=1.0*intensity.max(),extent=extents,origin="lower",cmap="nipy_spectral_r")
    plt.colorbar()
    plt.xlabel("X Position (m)")
    plt.ylabel("Y Position (m)")
    plt.title("2D diffraction with a circular aperture")
    plt.show()

def monte(aperture, z, k, screen_range, resolution, samples): #Function for part 4
    N = samples #Number of samples

    def doubleInteg(x, y, xp, yp, z, k, aperture): #Function to return calculated values for each sample
        values = []
        for i in range(len(xp)):
            if (xp[i]**2+yp[i]**2) > (aperture/2)**2: #Check if point is in the aperture
                values.append(0)
            else:
                value = np.exp(((1j*k)/(2*z))*((x-xp[i])**2+(y-yp[i])**2)) #Calculate value if in aperture
                values.append(value.imag)
        return np.array(values)
    
    def monteCarlo(x, y, z, k, aperture): #Perform onte carlo method
       xp = np.random.uniform(low=(-aperture/2), high=aperture/2, size=N) #Generate random samples 
       yp = np.random.uniform(low=(-aperture/2), high=aperture/2, size=N)
       values = doubleInteg(x, y, xp, yp, z, k , aperture) #Calculate value for each pair of samples
       mean = values.sum()/N
       meansq = (values*values).sum()/N
       integral = aperture*mean #Calculate the integral end error
       error = aperture*np.sqrt((meansq-mean*mean)/N)
       return integral, error
         
    def genData(aperture, z, k, resolution, screen_range): #Function to generate the data

        xs = np.linspace(-screen_range/2, screen_range/2, num=resolution) #Generate values to integrate for
        ys = np.linspace(-screen_range/2, screen_range/2, num=resolution)


        intensities = []
        
        constant = k/(2*np.pi*z)

        for y in tqdm(ys):
            xIntensities = []
            for x in tqdm(xs):
                integral, error = monteCarlo(x, y, z, k, aperture) #Find integral vaule using monte carlo method
                I = c*e0*constant*integral
                xIntensities.append(I)
            intensities.append(xIntensities)
        intensities = np.array(intensities)
        return intensities
    
    intensity = genData(aperture, z, k, resolution, screen_range) #Generate data
    extents = (-screen_range/2,screen_range/2,-screen_range/2,screen_range/2)

    for y in range(len(intensity)):
        for x in range(len(intensity[y])):
            if intensity[y][x] < 0.05*intensity.max():
                intensity[y][x] = 0

    plt.imshow(intensity,vmin=0,vmax=1.0*intensity.max(),extent=extents,origin="lower",cmap="nipy_spectral_r")
    plt.colorbar()
    plt.xlabel("X Position (m)")
    plt.ylabel("Y Position (m)")
    plt.title("2D diffraction through Monte Carlo")
    plt.show()

MyInput = '0' #Selection menu
while MyInput != 'q':
    MyInput = input('Enter a choice, "1", "2", "3", "4" or "q" to quit: ')
    print('You entered the choice: ',MyInput)
    if MyInput == '1':
        print('You have chosen part (1): 1D rectangular diffraction')
        aperture = input("Please input the desired aperture (m): ")
        z = input("Please enter the desired distance from the screen (m): ")
        wl = input("Please enter the desired wavelength of light (m): ")
        k = (2*np.pi)/float(wl)
        screen_range = input("Please enter the diameter of the screen (m): ")
        resolution = input("Please enter the resolution of the plot (pixels): ")
        plot1D(float(aperture), float(z), float(k), float(screen_range), int(resolution))
    elif MyInput == '2':
        print('You have chosen part (2): 2D rectangular diffraction')
        aperture = input("Please input the desired aperture (m): ")
        z = input("Please enter the desired distance from the screen (m): ")
        wl = input("Please enter the desired wavelength of light (m): ")
        k = (2*np.pi)/float(wl)
        screen_range = input("Please enter the diameter of the screen (m): ")
        resolution = input("Please enter the resolution of the plot (pixels): ")
        plot2Drectangular(float(aperture), float(z), float(k), float(screen_range), int(resolution))
    elif MyInput == '3':
        print('You have chosen part (3): 2D circular diffraction')
        aperture = input("Please input the desired aperture (m): ")
        z = input("Please enter the desired distance from the screen (m): ")
        wl = input("Please enter the desired wavelength of light (m): ")
        k = (2*np.pi)/float(wl)
        screen_range = input("Please enter the diameter of the screen (m): ")
        resolution = input("Please enter the resolution of the plot (pixels): ")
        plot2Dcircular(float(aperture), float(z), float(k), float(screen_range), int(resolution))
    elif MyInput == '4':
        print('You have chosen part (3): 2D circular diffraction using Monte Carlo')
        aperture = input("Please input the desired aperture (m): ")
        z = input("Please enter the desired distance from the screen (m): ")
        wl = input("Please enter the desired wavelength of light (m): ")
        k = (2*np.pi)/float(wl)
        screen_range = input("Please enter the diameter of the screen (m): ")
        resolution = input("Please enter the resolution of the plot (pixels): ")
        samples = input("Please enter the desired number of samples for the calculation: ")
        monte(float(aperture), float(z), float(k), float(screen_range), int(resolution), int(samples))
    elif MyInput != 'q':
        print('This is not a valid choice')
print('You have chosen to finish - goodbye.')