
from aotools.turbulence.phasescreen import ft_phase_screen
from astropy.modeling import models
import gnuplotlib as gp
import numpy as np

aperture = 400 # mm RADIUS
focal_distance = 6400 # mm
wavelength = 0.0005 # mm
defocus_distance = 1 # mm

# Pupil

padding = 4
y, x = np.mgrid[-aperture*padding:aperture*padding, -aperture*padding:aperture*padding]
ring = models.Ring2D(amplitude=1, x_0=0, y_0=0, r_in=0.2*aperture, r_out=aperture)
pupil = ring(x,y)

# Atmosphere

r_0 = 0.2 # m
L_0 = 100 # m
l_0 = wavelength*0.001 # m
mm_in_m = 0.001 # unit of our "pixels" in m 
phase_screen = ft_phase_screen(r_0, 2*aperture*padding, mm_in_m, L_0, l_0)
wavefront = pupil #* np.exp( 1j * phase_screen + 0j)

# Defocus

ptv = defocus_distance * 8*(focal_distance/(2*aperture))**2

r = np.sqrt(x**2 + y**2)
phase_angle = 12*(r**2 - 1) * ptv
kernel_def = np.exp(1j * phase_angle / wavelength *2*np.pi)

defocused_pupil = wavefront * kernel_def
defocused_psf = np.fft.fftshift(np.fft.fft2(defocused_pupil))
defocused_psf = np.abs(defocused_psf)**2

# gp.plotimage(pupil, set=["term x11 1", "size square"])
# gp.plotimage(defocused_psf, set=["term x11 2", "size square"])


import aotools
from aotools.turbulence import phasescreen
from astropy.convolution import convolve, convolve_fft, Gaussian2DKernel

def simulate_defocused_seeing_limited_psf(shape, pixel_scale, seeing, aperture, defocus):
    # create an Airy disk model with the given diameter
    airy = models.AiryDisk2D(radius=aperture, x_0=shape[0] // 2, y_0=shape[1] // 2)
    y, x = np.mgrid[-aperture:aperture, -aperture:aperture]
    airy_disk = airy(x,y)

    
    # convolve with a Gaussian kernel to simulate seeing
    seeing_kernel = Gaussian2DKernel(seeing / pixel_scale)
    seeing_limited = convolve_fft(airy_disk, seeing_kernel)

    # simulate atmosphere with a phase screen
    np.random.seed(0)
    screen = phasescreen.ft_phase_screen(
        r0=0.2,
        N=shape[0]*2,
        delta=1,
        L0=10,
        l0=0.01,
    )
    phase_screen = np.exp(1j * screen)
    seeing_limited *= np.abs(phase_screen)

    # simulate defocus
    defocus_kernel = Ring2D(r_in=0.2*aperture, r_out=aperture, amplitude=1)
    defocus_psf = defocus_kernel(x,y)
    defocus_psf /= np.sum(defocus_psf)
    defocus_psf = convolve(seeing_limited, np.pad(defocus_psf, ((0,1),(0,1)), 0), boundary='extend')
    

    return defocus_psf
