pipeline

#!/usr/bin/env python3
#
# This file is part of the Robotic Observatory Control Kit (rockit)
#
# rockit is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# rockit is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with rockit.  If not, see <http://www.gnu.org/licenses/>.

"""Commandline client for managing the reduction pipeline"""

# pylint: disable=invalid-name
# pylint: disable=broad-except

import glob
import os
from pathlib import Path
import random
import re
import subprocess
import sys
import time
import traceback
import Pyro4
from astropy.samp import SAMPIntegratedClient, lockfile_helpers
from rockit.common import print
from rockit.pipeline import CommandStatus, Config

SCRIPT_NAME = os.path.basename(sys.argv[0])
sys.excepthook = Pyro4.util.excepthook

# Mapping used for the 'type' command
TYPE_VALUES = {
    'bias': 'BIAS',
    'dark': 'DARK',
    'flat': 'FLAT',
    'science': 'SCIENCE',
    'junk': 'JUNK'
}


def run_command(command, args):
    """Runs a daemon command, handling error messages"""
    if 'PIPELINED_CONFIG_PATH' in os.environ:
        config = Config(os.environ['PIPELINED_CONFIG_PATH'])
    else:
        # Load the config file defined in the PIPELINED_CONFIG_PATH environment variable or from the
        # default system location (/etc/pipelined/). Exit with an error if zero or multiple are found.
        files = glob.glob("/etc/pipelined/*.json")
        if len(files) != 1:
            print('error: failed to guess the default config file. ' +
                  'Run as PIPELINED_CONFIG_PATH=/path/to/config.json pipeline <command>')
            return 1

        config = Config(files[0])

    try:
        ret = command(config, args)
    except Pyro4.errors.CommunicationError:
        ret = -101

    # Print message associated with error codes, except for -1 (error handled locally)
    if ret not in [-1, 0]:
        print(CommandStatus.message(ret))
    sys.exit(ret)


def get_ds9_instances():
    local_client = SAMPIntegratedClient()
    try:
        local_client.connect()
        client_name_map = {}
        extract_name = re.compile(r'.*XPA_NAME:\t(.*) XPA_METHOD:.*')
        for client_id in local_client.get_subscribed_clients('ds9.get'):
            try:
                # DS9 8.5 doesn't seem to offer a proper way to query the window name directly
                result = local_client.ecall_and_wait(client_id, 'ds9.get', "1", cmd="xpa info")
                if result['samp.status'] == 'samp.ok':
                    m = extract_name.match(result['samp.result']['value'])
                    client_name_map[m.group(1)] = client_id
            except Exception:
                pass
        return client_name_map
    except Exception:
        return {}
    finally:
        local_client.disconnect()


def camera_ids(config):
    """Returns a dictionary of lowercase camera id to camera id"""
    return {camera_id.lower(): camera_id for camera_id in config.cameras}


def list_cameras(config, _):
    print(' '.join(camera_ids(config)))
    return 0


def print_status(config, _):
    """Reports the current pipeline status"""
    with config.daemon.connect() as pipeline:
        data = pipeline.report_status()

    print('Frame Archiving:')
    print(f'   Archive subdirectory is [b]{data["archive_subdirectory"]}[/b]')
    print(f'   Frame prefix is [b]{data["frame_prefix"]}[/b]')
    print(f'   Frame type is [b]{data["frame_type"]}[/b]')
    if data['frame_type'] == 'SCIENCE':
        print(f'   Frame object is [b]{data["frame_object"]}[/b]')
    else:
        print('   Frame object not written (requires SCIENCE type)')

        for camera_id, enabled in data['archive_enabled'].items():
            label = '[green]ARCHIVED[/green]' if enabled else '[red]DISCARDED[/red]'
            if len(config.cameras) > 1:
                prefix = f'   {camera_id} camera frames are '
            else:
                prefix = '   Frames are '
            print(prefix + f'[b]{label}[/b]')

    def enabled_label(e):
        return '[green]ENABLED[/green]' if e else '[red]DISABLED[/red]'

    print('Pipeline flags:')
    print(f'   WCS solutions are [b]{enabled_label(data["wcs_enabled"])}[/b]')
    print(f'   Field rotation calculation is [b]{enabled_label(data["rotation_enabled"])}[/b]')
    print(f'   HFD calculation is [b]{enabled_label(data["hfd_enabled"])}[/b]')
    print(f'   HFD grid calculation is [b]{enabled_label(data["hfd_grid_enabled"])}[/b]')
    print(f'   Intensity statistics are [b]{enabled_label(data["intensity_stats_enabled"])}[/b]')

    guideprofiles = enabled_label(data['guide_camera_id'])
    if data['guide_camera_id'] and len(config.cameras) > 1:
        guideprofiles += f' ({data["guide_camera_id"]})'
    print(f'   Guide profiles are [b]{guideprofiles}[/b]')
    print(f'   Dashboard updates are [b]{enabled_label(data["dashboard_enabled"])}[/b]')
    return 0


def dashboard(config, args):
    """Configure dashboard preview generation"""
    if len(args) == 1 and (args[0] == 'enable' or args[0] == 'disable'):
        with config.daemon.connect() as pipeline:
            return pipeline.set_dashboard(args[0] == 'enable')

    print(f'usage: {SCRIPT_NAME} dashboard <enable|disable>')
    return -1


def set_archive_subdirectory(config, args):
    """Set the frame archive subdirectory"""
    if len(args) == 1:
        with config.daemon.connect() as pipeline:
            return pipeline.set_archive_subdirectory(args[0])
    print(f'usage: {SCRIPT_NAME} subdir <subdirectory name>')
    return -1


def set_prefix(config, args):
    """Set the output frame filename prefix"""
    if len(args) == 1:
        with config.daemon.connect() as pipeline:
            return pipeline.set_output_frame_prefix(args[0])
    print(f'usage: {SCRIPT_NAME} prefix <output frame prefix>')
    return -1


def set_object(config, args):
    """Set the OBJECT header keyword"""
    if len(args) == 1:
        with config.daemon.connect() as pipeline:
            return pipeline.set_frame_object(args[0])
    print(f'usage: {SCRIPT_NAME} object <object name>')
    return -1


def set_type(config, args):
    """Set the IMAGETYP header keyword"""
    if len(args) == 1 and args[0] in TYPE_VALUES:
        with config.daemon.connect() as pipeline:
            return pipeline.set_frame_type(TYPE_VALUES[args[0]])
    print(f'usage: {SCRIPT_NAME} type <bias|dark|flat|science|junk>')
    return -1


def set_archive(config, args):
    """Enable or disable archiving of acquired frames to disk"""
    if len(config.cameras) == 1:
        if len(args) == 1 and (args[0] == 'enable' or args[0] == 'disable'):
            with config.daemon.connect() as pipeline:
                return pipeline.set_archive(next(iter(config.cameras)), args[0] == 'enable')

        print(f"usage: {SCRIPT_NAME} archive <enable|disable>")
    else:
        cameras = camera_ids(config)
        if len(args) == 2 and args[0] in cameras and (args[1] == 'enable' or args[1] == 'disable'):
            with config.daemon.connect() as pipeline:
                return pipeline.set_archive(cameras[args[0]], args[1] == 'enable')

        print(f"usage: {SCRIPT_NAME} archive <{'|'.join(cameras)}> <enable|disable>")
    return -1


def set_wcs(config, args):
    """Enable or disable wcs solutions for acquired frames"""
    if len(args) == 1 and (args[0] == 'enable' or args[0] == 'disable'):
        with config.daemon.connect() as pipeline:
            return pipeline.set_wcs(args[0] == 'enable')

    print(f'usage: {SCRIPT_NAME} wcs <enable|disable>')
    return -1


def set_rotation(config, args):
    """Enable or disable field rotation calculation for acquired frames"""
    if len(args) == 1 and (args[0] == 'enable' or args[0] == 'disable'):
        with config.daemon.connect() as pipeline:
            return pipeline.set_rotation(args[0] == 'enable')

    print(f'usage: {SCRIPT_NAME} rotation <enable|disable>')
    return -1


def set_hfd(config, args):
    """Enable or disable half-flux diameter calculation for acquired frames"""
    if len(args) == 1 and (args[0] == 'enable' or args[0] == 'disable'):
        with config.daemon.connect() as pipeline:
            return pipeline.set_hfd(args[0] == 'enable')

    print(f'usage: {SCRIPT_NAME} hfd <enable|disable>')
    return -1


def set_hfd_grid(config, args):
    """Enable or disable half-flux diameter grid calculation for acquired frames"""
    if len(args) == 1 and (args[0] == 'enable' or args[0] == 'disable'):
        with config.daemon.connect() as pipeline:
            return pipeline.set_hfd_grid(args[0] == 'enable')

    print(f'usage: {SCRIPT_NAME} hfdgrid <enable|disable>')
    return -1


def set_intensity_stats(config, args):
    """Enable or disable intensity statistics for acquired frames"""
    if len(args) == 1 and (args[0] == 'enable' or args[0] == 'disable'):
        with config.daemon.connect() as pipeline:
            return pipeline.set_intensity_stats(args[0] == 'enable')

    print(f'usage: {SCRIPT_NAME} intstats <enable|disable>')
    return -1


def register_preview(config, args):
    """Open a new ds9 preview window and register it with the server"""
    if len(config.cameras) == 1:
        if len(args) != 0:
            print(f"usage: {SCRIPT_NAME} preview")
            return -1

        camera_id = next(iter(config.cameras))
    else:
        cameras = camera_ids(config)
        if len(args) != 1 or args[0] not in cameras:
            print(f"usage: {SCRIPT_NAME} preview <{'|'.join(cameras)}>")
            return -1

        camera_id = cameras[args[0]]

    # Increase display size on HiDPI displays
    window_scale = 1.0
    try:
        output = subprocess.check_output(['/usr/bin/xrdb', '-query']).decode('ascii').split('\n')
        for line in output:
            if line.startswith('Xft.dpi'):
                window_scale = float(line[8:]) / 96
    except Exception:
        traceback.print_exc(file=sys.stdout)
        pass

    # Generate a random name to avoid collisions
    ds9_instances = get_ds9_instances()
    while True:
        client_index = random.randint(0, 999999)
        name = f'preview-{camera_id}-{client_index:06d}'
        if name not in ds9_instances:
            break

    client = SAMPIntegratedClient()
    ds9 = None
    try:
        os.system(f'/usr/bin/ds9 -title {name} &')
        for i in range(10):
            time.sleep(1)
            ds9_instances = get_ds9_instances()
            if name in ds9_instances:
                ds9 = ds9_instances[name]
                break

        if ds9 is None:
            print('error: failed to open DS9 preview window')
            return -1

        client.connect()
        client.enotify(ds9, 'ds9.set', cmd='view filename no')
        client.enotify(ds9, 'ds9.set', cmd='view physical no')
        client.enotify(ds9, 'ds9.set', cmd='view frame no')
        client.enotify(ds9, 'ds9.set', cmd='view buttons no')
        client.enotify(ds9, 'ds9.set', cmd='view colorbar no')
        client.enotify(ds9, 'ds9.set', cmd='view layout vertical')
        client.enotify(ds9, 'ds9.set', cmd='scale mode zscale')
        client.enotify(ds9, 'ds9.set', cmd='background black')
        client.enotify(ds9, 'ds9.set', cmd=f'width {int(window_scale * config.cameras[camera_id]["preview_ds9_width"])}')
        client.enotify(ds9, 'ds9.set', cmd=f'height {int(window_scale * config.cameras[camera_id]["preview_ds9_height"])}')
        client.enotify(ds9, 'ds9.set', cmd=f'zoom {window_scale * config.cameras[camera_id]["preview_ds9_zoom"]}')
    except Exception:
        traceback.print_exc(file=sys.stdout)
        print('error: failed to open DS9 preview window')
        return -1
    finally:
        client.disconnect()

    with config.daemon.connect() as pipeline:
        _, registration = lockfile_helpers.check_running_hub(os.path.join(Path.home(), ".samp"))
        registration['rockit.client'] = ds9
        registration['rockit.window_scale'] = window_scale

        # Note: If the server is running on a different machine scp/ssh will be used to update the preview
        # This requires that key-based ssh has been configured from the server machine to this client machine.
        return pipeline.register_preview(camera_id, registration)


def display_preview(config, args):
    """Internal function to update the image displayed in a given ds9 preview window"""

    from astropy.io import fits
    try:
        samp_url, samp_secret, samp_client, path, window_scale = args
        window_scale = float(window_scale)
        header = fits.getheader(path)
        camera_config = config.cameras[header['CAMID']]

        client = SAMPIntegratedClient()
        try:
            client.connect(hub_params={
                'samp.hub.xmlrpc.url': samp_url,
                'samp.secret': samp_secret
            })

            result = client.ecall_and_wait(samp_client, "ds9.get", "1", cmd="width")
            width = int(result.get('samp.result', {}).get('value', camera_config['preview_ds9_width']))

            result = client.ecall_and_wait(samp_client,  "ds9.get", "1", cmd="height")
            height = int(result.get('samp.result', {}).get('value', camera_config['preview_ds9_height']))
            half_margin = window_scale * 10

            def text_label(pos, text):
                if pos[1] == 'l':
                    x = width / 4
                elif pos[1] == 'r':
                    x = 3 * width / 4
                else:
                    x = width / 2

                y = half_margin if pos[0] == 't' else height - half_margin

                return f'text {x} {y + 2} "{text}" # color = white fontweight = bold'

            timestamp = header['DATE-OBS']
            time_src = header.get('TIME-SRC', None)
            if time_src:
                timestamp += f' ({time_src})'

            prefix = []
            if camera_config['preview_display_camera_id']:
                prefix.append(header['CAMID'])
            if camera_config['preview_display_filter']:
                prefix.append(header.get('FILTER', 'UNKNOWN'))
            exptime = (', '.join(prefix) + ' @ ' if prefix else '') + f'{header["EXPTIME"]:.3f}s'

            cam_transfer = header.get('CAM-TFER', None)
            if cam_transfer:
                exptime += f' ({cam_transfer})'

            regions = []
            ilustrations = [
                f'box 0 {half_margin:.1f} {width} {half_margin:.1f} # color = black fill = yes',
                f'box 0 {height - half_margin:.1f} {width} {half_margin:.1f} # color = black fill = yes',
                text_label('tl', timestamp),
                text_label('tr', exptime),
                text_label('bm', 'SAVED' if 'FILENAME' in header else 'NOT SAVED'),
            ]

            if 'WCSSOLVE' in header:
                ilustrations.append(text_label('bl', 'WCS: SOLVED' if header['WCSSOLVE'] else 'WCS: FAILED'))

            if 'FIELDROT' in header:
                ilustrations.append(text_label('br', f'North is {header["FIELDROT"]:.1f} deg CCW of +X'))

            elif 'MEDHFD' in header:
                ilustrations.append(text_label('br', f'HFD: {header["MEDHFD"]:.2f} arcsec'))

            if 'HFD_ROWS' in header and 'HFD_COLS' in header:
                x_step = header['NAXIS1'] / header['HFD_COLS']
                y_step = header['NAXIS2'] / header['HFD_ROWS']

                for j in range(header['HFD_ROWS']):
                    for i in range(header['HFD_COLS']):
                        x = (i + 0.5) * x_step
                        y = (j + 0.5) * y_step
                        regions.append(f'box {x:.6f} {y:.6f} {x_step} {y_step} # color = yellow')
                        median = header.get(f'HFD_{j:02d}{i:02d}', None)
                        if median is not None:
                            regions.append(f'text {x:.6f} {y:.6f} "{median:.2f}" # color = green')

            client.enotify(samp_client, 'ds9.set', cmd=f'file "{path}"')
            client.enotify(samp_client, 'ds9.set', cmd='illustrate delete all')
            for illustration in ilustrations:
                command = illustration.replace('"', '\\\"')
                client.enotify(samp_client, 'ds9.set', cmd=f'illustrate command "{command}"')

            for region in regions:
                command = region.replace('"', '\\\"')
                client.enotify(samp_client, 'ds9.set', cmd=f'region command "{command}"')

        finally:
            client.disconnect()
        return CommandStatus.Succeeded
    except Exception:
        return CommandStatus.Failed


def reset(config, _):
    """Reset to the default pipeline configuration"""
    with config.daemon.connect() as pipeline:
        return pipeline.configure({}, quiet=True)


def print_usage():
    """Prints the utility help"""
    print(f'usage: {SCRIPT_NAME} <command> \\[<args>]')
    print()
    print('general commands:')
    print('   status       print a human-readable summary of the pipeline status')
    print('   preview      open a ds9 window to preview acquired frames')
    print('   reset        reset to the default pipeline configuration')
    print('output frame commands')
    print('   archive      enable or disable frame archiving')
    print('   subdir       set archive subdirectory')
    print('   prefix       set archive frame prefix')
    print('   object       set metadata object name')
    print('   type         set metadata frame type')
    print('reduction flags:')
    print('   wcs          enable or disable WCS solutions')
    print('   rotation     enable or disable field rotation preview')
    print('   hfd          enable or disable half-flux diameter calculations')
    print('   hfdgrid      enable or disable half-flux diameter grid preview')
    print('   intstats     enable or disable intensity statistics')
    print('   dashboard    enable or disable dashboard preview updates')
    print()
    return 0


if __name__ == '__main__':
    commands = {
        'status': print_status,
        'preview': register_preview,
        'display': display_preview,
        'reset': reset,
        'archive': set_archive,
        'subdir': set_archive_subdirectory,
        'prefix': set_prefix,
        'object': set_object,
        'type': set_type,
        'wcs': set_wcs,
        'rotation': set_rotation,
        'hfd': set_hfd,
        'hfdgrid': set_hfd_grid,
        'intstats': set_intensity_stats,
        'dashboard': dashboard,
        'list-cameras': list_cameras
    }

    if len(sys.argv) >= 2 and sys.argv[1] in commands:
        sys.exit(run_command(commands[sys.argv[1]], sys.argv[2:]))

    sys.exit(print_usage())