{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bc95662-a958-4fb2-b7bd-b09148f8f288",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bbi import *\n",
    "from bbi.experimental import SpawningStrategy, SpawnZoneV2\n",
    "\n",
    "world = World(\"ocean\")\n",
    "\n",
    "world.camera.set_transform(\n",
    "    Transform(\n",
    "        position=Vector3(0, 0, 20),\n",
    "        rotation=Vector3(0, 0, 0)\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "761acd69-ff63-4e8c-8c8d-30e1eddf3654",
   "metadata": {},
   "source": [
    "# Enable thermal mode\n",
    "\n",
    "- Set thermal mode to view thermal settings on the viewport.\n",
    "- Apply the `high_noise_thermal` sensor preset to get thermal artifacts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b984193-b4a4-498f-98d2-3b133c23ae48",
   "metadata": {},
   "outputs": [],
   "source": [
    "world.camera.set_thermal_mode(True)\n",
    "world.camera.set_preset(\"high_noise_thermal\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90bfa706-4781-46df-bcff-93ea059a4c84",
   "metadata": {},
   "source": [
    "# Camera and environment\n",
    "\n",
    "- Set the thermal conditions of environment elements.\n",
    "- Tune sensor artifacts if `high_noise_thermal` does not match your sensor well enough."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1171d905-936f-4341-b209-dac3eac5f1b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "world.camera.thermal.set_sky_temperature(0.3)\n",
    "world.camera.thermal.set_ocean_temperature(0.2)\n",
    "world.camera.thermal.set_sky_transition_point(0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e589f39-c2e1-4e35-946c-aa05b4d1d6a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "world.camera.post_processing.set_blur(0.2)\n",
    "world.camera.post_processing.set_film_grain_amount(0.3)\n",
    "world.camera.post_processing.set_film_grain_scale(0.4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f53a1650-64c3-4cf9-bb46-dc9a08f8c78a",
   "metadata": {},
   "source": [
    "# Assets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90f5ad15-cf63-4544-810e-7a0d3d2271b1",
   "metadata": {},
   "source": [
    "## Singular asset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05379da4-af1c-4631-901a-d194ea7cd16f",
   "metadata": {},
   "outputs": [],
   "source": [
    "asset1 = BuoyantAsset(\n",
    "    \"asset1\",\n",
    "    \"Barzan Bergen Unloaded\",\n",
    "    position=Vector3(382, -250, 0),\n",
    "    rotation=Vector3(0, 0, 210)\n",
    ")\n",
    "world.add(asset1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "907ef2a2-4ab1-4332-a200-de89efb7af46",
   "metadata": {},
   "outputs": [],
   "source": [
    "asset1.thermal.set_temperature(0.7)\n",
    "asset1.thermal.set_glass_brightness(0.7)\n",
    "asset1.thermal.set_reflectivity(0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e186ee4-778b-42a0-a79b-ac2dd7293da8",
   "metadata": {},
   "outputs": [],
   "source": [
    "asset2 = BuoyantAsset(\n",
    "    \"asset2\",\n",
    "    \"Manhattan\",\n",
    "    position=Vector3(5000, 4000, 0),\n",
    "    rotation=Vector3(0, 0, 0)\n",
    ")\n",
    "world.add(asset2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9627f706-7624-43be-a9e3-85a93992d1aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "asset2.thermal.set_temperature(0.1)\n",
    "asset2.thermal.set_reflectivity(0.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa8c7dc9-4337-4d52-a468-2f34897181a1",
   "metadata": {},
   "source": [
    "## Spawnzone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ede1bff-1766-4001-95ec-69b7f29c2f15",
   "metadata": {},
   "outputs": [],
   "source": [
    "spawn_area = Polygon(\n",
    "    \"spawn_area\",\n",
    "    points=[\n",
    "        Vector3(117.75, -104.51, 0),\n",
    "        Vector3(50.37, -43.15, 0),\n",
    "        Vector3(51.52, 39.25, 0),\n",
    "        Vector3(113.33, 97.34, 0),\n",
    "    ]\n",
    ")\n",
    "world.add(spawn_area)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c386e6de-4cde-4650-9648-29b9a4c7629e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MySpawningStrategy(SpawningStrategy):\n",
    "    def __init__(\n",
    "        self, \n",
    "        area: Polygon, \n",
    "        asset_list: list[str],\n",
    "        temperature_range: tuple[float, float],\n",
    "    ):\n",
    "        self.area = area\n",
    "        self.asset_list = asset_list\n",
    "        self.temperature_range = temperature_range\n",
    "        \n",
    "    def asset_factory(self):\n",
    "        return BuoyantAsset(\n",
    "            SpawningStrategy.NODE_ID_PLACEHOLDER,\n",
    "            asset_name=rng.choice(self.asset_list)\n",
    "        )\n",
    "\n",
    "    def placement_strategy(self) -> Transform:\n",
    "        random_position = self.area.uniform_random_points(1)[0]\n",
    "        random_rotation = Vector3(0, 0, rng.uniform(0, 360))\n",
    "        return Transform(random_position, random_rotation)\n",
    "\n",
    "    def post_spawn_hook(self, asset: Asset) -> None:\n",
    "        temperature = rng.uniform(*self.temperature_range)\n",
    "        asset.thermal.set_temperature(temperature)\n",
    "        asset.thermal.set_glass_brightness(temperature)\n",
    "        asset.thermal.set_reflectivity(0.3)\n",
    "\n",
    "spawnzone = SpawnZoneV2(\n",
    "    node_id=\"spawnzone\",\n",
    "    spawning_strategy=MySpawningStrategy(\n",
    "        area=spawn_area,\n",
    "        asset_list=[\"Offshore (Harbour Launch)\", \"Halifax\", \"Greendeck\"],\n",
    "        temperature_range=(0.4, 1.0)\n",
    "    ),\n",
    "    quantity=5,\n",
    ")\n",
    "world.add(spawnzone)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34ec2466-1350-48bc-8751-50820982b255",
   "metadata": {},
   "source": [
    "# Render\n",
    "\n",
    "- Render with argument `params=RenderParams.Thermal()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bbdfce8-73e9-42e7-b32f-631fd5a484c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "scenario = world.new_scenario(1)\n",
    "\n",
    "scenario.render(params=RenderParams.Thermal())"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}