from __future__ import print_function, division
import numpy as np
import scipy as sp
import matplotlib.cm as cm
import matplotlib.pyplot as plt
from IPython.display import Image
%matplotlib inline

Image('images/relationAFandOTF.png', embed=True, width=300) 

Image('images/ambiguityFnKernel.png', embed=True) 

x = np.linspace(-8, 8, 150)
y = np.sinc(x)
fig, ax = plt.subplots(1,1)
ax.plot(x, y, label='$sinc(x)$')
ax.set_ylim(-0.3, 1.02)
ax.set_xlim(-8, 8)
ax.set_xlabel('x')
ax.set_title('sinc(x)', y=1.02)
rootsx = range(-8, 0) + range(1,9)
ax.scatter(rootsx, np.zeros(len(rootsx)), 
           c='r', zorder=20)
ax.grid()
plt.show()

def plot_1dRectAF(w20LambdaBy2, N=15, umin=-2, umax=2, ymin=-8, ymax=8):
    """rudimentary function to show the AF
    
    Parameters
    ----------
    w20LambdaBy2 : list of real values
        specifies the amount defocus error W_{20} in
        terms of lambda/2. The slope of the OTF 
        associated with W_{20} in AF plane is 2*W20/lambda.
    N : integer
        number of zero value loci to plot 
    """
    u = np.linspace(umin, umax, 200)
    y = np.linspace(ymin, ymax, 400)
    uu, yy = np.meshgrid(u, y) # grid

    # Numpy's normalized sinc function = sin(pi*x)/(pi*x)
    af = (1 - np.abs(uu)/2)*np.sinc(yy*(2 - np.abs(uu)))
    # plot
    fig = plt.figure(figsize=(12, 7))
    ax1 = fig.add_axes([0.12, 0, 0.42, 1.0]) # [*left*, *bottom*, *width*,*height*]
    ax2 = fig.add_axes([0.6, 0.12, 0.38, 0.76])
    im = ax1.imshow(af, cmap=cm.bwr, origin='lower', 
                   extent=[umin, umax, ymin, ymax], 
                   vmin=-1.0, vmax=1.0, aspect=2./6)
    plt.colorbar(im, ax=ax1, shrink=0.77, aspect=35)
    # zero value loci
    for n in range(1, N+1):
        zvl = n/(2.0 - abs(u[1:-1]))
        ax1.plot(u[1:-1], zvl, color='#888888', 
                linestyle='dashed')
        ax1.plot(u[1:-1], -zvl, color='#888888', 
            linestyle='dashed')
    # OTF line on AF plane
    for elem in w20LambdaBy2:
        otfY = elem*u # OTF line in AF with slope 2w_{20}/lambda
        ax1.plot(u, otfY)
    # intersections
    def get_intersections(b): 
        # b is tan(phi) or 2w_{20}/lambda
        n = np.linspace(1, np.floor(b), np.floor(b))
        u1 = 1 + np.sqrt(1 - n/b)
        u2 = 1 - np.sqrt(1 - n/b)
        y1 = u1*b
        y2 = u2*b
        u = np.hstack((u1, u2))
        y = np.hstack((y1, y2))
        return u, y
    for elem in w20LambdaBy2:
        intersectionsU, intersectionsY = get_intersections(elem)
        ax1.scatter(intersectionsU, intersectionsY,
                    marker='x', c='k', zorder=20)

    # OTF plots
    for elem in w20LambdaBy2:
        otf = (1 - np.abs(u)/2)*np.sinc(elem*u*(2 - np.abs(u)))
        ax2.plot(u, otf, label='$W_{20}' + '= {}\lambda/2$'.format(elem))
    # axis settings
    ax1.set_xlim(umin, umax)
    ax1.set_ylim(ymin, ymax)
    ax1.set_title('2-D AF of 1-D rect pupil P(x)', y=1.01)
    ax1.set_xlabel('u', fontsize=14)
    ax1.set_ylabel('y', fontsize=14)
    ax2.axhline(y=0, xmin=-2, xmax=2, color='#888888',
                zorder=0, linestyle='dashed')
    ax2.grid(axis='x')
    ax2.legend(fontsize=12)
    ax2.set_xlim(-2, 2); ax2.set_ylim(-0.2, 1.005)
    ax2.set_title("Optical Transfer Function", y=1.02)
    ax2.set_xlabel("u (scaled saptial frequency)", fontsize=14)
    #fig.tight_layout()
    plt.show()

plot_1dRectAF(w20LambdaBy2=[0, 1, 2, 3])

plot_1dRectAF(w20LambdaBy2=[0, 0.5, 0.99, 1.5, 3.6])

from scipy.integrate import quad
import warnings
warnings.simplefilter(action='error', category=np.ComplexWarning)
# Turn on the warning to ensure that the numerical integration is "reliable"?
warnings.simplefilter(action='always', category=sp.integrate.IntegrationWarning)

ifft = np.fft.fft
fftshift = np.fft.fftshift
fftfreq = np.fft.fftfreq

# cubic phase mask parameters
alpha = 90
gamma = 3
umin, umax = -2, 2
ymin, ymax = -60, 60
w20LambdaBy2 = [0, 5, 15] # amounts of defocus in units of wavelength (by 2)

uVec = np.linspace(umin, umax, 300)
N = 512  # number of samples along "t" ... and for FFT
L = 1

def gut(t, alpha, gamma, u): 
    return 0.5*np.exp(1j*alpha*((t + u/2)**gamma - (t - u/2)**gamma))

guy = np.empty((N, len(uVec)))
#roi = np.empty((N, len(uVec))) # for debugging & seeing the region of integration
t = np.linspace(-2*L, 2*L, N)
dt = (4*L)/(N-1)

for i, u in enumerate(uVec):
    g = np.zeros_like(t, dtype='complex64')
    if -1 <= u/2.0 < 0:
        mask = (t > (-u/2 - 1))*(t < (u/2 + 1))
        #roi[:, i] = mask.astype('float32')
        g[mask] = gut(t[mask], alpha, gamma, u)
        guy[:, i] = np.abs(fftshift(ifft(g)))
    elif 0 <= u/2.0 <= 1:
        mask = (t > (u/2 - 1))*(t < (-u/2 + 1))
        #roi[:, i] = mask.astype('float32')
        g[mask] = gut(t[mask], alpha, gamma, u)
        guy[:, i] = np.abs(fftshift(ifft(g)))

# Normalize to make maximum value = 1
guyMax = np.max(np.abs(guy.flat))
guy = guy/guyMax

yindex = fftshift(fftfreq(N, dt))
ymin, ymax = yindex[0], yindex[-1]

fig = plt.figure(figsize=(12, 7))
ax1 = fig.add_axes([0.12, 0, 0.5, 1.0]) # [*left*, *bottom*, *width*,*height*]
ax2 = fig.add_axes([0.66, 0.23, 0.32, 0.54])
    
im = ax1.imshow(guy**0.8, cmap=cm.YlGnBu_r, origin='lower',
               extent=[umin, umax, ymin, ymax],
               vmin=0.0, vmax=1.0,aspect=1./40)  
plt.colorbar(im, ax=ax1, shrink=0.55, aspect=35)

# OTF line in AF 
for elem in w20LambdaBy2:
    otfY = elem*uVec # OTF line in AF with slope 2w_{20}/lambda
    ax1.plot(uVec, otfY, alpha = 0.6, linestyle='solid')

ax1.set_xlim(umin, umax)
ax1.set_ylim(ymin, ymax)
ax1.set_xlabel('u', fontsize=14)
ax1.set_ylabel('y', fontsize=14)
ax1.set_title('2-D AF of 1-D cpm', y=1.01)

# Magnitude plots of the OTF of the cpm

def otf_cpm(t, alpha, gamma, u, w20LamBy2): 
    return (0.5*np.exp(1j*alpha*((t + u/2)**gamma - (t - u/2)**gamma))
            *np.exp(1j*2*np.pi*u*w20LamBy2*t))

def complex_quad(func, a, b, **kwargs):
    """Compute numerical integration of complex function between
    limits a and b
    Adapted from the following SO post:
    stackoverflow.com/questions/5965583/use-scipy-integrate-quad-to-integrate-complex-numbers
    """
    def real_func(x, *args):
        if args:
            return sp.real(func(x, *args))
        else:
            return sp.real(func(x))
    def imag_func(x, *args):
        if args:
            return sp.imag(func(x, *args))
        else:
            return sp.imag(func(x))
    real_integral = quad(real_func, a, b, **kwargs)
    imag_integral = quad(imag_func, a, b, **kwargs)
    return (real_integral[0] + 1j*imag_integral[0], real_integral[1:], imag_integral[1:])

for elem in w20LambdaBy2:
    Huw = np.empty_like(uVec, dtype='complex64')
    for i, u in enumerate(uVec):
        if -1 <= u/2.0 < 0:
            Huw[i] = complex_quad(func=otf_cpm, a=-u/2 - 1, b=u/2 + 1, args=(alpha, gamma, u, elem))[0]
        elif 0 <= u/2.0 <= 1:
            Huw[i] = complex_quad(func=otf_cpm, a=u/2 - 1, b= -u/2 + 1, args=(alpha, gamma, u, elem))[0]
    HuwMax = np.max(np.abs(Huw))
    ax2.plot(uVec, np.abs(Huw)/HuwMax, label='$W_{20}' + '= {}\lambda/2$'.format(elem))

ax2.legend(fontsize=12)
ax2.set_ylabel("Magnitude of OTF", fontsize=14)
ax2.set_xlabel("Spatial frequency, u", fontsize=14)
ax2.set_title('Optical Transfer Function of cpm', y=1.01)

plt.show()