{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mixed Ground and Space Imaging\n",
    "\n",
    "This minimal example uses three spacecraft: a ground imager, a space-to-space imager, and one passive RSO target. Each observer has its own scenario and reward channel. `MixedReward` keeps the different data types separate while `FreeCommunication` shares knowledge between spacecraft.\n",
    "\n",
    "BSK-RL and its dependencies must be installed. The fixed geometry and permissive imaging tolerances keep this a short configuration demonstration; they are not mission settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import ClassVar\n",
    "\n",
    "import numpy as np\n",
    "from Basilisk.utilities import orbitalMotion\n",
    "\n",
    "from bsk_rl import ConstellationTasking, act, data, obs, sats, scene\n",
    "from bsk_rl.comm import FreeCommunication\n",
    "from bsk_rl.scene.targets import Target\n",
    "from bsk_rl.sim import dyn, fsw\n",
    "\n",
    "EPOCH = \"2026 OCT 01 00:00:00.000 (UTC)\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the spacecraft\n",
    "\n",
    "The observers use different dynamics, flight software, observations, and imaging actions. The passive spacecraft is a participant in the RSO scenario but produces no reward data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GroundObserver(sats.ImagingSatellite):\n",
    "    dyn_type = dyn.GroundStationDynModel\n",
    "    fsw_type = fsw.ImagingFSWModel\n",
    "    observation_spec: ClassVar[list[obs.Observation]] = [obs.Time()]\n",
    "    action_spec: ClassVar[list[act.Action]] = [\n",
    "        act.Image(1, max_duration=30),\n",
    "        act.Charge(duration=10),\n",
    "    ]\n",
    "\n",
    "\n",
    "class RSOObserver(sats.Satellite):\n",
    "    dyn_type = (dyn.SpaceToSpaceImagingDynModel, dyn.GroundStationDynModel)\n",
    "    fsw_type = fsw.SpaceToSpaceImagingFSWModel\n",
    "    observation_spec: ClassVar[list[obs.Observation]] = [\n",
    "        obs.Time(),\n",
    "        obs.RSOTargetProperties(n_ahead_observe=1),\n",
    "    ]\n",
    "    action_spec: ClassVar[list[act.Action]] = [\n",
    "        act.ImageRSO(\n",
    "            1,\n",
    "            max_duration=30,\n",
    "            min_pointing_hold_s=0,\n",
    "            require_illumination_during_hold=False,\n",
    "        ),\n",
    "        act.Charge(duration=10),\n",
    "    ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the targets\n",
    "\n",
    "Both ground (`FixedGroundTargets`) and space (`PassiveTarget`) targets are defined. `FixedGroundTargets` places one ground target below the ground observer after the Earth's orientation is initialized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FixedGroundTargets(scene.UniformTargets):\n",
    "    \"\"\"Put a ground target directly below the initial imaging satellite.\"\"\"\n",
    "\n",
    "    def regenerate_targets(self):\n",
    "        self.targets = [\n",
    "            Target(\"city\", [orbitalMotion.REQ_EARTH * 1e3, 0, 0], priority=3)\n",
    "        ]\n",
    "\n",
    "    def reset_post_sim_init(self):\n",
    "        satellite = self.satellites[0]\n",
    "        position_P = satellite.dynamics.world.PN @ satellite.dynamics.r_BN_N\n",
    "        # Access registration retains this array; update it before windows are\n",
    "        # generated, accounting for the realized Earth orientation at the epoch.\n",
    "        self.targets[0].r_LP_P[:] = (\n",
    "            self.radius * position_P / np.linalg.norm(position_P)\n",
    "        )\n",
    "\n",
    "\n",
    "class PassiveTarget(sats.Satellite):\n",
    "    dyn_type = dyn.RSOTargetDynModel\n",
    "    fsw_type = fsw.FSWModel\n",
    "    observation_spec: ClassVar[list[obs.Observation]] = [obs.Time()]\n",
    "    action_spec: ClassVar[list[act.Action]] = [act.Drift(duration=1e9)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bind scenarios and reward channels\n",
    "\n",
    "`MixedScenario.satellite_mapping` lists all participants, including the passive RSO. `MixedReward.satellite_mapping` lists only the satellites producing each reward channel, and `scenario_mapping` binds each channel to its scenario ID. Channel IDs are explicit names, independent of data classes.\n",
    "\n",
    "The ground target has priority 3 and the RSO has priority 2. The RSO reward callback awards its priority at acquisition for this example; the default RSO reward instead requires complete downlink."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use a nearby RSO and a ground target below the ground observer so that\n",
    "# this short demonstration can acquire both images.\n",
    "satellite_args = dict(\n",
    "    rN=[7e6, 0, 0],\n",
    "    vN=[0, 7500, 0],\n",
    "    oe=None,\n",
    "    sigma_init=[0, 0, 0],\n",
    "    omega_init=[0, 0, 0],\n",
    "    wheelSpeeds=[0.0, 0.0, 0.0],\n",
    "    imageAttErrorRequirement=2.0,  # Permissive pointing tolerance for this demo.\n",
    "    instrumentBaudRate=1000,\n",
    "    batteryStorageCapacity=1e9,\n",
    "    storedCharge_Init=1e9,\n",
    "    dataStorageCapacity=1e6,\n",
    "    storageInit=0,\n",
    ")\n",
    "satellites = [\n",
    "    GroundObserver(\"ground\", sat_args=satellite_args),\n",
    "    RSOObserver(\"observer\", sat_args=satellite_args),\n",
    "    PassiveTarget(\"debris\"),\n",
    "]\n",
    "catalog = scene.RSOTargetCatalog(\n",
    "    (\n",
    "        scene.RSOTarget(\n",
    "            \"object\", \"debris\", (7e6 + 1000, 0, 0), (0, 7500, 0), priority=2\n",
    "        ),\n",
    "    ),\n",
    "    EPOCH,\n",
    ")\n",
    "\n",
    "mixed_scene = scene.MixedScenario(\n",
    "    scenarios={\n",
    "        \"ground_scene\": FixedGroundTargets(1),\n",
    "        \"rso_scene\": scene.RSOTargets(catalog, [\"observer\"]),\n",
    "    },\n",
    "    satellite_mapping={\n",
    "        \"ground_scene\": [\"ground\"],\n",
    "        \"rso_scene\": [\"observer\", \"debris\"],\n",
    "    },\n",
    ")\n",
    "mixed_reward = data.MixedReward(\n",
    "    rewarders={\n",
    "        \"ground_images\": data.UniqueImageReward(),\n",
    "        \"rso_images\": data.RSOImageReward(\n",
    "            quality_threshold=0,\n",
    "            acquisition_reward_fn=lambda record, target: target.priority,\n",
    "        ),\n",
    "    },\n",
    "    satellite_mapping={\n",
    "        \"ground_images\": [\"ground\"],\n",
    "        \"rso_images\": [\"observer\"],\n",
    "    },\n",
    "    scenario_mapping={\n",
    "        \"ground_images\": \"ground_scene\",\n",
    "        \"rso_images\": \"rso_scene\",\n",
    "    },\n",
    ")\n",
    "env = ConstellationTasking(\n",
    "    satellites=satellites,\n",
    "    scenario=mixed_scene,\n",
    "    rewarder=mixed_reward,\n",
    "    communicator=FreeCommunication(),\n",
    "    world_args=dict(utc_init=EPOCH),\n",
    "    time_limit=60,\n",
    "    max_step_duration=1,\n",
    "    sim_rate=1,\n",
    "    log_level=\"INFO\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the environment\n",
    "\n",
    "Both observers image their first candidate. This fixed geometry ensures that the RSO candidate slot is populated; general policies must avoid selecting empty RSO slots.\n",
    "\n",
    "Expected cumulative rewards are `ground: 3`, `observer: 2`, and `debris: 0`. Every spacecraft learns both channels through communication, but each observer generates data only for its assigned channel. Access heterogeneous data explicitly through `data.data[channel_id]`.\n",
    "\n",
    "Mappings can overlap to assign several reward channels and scenarios to one satellite, provided that satellite implements the corresponding models and actions. Combining ground and RSO imaging on the same physical satellite requires additional model integration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    observations, info = env.reset(seed=0)\n",
    "    # Action 0 images a target on each observer, and drifts on the passive RSO.\n",
    "    # None continues the current action on subsequent steps.\n",
    "    for step in range(8):\n",
    "        actions = {sat.name: 0 if step == 0 else None for sat in env.satellites}\n",
    "        observations, rewards, terminated, truncated, info = env.step(actions)\n",
    "\n",
    "    print(\"Cumulative rewards:\", env.rewarder.cum_reward)\n",
    "    print(\"Ground images:\", len(env.rewarder.data.data[\"ground_images\"].imaged))\n",
    "    print(\"RSO captures:\", len(env.rewarder.data.data[\"rso_images\"].captures))\n",
    "    for sat in env.satellites:\n",
    "        print(\n",
    "            sat.name,\n",
    "            \"produces\",\n",
    "            list(sat.data_store.data_stores),\n",
    "            \"knows\",\n",
    "            list(sat.data_store.data.data),\n",
    "        )\n",
    "finally:\n",
    "    env.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv_issue_360 (3.12.7)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}