# NOTE: You also need a version of the "measures" tables (ephemerides and geodetic data).
#       To download the latest and configure casacore to use them, run:
#
#           mkdir $HOME/IERS && cd $HOME/IERS
#           wget ftp://ftp.astron.nl/outgoing/Measures/WSRT_Measures.ztar
#           tar xfz WSRT_Measures.ztar
#
#           echo "measures.directory: $HOME/IERS" > $HOME/.casarc 
import casacore.measures # pip install python-casacore
from math import sin, cos
import numpy
import datetime

def subtract(a, b):
    return numpy.array([x - y for x, y in zip(a, b)])

def get_measure(itrf: list[float]) -> casacore.measures.measures:
    """ Return a measure object, configured for the specified location. """
    measure = casacore.measures.measures()
    frame_location = measure.position("ITRF",*[f"{x}m" for x in itrf])

    result = measure.do_frame(frame_location)
    assert result == True

    return measure

def set_measure_time(measure: casacore.measures.measures, utc_time: datetime.datetime):
    """ Configure the measure object for the specified time. """
    frame_time = measure.epoch("UTC", utc_time.isoformat(' '))

    result = measure.do_frame(frame_time)
    assert result == True

def get_direction_vector(measure: casacore.measures.measures, pointing: dict) -> numpy.array:
    """ Compute a direction vector for the given pointing, relative to the measure. """
    angles = meas.measure(pointing, "ITRF")
    angle0, angle1 = angles["m0"]["value"], angles["m1"]["value"]

    # Convert polar to carthesian coordinates
    # see also https://github.com/casacore/casacore/blob/e793b3d5339d828a60339d16476bf688a19df3ec/casa/Quanta/MVDirection.cc#L67
    direction_vector = numpy.array(
                       (cos(angle0)*cos(angle1),
                        sin(angle0)*cos(angle1),
                        sin(angle1)))

    return direction_vector

def get_delay(reference_direction_vector: numpy.array, relative_itrf: numpy.array) -> float:
    """ Compute the delay to apply, in seconds, to align signals for a distance of `relative_itrf` in the
        direction of `reference_direction_vector`. """
    speed_of_light = 299792458.0

    return numpy.dot(reference_direction_vector, relative_itrf) / speed_of_light

# create a frame tied to the reference position
reference_itrf = [3826577.066, 461022.948, 5064892.786] # CS002LBA, in ITRF2005 epoch 2012.5
meas = get_measure(reference_itrf)

# set the timestamp to solve for
timestamp = datetime.datetime(2021,1,1,0,0,5)
set_measure_time(meas, timestamp)

# obtain the direction vector for a specific pointing
pointing = meas.direction("J2000","0deg","0deg")
reference_dir_vector = get_direction_vector(meas, pointing)

# compute the delays for an antennas w.r.t. the reference position
antenna_itrf = [[3826923.546, 460915.441, 5064643.489]] # CS001LBA, in ITRF2005 epoch 2012.5
antenna_relative_itrf = [subtract(pos, reference_itrf) for pos in antenna_itrf]
delays = [get_delay(reference_dir_vector, relative_itrf) for relative_itrf in antenna_relative_itrf]

# print the delays
import pprint
pprint.pprint(delays)