%matplotlib inline

import sys
sys.path.append('../..')

from matplotlib import pylab

pylab.rcParams['figure.figsize'] = 12, 10

import functools
import numpy
import scipy
import scipy.special

from crocodile.clean import *
from crocodile.synthesis import *
from crocodile.simulate import *
from util.visualize import *
from arl.test_support import create_named_configuration

"""
Explanation: Illustration of w-imaging
End of explanation
"""


vlas = create_named_configuration('VLAA')
ha_range = numpy.arange(numpy.radians(0),
                        numpy.radians(90),
                        numpy.radians(90 / 36))
dec = numpy.radians(45)
vobs = xyz_to_baselines(vlas.data['xyz'], ha_range, dec)

# Wavelength: 5 metres 
wvl=5
uvw = vobs / wvl

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
ax = plt.figure().add_subplot(121, projection='3d')
ax.scatter(uvw[:,0], uvw[:,1] , uvw[:,2])
max_uvw = numpy.amax(uvw)
ax.set_xlabel('U [$\lambda$]'); ax.set_xlim((-max_uvw, max_uvw))
ax.set_ylabel('V [$\lambda$]'); ax.set_ylim((-max_uvw, max_uvw))
ax.set_zlabel('W [$\lambda$]'); ax.set_zlim((-max_uvw, max_uvw))
ax.view_init(20, 20)
pylab.show()

"""
Explanation: Generate baseline coordinates for an observation with the VLA over 6 hours, with a visibility recorded every 10 minutes. The phase center is fixed at a declination of 45 degrees. We assume that the imaged sky says at that position over the course of the observation.
Note how this gives rise to fairly large $w$-values.
End of explanation
"""


import itertools
vis = numpy.zeros(len(uvw), dtype=complex)
for u,v in itertools.product(range(-3, 4), range(-3, 4)):
    vis += 1.0*simulate_point(uvw, 0.010*u, 0.010*v)
plt.clf()
uvdist=numpy.sqrt(uvw[:,0]**2+uvw[:,1]**2)
plt.plot(uvdist, numpy.abs(vis), '.', color='r')

"""
Explanation: We can now generate visibilities for these baselines by simulation. We place three sources.
End of explanation
"""


theta = 2*0.05
lam = 18000
d,p,_=do_imaging(theta, lam, uvw, None, vis, simple_imaging)
show_image(d, "image", theta)

"""
Explanation: Using imaging, we can now reconstruct the image. The easiest option is to use simple imaging without a convolution function:
End of explanation
"""


step=int(theta*lam/10)
def zoom(x, y=step): pylab.matshow(d[y:y+2*step,x:x+2*step]) ; pylab.colorbar(shrink=.4,pad=0.025);  pylab.show()
from ipywidgets import interact
interact(zoom, x=(0,d.shape[0]-2*step,step), y=(0,d.shape[1]-2*step,step));

"""
Explanation: Zooming in shows the source structure in detail
End of explanation
"""


wstep=100
wcachesize=int(numpy.ceil(numpy.abs(uvw[:,2]).max()/wstep))
print("Making w-kernel cache of %d kernels" % wcachesize)
wcache=pylru.FunctionCacheManager(w_kernel, wcachesize)
imgfn = functools.partial(w_cache_imaging, kernel_cache=w_conj_kernel_fn(wcache),
                          wstep=wstep, Qpx=2, NpixFF=256, NpixKern=31)

d_w,p_w,_=do_imaging(theta, lam, uvw, None, vis, imgfn)
show_image(d_w, "image", theta)

step=int(theta*lam/10)
def zoom_w(x=720,y=step): pylab.matshow(d_w[y:y+2*step,x:x+2*step]); pylab.colorbar(shrink=.4,pad=0.025); pylab.show()
interact(zoom_w, x=(0,d.shape[0]-2*step,step), y=(0,d.shape[1]-2*step,step))

"""
Explanation: If we use convolution kernels for $w$-reprojection, we can improve the sharpness of imaging. First we make a cache to hold the convolution kernels.
End of explanation
"""


