%matplotlib inline
import matplotlib.pyplot as plt
from matplotlib import cm, colors
import numpy as np
from ipywidgets import interact
plt.rcParams['figure.figsize'] = 16, 8
import sys
sys.path.append('../..')
from crocodile.simulate import *
from util.visualize import *

"""
Explanation: Coordinates used in interferometry
Astronomy deals with a lot of graographic and celestial coordinate systems - and approximations thereof. This document is a quick introduction into the coordinate systems used by Crocodile:

hour angle/declination: positions on the sky relative to observer position
$(u,v,w)$: station / baseline positions (used for $uv$-grid coordinates)
$(l,m,n)$: sky directional cosines relative to the delay tracking centre (used for image coordinates)
End of explanation
"""


interact(visualise_uvw, latitude=(0,90,10), hour_angle=(-90,90,10), declination=(-90,90,10));

"""
Explanation: 1. Hour angle / declination
The only spherical coordinate system we are going to use for celestial coordinates. It is defined by:

declination == 90° points north towards the celestial north pole (parallel to the earth axis)
declination == -90° points south towards the celestial south pole (parallel to the earth axis)
hour_angle == 90°, declination == 0° points west locally
hour_angle == -90°, declination == 0° points east locally

(Note: it is more common to express celestial cooordinates in terms of right ascension and declination. The conversion rule is simply RA = LST - HA, where LST is the local sidereal time, correcting for the earth's rotation. Note that RA increases eastwards, while HA increases westwards!)
See this interactive visualisation - the $w$ arrow points towards the direction specified:
End of explanation
"""


vlas = numpy.genfromtxt("../../data/configurations/VLA_A_hor_xyz.csv", delimiter=",")
size = np.amax(np.abs(np.ravel(vlas)))

def draw_transformed(latitude, hour_angle, declination):
    xyz = xyz_at_latitude(vlas, np.radians(latitude))
    uvws = xyz_to_uvw(xyz, np.radians(hour_angle), np.radians(declination))
    ncp = xyz_to_uvw(np.array([[0,0,10000]]), np.radians(hour_angle), np.radians(declination))
    
    fig = plt.figure()
    ax = fig.add_subplot(121, projection='3d')
    ax.scatter(*np.transpose(xyz))
    ax.set_xlabel('X [$m$]'); ax.set_ylabel('Y [$m$]'); ax.set_zlabel('Z [$m$]')
    make_arrow(ax, [0,0,0], [10000,0,0], "black", "geographic east")
    make_arrow(ax, [0,0,0], [0,0,10000], "black", "celestial north")
    ax.set_xlim(-size,size); ax.set_ylim(-size,size); ax.set_zlim(-size,size)
    ax = fig.add_subplot(122, projection='3d')
    ax.scatter(*np.transpose(uvws))
    make_arrow(ax, [0,0,0], [0,0,10000], "black", "phase centre")
    make_arrow(ax, [0,0,0], ncp[0], "gray", "celestial north")
    ax.set_xlabel('u [$m$]'); ax.set_ylabel('v [$m$]'); ax.set_zlabel('w [$m$]')
    ax.set_xlim(-size,size); ax.set_ylim(-size,size); ax.set_zlim(-size,size)
    plt.show()
interact(draw_transformed, latitude=(-90,90,10), hour_angle=(-90,90,10), declination=(-90,90,10));

"""
Explanation: 2. Local coordinates
For imaging, it is a good idea to express celestial directions relative to a certain "middle" position on the sky - conceptually the position that will appear in the middle of our image. This is called the phase tracking centre.
From this we derive the $(u,v,w)$ coordinate system, which is a local coordinate system with its origin at the observer's location (centre of the telescope). The axes are determined as follows:


$w$ points towards the phase tracking centre,


$v, w$ and the north celestial pole are always on a common plane and


$u$ completes a right-handed coordinate system (cross product of $v$ and $w$)


This has the following properties:


when the direction of observation is the north celestial pole (hour_angle=0, declination=90), the UVW coordinates are aligned with XYZ


when $w$ is on the local meridian (hour_angle=0), $u$ points East


when the direction of observation is at declination=0, hour_angle=-90 makes $w$ point East.


(adapted from http://casa.nrao.edu/Memos/CoordConvention.pdf)
Here is an example of how we would translate the station coordinates of the VLA (http://www.vla.nrao.edu/) into the $(u,v,w)$ coordinate system appropriate for an observation:
End of explanation
"""


cmap = cm.ScalarMappable(norm=colors.Normalize(vmin=-45,vmax=45), cmap=cm.ocean)
def draw_range(latitude, declination):
    xyz = xyz_at_latitude(vlas, np.radians(latitude))
    ax = plt.figure().add_subplot(121, projection='3d')
    for hour_angle in numpy.linspace(-45, 45, 10):
        
        # Determine rotated station UVWs
        uvws = xyz_to_uvw(xyz, np.radians(hour_angle), np.radians(declination))

        ax.scatter(*np.transpose(uvws), c=[cmap.to_rgba(hour_angle)]*len(uvws), edgecolor=cmap.to_rgba(hour_angle))
    make_arrow(ax, [0,0,0], [0,0,10000], "black", "phase centre")
    ax.set_xlabel('u [$m$]'); ax.set_ylabel('v [$m$]'); ax.set_zlabel('w [$m$]')
    ax.set_xlim(-size,size); ax.set_ylim(-size,size); ax.set_zlim(-size,size)
    plt.show()
interact(draw_range, latitude=(-90,90,10), hour_angle=(-90,90,10), declination=(-90,90,10));

"""
Explanation: Note how the only observation where the target does not move with the hour angle is latitude = +-90 and declination = +-90. This means we are observing from the earth's north/south pole straight up/down, which makes stations rotate around us in a common plane.
Of course, the hour angle is not constant along an observation due to the earth's rotation. For example, let us assume a 6 hour observation in which the hour angle of the phase centre goes from $-45°$ to $45°$. Then we get a distribution of station $(u,v,w)$ coordinates, depending on telescope latitude and phase centre declination:
End of explanation
"""


def draw_baselines(latitude, declination):
    xyz = xyz_at_latitude(vlas, np.radians(latitude))
    ax = plt.figure().add_subplot(111, projection='3d')
    ax.set_xlabel('u [$m$]'); ax.set_ylabel('v [$m$]'); ax.set_zlabel('w [$m$]')
    ax.set_xlim(-size,size); ax.set_ylim(-size,size); ax.set_zlim(-size,size)
    uvws = xyz_to_uvw(xyz, np.radians(0), np.radians(declination))
    # Draw an arrow for every combination of station coordinates
    import itertools
    for source, target in itertools.combinations(uvws, 2):
        make_arrow(ax, source, target, 'b', None, scale=10, lw=1)
    ax.view_init(elev=80., azim=-45.)
    plt.show()
interact(draw_baselines, latitude=(-90,90,10), hour_angle=(-90,90,10), declination=(-90,90,10));

"""
Explanation: Observe that with latitude = 0 multiple stations do not move at all. This is because the VLA layout would have a line of stations going in a straight south-north line if placed on the equator. Fortunately, that is not where it is actually located.
3. Baselines
For interferometry we derive data about the sky from the interference between station/antenna signals. Every such pair forms a baseline. The interference pattern measured for it largely depends on its length and orientation. The number of possible station combinations grows quadratically with the number stations, which means that even for relatively few stations we get very many baselines:
End of explanation
"""


def draw_baseline_orientations(latitude, declination):
    xyz = xyz_at_latitude(vlas, np.radians(latitude))
    def draw_orientations(ax, xyz):
        ax.set_xlabel('u [$m$]'); ax.set_ylabel('v [$m$]'); ax.set_zlabel('w [$m$]')
        ax.set_xlim(-2*size,2*size); ax.set_ylim(-2*size,2*size); ax.set_zlim(-2*size,2*size)
        for hour_angle in numpy.linspace(-45, 45, 10):
            # Determine rotated station UVWs, generate baselines from pairs
            uvws = xyz_to_uvw(xyz, np.radians(hour_angle), np.radians(declination))
            bls = baselines(uvws)
            # Draw baseline orientation, both directions
            color = len(bls)*[cmap.to_rgba(hour_angle)]
            ax.scatter(*np.transpose(bls), c=color, edgecolor=cmap.to_rgba(hour_angle))
            ax.scatter(*np.transpose(-bls), c=color, edgecolor=cmap.to_rgba(hour_angle))
    fig = plt.figure()
    ax = fig.add_subplot(121, projection='3d')
    draw_orientations(ax, xyz)
    ax = fig.add_subplot(122, projection='3d')
    draw_orientations(ax, xyz)
    ax.view_init(elev=10., azim=-10.)
    plt.show()
interact(draw_baseline_orientations, latitude=(-90,90,10), hour_angle=(-90,90,10), declination=(-90,90,10));

"""
Explanation: In practice, we do not care about the baselines' spatial location, but just their direction and length. The difference between the station's $(u,v,w)$ coordinates yields the $(u,v,w)$ coordinates of the baseline.
End of explanation
"""


interact(visualise_lmn, hour_angle=(-90,0,10), declination=(0,90,10));

"""
Explanation: Important things to note here:

Strictly speaking every baseline has two coordinates $(u,v,w)$ and $(-u,-v,-w)$ depending on the order of the subtraction. We typically (arbitrarily) select only one direction to consider for every baseline. In the code above we draw all data twice for that reason.
In imaging, baseline $(u,v,w)$ coordinates are often measured in wavelenghts, as that makes it easier to reason about interference patterns. For the moment, we are still sticking with simple meter units.

4. Image/sky coordinates
To complete our handling of coordinates, we need to revisit sky coordinates. While we can identify sky positions using the hour angle/declination system, it is not suitable for imaging. Just consider what would happen if we tried to image close to the celestial north pole. Instead, we will be using directional cosines, which are unit vectors pointing into the direction in question:
$$\begin{pmatrix}x\y\z\end{pmatrix} = \begin{pmatrix}
    \text{sin}(ha) * \text{cos}(dec) \
    \text{cos}(ha) * \text{cos}(dec) \
    \text{sin}(dec)
\end{pmatrix}$$
For example, we can calculate the directional cosines of the phase centre in the $(x,y,z)$ coordinate system:
End of explanation
"""


ax = pl.figure().add_subplot(111, projection='3d')
ax.set_xlabel('m [$1$]'); ax.set_ylabel('l [$1$]'); ax.set_zlabel('n [$1$]')
make_arrow(ax, (0,0,-1),(0,0,0), "black", "phase centre")
lons = numpy.linspace(-numpy.pi, numpy.pi, 80)
lats = numpy.linspace(numpy.pi/2, 0, 20)
l, m, n = circular_to_xyz(numpy.outer(lons, numpy.ones(len(lats))),
                          numpy.outer(numpy.ones(len(lons)), lats))
ax.plot_surface(l, m, n-1, rstride=1, cstride=1, linewidth=0, alpha=0.4, color='white')
ax.plot_surface(l, m, numpy.array([[0]]), rstride=1, cstride=1, linewidth=0, alpha=0.4, color='white')
ax.view_init(elev=-35, azim=25)
pl.show()

"""
Explanation: Note that in contrast to before, the sphere here stands for the sky instead of the earth. The sky is assumed to be very far away ("far field") to the point where we only care about directions, not distance. This means that we only have two degrees of freedom despite the fact that we are using three components.
However, the local $(x,y,z)$ coordinate system is not particularly useful for imaging either. Just as with $(u,v,w)$ coordinates, we want to eliminate the influence of the observation direction. Therefore we define the $(l,m,n)$ coordinate system as parallel to the $(u,v,w)$ coordinate system, but with its origin at the phase tracking centre on the sky sphere:
End of explanation
"""


