import sys
import cv2
import numpy as np
from matplotlib import pyplot as plt
from numpy import cos, sin, pi
from scipy.fft import fft2, fftshift
from scipy.special import jn

blurDimension = 3  # px


def blurImg(im):
    blur = cv2.blur(im, (blurDimension, blurDimension))
    # blur = cv2.GaussianBlur(im, (blurDimension, blurDimension), 0)
    return blur


def gamma(x):
    val = 1.055*fractionalPowerBase(x, (1/2.4)) - 0.055
    return val


def fractionalPowerBase(x, power):
    return np.sign(x)*(np.abs(x))**power


def getPointAtLineAngle(x0, y0, angle, length):
    x = round(x0 + length*cos(angle*pi/180.0));
    y = round(y0 + length*sin(angle*pi/180.0));

    return x, y


def makeApertureEmpty(imageSizePx, diameterAperture):
    im = np.zeros(imageSizePx, np.float64)
    centerX = int(imageSizePx[0]/2)
    centerY = int(imageSizePx[1]/2)

    im = cv2.circle(im, (centerX, centerY), int(diameterAperture/2), 255, -1)

    im = blurImg(im)

    return im


def makeApertureWithCircle(imageSizePx, diameterAperture, diamCenter):
    im = np.zeros(imageSizePx, np.float64)
    centerX = int(imageSizePx[0]/2)
    centerY = int(imageSizePx[1]/2)

    im = cv2.circle(im, (centerX, centerY), int(diameterAperture/2), 255, -1)
    im = cv2.circle(im, (centerX, centerY), int(diamCenter/2), 0, -1)

    im = blurImg(im)

    return im


def makeAperture3Vanes(imageSizePx, diameterAperture, diamCenter, thickness):
    im = np.zeros(imageSizePx, np.float64)
    centerX = int(imageSizePx[0]/2)
    centerY = int(imageSizePx[1]/2)

    im = cv2.circle(im, (centerX, centerY), int(diameterAperture/2), 255, -1)
    im = cv2.circle(im, (centerX, centerY), int(diamCenter/2), 0, -1)

    im = cv2.line(im, (centerX, centerY), getPointAtLineAngle(centerX, centerY, 30, diameterAperture/2), 0, int(thickness))
    im = cv2.line(im, (centerX, centerY), getPointAtLineAngle(centerX, centerY, 150, diameterAperture/2), 0, int(thickness))
    im = blurImg(im)
    im = cv2.line(im, (centerX, centerY), getPointAtLineAngle(centerX, centerY, -90, diameterAperture/2), 0, int(thickness))

    return im


def makeAperture3VanesArched(imageSizePx, diameterAperture, diamCenter, thickness):
    im = np.zeros(imageSizePx, np.float64)
    centerX = int(imageSizePx[0]/2)
    centerY = int(imageSizePx[1]/2)

    im = cv2.circle(im, (centerX, centerY), int(diameterAperture/2), 255, -1)
    im = cv2.circle(im, (centerX, centerY), int(diamCenter/2), 0, -1)

    r = int((diameterAperture/4) + thickness/4)
    im = cv2.ellipse(im, (centerX, int(centerY - r)), (r, r), 0, 90, 270, 0, thickness)
    im = cv2.ellipse(im, (getPointAtLineAngle(centerX, centerY, 30, r)), (r, r), 0, 210, 390, 0, thickness)
    im = cv2.ellipse(im, (getPointAtLineAngle(centerX, centerY, 150, r)), (r, r), 0, -30, 150, 0, thickness)

    im = blurImg(im)

    return im


def makeAperture4Vanes(imageSizePx, diameterAperture, diamCenter, thickness):
    im = np.zeros(imageSizePx, np.float64)
    centerX = int(imageSizePx[0]/2)
    centerY = int(imageSizePx[1]/2)

    im = cv2.circle(im, (centerX, centerY), int(diameterAperture/2), 255, -1)
    im = cv2.circle(im, (centerX, centerY), int(diamCenter/2), 0, -1)

    im = blurImg(im)

    im = cv2.line(im, (centerX, 0), (centerX, imageSizePx[1]), 0, int(thickness))
    im = cv2.line(im, (0, centerY), (imageSizePx[0], centerY), 0, int(thickness))

    return im


def makeAperture4VanesInv(imageSizePx, diameterAperture, diamCenter, thickness):
    im = np.zeros(imageSizePx, np.float64)
    centerX = int(imageSizePx[0]/2)
    centerY = int(imageSizePx[1]/2)

    #im = cv2.circle(im, (centerX, centerY), int(diameterAperture/2), 255, -1)
    #im = cv2.circle(im, (centerX, centerY), int(diamCenter/2), 0, -1)

    im = blurImg(im)

    #im = cv2.line(im, (centerX, 0), (centerX, imageSizePx[1]), 255, int(thickness))
    im = cv2.line(im, (0, centerY), (imageSizePx[0], centerY), 255, int(thickness))

    return im


def radToArcMin(rad):
    return rad*180/np.pi*3600


def getPhi(lam, D, i):
    return lam/D*i


def normalize(im):
    maxIntensity = np.max(im)
    return im/maxIntensity


def showAperture(D_img, im):
    e = int(D_img/2)
    plt.imshow(im, cmap='gray', extent=[-e, e, -e, e])
    plt.xlabel(r'$L_x$ [m]')
    plt.ylabel(r'$L_y$ [m]')
    plt.xlim(-2,2)
    plt.ylim(-2,2)
    plt.show()


def Bessel(i, D_img, D):
    x = np.pi*D*i/D_img
    jinc = lambda x: jn(1, x) / x  # Sombrero function
    airy = (2 * jinc(x))**2
    # plt.plot(x,airy)
    # plt.xlabel('$x$')
    # plt.ylabel('$I(x)/I_0$')
    # plt.show()
    return airy


def plotMiddleCut(im, lam, D_img, D):
    yMiddleCut = im[int(len(im)/2)]
    yMiddleCut = np.clip(gamma(yMiddleCut*150), 0, 1)

    x = np.arange(-len(im)/2, len(im)/2)
    bessel = Bessel(x, D_img, D)

    x = radToArcMin(getPhi(lam, D_img, x))


    plt.figure(figsize =(10,5))
    #plt.plot(x, bessel, "m", label = "Analytické riešenie")
    plt.plot(x, yMiddleCut, "b", label = "Numerické riešenie")
    plt.xlabel(r'$\varphi_x$ [″]')
    plt.ylabel(r'$\frac{I(\varphi)}{I_0}$')
    plt.xlim(-1.6,1.6)
    #plt.legend()
    plt.show()


def main():
    np.set_printoptions(threshold=sys.maxsize)  # print numpy whole arrays
    ### set params:
    diameterAperture = 180  #aperture in pixels
    diamCenter = 144  #central obstruction in pixels
    thickness = 4  #vane thickness in pixels, should be even number for arched vanes

    Dpx = 5000
    lam = 500e-9  # lambda
    D = 2.4  # aperture [m] Hubble
    D_img = Dpx * D / diameterAperture  # [m]
    D_secondary_Hubble = 30.5  # cm  TODO + vanes thickness
    ###############
    imageSizePx = (Dpx, Dpx)  # must be much bigger than aperture to avoid artifacts


    print("ourAiry2     ", radToArcMin(getPhi(lam, D_img, 33)))
    airy = radToArcMin(1.22*lam/D)
    print("Airy radius: ", airy)

    ### Make aperture, very small in an ocean of black
    #im = makeApertureEmpty(imageSizePx, diameterAperture)
    #showAperture(D_img, im)
    im = makeApertureWithCircle(imageSizePx, diameterAperture, diamCenter)
    showAperture(D_img, im)
    #im = makeAperture3Vanes(imageSizePx, diameterAperture, diamCenter, thickness)
    #showAperture(D_img, im)
    #im = makeAperture3VanesArched(imageSizePx, diameterAperture, diamCenter, thickness)
    #showAperture(D_img, im)
    #im = makeAperture4Vanes(imageSizePx, diameterAperture, diamCenter, thickness)
    #showAperture(D_img, im)
    #im = makeAperture4VanesInv(imageSizePx, diameterAperture, diamCenter, thickness)
    #showAperture(D_img, im)

    ### comment "exit" to call Jean-Baptiste Joseph Fourier
    # exit()

    ft = fftshift(fft2(im))

    real = (np.abs(ft))
    I = np.square(real)
    I = normalize(I)
    # print(I)

    plotMiddleCut(I, lam, D_img, D)

    e = radToArcMin(getPhi(lam, D_img, int(Dpx/2)))
    # plt.imshow(np.clip(gamma(I*100), 0, 1), extent=[-e, e, -e, e], cmap='gray')  # black&white
    plt.imshow(np.clip(gamma(I*150), 0, 1), extent=[-e, e, -e, e], cmap='gray_r')  # black&white inverted
    plt.xlabel(r'$\varphi_x$ [″]')
    plt.ylabel(r'$\varphi_y$ [″]')
    plt.xlim(-2.3,2.3)
    plt.ylim(-2.3,2.3)
    plt.show()


if __name__ == '__main__':
    main()