#!/usr/bin/env python3
"""
Genesis quick start — super-resolution + outpainting on a live Esri tile.

Downloads the main JiT-B/16 checkpoints from https://huggingface.co/MVRL/genesis,
fetches one real Esri World Imagery tile (Sydney Opera House / Circular Quay,
zoom 16), then:

  1. SR: the 256x256 parent tile (z16) -> 512x512 mosaic of its four z17
     children in a single forward pass            -> genesis_sr_out.png
  2. OP: the tile viewed as a 2x2 quadrant mosaic (each quadrant is the
     footprint of one real z17 child); keep the upper-left quadrant real,
     outpaint the other three                     -> genesis_op_out.png

Run from the root of a https://github.com/mvrl/genesis checkout,
with its environment active (`uv sync && source .venv/bin/activate`):

    python example_sr_op.py

Sampling uses the paper eval setting: 50 steps, cfg = 1.0 (the loader
defaults). EMA weights are loaded automatically.

Imagery: Esri World Imagery — Source: Esri, Maxar, Earthstar Geographics,
and the GIS User Community.
"""
import os
import sys

# Repo root: this file at the genesis repo root, or run from a checkout.
REPO = os.path.dirname(os.path.abspath(__file__))
if not os.path.isdir(os.path.join(REPO, "src")):
    REPO = os.getcwd()
    assert os.path.isdir(os.path.join(REPO, "src")), \
        "Run this script from the root of a genesis checkout (github.com/mvrl/genesis)."

# Same sys.path recipe as demos/*.py — demos/ goes LAST (highest priority)
# so its `utils/` package is not shadowed by src/*/utils.py.
sys.path.insert(0, os.path.join(REPO, "src"))
sys.path.insert(0, os.path.join(REPO, "demos"))

import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from utils.genesis_common import (  # demos/utils/genesis_common.py
    fetch_tile_at_zoom,
    inference_device,
    load_op_denoiser,
    load_sr_denoiser,
    outpaint_context_white_holes_preview,
    run_outpainting,
    run_superresolution,
)

HF_REPO = "MVRL/genesis"
ARCH = "JiT-B/16"  # main B/16 checkpoints (H/16 SR alone is 23 GB)

# Fixed scenic spot: Sydney Opera House / Circular Quay waterfront.
LAT, LON, ZOOM = -33.8568, 151.2153, 16  # training zoom levels are 10..19


def main() -> None:
    device = inference_device()  # cuda if available, else cpu
    print(f"Device: {device}")

    # 1. Checkpoints from the HF Hub (cached under ~/.cache/huggingface).
    #    The main SR model is DINOv3-conditioned; the main OP model is no-DINO.
    sr_ckpt = hf_hub_download(HF_REPO, "main/superresolution/main_B16/sr-full-tile-stage3-step0800000.ckpt")
    op_ckpt = hf_hub_download(HF_REPO, "main/outpainting/main_B16/op-new-stage3-step0800000.ckpt")
    dino = hf_hub_download(HF_REPO, "dinov3_vitl16_pretrain_sat493m-eadcf0ff.pth")

    # 2. One real 256x256 Esri World Imagery tile as the input.
    tile = fetch_tile_at_zoom(LON, LAT, ZOOM)  # PIL RGB, 256x256
    tile.save("genesis_sr_in.png")

    # 3. Super-resolution: z16 parent -> 512x512 mosaic of its 4 z17 children.
    #    target_zoom is the CHILD zoom level (feeds the model's GSD embedder).
    sr_model, _ = load_sr_denoiser(sr_ckpt, ARCH, dino_weights=dino, device=device)
    sr_out = run_superresolution(sr_model, tile, device, target_zoom=ZOOM + 1)
    sr_out.save("genesis_sr_out.png")
    print("Saved genesis_sr_in.png (256 parent) and genesis_sr_out.png (512 SR mosaic)")

    del sr_model  # free memory before loading the OP model
    if device.type == "cuda":
        torch.cuda.empty_cache()

    # 4. Outpainting: 2x2 quadrant mosaic — 1 real quadrant, 3 masked holes.
    #    Mask convention: white (255) = hole to generate, black (0) = known.
    mask = Image.new("L", (256, 256), 255)  # everything is a hole ...
    mask.paste(0, (0, 0, 128, 128))         # ... except the upper-left quadrant

    op_model, _ = load_op_denoiser(op_ckpt, ARCH, dino_weights="", device=device)
    outpaint_context_white_holes_preview(tile, mask).save("genesis_op_in.png")
    op_out = run_outpainting(op_model, tile, mask, device, tile_zoom=ZOOM)
    op_out.save("genesis_op_out.png")
    print("Saved genesis_op_in.png (masked input) and genesis_op_out.png (outpainted)")


if __name__ == "__main__":
    main()