update_config.py

#!/usr/bin/env python3


from pathlib import Path
import argparse
import json
import sys

import yaml
import zmq

# Defining the exposed configurations by the script.
AGIPD_CONFIGURATIONS = {
    "correct":
    {
        "common-mode": {'type': bool},
        "force-hg-if-below": {'type': int},
        "rel-gain": {'type': str},
        "xray-gain": {'type': bool},
        "blc-noise": {'type': bool},
        "blc-set-min": {'type': bool},
        "blc-stripes": {'type': bool},
        "zero-nans": {'type': bool},
        "zero-orange": {'type': bool},
        "max-pulses": {'type': list,
                       'msg': "Range list of maximum pulse indices "
                              "(--max-pulses start end step). "
                              "3 max input elements. "},
        'use-litframe-finder': {'type': str},
        'litframe-device-id': {'type': str},
        'energy-threshold': {'type': int}
    },
    "dark":
    {
        "thresholds-offset-hard-hg": {'type': list},
        "thresholds-offset-hard-mg": {'type': list},
        "thresholds-offset-hard-lg": {'type': list},
        "thresholds-offset-hard-hg-fixed": {'type': list},
        "thresholds-offset-hard-mg-fixed": {'type': list},
        "thresholds-offset-hard-lg-fixed": {'type': list},
        "thresholds-offset-sigma": {'type': list},
        "thresholds-noise-hard-hg": {'type': list},
        "thresholds-noise-hard-mg": {'type': list},
        "thresholds-noise-hard-lg": {'type': list},
        "thresholds-noise-sigma": {'type': list},
    }
}

REMI_CONFIGURATIONS = {
    'correct':
    {
        'first-pulse-offset': {'type': int}
    }
}

TIMEPIX_CONFIGURATIONS = {
    'correct':
    {
        'max-num-centroids': {'type': int},
        'clustering-epsilon': {'type': float},
        'clustering-tof-scale': {'type': float},
        'clustering-min-samples': {'type': int},
        'clustering-n-jobs': {'type': int},
        'threshold-tot': {'type': int},
        'raw-timewalk-lut-filepath': {'type': str},
        'centroiding-timewalk-lut-filepath': {'type': str}
    }
}

AGIPD_DATA_MAPPING = {
    "karabo-da": {
        'type': list,
        'choices': [f"AGIPD{i:02d}" for i in range(16)],
        'msg': "Choices: [AGIPD00 ... AGIPD15]. "
    }
}

REMI_DATA_MAPPING = {
    "karabo-da": {
        "type": list,
        "choices": ["DIGI02"],
        "msg": "Choices: [DIGI02]. "
    }
}

TIMEPIX_DATA_MAPPING = {
    "karabo-da": {
        "type": list,
        "choices": ["DA02"],
        "msg": "Choices: [DA02]. "
    }
}

AVAILABLE_DETECTORS = {
    "SPB_DET_AGIPD1M-1": [AGIPD_CONFIGURATIONS, AGIPD_DATA_MAPPING],
    "MID_DET_AGIPD1M-1": [AGIPD_CONFIGURATIONS, AGIPD_DATA_MAPPING],
    "SQS_REMI_DLD6": [REMI_CONFIGURATIONS, REMI_DATA_MAPPING],
    "SQS_AQS_CAM": [TIMEPIX_CONFIGURATIONS, TIMEPIX_DATA_MAPPING]
}


def formatter(prog):
    return argparse.HelpFormatter(prog, max_help_position=52)


parser = argparse.ArgumentParser(
    description='Request update of configuration',
    formatter_class=formatter,
    conflict_handler="resolve",
)
required_args = parser.add_argument_group('required arguments')

required_args.add_argument(
    '--karabo-id', type=str,
    choices=list(AVAILABLE_DETECTORS.keys()))
required_args.add_argument(
    '--proposal', type=str,
    help='The proposal number, without leading p, but with leading zeros. ')
required_args.add_argument('--cycle', type=str, help='The facility cycle, '
                           'detected automatically if omitted')

action_group = required_args.add_mutually_exclusive_group()
action_group.add_argument(
    '--correct', '-c', action='store_true')
action_group.add_argument(
    '--dark', '-d', action='store_true')

parser.add_argument(
    '--verbose', '-v', action='store_true',
    help='More verbose output, i.a. print the entire configuration written.')
parser.add_argument(
    '--apply', action='store_true',
    help='Apply and push the requested configuration update to the git.')
parser.add_argument(
    '--webservice-address',
    type=str,
    default="tcp://max-exfl-cal001:5555",
    help=('The port of the webservice to update '
          'calibration configurations repository.')
)

parser.add_argument(
    '--instrument',
    type=str, choices=["CALLAB"],
    help='This is only used for testing purposes.'
)


def _find_cycle(proposal: str, exp_root: Path = Path('/gpfs/exfel/exp')) -> str:
    try:
        proposal_no = int(proposal)
    except ValueError:
        raise ValueError('proposal number cannot be converted to a number')

    # /gpfs/exfel/exp/<instrument>/<cycle>/p<proposal>/
    proposal_path = next(exp_root.glob(f'*/*/p{proposal_no:06d}'), None)

    if proposal_path is None:
        raise ValueError('could not locate proposal on GPFS')

    return proposal_path.parts[-2]


def _add_available_configs_to_arg_parser(karabo_id: str, action: str):
    """Add the available configuration for the selected detector
    to the argument parser.

    Additionaly, negative booleans (-no-<bool>) are added
    along with the arguments.
    """

    available_conf = [{}, AVAILABLE_DETECTORS[karabo_id][1]]
    # adding "no" bools to available configurations

    # Loop over action configurations in available_detectors dictionary.
    for key, val in AVAILABLE_DETECTORS[karabo_id][0][action].items():
        available_conf[0][key] = val
        if val['type'] == bool:
            available_conf[0][f'no-{key}'] = {'type': bool}

    for conf in available_conf:
        for option, info in conf.items():
            type_ = info['type']
            choices = info.get('choices')
            if info['type'] == list:
                arguments = {
                    "action": 'append',
                    "nargs": '+',
                }
                type_ = str
                # Avoid having a big line of choices in the help message.
                if choices:
                    arguments.update({
                        "metavar": option.upper(),
                        "choices": choices,
                    })
            else:
                arguments = {"choices": choices}

            # Add help messages
            help_msg = ""
            if 'msg' in info.keys():
                help_msg += info['msg']

            help_msg += f"Type: {info['type'].__name__} ".upper()
            parser.add_argument(
                f"--{option}",
                type=type_,
                help=help_msg,
                **arguments,
            )

    return available_conf


def _create_new_config_from_args_input(
    instrument: str,
    args,
    available_conf,
):
    """Create an updated configuration from CLI args
    with data-mapping and correct configs."""
    karabo_id = args["karabo_id"]
    action = "dark" if args.get("dark") else "correct"
    new_conf = {action: {instrument: {karabo_id: {}}}}

    for key, value in args.items():
        key = key.replace("_", "-")
        for conf in available_conf:
            if key in conf and value is not None:

                if isinstance(value, list):
                    value = value[0]
                # convert no arguments to bool false
                if 'no-' in key and isinstance(value, bool):
                    if key not in AVAILABLE_DETECTORS[karabo_id][0].keys():
                        new_conf[action][instrument][karabo_id][key.replace('no-', '')] = False  # noqa
                    # avoid saving the "no-"key into the updated config
                    continue

                # checking if data-mapping was updated.
                if key in AVAILABLE_DETECTORS[karabo_id][1].keys():
                    if 'data-mapping' not in new_conf.keys():
                        new_conf['data-mapping'] = {karabo_id: {key: {}}}
                    new_conf['data-mapping'][karabo_id][key] = value
                else:
                    new_conf[action][instrument][karabo_id][key] = value
    return new_conf


def main():
    # remove help calls, to avoid exiting the argument parser.
    argv = sys.argv[1:]
    add_help = False
    if "-h" in argv:
        argv.remove("-h")
        add_help = True
    if "--help" in argv:
        argv.remove("--help")
        add_help = True
    known, _ = parser.parse_known_args(argv)
    args = vars(known)
    karabo_id = args["karabo_id"]
    webservice_address = args["webservice_address"]
    instrument = args['instrument']
    proposal = args['proposal']
    cycle = args['cycle']
    action = "dark" if args.get("dark") else "correct"

    # Avoid errors when karabo_id and action are not given.
    if karabo_id and action:
        available_conf = _add_available_configs_to_arg_parser(
            karabo_id, action)
        # check if instrument is not given from CLI (e.g. CALLAB)
        if instrument is None:
            # extract instrument from karabo_id
            instrument = karabo_id.split("_")[0]
        else:
            instrument = args['instrument']

    if add_help:
        argv.append("--help")

    args = vars(parser.parse_args(argv))

    if instrument is None or proposal is None:
        print("Need to define all required fields")
        sys.exit(1)
    elif cycle is None:
        cycle = _find_cycle(proposal)

    new_conf = _create_new_config_from_args_input(
        instrument=instrument,
        args=args,
        available_conf=available_conf,
    )

    if not args["apply"]:
        print("\n")
        print("-" * 80)
        print("THIS IS A DRY RUN ONLY, NO CHANGES ARE MADE")
        print("\n")
        print("-" * 80)

    pyaml = yaml.dump(new_conf, default_flow_style=False)
    print(f"# Sending the following update:\n{pyaml}")
    print("-" * 80)
    con = zmq.Context()
    socket = con.socket(zmq.REQ)
    socket.connect(webservice_address)
    msg = "','".join([
        "update_conf",
        "SASEX",
        args["karabo_id"],
        instrument,
        cycle,
        args["proposal"],
        json.dumps(new_conf),
        str(args["apply"]),
    ])
    socket.send(f"['{msg}']".encode())
    resp = socket.recv_multipart()[0]
    print("# Configuration now in place is:")

    if args['verbose']:
        print(resp.decode())
    else:
        total_config = yaml.safe_load(resp.decode())
        print(yaml.dump({
            action: {instrument: {
                karabo_id: total_config[action][instrument][karabo_id]
            }}
        }, default_flow_style=False))


if __name__ == '__main__':
    sys.exit(main())