mbari_wec_batch

#!/usr/bin/python3
# Copyright 2022 Open Source Robotics Foundation, Inc. and Monterey Bay Aquarium Research Institute
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Launch batch of buoy simulations."""

from em import invoke as empy
import numpy as np
import yaml

import os
import shutil
import subprocess
import sys
import time

from ament_index_python.packages import get_package_share_directory

from launch import LaunchDescription
from launch import LaunchService
from launch.actions import ExecuteProcess, IncludeLaunchDescription
from launch.actions import RegisterEventHandler
from launch.event_handlers import OnProcessExit, OnProcessStart

from launch.launch_description_sources import PythonLaunchDescriptionSource

from launch_testing.proc_info_handler import ActiveProcInfoHandler
from launch_testing.util import resolveProcesses

import rclpy
from rclpy.node import Node


DEFAULT_PHYSICS_MAX_STEP_SIZE = 0.001  # seconds


class MonoChromatic(object):
    def __init__(self, A, T):
        self.A, self.T = A, T
        self.name = 'MonoChromatic'

    def __str__(self):
        if self.A is not None:
            return f'{self.name};A:{self.A};T:{self.T}'
        else:
            return f'{self.name};A:default;T:default'


class Bretschneider(object):
    def __init__(self, Hs, Tp):
        self.Hs, self.Tp = Hs, Tp
        self.name = 'Bretschneider'

    def __str__(self):
        if self.Hs is not None:
            return f'{self.name};Hs:{self.Hs};Tp:{self.Tp}'
        else:
            return f'{self.name};Hs:default;Tp:default'


class Custom(object):
    def __init__(self, f, Szz):
        self.f, self.Szz = f, Szz
        self.name = 'Custom'

    def __str__(self):
        if self.f is not None:
            return f'{self.name};' + \
                   f'f:{":".join([str(f) for f in self.f])};' + \
                   f'Szz:{":".join([str(Szz) for Szz in self.Szz])}'
        else:
            return f'{self.name};f:defaults;Szz:defaults'


def generate_simulations(sim_params_yaml):

    rclpy.init()
    node = Node('mbari_wec_batch')

    with open(sim_params_yaml, 'r') as fd:
        sim_params = yaml.load(fd, Loader=yaml.CLoader)

    print_optional = lambda param: str(param) if param is not None else ''  # noqa: E731

    # Grab params from yaml
    # physics time step (seconds)
    if 'physics_step' in sim_params:
        physics_step = np.array(sim_params['physics_step'], dtype=float)
        physics_step = np.atleast_1d(physics_step)
    else:
        physics_step = [None]
        node.get_logger().warn(
            'sim_params_yaml: optional physics_step parameter not specified -- defaulting'
        )

    # physics real-time factor (RTF)
    try:
        if 'physics_rtf' in sim_params:
            physics_rtf = float(sim_params['physics_rtf'])
        else:
            physics_rtf = None
            node.get_logger().warn(
                'sim_params_yaml: optional physics_rtf parameter not specified -- defaulting'
            )
    except TypeError:
        node.get_logger().error('sim_params_yaml: physics_rtf parameter' +
                                ' must be a single value')
        sys.exit(-1)

    # Enable GUI
    try:
        if 'enable_gui' in sim_params:
            enable_gui = bool(sim_params['enable_gui'])
        else:
            enable_gui = False
            node.get_logger().debug(
                'sim_params_yaml: optional enable_gui parameter not specified' +
                ' -- defaulting to headless'
            )
    except TypeError:
        node.get_logger().error('sim_params_yaml: enable_gui parameter' +
                                ' must be a single value')
        sys.exit(-1)

    # seed (random seed for incident waves plugin)
    try:
        if 'seed' in sim_params:
            seed_ = int(sim_params['seed'])
        else:
            seed_ = None
            node.get_logger().warn('sim_params_yaml: optional seed parameter not specified')
    except TypeError:
        node.get_logger().error('sim_params_yaml: seed parameter' +
                                ' must be a single value')
        sys.exit(-1)

    # duration (seconds)
    try:
        if 'duration' not in sim_params:
            node.get_logger().error('sim_params_yaml: required duration parameter not specified')
            sys.exit(-1)
        duration = float(sim_params['duration'])
    except TypeError:
        node.get_logger().error('sim_params_yaml: duration parameter' +
                                ' must be a single value in seconds')
        sys.exit(-1)

    # door state ('open', 'closed')
    if 'door_state' in sim_params:
        door_state = sim_params['door_state']
        door_state = np.atleast_1d(door_state)
        if any(['open' not in ds and 'closed' not in ds for ds in door_state]):
            node.get_logger().error(
                "sim_params_yaml: all door_state parameters must be either 'open' or 'closed'"
            )
            sys.exit(-1)
    else:
        node.get_logger().warn(
            "sim_params_yaml: door_state parameter not specified -- defaulting to 'closed'"
        )
        door_state = ['closed']

    # scale factor (0.5 - 1.4)
    if 'scale_factor' in sim_params:
        scale_factor = np.array(sim_params['scale_factor'], dtype=float)
        scale_factor = np.atleast_1d(scale_factor)
        if not np.all((0.5 <= scale_factor) & (scale_factor <= 1.4)):
            node.get_logger().error(
                'sim_params_yaml: all scale_factor parameters must be between 0.5 to 1.4'
            )
            sys.exit(-1)
    else:
        node.get_logger().warn(
            'sim_params_yaml: scale_factor parameter not specified -- defaulting to 1.0'
        )
        scale_factor = [1.0]

    # Battery State
    # State of Charge (0.0 [270V] to 1.0 [320V])
    if 'battery_soc' in sim_params:
        battery_state = np.array(sim_params['battery_soc'], dtype=float)
        battery_state = np.atleast_1d(battery_state)
        if np.all((0.0 <= battery_state) & (battery_state <= 1.0)):
            batt_state_type = 'SoC'
        else:
            node.get_logger().error(
                'sim_params_yaml: all battery_soc parameters must be between 0.0 to 1.0'
            )
            sys.exit(-1)
    # EMF (270V to 320V)
    elif 'battery_emf' in sim_params:
        battery_state = np.array(sim_params['battery_emf'], dtype=float)
        battery_state = np.atleast_1d(battery_state)
        if np.all((270.0 <= battery_state) & (battery_state <= 300.0)):
            batt_state_type = 'EMF'
        else:
            node.get_logger().error(
                'sim_params_yaml: all battery_emf parameters must be between 300V to 270V'
            )
            sys.exit(-1)
    else:
        node.get_logger().warn(
            'sim_params_yaml: optional battery_emf or battery_soc parameter not specified'
        )
        battery_state = [None]
        batt_state_type = ''

    # mean_piston_position (meters)
    if 'mean_piston_position' in sim_params:
        mean_piston_position = np.array(sim_params['mean_piston_position'], dtype=float)
        mean_piston_position = np.atleast_1d(mean_piston_position)
    else:
        mean_piston_position = [None]
        node.get_logger().warn(
            'sim_params_yaml: optional mean_piston_position parameter not specified' + \
            ' -- defaulting'
        )

    # IncWaveSpectrumType
    if 'IncidentWaveSpectrumType' in sim_params:
        inc_wave_spectrum_types = sim_params['IncidentWaveSpectrumType']
        incident_waves = []
        for inc_wave_spectrum_type in inc_wave_spectrum_types:
            if 'MonoChromatic' in inc_wave_spectrum_type:
                try:
                    monochromatic_spectrum = inc_wave_spectrum_type['MonoChromatic']
                except TypeError:  # just default params
                    monochromatic_spectrum = {}
                if 'A' in monochromatic_spectrum:
                    A = np.array(monochromatic_spectrum['A'], dtype=float)
                    A = np.atleast_1d(A)
                else:
                    A = None
                if 'T' in monochromatic_spectrum:
                    T = np.array(monochromatic_spectrum['T'], dtype=float)
                    T = np.atleast_1d(T)
                else:
                    T = None

                if (A is None) ^ (T is None):
                    node.get_logger().error(
                        'sim_params_yaml: MonoChromatic: A or T unspecified.' +
                        ' Please specify both or neither'
                    )
                    sys.exit(-1)
                if A.shape != T.shape:
                    node.get_logger().error(
                        'sim_params_yaml: MonoChromatic: A and T have different length'
                    )
                    sys.exit(-1)
                incident_waves.extend([MonoChromatic(a, t) for a, t in zip(A, T)])
            elif 'Bretschneider' in inc_wave_spectrum_type:
                try:
                    bretschneider_spectrum = inc_wave_spectrum_type['Bretschneider']
                except TypeError:  # just default params
                    bretschneider_spectrum = {}
                if 'Hs' in bretschneider_spectrum:
                    Hs = np.array(bretschneider_spectrum['Hs'], dtype=float)
                    Hs = np.atleast_1d(Hs)
                else:
                    Hs = None
                if 'Tp' in bretschneider_spectrum:
                    Tp = np.array(bretschneider_spectrum['Tp'], dtype=float)
                    Tp = np.atleast_1d(Tp)
                else:
                    Tp = None

                if (Hs is None) ^ (Tp is None):
                    node.get_logger().error(
                        'sim_params_yaml: Bretschneider: Hs or Tp unspecified.' +
                        ' Please specify both or neither'
                    )
                    sys.exit(-1)
                if Hs.shape != Tp.shape:
                    node.get_logger().error(
                        'sim_params_yaml: Bretschneider: Hs and Tp have different length.'
                    )
                    sys.exit(-1)
                incident_waves.extend([Bretschneider(hs, tp) for hs, tp in zip(Hs, Tp)])
            elif 'Custom' in inc_wave_spectrum_type:
                try:
                    custom_spectrum = inc_wave_spectrum_type['Custom']
                except TypeError:  # just default params
                    custom_spectrum = {}
                if 'f' in custom_spectrum:
                    f = np.array(custom_spectrum['f'], dtype=float)
                    f = np.atleast_1d(f)
                else:
                    f = None
                if 'Szz' in custom_spectrum:
                    Szz = np.array(custom_spectrum['Szz'], dtype=float)
                    Szz = np.atleast_1d(Szz)
                else:
                    Szz = None

                if (f is None) ^ (Szz is None):
                    node.get_logger().error(
                        'sim_params_yaml: Custom: f or Szz unspecified.' +
                        ' Please specify both or neither'
                    )
                    sys.exit(-1)
                if f is not None and Szz is not None:
                    if f.shape != Szz.shape:
                        node.get_logger().error(
                            'sim_params_yaml: Custom: f.shape != Szz.shape'
                        )
                        sys.exit(-1)
                incident_waves.append(Custom(f, Szz))
            else:
                node.get_logger().error(
                    f'sim_params_yaml: IncWaveSpectrumType [{inc_wave_spectrum_type}]' +
                    ' did not match supported types:\n    ' +
                    '\n    '.join(['MonoChromatic', 'Bretschneider', 'Custom'])
                )
                sys.exit(-1)
    else:
        node.get_logger().warn(
            'sim_params_yaml: optional IncWaveSpectrumType parameter not specified' +
            ' -- defaulting to No Waves'
        )
        incident_waves = [None]

    # create batch results directory
    timestr = time.strftime("%Y%m%d%H%M%S")
    batch_results_dir = f'batch_results_{timestr}'
    node.get_logger().info(f'Creating batch results directory: {batch_results_dir}')
    os.makedirs(batch_results_dir)
    # python workaround for 'ln -sf'
    os.symlink(batch_results_dir, f'latest_batch_results_{timestr}', target_is_directory=True)
    os.replace(f'latest_batch_results_{timestr}', 'latest_batch_results')

    sim_params_name, dot_yaml = os.path.splitext(os.path.basename(sim_params_yaml))
    sim_params_date_yaml = sim_params_name + f'_{timestr}' + dot_yaml
    node.get_logger().debug('Copying sim params yaml to batch results directory: ' +
                            os.path.join(batch_results_dir,
                                         sim_params_date_yaml))
    shutil.copy(sim_params_yaml,
                os.path.join(batch_results_dir,
                             sim_params_date_yaml))

    # generate test matrix
    batch_params = list(zip(*[param.ravel() for param in np.meshgrid(physics_step,
                                                                     door_state,
                                                                     scale_factor,
                                                                     battery_state,
                                                                     mean_piston_position,
                                                                     incident_waves)]))

    node.get_logger().info(f'Generated {len(batch_params)} simulation runs')
    node.get_logger().debug('PhysicsStep, PhysicsRTF, Seed, Duration, DoorState, ScaleFactor' +
                            ', BatteryState, MeanPistonPosition' +
                            ', IncWaveSpectrumType;IncWaveSpectrumParams')
    [node.get_logger().debug(f'{print_optional(ps)}, {print_optional(physics_rtf)}' +
                             f', {print_optional(seed_)}, {duration}, {ds}, {sf}' +
                             f', {print_optional(bs)}' +
                             f', {print_optional(mpp)}' +
                             f', {print_optional(iw)}')
        for ps, ds, sf, bs, mpp, iw in batch_params]

    node.get_logger().info('Creating log file: ' +
                           os.path.join(batch_results_dir, 'batch_runs.log'))
    with open(os.path.join(batch_results_dir, 'batch_runs.log'), 'w') as fd:
        fd.write(f'# Generated {len(batch_params)} simulation runs\n')
        fd.write('RunIndex, SimReturnCode, StartTime, rosbag2Filename, pblogFilename' +
                 ', PhysicsStep, PhysicsRTF' +
                 ', Seed, Duration, DoorState, ScaleFactor, BatteryState' +
                 ', MeanPistonPosition, IncWaveSpectrumType;IncWaveSpectrumParams\n')

    # Find packages
    pkg_buoy_gazebo = get_package_share_directory('buoy_gazebo')
    pkg_buoy_description = get_package_share_directory('buoy_description')

    # Find model templates
    model_dir = 'mbari_wec_base'
    empy_base_sdf_file = os.path.join(pkg_buoy_description, 'models', model_dir, 'model.sdf.em')
    base_sdf_file = os.path.join(pkg_buoy_description, 'models', model_dir, 'model.sdf')

    model_dir = 'mbari_wec'
    empy_sdf_file = os.path.join(pkg_buoy_description, 'models', model_dir, 'model.sdf.em')
    sdf_file = os.path.join(pkg_buoy_description, 'models', model_dir, 'model.sdf')

    # Find world file template
    empy_world_file = os.path.join(pkg_buoy_gazebo, 'worlds', 'mbari_wec.sdf.em')
    world_file = os.path.join(pkg_buoy_gazebo, 'worlds', 'mbari_wec_batch.sdf')

    # Start batch runs
    rng = np.random.default_rng()
    for idx, (ps, ds, sf, bs, mpp, iw) in enumerate(batch_params):
        if not rclpy.ok():
            break  # Shutting down
        seed = rng.integers(2**32-1) if seed_==0 else seed_
        node.get_logger().info(f'\n\nSim run [{idx}] for {duration} seconds:' +
                               f" door state='{ds}', scale factor={sf}" +
                               f", battery state={bs if bs is not None else 'None'}" +
                               f", mean piston position={mpp if mpp is not None else 'None'}" +
                               ', IncidentWaveSpectrumType=' +
                               f"{str(iw) if iw is not None else 'None'}\n")

        # fill mbari_wec_base model template with params
        empy(['-D', f"door_state = '{ds}'",
              '-o', base_sdf_file,
              empy_base_sdf_file])

        # fill mbari_wec world template with params
        mbari_wec_world_params = []
        if ps is not None:
            mbari_wec_world_params.extend(['-D', f'physics_step = {ps}'])
        if physics_rtf is not None:
            mbari_wec_world_params.extend(['-D', f'physics_rtf = {physics_rtf}'])
        mbari_wec_world_params.extend(['-o', world_file,
                                       empy_world_file])
        empy(mbari_wec_world_params)

        # fill mbari_wec model template with params
        mbari_wec_model_params = ['-D', f'scale_factor = {sf}']
        if seed is not None:
            mbari_wec_model_params.extend(['-D', f'inc_wave_seed = {seed}'])
        if bs is not None:
            if 'SoC' in batt_state_type:
                mbari_wec_model_params.extend(['-D', f'battery_soc = {bs}'])
            elif 'EMF' in batt_state_type:
                mbari_wec_model_params.extend(['-D', f'battery_emf = {bs}'])
        if mpp is not None:
            mbari_wec_model_params.extend(['-D', f'x_mean_pos = {mpp}'])
        if iw is not None:
            mbari_wec_model_params.extend(['-D', f'inc_wave_spectrum_type = "{iw.name}"'])
            if 'MonoChromatic' in iw.name:
                if iw.A is not None:
                    mbari_wec_model_params.extend(['-D', f'A = {iw.A}'])
                    mbari_wec_model_params.extend(['-D', f'T = {iw.T}'])
            if 'Bretschneider' in iw.name:
                if iw.Hs is not None:
                    mbari_wec_model_params.extend(['-D', f'Hs = {iw.Hs}'])
                    mbari_wec_model_params.extend(['-D', f'Tp = {iw.Tp}'])
            if 'Custom' in iw.name:
                if iw.f is not None:
                    mbari_wec_model_params.extend(['-D', f'f = {iw.f.tolist()}'])
                    mbari_wec_model_params.extend(['-D', f'Szz = {iw.Szz.tolist()}'])
        mbari_wec_model_params.extend(['-o', sdf_file,
                                       empy_sdf_file])
        node.get_logger().debug(mbari_wec_model_params)
        empy(mbari_wec_model_params)

        # convert duration to iterations
        if ps is None:
            step_size = DEFAULT_PHYSICS_MAX_STEP_SIZE
        else:
            step_size = ps
        iterations = int(round(duration / step_size))

        node.get_logger().debug('running gz-sim with gz_args:=' +
                                '-r' if enable_gui else '-rs' +
                                f' --iterations {iterations}')

        # Dynamically create launch file
        start_time = time.strftime("%Y%m%d%H%M%S")
        rosbag2_dir = 'rosbag2'
        pblog_dir = 'pblog'
        run_results_dir = f'results_run_{idx}_{start_time}'
        os.makedirs(os.path.join(batch_results_dir, run_results_dir))

        # get info from processes
        proc_start_info = ActiveProcInfoHandler()
        proc_exit_info = ActiveProcInfoHandler()

        # python workaround for 'ln -sf'
        os.symlink(os.path.join(run_results_dir, rosbag2_dir),
                   os.path.join(batch_results_dir, f'latest_rosbag_{start_time}'), target_is_directory=True)
        os.replace(os.path.join(batch_results_dir, f'latest_rosbag_{start_time}'),
                   os.path.join(batch_results_dir, f'latest_rosbag'))

        def generate_launch_description():
            mbari_wec = IncludeLaunchDescription(
                PythonLaunchDescriptionSource(
                    os.path.join(pkg_buoy_gazebo, 'launch', 'mbari_wec.launch.py'),
                ),
                launch_arguments={'extra_gz_args': '-r' if enable_gui else '-rs' +
                                                   f' --iterations {iterations}',
                                  'world_file': 'mbari_wec_batch.sdf',
                                  'regenerate_models': 'false',
                                  'pbloghome': batch_results_dir,
                                  'pblogdir': os.path.join(run_results_dir, pblog_dir)}.items(),
            )

            # record all topics with rosbag2
            rosbag2 = ExecuteProcess(
                cmd=['ros2', 'bag', 'record',
                     '-o', os.path.join(batch_results_dir, run_results_dir, rosbag2_dir), '-a'],
                output='screen'
            )

            def capture_start_pids(info_, unused):
                proc_start_info.append(info_)
                start_procs = resolveProcesses(proc_start_info, process='ruby $(which gz) sim')
                for info in [proc_start_info[item] for item in start_procs]:
                    print(f'{info.pid}')
                    print(f'{info.cmd}')
                    import subprocess
                    output = subprocess.run(['ps', 'h', '--ppid', f'{info.pid}', '-o', 'pid'])
                    print(f'{output.stdout}')
                    print(f'{output.stderr}')

            return LaunchDescription([
                rosbag2,
                mbari_wec,
                RegisterEventHandler(  # get info from processes
                    OnProcessStart(on_start=lambda info, unused: proc_start_info.append(info))
                ),
                RegisterEventHandler(  # get return codes from processes
                    OnProcessExit(on_exit=lambda info, unused: proc_exit_info.append(info))
                ),
            ])

        # Launch sim instance
        ls = LaunchService()
        ls.include_launch_description(generate_launch_description())
        result = ls.run()

        if result != 0:
            node.get_logger().error(f'Simulation run [{idx}] was not successful: ' +
                                    f'batch launch return code [{result}]')

        # get return codes from processes
        procs = resolveProcesses(proc_exit_info, process='ruby $(which gz) sim')
        gzsim_result = 0
        for info in [proc_exit_info[item] for item in procs]:
            gzsim_result = info.returncode
            if gzsim_result != 0:
                node.get_logger().error(f'Simulation run [{idx}] was not successful: ' +
                                        f'gz-sim return code [{gzsim_result}]')
                result = gzsim_result
                if info.returncode < -2:
                    node.get_logger().error(f'gz sim process may still be running...')
                    start_procs = resolveProcesses(proc_start_info, process='ruby $(which gz) sim')
                    for info in [proc_start_info[item] for item in start_procs]:
                        for proc in ['ruby', 'sim_pblog']:
                            # kill whole process tree
                            # kill_proc = 'ps x -o "%r %c"' \
                            #             + f' | grep {proc}' \
                            #             + ' | cut -d " " -f 2' \
                            #             + ' | xargs -I{} kill -9 -- -{}'

                            # kill just process
                            kill_proc = 'ps x -o "%p %c"' \
                                        + f' | grep {proc}' \
                                        + " | awk '{ print $1 }'" \
                                        + ' | xargs -I{} kill -9 {}'

                            output = subprocess.run(kill_proc, shell=True, check=True)
                            print(f'{output.args}')

        # Write to log
        with open(os.path.join(batch_results_dir, 'batch_runs.log'), 'a') as fd:
            fd.write(f'{idx}, {result}, {start_time}' +
                     f', {os.path.join(run_results_dir, rosbag2_dir)}' +
                     f', {os.path.join(run_results_dir, pblog_dir)}' +
                     f', {print_optional(ps)}, {print_optional(physics_rtf)}' +
                     f', {print_optional(seed)}, {duration}, {ds}, {sf}' +
                     f', {print_optional(bs)}' +
                     f', {print_optional(mpp)}' +
                     f', {print_optional(iw)}\n')

    # Return all files to defaults
    empy(['-o', base_sdf_file,
          empy_base_sdf_file])
    empy(['-o', world_file,
          empy_world_file])
    empy(['-o', sdf_file,
          empy_sdf_file])


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Batch running for mbari_wec_gz')
    parser.add_argument('sim_params_yaml',
                        help='yaml file with batch parameters')
    args, unknown = parser.parse_known_args()
    generate_simulations(args.sim_params_yaml)