from matplotlib import pyplot as plt
from matplotlib import animation
import numpy as np
from scipy import special
bessel = special.jv
i = complex(0, 1)


def q(z):
    return z + i*10


def calc_diffraction(r, z):
    return (1 ** 2) * (0.0015 ** 2 / 0.001 ** 2) * np.exp(
        (-2 * z ** 2 * (np.sin(0.1)) ** 2) / (0.0015 ** 2)) * np.exp(
        (-2 * r ** 2) / (0.0015 ** 2)) * bessel(0, np.conjugate(q(0)) * 2 *
            np.pi * r * np.sin(0.1) / (np.conjugate(q(z)) * 0.000000750)) *
            bessel(0, (q(0)) * 2 * np.pi * r * np.sin(0.01) / ((q(z)) * 0.000000750))


def plotradialevolutionArtist(begin, end, number):
    xDim = 0.0000055 * 2048
    yDim = 0.0000055 * 1088

    x = np.linspace(-xDim / 2, xDim / 2, 2048)
    y = np.linspace(-yDim / 2, yDim / 2, 1088)

    X, Y = np.meshgrid(x, y)
    R = np.sqrt(X ** 2 + Y ** 2)

    fig, ax = plt.subplots(1, number)

    number_list = np.linspace(begin, end, number)

    ims = []

    for i in range(len(number_list)):
        print(i)
        Z = np.real(calc_diffraction(R, number_list[i]))
        ax[i].set_aspect('equal')
        ax[i].contourf(X, Y, Z, 3)
        ax[i].tick_params(direction='inout', length=8, width=1, pad=8, labelsize=18, top=False, right=False)
        ims.append(ax[i])

    anim = animation.ArtistAnimation(fig, ims, interval=100, blit=True,
repeat=True, repeat_delay=100)
    plt.show()


plotradialevolutionArtist(0, 1, 10)

from matplotlib import pyplot as plt
from matplotlib import animation
import numpy as np
from scipy import special
bessel = special.jv
i = complex(0, 1)


def q(z):
    return z + i*10


def calc_diffraction(r, z):
    return (1 ** 2) * (0.0015 ** 2 / 0.001 ** 2) * np.exp(
        (-2 * z ** 2 * (np.sin(0.1)) ** 2) / (0.0015 ** 2)) * np.exp(
        (-2 * r ** 2) / (0.0015 ** 2)) * bessel(0, np.conjugate(q(0)) * 2 *
            np.pi * r * np.sin(0.1) / (np.conjugate(q(z)) * 0.000000750)) *
            bessel(0, (q(0)) * 2 * np.pi * r * np.sin(0.01) / ((q(z)) * 0.000000750))


def plotradialevolutionArtist(begin, end, number):
    xDim = 0.0000055 * 2048
    yDim = 0.0000055 * 1088

    x = np.linspace(-xDim / 2, xDim / 2, 2048)
    y = np.linspace(-yDim / 2, yDim / 2, 1088)

    X, Y = np.meshgrid(x, y)
    R = np.sqrt(X ** 2 + Y ** 2)

    fig, ax = plt.subplots()

    number_list = np.linspace(begin, end, number)

    ims = []

    for i in range(len(number_list)):
        print(i)
        ax.clear()
        Z = np.real(calc_diffraction(R, number_list[i]))
        ax.set_aspect('equal')
        ax.contourf(X, Y, Z, 3)
        ax.tick_params(direction='inout', length=8, width=1, pad=8, labelsize=18, top=False, right=False)
        ims.append(ax)

    anim = animation.ArtistAnimation(fig, ims, interval=100, blit=True,
repeat=True, repeat_delay=100)
    plt.show()


plotradialevolutionArtist(0, 1, 10)