# Script to generate random targets, observatory locations, and times, and
# run these using the Starlink rv command to generate reference values for the
# velocity frame corrections. This requires that Starlink is installed and that
# the rv command is in your PATH. More information about Starlink can be found
# at http://starlink.eao.hawaii.edu/starlink

if __name__ == "__main__":
    from random import choice
    from subprocess import check_output

    import numpy as np

    from astropy import units as u
    from astropy.coordinates import Angle, SkyCoord
    from astropy.table import QTable
    from astropy.time import Time

    np.random.seed(12345)

    N = 100

    tab = QTable()
    target_lon = np.random.uniform(0, 360, N) * u.deg
    target_lat = np.degrees(np.arcsin(np.random.uniform(-1, 1, N))) * u.deg
    tab["target"] = SkyCoord(target_lon, target_lat, frame="fk5")
    tab["obstime"] = Time(
        np.random.uniform(Time("1997-01-01").mjd, Time("2017-12-31").mjd, N),
        format="mjd",
        scale="utc",
    )
    tab["obslon"] = Angle(np.random.uniform(-180, 180, N) * u.deg)
    tab["obslat"] = Angle(np.arcsin(np.random.uniform(-1, 1, N)) * u.deg)
    tab["geocent"] = 0.0
    tab["heliocent"] = 0.0
    tab["lsrk"] = 0.0
    tab["lsrd"] = 0.0
    tab["galactoc"] = 0.0
    tab["localgrp"] = 0.0

    for row in tab:
        # Produce input file for rv command
        with open("rv.input", "w") as f:
            f.write(
                f"{row['obslon'].to_string('deg', sep=' ')}"
                f" {row['obslat'].to_string('deg', sep=' ')}\n"
            )
            f.write(
                f"{row['obstime'].datetime.year} {row['obstime'].datetime.month}"
                f" {row['obstime'].datetime.day} 1\n"
            )
            f.write(row["target"].to_string("hmsdms", sep=" ") + " J2000\n")
            f.write("END\n")

        # Run Starlink rv command
        check_output(["rv", "rv.input"])

        # Parse values from output file
        lis_lines = []
        started = False
        for lis_line in open("rv.lis"):
            if started and lis_line.strip() != "":
                lis_lines.append(lis_line.strip())
            elif "LOCAL GROUP" in lis_line:
                started = True

        # Some sources are not observable at the specified time and therefore don't
        # have entries in the rv output file
        if len(lis_lines) == 0:
            continue

        # If there are lines, we pick one at random. Note that we can't get rv to
        # run at the exact time we specified in the input, so we will re-parse the
        # actual date/time used and replace it in the table
        lis_line = choice(lis_lines)

        # The column for 'SUN' has an entry also for the light travel time, which
        # we want to ignore. It sometimes includes '(' followed by a space which
        # can cause issues with splitting, hence why we get rid of the space.
        lis_line = lis_line.replace("(  ", "(").replace("( ", "(")
        (
            year,
            month,
            day,
            time,
            zd,
            row["geocent"],
            row["heliocent"],
            _,
            row["lsrk"],
            row["lsrd"],
            row["galactoc"],
            row["localgrp"],
        ) = lis_line.split()
        row["obstime"] = Time(
            f"{year}-{month}-{day}T{time}:00", format="isot", scale="utc"
        )

    # We sampled 100 coordinates above since some may not have results - we now
    # truncate to 50 sources since this is sufficient.
    tab[:50].write("reference_rv.ecsv", format="ascii.ecsv")
