scenario_example.py

"""Example mosaik scenario: OMOD EV agents coupled with midas-powergrid.

This script demonstrates how to:

1. Start the OmodSimulator with an OMOD JSON file (or auto-generate
   one by running OMOD from Python).
2. Create N EV agents with dynamic bus mapping.
3. Route EV charging power to the correct grid bus via LoadRouter
   entities, which handle the dynamic bus-to-load mapping that mosaik's
   static topology cannot express natively.
4. Run the simulation for 24 hours with step_size = 60 seconds.
5. Persist results via midas-store.

Prerequisites
-------------
- ``midas-mosaik``, ``midas-powergrid``, and ``midas-omod`` must be
  installed.
- Either an OMOD JSON file, or Java + OMOD JAR for auto-generation.

Usage
-----

::

    # Use an existing OMOD JSON file:
    python scenario_example.py --json-path path/to/omod_output.json

    # Auto-run OMOD from Python:
    python scenario_example.py \
        --omod-jar /path/to/omod.jar \
        --osm-file region.osm.pbf \
        --area-geojson area.geojson \
        --n-agents-omod 100 --n-days-omod 1

::

    # Use the test fixture (no OMOD or Java required):
    python scenario_example.py \
        --json-path midas_omod/tests/data/test_agents.json \
        --num-agents 2
"""

from __future__ import annotations

import argparse
import importlib.util
import sys
from pathlib import Path

import mosaik  # noqa: I001

# ------------------------------------------------------------------
# Configuration
# ------------------------------------------------------------------

SIM_CONFIG = {
    "Omod": {
        "python": "midas_omod.simulator:OmodSimulator",
    },
    "LoadRouter": {
        "python": "midas_omod.load_router:LoadRouterSimulator",
    },
    "Powergrid": {
        "python": "midas_powergrid.simulator:PandapowerSimulator",
    },
    "Store": {
        "python": "midas_store.simulator:StoreSimulator",
    },
}

# Simulation parameters
# time is in seconds; step_size = 60 s; 24 h = 86400 s -> 1440 steps
END = 24 * 60 * 60
STEP_SIZE = 60  # seconds


# ------------------------------------------------------------------
# Mosaik transform: kW -> MW
# ------------------------------------------------------------------


def kw_to_mw(value: float) -> float:
    """Convert kilowatts to megawatts."""
    if value is None:
        return 0.0
    return value / 1000.0


# ------------------------------------------------------------------
# Main scenario
# ------------------------------------------------------------------


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Run an OMOD + powergrid co-simulation scenario."
    )
    parser.add_argument(
        "--json-path",
        type=str,
        default=None,
        help="Path to the OMOD JSON output file. "
        "Takes priority over --omod-jar auto-run.",
    )
    parser.add_argument(
        "--num-agents",
        type=int,
        default=2,
        help="Number of EV agents to create.",
    )
    parser.add_argument(
        "--step-size",
        type=int,
        default=STEP_SIZE,
        help="Step size in seconds.",
    )
    parser.add_argument(
        "--end",
        type=int,
        default=END,
        help="End time in seconds.",
    )
    # OMOD auto-run arguments.
    parser.add_argument(
        "--omod-jar",
        type=str,
        default=None,
        help="Path to OMOD JAR file for auto-generation.",
    )
    parser.add_argument(
        "--osm-file",
        type=str,
        default=None,
        help="Path to OpenStreetMap .osm.pbf file.",
    )
    parser.add_argument(
        "--area-geojson",
        type=str,
        default=None,
        help="Path to focus area GeoJSON file.",
    )
    parser.add_argument(
        "--n-agents-omod",
        type=int,
        default=100,
        help="Number of agents for OMOD auto-run.",
    )
    parser.add_argument(
        "--n-days-omod",
        type=int,
        default=1,
        help="Number of days for OMOD auto-run.",
    )
    args = parser.parse_args()

    # ------------------------------------------------------------------
    # Determine JSON path: explicit > auto-run > test fixture
    # ------------------------------------------------------------------

    json_path = None

    if args.json_path:
        json_path = Path(args.json_path)
        if not json_path.exists():
            print(
                f"Error: JSON file not found: {json_path}",
                file=sys.stderr,
            )
            return 1
    elif args.omod_jar and args.osm_file and args.area_geojson:
        # Auto-run OMOD from Python.
        from midas_omod.omod_runner import (
            cached_omod_path,
            run_omod,
        )

        cache_dir = ".cache/omod"
        cached = cached_omod_path(
            cache_dir,
            args.omod_jar,
            args.osm_file,
            args.area_geojson,
            args.n_agents_omod,
            args.n_days_omod,
        )
        if cached is not None:
            json_path = cached
        else:
            print("Running OMOD to generate mobility data...")
            out_path = (
                f"{cache_dir}/omod_output_{args.n_agents_omod}.json"
            )
            Path(cache_dir).mkdir(parents=True, exist_ok=True)
            json_path = run_omod(
                jar_path=args.omod_jar,
                osm_file=args.osm_file,
                area_geojson=args.area_geojson,
                n_agents=args.n_agents_omod,
                n_days=args.n_days_omod,
                out_path=out_path,
            )
            print(f"OMOD output: {json_path}")
    else:
        # Default: use test fixture.
        json_path = Path("midas_omod/tests/data/test_agents.json")
        if not json_path.exists():
            print(
                "Error: No JSON path or OMOD parameters provided.\n"
                "Use --json-path, or --omod-jar + --osm-file + "
                "--area-geojson.",
                file=sys.stderr,
            )
            return 1

    # Check for optional dependencies before starting.
    missing = []
    if importlib.util.find_spec("midas_powergrid") is None:
        missing.append("midas-powergrid")
    if importlib.util.find_spec("midas_store") is None:
        missing.append("midas-store")
    if missing:
        print(
            "The following optional packages are required for the full "
            "scenario but are not installed:",
            file=sys.stderr,
        )
        for pkg in missing:
            print(f"  pip install {pkg}", file=sys.stderr)
        print(
            "\nInstall them and re-run, or use the OMOD simulator "
            "standalone for mobility-only simulation.",
            file=sys.stderr,
        )
        return 1

    world = mosaik.World(SIM_CONFIG)

    # ------------------------------------------------------------------
    # Start simulators
    # ------------------------------------------------------------------

    # OMOD mobility simulator with bus mapper support.
    # Use location_bus_mapping as fallback for grids without geodata.
    location_bus_mapping = {
        "HOME": "Bus-1-0",
        "WORK": "Bus-1-1",
        "OTHER": "Bus-1-2",
    }

    omod_sim = world.start(
        "Omod",
        time_resolution=args.step_size,
        step_size=args.step_size,
        json_path=str(json_path),
        battery_capacity_kwh=60.0,
        charging_power_kw=11.0,
        consumption_kwh_per_km=0.18,
        location_bus_mapping=location_bus_mapping,
    )

    # Power grid simulator (pandapower)
    pg_sim = world.start(
        "Powergrid",
        time_resolution=args.step_size,
        step_size=args.step_size,
        gridfile="midasmv",
    )

    # LoadRouter simulator for dynamic bus routing.
    router_sim = world.start(
        "LoadRouter",
        time_resolution=args.step_size,
        step_size=args.step_size,
    )

    # Store (database for persistence)
    store_sim = world.start(
        "Store",
        step_size=args.step_size,
        filename="omod_powergrid_results.hdf5",
    )

    # ------------------------------------------------------------------
    # Create models
    # ------------------------------------------------------------------

    # Create EV agents
    ev_agents = omod_sim.EVAgent.create(num=args.num_agents)

    # Create the power grid
    grid = pg_sim.Grid()

    # Get the load nodes from the grid.
    loads = [e for e in grid.children if e.type == "Load"]

    # Create a database
    db = store_sim.Database()

    # ------------------------------------------------------------------
    # Dynamic bus routing via LoadRouter
    # ------------------------------------------------------------------
    #
    # Mosaik connections are static — an EV agent cannot dynamically
    # reconnect to a different Load entity at runtime. The LoadRouter
    # solves this by acting as a filter: each LoadRoute entity is
    # bound to a fixed target_bus_id, receives grid_power_mw and
    # bus_id from all EV agents, and forwards only the power from
    # agents currently at that bus.
    #
    # In a real scenario, we would extract bus coordinates from the
    # pandapower grid and use BusMapper.nearest_bus(). Here we use a
    # simplified mapping for demonstration.

    # Build LoadRouter entities — one per available load.
    num_loads = min(len(loads), 3)  # HOME, WORK, OTHER
    load_routers = []
    bus_ids = ["Bus-1-0", "Bus-1-1", "Bus-1-2"]

    for i in range(num_loads):
        router = router_sim.LoadRoute.create(
            num=1, target_bus_id=bus_ids[i]
        )[0]
        load_routers.append((router, loads[i], bus_ids[i]))

    # Connect each EV agent's grid_power_mw and bus_id to each
    # LoadRouter, and each LoadRouter's p_mw to the corresponding
    # Load's p_mw.
    for router, load, _bus_id in load_routers:
        for ev in ev_agents:
            world.connect(ev, router, "grid_power_mw", "bus_id")
        world.connect(router, load, "p_mw")

    # ------------------------------------------------------------------
    # Connect to database (store)
    # ------------------------------------------------------------------

    for ev in ev_agents:
        world.connect(
            ev, db,
            "lat", "lon", "soc_kwh",
            "charging_power_kw", "status",
            "bus_id", "grid_power_mw",
        )

    # Connect grid entities to database
    for child in grid.children:
        world.connect(child, db, "p_mw")

    # ------------------------------------------------------------------
    # Run the simulation
    # ------------------------------------------------------------------

    print(
        f"Starting simulation: end={args.end}s, "
        f"step_size={args.step_size}s, "
        f"agents={args.num_agents}, "
        f"loads_routed={num_loads}"
    )
    world.run(until=args.end)

    print(
        "Simulation complete. Results saved to "
        "omod_powergrid_results.hdf5"
    )
    return 0


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