%matplotlib inline

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

from matplotlib import pylab

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

import numpy
import scipy
import scipy.special

from crocodile.clean import *
from crocodile.synthesis import *
from crocodile.simulate import *
from crocodile.antialias import *

vlas=numpy.genfromtxt("../../data/configurations/VLA_A_hor_xyz.csv", delimiter=",")
vobs=xyz_to_baselines(vlas, numpy.arange(0,numpy.pi,0.04) ,  numpy.pi/4)
yyone=simulate_point(vobs/5, 0.001, 0.001)
yytwo=yyone+ simulate_point(vobs/5, 0.0025, 0.0025)*5

"""
Explanation: Anti-Aliasing Functions in Interferometry
Setup the data problem. Will use field of view of 0.002 radian. One
source withing the field of view and a 5 times stronger just outside.
End of explanation
"""


d,p,_=do_imaging(2*0.002, 2*15000, vobs/5, None, yyone, simple_imaging)
pylab.matshow(d)  ; pylab.colorbar();  pylab.show()

"""
Explanation: Simple imaging (no convolution) with single source within field of view:
End of explanation
"""


d,p,_=do_imaging(2*0.002, 2*15000, vobs/5, None, yytwo, simple_imaging)
pylab.matshow(d)  ; pylab.colorbar();  pylab.show()
print(d[40:60,40:60].std())

"""
Explanation: Simple imaging (no convolution) with both sources, shows the strong
source being "aliased" back into the field of view:
End of explanation
"""


over = 64
support = 10
x0 = 0.27
aa=anti_aliasing_function(120, 0, 4)
pylab.plot(coordinates(120), aa);
aa2 = sze_tan_grid_correction(support//2, x0, 120)
pylab.plot(coordinates(120), aa2);
aa3, kv3o = anti_aliasing_calculate(120, over, support); aa3 /= aa3[60]
pylab.plot(coordinates(120), aa3);
legend = ["PSWF", "Sze Tan", "PSWF 'Schwab?'"]
pylab.legend(legend)

"""
Explanation: Anti-aliasing function
This is an example anti-aliasing function to use:
End of explanation
"""


r = numpy.arange(-over*(support//2), over*((support+1)//2)) / over
print(over * support, r.shape)
#slc = slice(over * support // 2 - 10, over * support // 2 + 10) # 120,136)
slc = slice(0, over * support)
kv=kernel_oversample(aa, over, support)
kv/=numpy.sum(kv[0])
pylab.plot(r, numpy.transpose(kv).flatten().real)
kv2=kernel_oversample(aa2, over, support)
kv2/=numpy.sum(kv2)/over
pylab.plot(r[slc], numpy.transpose(kv2).flatten().real[slc])
kv3=kernel_oversample(aa3, over, support)
kv3/=numpy.sum(kv3)/over
pylab.plot(r[slc], numpy.transpose(kv3).flatten().real[slc])

kv4 = sze_tan_gridder(support//2, x0, over)
pylab.plot(r[slc], numpy.transpose(kv4).flatten().real[slc])
kv2=kv4

pylab.legend(legend + ['Sze Tan Direct']);

"""
Explanation: After FFT-ing and extracting the middle this is what the anti-aliasing
function looks like:
End of explanation
"""


pylab.plot(r, numpy.transpose(kv)[::-1].flatten().imag)
pylab.plot(r, numpy.transpose(kv2)[::-1].flatten().imag)
pylab.plot(r, numpy.transpose(kv3)[::-1].flatten().imag)
pylab.legend(legend);

"""
Explanation: Imaginary part is close to nil:
End of explanation
"""


d,p,_=do_imaging(2*0.002, 2*15000, vobs/5, None, yyone, conv_imaging, kv=kv)
pylab.matshow(-d); pylab.colorbar();  pylab.show()

"""
Explanation: Gridding with the anti-aliasing function:
This is the image of single source within field of view without correcting the taper . Note that the strength of the source is
surpressed.
End of explanation
"""


d,p,_=do_imaging(2*0.002, 2*15000, vobs/5, None, yyone, conv_imaging, kv=kv)
pylab.matshow(-d/numpy.outer(aa, aa))  ; pylab.colorbar();  pylab.show()

"""
Explanation: This is the image of single source within field of view. Note that the
image is divided by the anti-aliasing funciton to correct for the
taper introduced by it:
End of explanation
"""


d,_,_=do_imaging(2*0.002, 2*15000, vobs/5, None, yytwo, conv_imaging, kv=kv)
pylab.matshow(d/numpy.outer(aa, aa)); pylab.colorbar(); pylab.title(legend[0]); pylab.show()
print((d/aa)[40:60,40:60].std())
d,_,_=do_imaging(2*0.002, 2*15000, vobs/5, None, yytwo, conv_imaging, kv=kv2)
pylab.matshow(d/numpy.outer(aa2, aa2)); pylab.colorbar(); pylab.title(legend[1]); pylab.show()
print((d/aa2)[40:60,40:60].std())
d,_,_=do_imaging(2*0.002, 2*15000, vobs/5, None, yytwo, conv_imaging, kv=kv3)
pylab.matshow(d/numpy.outer(aa3, aa3))  ; pylab.colorbar(); pylab.title(legend[2]); pylab.show()
print((d/aa3)[40:60,40:60].std())

"""
Explanation: This is the image with source outside of field of view also: note that
is entirely surpressed. All that's left is the fringes around the
source's off-screen position.
End of explanation
"""


d,p,_=do_imaging(2*0.002, 2*15000, vobs/5, None, yytwo, conv_imaging, kv=kv)
pylab.matshow(d)  ; pylab.colorbar();  pylab.show()
print((d)[40:60,40:60].std())
d,p,_=do_imaging(2*0.002, 2*15000, vobs/5, None, yytwo, conv_imaging, kv=kv2)
pylab.matshow(d)  ; pylab.colorbar();  pylab.show()
print((d)[40:60,40:60].std())
d,p,_=do_imaging(2*0.002, 2*15000, vobs/5, None, yytwo, conv_imaging, kv=kv3)
pylab.matshow(d)  ; pylab.colorbar();  pylab.show()
print((d)[40:60,40:60].std())

"""
Explanation: As above but without correcting the taper.
End of explanation
"""


