# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "marimo>=0.21.1",
#     "pyzmq>=27.1.0",
# ]
# ///

# pyright: reportMissingImports=false, reportMissingParameterType=false, reportUnknownParameterType=false, reportUnknownVariableType=false, reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnusedParameter=false, reportUnusedCallResult=false, reportUnusedExpression=false

import marimo

__generated_with = "0.22.4"
app = marimo.App(width="medium", app_title="Test Scene Image Collector")


@app.cell
def _():
    from pathlib import Path

    import marimo as mo

    return Path, mo


@app.cell
def _(Path, mo):
    SCENES = {
        "skokloster-castle": "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb",
        "apartment_1": "data/scene_datasets/habitat-test-scenes/apartment_1.glb",
        "van-gogh-room": "data/scene_datasets/habitat-test-scenes/van-gogh-room.glb",
    }
    IMAGE_SIZE = 1024
    VIEWS_PER_POINT = 12
    POINTS_PER_SCENE = 5
    SEED = 42
    OUTPUT_DIR = Path("outputs/test_image")

    mo.md(
        f"## Configuration\n- Scenes: {len(SCENES)}\n- Points/scene: {POINTS_PER_SCENE}\n- Views/point: {VIEWS_PER_POINT}\n- Resolution: {IMAGE_SIZE}x{IMAGE_SIZE}\n- Total images: {len(SCENES) * POINTS_PER_SCENE * VIEWS_PER_POINT}"
    )
    return (
        IMAGE_SIZE,
        OUTPUT_DIR,
        POINTS_PER_SCENE,
        SCENES,
        SEED,
        VIEWS_PER_POINT,
    )


@app.cell
def _(IMAGE_SIZE, OUTPUT_DIR, POINTS_PER_SCENE, SCENES, SEED, VIEWS_PER_POINT):
    from simulator import collect_scene_images

    scene_stats = {}
    total_images = 0

    for scene_name, scene_path in SCENES.items():
        print(f"Collecting {scene_name}...")
        collected = collect_scene_images(
            scene_name=scene_name,
            scene_path=scene_path,
            output_dir=OUTPUT_DIR,
            image_size=IMAGE_SIZE,
            views_per_point=VIEWS_PER_POINT,
            points_per_scene=POINTS_PER_SCENE,
            seed=SEED,
        )
        scene_stats[scene_name] = collected
        total_images += collected
        print(f"Collected {collected} images for {scene_name}")
    return scene_stats, total_images


@app.cell
def _(OUTPUT_DIR, mo, scene_stats, total_images):
    mo.md(f"""
    ## Summary\n- Output directory: `{OUTPUT_DIR}`\n- Total images: {total_images}\n- Per-scene: {scene_stats}
    """)
    return


@app.cell
def _(OUTPUT_DIR, SCENES, mo):
    from PIL import Image

    preview_grid = []
    for _scene_name in SCENES:
        sample_path = OUTPUT_DIR / _scene_name / "000" / "view_000.png"
        if not sample_path.exists():
            continue
        preview_grid.append(
            mo.vstack(
                [mo.md(f"**{_scene_name}**"), mo.image(Image.open(sample_path))],
                align="center",
            )
        )

    preview = (
        mo.vstack(preview_grid, align="center")
        if preview_grid
        else mo.md("No preview images yet.")
    )
    preview
    return


if __name__ == "__main__":
    app.run()