import matplotlib.pyplot as plt
import cupy as cp
import numpy as np
import math
from numpy.fft import fft2, fftshift, ifftshift
from cupyx.scipy import fft as fft_gpu

from zern_generator import Zernike


def denoise(screen: np.ndarray):
    return np.where(np.abs(screen) < 1e-4, 0, screen)


class FraunhoferGPU:
    def __init__(self, zernike: Zernike):
        self.focal_distance = 100 * 10 ** -3
        self.distance = 100 * 10 ** -3
        self.radius = 5 * 10 ** -3
        self.span = (-(10 ** -2), 10 ** -2)
        self.wavelength = 500 * 10 ** -9
        self.k = 2*math.pi/self.wavelength

        self.zernike = zernike


    def get_focal_point(self,input_screen):
        input_screen_gpu = cp.asarray(input_screen)
        xv = cp.asarray(self.zernike.xv)
        yv = cp.asarray(self.zernike.yv)
        phase_mul = ((cp.exp(1j*self.k*self.distance) * cp.exp(1j*self.k*(xv**2 + yv**2)/(2*self.distance)))/(1j*self.wavelength*self.distance))
        U_fraun = fft_gpu.fft2(fft_gpu.ifftshift(input_screen_gpu), overwrite_x=True)
        U_fraun = fftshift(U_fraun * phase_mul)
        
        I_fraun = cp.abs(U_fraun)**2
        return I_fraun / cp.max(I_fraun)


    def generate_diff(self, defocus_amount, *zernikes):
        zernikes = list(zernikes)
        if len(zernikes) < 5:
            zernikes += [0] * (5-len(zernikes))

        wf = self.zernike.generate_zern_screen([1], self.radius)

        f_plus = self.zernike.generate_zern_screen(self.zernike.add_defocus(defocus_amount,zernikes), norm_radius=self.radius)
        f_minus = self.zernike.generate_zern_screen(self.zernike.add_defocus(-defocus_amount,zernikes), norm_radius=self.radius)

        wf_c = cp.asarray(wf)
        f_plus_c = cp.asarray(f_plus)
        f_minus_c = cp.asarray(f_minus)
        focal_image_plus = denoise(self.get_focal_point(wf_c*cp.exp(f_plus_c * 1j)))
        focal_image_minus = denoise(self.get_focal_point(wf_c*cp.exp(f_minus_c * 1j)))

        
        return (
            cp.asnumpy(focal_image_plus),
            cp.asnumpy(focal_image_minus),
            cp.asnumpy(focal_image_plus - focal_image_minus)
        )
        
class Fraunhofer:
    def __init__(self, zernike: Zernike):
        self.focal_distance = 100 * 10 ** -3
        self.distance = 100 * 10 ** -3
        self.radius = 5 * 10 ** -3
        self.span = (-(10 ** -2), 10 ** -2)
        self.wavelength = 500 * 10 ** -9
        self.k = 2*math.pi/self.wavelength

        self.zernike = zernike

    # def _denoise(screen: np.ndarray):
    #     return np.where(np.abs(screen) < 1e-4, 0, screen)

    def get_focal_point(self,input_screen):
        input_screen_gpu = np.asarray(input_screen)
        phase_mul = ((np.exp(1j*self.k*self.distance) * np.exp(1j*self.k*(self.zernike.xv**2 + self.zernike.yv**2)/(2*self.distance)))/(1j*self.wavelength*self.distance))
        U_fraun = fft2(ifftshift(input_screen_gpu))
        U_fraun = fftshift(U_fraun * phase_mul)
        #U_fraun = fftshift(U_fraun)


        I_fraun = np.abs(U_fraun)**2
        return I_fraun / np.max(I_fraun)


    def generate_diff(self, defocus_amount, *zernikes):
        zernikes = list(zernikes)
        if len(zernikes) < 5:
            zernikes += [0] * (5-len(zernikes))

        wf = self.zernike.generate_zern_screen([1], self.radius)

        f_plus = self.zernike.generate_zern_screen(
            self.zernike.add_defocus(defocus_amount,zernikes),
            norm_radius=self.radius
        )
        f_minus = self.zernike.generate_zern_screen(
            self.zernike.add_defocus(-defocus_amount,zernikes), 
            norm_radius=self.radius
        )

        wf_c = np.asarray(wf)
        f_plus_c = np.asarray(f_plus)
        f_minus_c = np.asarray(f_minus)
        focal_image_plus = denoise(
            self.get_focal_point(wf_c*np.exp(f_plus_c * 1j))
        )
        focal_image_minus = denoise(
            self.get_focal_point(wf_c*np.exp(f_minus_c * 1j))
        )
        diff = focal_image_plus - focal_image_minus
        normalized_diff = diff / cp.max(diff)

        return (
            cp.asnumpy(focal_image_plus),
            cp.asnumpy(focal_image_minus),
            cp.asnumpy(normalized_diff)
        )



if __name__ == "__main__":
    zernike = Zernike()
    zernike.set_image_params(npix=2 ** 8)
    zoom_factor = 20
    f = Fraunhofer(zernike=zernike)
    ratio = (f.wavelength*f.focal_distance)/f.radius
    zernikes = [[1 if i == j else 0 for i in range(0,24)] for j in range(20,24)]
    print(zernikes)
    fig, ax = plt.subplots(4,1)
    for i, z in enumerate(zernikes):
        print(zernike._osa_to_noll_list(z))
        wf_a = zernike.generate_zern_screen([1], 0.01/zoom_factor)
        wf_c = zernike.generate_zern_screen(z, 0.01/zoom_factor)
        wf_c_nz = zernike.generate_zern_screen(z, 0.01)
        wf = wf_a * np.exp(wf_c*1j)


        focal = f.get_focal_point(wf)

        

        zoom = ratio*zoom_factor
        #ax[i][0].imshow(np.abs(wf_c_nz))

        ax[i].imshow(np.abs(focal))
        ax[i].imshow(np.abs(focal))
    ax[0].set_title("Zernike 5 - DEFOCUS")
    ax[1].set_title("Zernike 6 - V-PR-ASTIGMATISM")
    ax[2].set_title("Zernike 7 - V-TREFOIL")
    ax[3].set_title("Zernike 8 - V-PR_COMA")
    ax[0].axis('off')
    ax[1].axis('off')
    ax[2].axis('off')
    ax[3].axis('off')
    plt.show()