cli.py

# SWARM-IS
# Copyright (c) 2026 Hugh Parsons
# SPDX-License-Identifier: AGPL-3.0-or-later

"""Command-line interface for the swarm surveillance predator control simulation.

Usage
-----
    python cli.py -c configs/adaptive.json
    python cli.py -c configs/adaptive.json --make_video
    python cli.py -c configs/adaptive.json --tune --n_samples 100

All simulation logic lives in src/experiment/. This file handles only:
  - Argument parsing
  - Config loading from JSON
  - Output directory creation
  - Result saving (CSV, plots, video)
  - Dispatching to run_experiment() or main_tune()
"""
import itertools
import json
import os
import re
from datetime import datetime
from pathlib import Path

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt

from src.experiment import run_experiment, main_tune, build_components
from src.experiment.config import (
    apply_transforms,
    flat_iter_dict,
    update_nested_dict,
    get_sample_space_function_and_transform,
    get_risk_maps,
    RE_MATCH_SAMPLE_FUNC,
    RE_MATCH_SELF_REFERENCE,
    EXCLUDED_ID_KEYS,
)
from src.utils import render_frames, save_video_opencv, extract_distribution_frames
from src.constants import SQM_PER_HA


def init_args():
    from argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument('-c', '--config', type=str, required=False)

    parser.add_argument('--tune', action='store_true', default=False)
    parser.add_argument('--sensitivity_combinations', type=str, default=None)
    parser.add_argument("--n_samples", type=int, default=15)
    parser.add_argument("--n_initialization_trials", type=int, default=None)

    parser.add_argument("--n_repeats", type=int, default=10)
    parser.add_argument("--n_workers", type=int, default=None,
                        help="Workers for parallel repeats (default: n_repeats, capped at CPU count)")
    parser.add_argument("--max_concurrent", type=int, default=1,
                        help="Max Ax configs evaluated concurrently (1 = fully sequential Bayesian)")
    parser.add_argument("--cpus_per_trial", type=int, default=None,
                        help="Ray per-trial CPU budget. Enables concurrent-config "
                             "scheduling when combined with --max_concurrent > 1. "
                             "Default: unset (current behaviour, no resource cap).")
    parser.add_argument("--contextual", action='store_true', default=False,
                        help="Contextual BO: one joint GP over tuning + sensitivity "
                             "params instead of independent per-cell searches")
    parser.add_argument("--extract_bests", action='store_true', default=False,
                        help="Extract empirical + model-predicted best configs per cell "
                             "from a completed/partial contextual run. Requires --resume.")
    parser.add_argument("--n_prediction_candidates", type=int, default=500,
                        help="Candidates per cell for model-predicted best extraction")
    parser.add_argument("--track_simulation_progression", "-k", action='store_true', default=False)
    parser.add_argument("--tag", type=str, default="")

    parser.add_argument("--resume", type=str, default=None,
                        help="Path to existing results dir to resume (skips completed combos)")

    parser.add_argument("--make_video", action='store_true', default=False)
    parser.add_argument("--make_plots", action='store_true', default=False)
    parser.add_argument("--render_distributions", action='store_true', default=False)
    args = parser.parse_args()
    return args


def main(args):
    if args.extract_bests:
        if not args.resume:
            raise ValueError("--extract_bests requires --resume <path>")
        output_dir = Path(args.resume)
        if not output_dir.exists():
            raise ValueError(f"Resume directory does not exist: {output_dir}")
        from src.experiment import extract_bests
        extract_bests(output_dir, n_candidates=args.n_prediction_candidates)
        return

    if not args.config:
        raise ValueError("--config is required (unless using --extract_bests)")

    with open(args.config, "r") as f:
        config = json.load(f)

    if args.resume:
        output_dir = Path(args.resume)
        if not output_dir.exists():
            raise ValueError(f"Resume directory does not exist: {output_dir}")
        print(f"Resuming from {output_dir}")
    else:
        tag = args.tag
        if len(tag) != 0:
            tag += "__"
        tag += datetime.now().strftime(r'%y-%m-%d_%H-%M-%S')
        output_dir = Path(os.getcwd()) / "results" / config["name"] / tag
        os.makedirs(output_dir, exist_ok=True)
        with open(output_dir / "config.json", "w+") as f:
            json.dump(config, f, indent=2)

    if args.tune:
        # Suppress Ax's verbose logging before importing it
        import logging
        for _name in ("ax", "ax.service", "ax.modelbridge", "botorch"):
            _logger = logging.getLogger(_name)
            _logger.setLevel(logging.ERROR)
            _logger.handlers.clear()
            _logger.addHandler(logging.NullHandler())
            _logger.propagate = False
        from src.tuning import AxSearch, Repeater  # noqa: F401 — triggers Ray/Ax setup

        base_search_space: dict = config
        sensitivity_combinations: dict = {}
        sensitivity_combination_grid = None

        if args.sensitivity_combinations is not None:
            with open(args.sensitivity_combinations, "r") as f:
                sensitivity_combinations = json.load(f)
            if not args.resume:
                with open(output_dir / "sensitivity_combinations.json", "w+") as f:
                    json.dump(sensitivity_combinations, f, indent=2)

        sensitivity_combinations = {k: v for k, v in flat_iter_dict(sensitivity_combinations)}
        sensitivity_combination_grid = list(itertools.product(*sensitivity_combinations.values()))

        # Only include parameters that actually vary across the grid
        varying_keys = [
            k for k, vals in sensitivity_combinations.items()
            if len(set(str(v) for v in vals)) > 1
        ]

        _KEY_ABBREV = {
            "initial_density_pha": "d",
            "prior_risk_estimate": "prior",
            "g0": "g0",
            "g0_decay_rate": "g0dr",
            "minimum_home_range_area_ha": "minhr",
            "false_positive_probability": "fpr",
            "predator_elimination_probability": "pelim",
            "n_nodes": "n",
            "site_height": "h",
            "site_width": "w",
            "tile_size": "ts",
        }

        def _abbrev_val(v):
            s = str(v)
            if "/" in s:
                s = s.split("/")[-1]
            if s.endswith(".csv"):
                s = s[:-4]
            return s

        def get_run_id(sensitivity_combination):
            components = []
            for k, v in zip(sensitivity_combinations.keys(), sensitivity_combination):
                sub_key = k.split(".")[-1]
                if sub_key in EXCLUDED_ID_KEYS:
                    continue
                if varying_keys and k not in varying_keys:
                    continue
                abbrev_key = _KEY_ABBREV.get(sub_key, sub_key)
                components.append(f"{abbrev_key}={_abbrev_val(v)}")
            return base_search_space["name"] + ("_" if components else "") + "_".join(components)

        def _build_search_space(search_space):
            """Process sample() directives and self-references in a search space."""
            flat_search_space = {k: v for (k, v) in flat_iter_dict(search_space)}
            for k, v in flat_iter_dict(search_space):
                if isinstance(v, str) and RE_MATCH_SAMPLE_FUNC.match(v):
                    self_references = RE_MATCH_SELF_REFERENCE.findall(v)
                    for match in self_references:
                        full_match, self_ref_key = match
                        keys = self_ref_key.split(".")
                        val = search_space.copy()
                        for sub_key in keys:
                            val = val[sub_key]
                        v = v.replace(full_match, str(val), 1)

                    sampler, transform = get_sample_space_function_and_transform(v)
                    update_nested_dict(search_space, k, sampler)
                    if transform is not None:
                        transform_key = f"{k}__transforms"
                        existing_transforms = []
                        if transform_key in flat_search_space:
                            existing_transforms = flat_search_space[transform_key]
                        update_nested_dict(search_space, transform_key, [transform] + existing_transforms)
            return search_space

        n_workers = args.n_workers
        if n_workers is None:
            if args.cpus_per_trial is not None:
                n_workers = min(args.n_repeats, args.cpus_per_trial)
            else:
                n_workers = min(args.n_repeats, os.cpu_count() or 1)

        if args.contextual:
            if args.sensitivity_combinations is None:
                raise ValueError("--contextual requires --sensitivity_combinations")

            from ray import tune as ray_tune
            from src.experiment import run_contextual_tune

            joint_search_space = base_search_space.copy()
            joint_search_space["id"] = base_search_space["name"] + "_contextual"

            for k, vals in sensitivity_combinations.items():
                unique_vals = list(dict.fromkeys(str(v) for v in vals))
                if len(unique_vals) > 1:
                    typed_vals = list(vals)
                    if any(isinstance(v, float) for v in typed_vals):
                        typed_vals = [float(v) for v in typed_vals]
                    update_nested_dict(joint_search_space, k, ray_tune.choice(typed_vals))
                else:
                    update_nested_dict(joint_search_space, k, vals[0])

            _build_search_space(joint_search_space)

            sensitivity_cells = []
            cell_ids = []
            for combo in sensitivity_combination_grid:
                cell = dict(zip(sensitivity_combinations.keys(), combo))
                sensitivity_cells.append(cell)
                cell_ids.append(get_run_id(combo))

            run_contextual_tune(
                search_space=joint_search_space,
                sensitivity_keys=varying_keys,
                sensitivity_cells=sensitivity_cells,
                cell_ids=cell_ids,
                n_samples=args.n_samples,
                n_repeats=args.n_repeats,
                n_workers=n_workers,
                max_concurrent=args.max_concurrent,
                cpus_per_trial=args.cpus_per_trial,
                n_initialization_trials=args.n_initialization_trials,
                output_dir=output_dir,
                resume=args.resume is not None,
                track_metrics_over_simulation=args.track_simulation_progression,
            )
        else:
            search_spaces = []
            for sensitivity_combination in sensitivity_combination_grid:
                id_ = get_run_id(sensitivity_combination)
                search_space = base_search_space.copy()
                search_space["id"] = id_

                for k, v in zip(sensitivity_combinations.keys(), sensitivity_combination):
                    update_nested_dict(search_space, k, v)

                _build_search_space(search_space)
                search_spaces.append(search_space)

            main_tune(
                search_spaces,
                n_samples=args.n_samples,
                n_repeats=args.n_repeats,
                n_workers=n_workers,
                max_concurrent=args.max_concurrent,
                cpus_per_trial=args.cpus_per_trial,
                n_initialization_trials=args.n_initialization_trials,
                output_dir=output_dir,
                track_metrics_over_simulation=args.track_simulation_progression,
            )

    else:
        # Apply any `__transforms` sentinels left over from a tuning run
        # (e.g. ``beta__transforms: ["exponential"]``). The tune path does
        # this per-trial via preprocess_config_for_trial; the single-run
        # path needs it at config-load time so a replayed best-config
        # uses real-space values, not the log-space samples from the CSV.
        config = apply_transforms(config)

        config["track_history"] = (
            config.get("track_history", False)
            or args.make_plots
            or args.make_video
            or args.render_distributions
        )

        # Resolve risk-map specifier strings (e.g. "Patchy", "Homogeneous") to
        # actual file paths. In tuning mode this is done by preprocess_config_for_trial;
        # for direct runs we pick the map that matches the seed.
        prior = config.get("prior_risk_estimate", "")
        if isinstance(prior, str) and not prior.endswith(".csv"):
            risk_maps = get_risk_maps(prior)
            seed = config.get("seed", 0)
            config["prior_risk_estimate"] = str(risk_maps[seed % len(risk_maps)])

        # Run the simulation
        history = run_experiment(config)

        # Build components only when needed for visualization (avoids a double build)
        site = None
        if args.make_video or args.render_distributions:
            site = build_components(config)["site"]

        print(f"Loss: {history.loss[-1]:.4f}")
        print(f"Distance moved: {np.sum(history.distance_moved):.1f}")

        pd.DataFrame({
            "loss": history.loss,
            "distance_moved": history.distance_moved,
        }).to_csv(output_dir / "results.csv")

        if args.make_plots:
            fig = plt.figure(figsize=(16, 8))

            ax1 = fig.add_subplot(1, 2, 1)
            ax1.plot(history.loss)
            ax1.set_ylim(0, 1)
            ax1.set_xlabel("Iteration")
            ax1.set_ylabel("Loss")
            ax1.set_title("Loss over time")

            ax2 = fig.add_subplot(1, 2, 2)
            ax2.plot(np.cumsum(history.distance_moved))
            ax2.set_xlabel("Iteration")
            ax2.set_ylabel("Distance moved")
            ax2.set_title("Cumulative distance moved over time")

            fig.tight_layout()
            fig.savefig(output_dir / "plot.png")

        if args.make_video:
            frames = render_frames(
                site.site_X, site.site_Y,
                np.max(site.site_X) * np.max(site.site_Y) / SQM_PER_HA,
                history.predator_locations,
                history.predators_alive,
                history.node_locations,
                history.detection_surface,
            )
            save_video_opencv(frames, str(output_dir / "output.mp4"))

        if args.render_distributions:
            start, mid, end = extract_distribution_frames(
                site.site_X, site.site_Y,
                np.max(site.site_X) * np.max(site.site_Y) / SQM_PER_HA,
                history.predator_locations,
                history.predators_alive,
                history.node_locations,
                history.detection_surface,
            )
            start.save(output_dir / "start.png")
            mid.save(output_dir / "mid.png")
            end.save(output_dir / "end.png")

        if args.make_plots:
            print(f"plots saved to {str(output_dir / 'plot.png')}")
        if args.make_video:
            print(f"video saved to {str(output_dir / 'output.mp4')}")


if __name__ == "__main__":
    args = init_args()
    main(args)