t0 = Time("2019-01-01T01:00")
stationname = "LV614LBA"

def station_to_itrs(station_name):
    """Returns an astropy ITRS coordinate for a given station (e.g. 'CS001LBA')"""
    x,y,z = db.phase_centres[station_name]
    return ITRS(x=x*u.m, y=y*u.m, z=z*u.m)

def station_to_earthlocation(station_name):
    """Returns an astropy EarthLocation for a given station (e.g. 'CS001LBA')"""
    x,y,z = db.phase_centres[station_name]
    return EarthLocation(x=x*u.m, y=y*u.m, z=z*u.m)

"""
Explanation: Test / fiddle around
End of explanation
"""


test_coord_skycoord = AltAz(location=station_to_earthlocation(stationname),
                            alt=89.0*u.deg, az=0*u.deg,
                            obstime=t0).transform_to(GCRS)

test_coord_gcrs = test_coord_skycoord.cartesian.get_xyz().value

"""
Explanation: We will test at zenith:
End of explanation
"""


def gcrs_to_itrs(t):
    return SkyCoord(x=np.array([1,0,0]),
                    y=np.array([0,1,0]),
                    z=np.array([0,0,1]),
                    representation_type='cartesian'
                   ).transform_to(ITRS(obstime=t)).cartesian.get_xyz().value

"""
Explanation: Let's make a matrix to go from sky (GCRS) to ITRS:
End of explanation
"""


m31_skycoord = SkyCoord.from_name("M31")

m31_gcrs = m31_skycoord.cartesian.get_xyz().value

gcrs_to_itrs(t0).dot(m31_gcrs)

m31_skycoord.transform_to(ITRS(obstime=t0))

"""
Explanation: Test that this matrix does the same thing as an astropy transformation:
End of explanation
"""


station_itrs = station_to_itrs(stationname).cartesian.get_xyz().value
station_itrs /= norm(station_itrs)

np.rad2deg(np.arccos((gcrs_to_itrs(t0)@test_coord_gcrs)@station_itrs))

"""
Explanation: Yay.
To test whether we computed zenith, let's test whether the ITRS-vector of zenith corresponds to the ITRS-vector that points towards the station phase center.
End of explanation
"""


a = np.sqrt(.5)
dipoles_pqr = np.array([[ a,  a, 0],
                        [ a, -a, 0],
                        [-a, -a, 0],
                        [-a,  a, 0]]).T

"""
Explanation: Ok, somehow there is an error of about one degree. That's probably because the station plane is not exactly tangent to a spherical earth. Let's ignore for now.
Define the dipoles, as unit vectors in the PQR system, and convert them to ETRS.
End of explanation
"""


db.pqr_to_etrs[stationname]@(db.pqr_to_etrs[stationname].T)

"""
Explanation: Note that we assume here that ITRS = ETRS, which is true to the meter level. (We could use Michiel's etrsitrs package here, but I think we've already got enough coordinate frames.)
So, now we have three coordinate frames:
 * GCRS (sky)
 * ITRS (Earth)
 * pqr (station)
We have two matrices to convert between these:
 * pqr_to_etrs
 * gcrs_to_itrs
To convert the other way around, these matrices need to be inverted, which is the same as transposing them because they are orthonormal.
End of explanation
"""


test_coord_itrs = gcrs_to_itrs(t0)@test_coord_gcrs

dipoles_itrs = db.pqr_to_etrs[stationname]@dipoles_pqr

dipoles_projected_itrs = dipoles_itrs - \
                         test_coord_itrs.dot(dipoles_itrs) * test_coord_itrs[:, np.newaxis]

"""
Explanation: Project the dipoles to the plane orthogonal to the test_coord. Do this in ITRS frame.
End of explanation
"""


dipoles_projected_itrs.T@test_coord_itrs

"""
Explanation: Check that the projected dipoles are indeed orthogonal to the test coordinate:
End of explanation
"""


dipoles_projected_gcrs = (gcrs_to_itrs(t0).T)@dipoles_projected_itrs

"""
Explanation: Yay.
End of explanation
"""


def rot_ra(phi):
    """Rotates back along the z-axis with an angle phi"""
    return  np.array([[np.cos(-phi), - np.sin(-phi), 0],
                      [np.sin(-phi),   np.cos(-phi), 0],
                      [0, 0, 1]])

"""
Explanation: These are the dipoles projected to GCRS. All we need to do now is plot them on the plane orthogonal to the vector test_coord_gcrs.
Unfortunately I don't know how to do this properly. Idea for now: rotate along the z-axis to set RA to zero, then rotate along the y-axis to set DEC to zero. Then the vectors should point straight up.
End of explanation
"""


test_coord_gcrs_rotra = rot_ra(test_coord_skycoord.ra)@test_coord_gcrs

test_coord_skycoord_rotra = SkyCoord(x=test_coord_gcrs_rotra[0],
                                     y=test_coord_gcrs_rotra[1],
                                     z=test_coord_gcrs_rotra[2],
                                     representation_type='cartesian',frame=GCRS).transform_to(GCRS)

test_coord_skycoord.ra.degree

test_coord_skycoord_rotra.ra.degree

"""
Explanation: Let's check this matrix by rotating test_coord back to ra=0:
End of explanation
"""


test_coord_skycoord.dec.degree

test_coord_skycoord_rotra.dec.degree

"""
Explanation: Yay. Check that the declination does not change:
End of explanation
"""


def rot_dec(theta):
    """Rotates back along the y-axis with an angle theta"""
    return  np.array([[np.cos(-theta), 0, - np.sin(-theta)],
                      [0,              1,                0],
                      [np.sin(-theta), 0,   np.cos(-theta)]])

"""
Explanation: Yay.
Now rotate the declination.
End of explanation
"""


test_coord_gcrs_rotdec = rot_dec(test_coord_skycoord_rotra.dec)@test_coord_gcrs_rotra

test_coord_skycoord_rotdec = SkyCoord(x=test_coord_gcrs_rotdec[0],
                                      y=test_coord_gcrs_rotdec[1],
                                      z=test_coord_gcrs_rotdec[2],
                                      representation_type='cartesian',frame=GCRS).transform_to(GCRS)

test_coord_skycoord_rotdec.dec.degree

test_coord_skycoord_rotra.ra.degree

test_coord_skycoord_rotdec.ra.degree

"""
Explanation: Test that this works (note we need to rotate back the RA first):
End of explanation
"""


def gcrs_to_harry(pointing_skycoord: SkyCoord):
    """Rotate back so that at coord_skycoord is at ra=0, dec=0"""
    return rot_dec(pointing_skycoord.dec)@rot_ra(pointing_skycoord.ra)

gcrs_to_harry(test_coord_skycoord)@test_coord_gcrs

gcrs_to_harry(m31_skycoord)@m31_gcrs

dipoles_projected_gcrs

gcrs_to_harry(test_coord_skycoord)@dipoles_projected_gcrs

"""
Explanation: Yay.
Now let's introduce another coordinate system, which is the rotated system so that the viewing direction looks straight at the origin (and has the same north as GCRS, because of the order in which the rotations from GCRS are done). Let's call this system harry.
End of explanation
"""


dipoles_gcrs = (gcrs_to_itrs(t0).T)@db.pqr_to_etrs[stationname]@dipoles_pqr

gcrs_to_harry(test_coord_skycoord)@dipoles_gcrs

dipoles_harry = gcrs_to_harry(test_coord_skycoord)@dipoles_gcrs

"""
Explanation: Ok, so we didn't need to project the dipoles, since in the harry coordinate frame, projection corresponds to setting the first coordinate to zero.
End of explanation
"""


def station_to_earthlocation(station_name):
    """Returns an astropy EarthLocation for a given station (e.g. 'CS001LBA')"""
    x,y,z = db.phase_centres[station_name]
    return EarthLocation(x=x*u.m, y=y*u.m, z=z*u.m)

def gcrs_to_itrs(t):
    return SkyCoord(x=np.array([1,0,0]),
                    y=np.array([0,1,0]),
                    z=np.array([0,0,1]),
                    representation_type='cartesian'
                   ).transform_to(ITRS(obstime=t)).cartesian.get_xyz().value

def rot_ra(phi):
    """Rotates back along the z-axis with an angle phi"""
    return  np.array([[np.cos(-phi), - np.sin(-phi), 0],
                      [np.sin(-phi),   np.cos(-phi), 0],
                      [0, 0, 1]])



def rot_dec(theta):
    """Rotates back along the y-axis with an angle theta"""
    return  np.array([[np.cos(-theta), 0, - np.sin(-theta)],
                      [0,              1,                0],
                      [np.sin(-theta), 0,   np.cos(-theta)]])

def gcrs_to_harry(pointing_skycoord: SkyCoord):
    """Rotate back so that at coord_skycoord is at ra=0, dec=0"""
    return rot_dec(pointing_skycoord.dec)@rot_ra(pointing_skycoord.ra)

def get_dipoles_harry(pointing_skycoord, time, stationname):
    a = np.sqrt(.5)
    dipoles_pqr = np.array([[ a,  a, 0],
                            [ a, -a, 0],
                            [-a, -a, 0],
                            [-a,  a, 0]]).T
    dipoles_gcrs = (gcrs_to_itrs(time).T) @ db.pqr_to_etrs[stationname] @ dipoles_pqr
    dipoles_harry = gcrs_to_harry(pointing_skycoord) @ dipoles_gcrs
    return dipoles_harry

dipoles_harry = get_dipoles_harry(SkyCoord.from_name("M31"), Time.now(), "IE613LBA")

def plot_dipoles_harry(pointing_str: str, timedelta: float, stationname: str):
    """Plot projected dipoles"""
    fig, ax = plt.subplots(1)
    
    t0 = Time("2019-01-01T01:00")
    
    if pointing_str == "zenith":
        pointing = AltAz(location=station_to_earthlocation(stationname),
                         alt=89.0*u.deg, az=0*u.deg,
                         obstime=t0).transform_to(GCRS)
    else:
        pointing = SkyCoord.from_name(pointing_str)

    time = t0 + timedelta * u.day
    
    dipoles_harry = get_dipoles_harry(pointing,
                                      time,
                                      stationname)
    
    x, y = dipoles_harry[1:3]
    ax.grid()
    ax.arrow(x[2], y[2], x[0] - x[2], y[0] - y[2], head_width=0.1, color='r',
             length_includes_head=True)
    ax.arrow(x[1], y[1], x[3] - x[1], y[3] - y[1], head_width=0.1, color='b',
             length_includes_head=True)
    ax.set_aspect(1)
    ax.set_xlim(-1, 1)
    ax.set_ylim(-1, 1)
    ax.set_xticks([-1,0,1])
    ax.set_yticks([-1,0,1])
    ax.tick_params(axis="x", bottom=False, top=False, labelbottom=False)
    ax.tick_params(axis="y", left=False, right=False, labelleft=False)
    ax.set_xlabel(r"$\Delta  \alpha$", fontsize=16)
    ax.set_ylabel(r"$\Delta  \delta$", fontsize=16)
    ax.set_title(f"Dipoles of {stationname} as seen from \n{pointing_str} at {time.iso[:16]}")
    _ = ax.plot(0, 0, 'kx');
    return fig

plot_dipoles_harry("M31", 0, "IE613LBA");

plot_dipoles_harry("zenith", 0, "IE613LBA");

"""
Explanation: Ok, let's put all of this in a function:
Real code
End of explanation
"""


plot_dipoles_harry("zenith", 0, "CS002LBA");

plot_dipoles_harry("zenith", 0, "LV614LBA");

"""
Explanation: Try from the superterp (where up should be North):
End of explanation
"""


from ipywidgets import interact, interactive, fixed, interact_manual

interactive_plot = interactive(plot_dipoles_harry,
                               stationname=list(db.phase_centres.keys()),
                               timedelta=(0., 1., 0.05),
                               pointing_str=["M31", "Cas A", "zenith"])
output = interactive_plot.children[-1]
output.layout.height = '350px'
interactive_plot

for framenum, time_delta in enumerate(np.linspace(0, 1, 120)):
    f = plot_dipoles_harry("zenith", time_delta, "CS002LBA");
    _ = f.savefig(f"frame{framenum:03d}.png");

"""
Explanation: Now make it interactive:
End of explanation
"""


