#!/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']}")