L2TS test observation

#!/usr/bin/env python3
import os
import sys
import yaml
import argparse
from datetime import datetime, timedelta

import astropy.units as u
from astropy.coordinates import SkyCoord, get_icrs_coordinates

from tmss_http_rest_client import TMSSsession

from print_dict import print_dict

if __name__ == "__main__":
    # Read command line arguments
    parser = argparse.ArgumentParser(description="Schedule pulsar timing observations in TMSS")
    parser.add_argument("-s", "--source", help="Source to observe [str, default: PSR B0329+54]",
                        default="PSR B0329+54", type=str)
    parser.add_argument("-a", "--antennafield", help="Antenna field to select [str, default: HBA_DUAL]",
                        default="HBA_DUAL", type=str)
    parser.add_argument("-f", "--filter", help="Filter to select [str, default: HBA_110_190]",
                        default="HBA_110_190", type=str)
    parser.add_argument("-S", "--subband", help="Subbands to select [str, default: 201..400]",
                        default="201..400", type=str)
    parser.add_argument("-t", "--tstart", help="Start time (isot) [str, default: now]",
                        default=None, type=str)
    parser.add_argument("-l", "--length", help="Observation length (s) [int, default: 60]",
                        default=60, type=int)
    parser.add_argument("-d", "--delay", help="Delay in scheduling (s) [int, default: 300]",
                        default=300, type=int)
    parser.add_argument("-u", "--upload", help="Upload specification document and create scheduling unit draft",
                        action="store_true")
    parser.add_argument("-b", "--blueprint", help="Blueprint uploaded scheduling unit draft",
                        action="store_true")
    parser.add_argument("-L", "--stationlist", help="Station to use [str, default: CS001, can be multiple, e.g. \"CS001,CS002,CS003\"]",
                        default="CS001", type=str)
    parser.add_argument("-i", "--scheduling_set_id", help="Scheduling set ID [int, default: 305]",
                        default=305, type=int)

    args = parser.parse_args()

    # Settings
    run_spec = {"strategy_name": "FE Stokes I - raw",
                "description": "L2TS test observation",
                "scheduling_set_id": args.scheduling_set_id}
    p = get_icrs_coordinates(args.source)
    src_spec = {"name": args.source,
                "angle1": p.ra.rad,
                "angle2": p.dec.rad,
                "duration_s": args.length,
                "elev_min_deg": 0,
                "lst_min_s": -21600,
                "lst_max_s": 21600,
                "priority_queue": "A",
                "rank": 1.00}

    # Decode subbands
    smin = int(args.subband.split("..")[0])
    smax = int(args.subband.split("..")[-1])

    # Start time
    if args.tstart is None:
        # Round to start of minute
        tstart = datetime.utcnow() + timedelta(seconds=args.delay)
        tstart = tstart.strftime("%Y-%m-%dT%H:%M:00")
    else:
        tstart = datetime.strptime(args.tstart, "%Y-%m-%dT%H:%M:%S")
        tstart = tstart.strftime("%Y-%m-%dT%H:%M:%S")

    # Station groups

    # Station groups for Fly's Eye observations
    stationlist = [item for item in args.stationlist.split(",")]
    station_groups = [{"stations": [item],
                       "max_nr_missing": 0} for item in stationlist]

    # Check antenna field against LOFAR2.0 stations
    if args.antennafield == "LBA_ALL":
        if not set(stationlist) <= set(["CS001"]):
            print("Can not use LBA_ALL on LOFAR stations")
            sys.exit()

    # Read credentials
    with open("login.yaml", "r") as fp:
        settings = yaml.full_load(fp)

    # Open session
    with TMSSsession(host=settings['host'],
                     port=settings['port'],
                     username=settings['username'],
                     password=settings['password']) as client:
        print("Opened TMSS connection")
        
        # Get the latest satellite monitoring template
        template = client.get_scheduling_unit_observing_strategy_template(run_spec['strategy_name'])
        print(f"Using strategy template {template['url']}")
        
        # Get the specifications document
        original_spec_doc = client.get_scheduling_unit_observing_strategy_template_specification_with_just_the_parameters(template['name'], template['version'])

        # Copy original specification document
        spec_doc = original_spec_doc.copy()

        pointing = {"target": src_spec["name"],
                    "angle1": src_spec["angle1"],
                    "angle2": src_spec["angle2"],
                    "direction_type": "J2000"}

        
        spec_doc['scheduling_constraints_doc']['time']['at'] = f'{tstart}'
        spec_doc['scheduling_constraints_doc']['scheduler'] = 'fixed_time'
        spec_doc['tasks']['Observation']['specifications_doc']['duration'] = src_spec['duration_s']
        spec_doc['tasks']['Observation']['short_description'] = src_spec["name"]
        spec_doc['tasks']['Observation']['specifications_doc']['station_configuration']['SAPs'][0]['digital_pointing'] = pointing
        spec_doc['tasks']['Observation']['specifications_doc']['station_configuration']['tile_beam'] = pointing
        spec_doc['tasks']['Observation']['specifications_doc']['station_configuration']['antenna_set'] = args.antennafield
        spec_doc['tasks']['Observation']['specifications_doc']['station_configuration']['filter'] = args.filter

        spec_doc['tasks']['Observation']['specifications_doc']['station_configuration']['SAPs'][0]['subbands'] = list(range(smin, smax + 1))
        spec_doc['tasks']['Observation']['specifications_doc']['station_configuration']['station_groups'] = station_groups
        spec_doc['tasks']['Observation']['specifications_doc']['beamformer']['pipelines'][0]['station_groups'] = station_groups
        
        # Show spec doc
        #        print(spec_doc)
        print_dict(spec_doc, prefix="spec_doc")

        # Skip if not uploaded
        if not args.upload:
            print("No scheduling units uploaded to TMSS.")
            print()
            sys.exit()

        # Create scheduling unit
        scheduling_unit_draft = client.create_scheduling_unit_draft_from_strategy_template(template['id'], run_spec['scheduling_set_id'], spec_doc, src_spec["name"], run_spec['description'], src_spec['priority_queue'], src_spec['rank'])

        # Patch scheduling unit with keys that are not included in the strategy
#        url = scheduling_unit_draft['url']
#        json_data = {"name": src_spec['name'],
#                     "rank": src_spec['rank'],
#                     "description": run_spec['description']}
#        print(json_data)
#        response = client.do_request_and_get_result_as_json_object(method="PATCH", full_url=url, json_data=json_data)
        print(f"Created SU draft {scheduling_unit_draft['id']} for {src_spec['name']} at {scheduling_unit_draft['url']}")

        # Skip if not blueprinted
        if not args.blueprint:
            print("No scheduling units blueprinted.")
            print()
            sys.exit()

        # Blueprint SU
        scheduling_unit_blueprint = client.create_scheduling_unit_blueprint_and_tasks_and_subtasks_tree(scheduling_unit_draft['id'])
        print(f"Created SU blueprint {scheduling_unit_blueprint['id']} for {src_spec['name']} at {scheduling_unit_blueprint['url']}")