#!/usr/bin/env python3
"""Generate review-only E2 prelabels from the exact E0 model generations."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import platform
import resource
import time
from collections import Counter
from datetime import UTC, datetime
from pathlib import Path, PurePosixPath
from typing import Any

from run_recorded_perception_epoch import (
    INSTANCE_MASK_THRESHOLD,
    INSTANCE_SCORE_THRESHOLD,
    SEMANTIC_MODEL_ID,
    SEMANTIC_REVISION,
    _canonical_json,
    _GpuTelemetry,
    _model_files,
    _palette,
    _percentiles,
    _read_object,
    _sha256,
    _valid_sha256,
    _write_json,
    _write_png,
)

PACK_SCHEMA = "missioncore.perception-evaluation-pack/v1"
PACK_IDENTITY_SCHEMA = "missioncore.perception-evaluation-pack-identity/v1"
CONTRACT_SCHEMA = "missioncore.perception-annotation-contract/v1"
VALID_FOV_SCHEMA = "missioncore.k1-valid-fov-mask/v1"
RESULT_SCHEMA = "missioncore.perception-evaluation-prelabels/v1"
RESULT_IDENTITY_SCHEMA = "missioncore.perception-evaluation-prelabels-identity/v1"
REPORT_SCHEMA = "missioncore.perception-evaluation-prelabels-report/v1"
FRAME_SCHEMA = "missioncore.perception-evaluation-prelabel-frame/v1"
MAPPING_PROFILE = "e0-coco-ade-to-e2-taxonomy/v1"
PREVIEW_COUNT = 16

TARGET_NAMES = {
    0: "outside_valid_fov",
    1: "person",
    2: "bicycle",
    3: "motorcycle",
    4: "car",
    5: "heavy_vehicle",
    6: "building_structure",
    7: "paved_road",
    8: "sidewalk_curb",
    9: "ground_dirt",
    10: "grass_low_vegetation",
    11: "tree_woody_vegetation",
    12: "sky",
    13: "static_obstacle",
    14: "animal",
    15: "other_background",
}

INSTANCE_MAPPING = {
    "person": 1,
    "bicycle": 2,
    "motorcycle": 3,
    "car": 4,
    "bus": 5,
    "truck": 5,
    "cat": 14,
    "dog": 14,
    "horse": 14,
    "sheep": 14,
    "cow": 14,
    "elephant": 14,
    "bear": 14,
    "zebra": 14,
    "giraffe": 14,
    "bench": 13,
    "fire hydrant": 13,
    "stop sign": 13,
    "parking meter": 13,
    "chair": 13,
    "potted plant": 13,
}

SEMANTIC_MAPPING = {
    "person": 1,
    "bicycle": 2,
    "minibike": 3,
    "car": 4,
    "bus": 5,
    "truck": 5,
    "van": 5,
    "building": 6,
    "house": 6,
    "skyscraper": 6,
    "wall": 6,
    "door": 6,
    "windowpane": 6,
    "hovel": 6,
    "awning": 6,
    "road": 7,
    "runway": 7,
    "path": 7,
    "sidewalk": 8,
    "stairs": 8,
    "stairway": 8,
    "step": 8,
    "earth": 9,
    "sand": 9,
    "field": 9,
    "land": 9,
    "dirt track": 9,
    "grass": 10,
    "plant": 10,
    "flower": 10,
    "tree": 11,
    "palm": 11,
    "sky": 12,
    "animal": 14,
}


def _arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--evaluation-pack", type=Path, required=True)
    parser.add_argument("--valid-fov-root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--cache", type=Path, required=True)
    parser.add_argument("--telemetry-interval-seconds", type=float, default=1.0)
    return parser.parse_args()


def _safe_artifact(root: Path, encoded: object) -> Path:
    if not isinstance(encoded, str):
        raise RuntimeError("evaluation artifact path is not a string")
    relative = PurePosixPath(encoded)
    if relative.is_absolute() or ".." in relative.parts or not relative.parts:
        raise RuntimeError("evaluation artifact path is unsafe")
    path = root.joinpath(*relative.parts).resolve(strict=True)
    if not path.is_file() or not path.is_relative_to(root):
        raise RuntimeError("evaluation artifact escapes the pack")
    return path


def _validate_pack(root: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]:
    resolved = root.resolve(strict=True)
    manifest = _read_object(resolved / "manifest.json")
    identity = manifest.get("identity")
    identity_sha256 = manifest.get("identity_sha256")
    if (
        manifest.get("schema_version") != PACK_SCHEMA
        or not isinstance(identity, dict)
        or identity.get("schema_version") != PACK_IDENTITY_SCHEMA
        or not _valid_sha256(identity_sha256)
        or hashlib.sha256(_canonical_json(identity)).hexdigest() != identity_sha256
        or manifest.get("generation_id") != f"evaluation-pack-{identity_sha256}"
        or identity.get("preprocessing_profile") != "fixed-valid-fov-fill/v1"
    ):
        raise RuntimeError("evaluation pack identity is invalid")
    artifacts = manifest.get("artifacts")
    if not isinstance(artifacts, list) or not artifacts:
        raise RuntimeError("evaluation pack artifacts are unavailable")
    seen: set[str] = set()
    for artifact in artifacts:
        if not isinstance(artifact, dict) or not isinstance(artifact.get("path"), str):
            raise RuntimeError("evaluation artifact descriptor is invalid")
        encoded = str(artifact["path"])
        if encoded in seen:
            raise RuntimeError("evaluation artifact descriptor is duplicated")
        path = _safe_artifact(resolved, encoded)
        if (
            artifact.get("byte_length") != path.stat().st_size
            or not _valid_sha256(artifact.get("sha256"))
            or _sha256(path) != artifact["sha256"]
        ):
            raise RuntimeError(f"evaluation artifact changed: {encoded}")
        seen.add(encoded)
    contract = _read_object(resolved / "annotation-contract.json")
    if (
        contract.get("schema_version") != CONTRACT_SCHEMA
        or contract.get("evaluation_identity_sha256") != identity_sha256
        or [row.get("id") for row in contract.get("categories", [])]
        != list(range(1, 16))
    ):
        raise RuntimeError("evaluation annotation contract is invalid")
    frames = identity.get("frames")
    if not isinstance(frames, list) or len(frames) != 64:
        raise RuntimeError("evaluation frame set is invalid")
    for expected_image_id, frame in enumerate(frames, start=1):
        if not isinstance(frame, dict):
            raise RuntimeError("evaluation frame descriptor is invalid")
        image_id = frame.get("image_id")
        frame_index = frame.get("frame_index")
        if (
            image_id != expected_image_id
            or not isinstance(frame_index, int)
            or isinstance(frame_index, bool)
        ):
            raise RuntimeError("evaluation frame order changed")
        name = f"image-{image_id:03d}-frame-{frame_index:06d}.png"
        expected_paths = {
            f"images/raw/{name}",
            f"images/valid-fov-fill/{name}",
        }
        if not expected_paths <= seen:
            raise RuntimeError("evaluation frame artifacts are incomplete")
    return manifest, frames


def _load_valid_fov(root: Path, pack: dict[str, Any]) -> Any:
    import numpy as np
    from PIL import Image

    resolved = root.resolve(strict=True)
    manifest = _read_object(resolved / "manifest.json")
    identity = manifest.get("identity")
    identity_sha256 = manifest.get("identity_sha256")
    artifact = manifest.get("artifact")
    pack_identity = pack["identity"]
    if (
        manifest.get("schema_version") != VALID_FOV_SCHEMA
        or not isinstance(identity, dict)
        or not _valid_sha256(identity_sha256)
        or hashlib.sha256(_canonical_json(identity)).hexdigest() != identity_sha256
        or manifest.get("generation_id") != f"valid-fov-mask-{identity_sha256}"
        or manifest.get("generation_id") != pack_identity["valid_fov_generation_id"]
        or identity.get("calibration_sha256") != pack_identity["calibration_sha256"]
        or identity.get("calibration_slot") != pack_identity["calibration_slot"]
        or identity.get("source_id") != pack_identity["source_id"]
        or identity.get("admitted_resolution") != [800, 600]
        or not isinstance(artifact, dict)
        or artifact.get("path") != "mask.png"
        or not _valid_sha256(artifact.get("sha256"))
    ):
        raise RuntimeError("valid-FOV binding is invalid")
    mask_path = (resolved / "mask.png").resolve(strict=True)
    if mask_path.parent != resolved or _sha256(mask_path) != artifact["sha256"]:
        raise RuntimeError("valid-FOV artifact changed")
    with Image.open(mask_path) as opened:
        mask = np.asarray(opened, dtype=np.uint8)
    if mask.shape != (600, 800) or not np.isin(mask, (0, 255)).all():
        raise RuntimeError("valid-FOV mask pixels are invalid")
    return mask > 0


def _resource_snapshot() -> dict[str, float]:
    usage = resource.getrusage(resource.RUSAGE_SELF)
    result = {
        "user_cpu_seconds": float(usage.ru_utime),
        "system_cpu_seconds": float(usage.ru_stime),
        "minor_page_faults": float(usage.ru_minflt),
        "major_page_faults": float(usage.ru_majflt),
        "voluntary_context_switches": float(usage.ru_nvcsw),
        "involuntary_context_switches": float(usage.ru_nivcsw),
    }
    try:
        for line in Path("/proc/self/io").read_text(encoding="ascii").splitlines():
            name, value = line.split(":", 1)
            if name in {"read_bytes", "write_bytes", "rchar", "wchar"}:
                result[f"proc_io_{name}"] = float(value.strip())
    except (OSError, UnicodeDecodeError, ValueError):
        pass
    return result


def _resource_delta(before: dict[str, float], after: dict[str, float]) -> dict[str, float]:
    return {
        name: round(value - before[name], 6)
        for name, value in after.items()
        if name in before
    }


def _semantic_target_map(labels: dict[int, str]) -> Any:
    import numpy as np

    mapping = np.full(256, 15, dtype=np.uint8)
    for source_id, name in labels.items():
        mapping[source_id] = SEMANTIC_MAPPING.get(name.strip().lower(), 15)
    return mapping


def _preview_indices(count: int) -> set[int]:
    admitted = min(PREVIEW_COUNT, count)
    if admitted == 1:
        return {0}
    return {
        (position * (count - 1) + (admitted - 1) // 2) // (admitted - 1)
        for position in range(admitted)
    }


def _preview(image: Any, semantic: Any, instance_map: Any, instances: list[dict[str, Any]]) -> Any:
    import numpy as np
    from PIL import Image, ImageDraw

    colors = np.zeros_like(image)
    for category_id in range(1, 16):
        colors[semantic == category_id] = _palette(category_id)
    blended = np.clip(
        image.astype(np.float32) * 0.52 + colors.astype(np.float32) * 0.48,
        0,
        255,
    ).astype(np.uint8)
    canvas = Image.fromarray(blended)
    draw = ImageDraw.Draw(canvas)
    for instance in instances:
        instance_id = int(instance["instance_id"])
        color = _palette(500 + instance_id)
        x1, y1, x2, y2 = (int(round(value)) for value in instance["box_xyxy"])
        draw.rectangle((x1, y1, x2, y2), outline=color, width=2)
        label = f"{instance['draft_category']} {float(instance['score']):.0%}"
        draw.text((x1, max(0, y1 - 13)), label, fill=color)
        if not np.any(instance_map == instance_id):
            raise RuntimeError("preview instance lost its mask")
    return np.asarray(canvas, dtype=np.uint8)


def _artifact(path: Path, root: Path) -> dict[str, Any]:
    return {
        "path": path.relative_to(root).as_posix(),
        "byte_length": path.stat().st_size,
        "sha256": _sha256(path),
    }


def _run(args: argparse.Namespace) -> int:
    pack_root = args.evaluation_pack.resolve(strict=True)
    output_root = args.output.resolve()
    cache_root = args.cache.resolve(strict=True)
    if output_root.exists():
        raise RuntimeError("prelabel output must be absent")
    pack, frames = _validate_pack(pack_root)
    valid_mask = _load_valid_fov(args.valid_fov_root, pack)

    import numpy as np
    import torch
    import torch.nn.functional as functional
    import torchvision
    import transformers
    from huggingface_hub import snapshot_download
    from PIL import Image
    from torchvision.models.detection import (
        MaskRCNN_ResNet50_FPN_V2_Weights,
        maskrcnn_resnet50_fpn_v2,
    )
    from transformers import AutoImageProcessor, BeitForSemanticSegmentation

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for E2 prelabels")
    output_root.mkdir(mode=0o700, parents=True, exist_ok=False)
    instance_root = output_root / "instance-prelabels"
    semantic_root = output_root / "semantic-prelabels"
    preview_root = output_root / "previews"
    instance_root.mkdir(mode=0o700)
    semantic_root.mkdir(mode=0o700)
    preview_root.mkdir(mode=0o700)
    telemetry_path = output_root / "gpu-telemetry.jsonl"
    metadata_path = output_root / "frames.jsonl"
    device = torch.device("cuda:0")
    frame_paths = []
    for frame in frames:
        name = f"image-{int(frame['image_id']):03d}-frame-{int(frame['frame_index']):06d}.png"
        frame_paths.append(pack_root / "images" / "valid-fov-fill" / name)

    resource_before = _resource_snapshot()
    wall_started = time.perf_counter()
    instance_latencies: list[float] = []
    semantic_latencies: list[float] = []
    instances_by_frame: list[list[dict[str, Any]]] = []
    instance_counts: Counter[str] = Counter()
    ignored_instance_counts: Counter[str] = Counter()
    preview_indices = _preview_indices(len(frames))
    torch.cuda.reset_peak_memory_stats()

    with telemetry_path.open("x", encoding="utf-8", newline="\n") as telemetry_stream:
        with _GpuTelemetry(telemetry_stream, args.telemetry_interval_seconds) as telemetry:
            instance_weights = MaskRCNN_ResNet50_FPN_V2_Weights.DEFAULT
            instance_model = maskrcnn_resnet50_fpn_v2(weights=instance_weights).to(device).eval()
            instance_labels = list(instance_weights.meta["categories"])
            with Image.open(frame_paths[0]) as opened:
                warm_image = opened.convert("RGB")
            with torch.inference_mode():
                warm_tensor = instance_weights.transforms()(warm_image).to(device)
                _ = instance_model([warm_tensor])[0]
                torch.cuda.synchronize()
            del warm_tensor

            with torch.inference_mode():
                for order, path in enumerate(frame_paths, start=1):
                    with Image.open(path) as opened:
                        image = np.asarray(opened.convert("RGB"), dtype=np.uint8)
                    tensor = instance_weights.transforms()(Image.fromarray(image)).to(device)
                    torch.cuda.synchronize()
                    started = time.perf_counter()
                    prediction = instance_model([tensor])[0]
                    torch.cuda.synchronize()
                    instance_latencies.append((time.perf_counter() - started) * 1000.0)
                    scores = prediction["scores"].detach().cpu().numpy()
                    keep = np.flatnonzero(scores >= INSTANCE_SCORE_THRESHOLD)
                    index_map = np.zeros(valid_mask.shape, dtype=np.uint16)
                    instances: list[dict[str, Any]] = []
                    for output_index in keep:
                        source_id = int(prediction["labels"][output_index].item())
                        source_name = instance_labels[source_id].strip().lower()
                        target_id = INSTANCE_MAPPING.get(source_name)
                        if target_id is None:
                            ignored_instance_counts[source_name] += 1
                            continue
                        mask = prediction["masks"][output_index, 0].detach().cpu().numpy()
                        mask = np.logical_and(mask >= INSTANCE_MASK_THRESHOLD, valid_mask)
                        admitted_mask = np.logical_and(mask, index_map == 0)
                        mask_pixels = int(np.count_nonzero(admitted_mask))
                        if mask_pixels < 8:
                            continue
                        instance_id = len(instances) + 1
                        index_map[admitted_mask] = instance_id
                        box = [
                            round(float(value), 6)
                            for value in prediction["boxes"][output_index]
                            .detach()
                            .cpu()
                            .tolist()
                        ]
                        instance = {
                            "instance_id": instance_id,
                            "draft_category_id": target_id,
                            "draft_category": TARGET_NAMES[target_id],
                            "source_model_category_id": source_id,
                            "source_model_category": source_name,
                            "score": round(float(scores[output_index]), 9),
                            "box_xyxy": box,
                            "mask_pixels": mask_pixels,
                            "review_state": "unreviewed-model-draft",
                        }
                        instances.append(instance)
                        instance_counts[TARGET_NAMES[target_id]] += 1
                    _write_png(instance_root / f"image-{order:03d}.png", index_map)
                    instances_by_frame.append(instances)
                    del tensor, prediction
                    if order % 16 == 0 or order == len(frames):
                        print(
                            json.dumps(
                                {"phase": "instance", "processed": order, "total": len(frames)},
                                sort_keys=True,
                            ),
                            flush=True,
                        )
            checkpoint = Path(instance_weights.url).name
            checkpoint_path = Path(torch.hub.get_dir()) / "checkpoints" / checkpoint
            del instance_model
            torch.cuda.empty_cache()

            semantic_snapshot = Path(
                snapshot_download(
                    repo_id=SEMANTIC_MODEL_ID,
                    revision=SEMANTIC_REVISION,
                    cache_dir=cache_root / "huggingface",
                    allow_patterns=("config.json", "preprocessor_config.json", "pytorch_model.bin"),
                )
            )
            processor = AutoImageProcessor.from_pretrained(
                semantic_snapshot,
                local_files_only=True,
                use_fast=False,
            )
            semantic_model, loading = BeitForSemanticSegmentation.from_pretrained(
                semantic_snapshot,
                local_files_only=True,
                output_loading_info=True,
            )
            load_problems = {
                name: loading.get(name, [])
                for name in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")
                if loading.get(name)
            }
            if load_problems:
                raise RuntimeError(
                    "semantic checkpoint did not load exactly: " + json.dumps(load_problems)
                )
            semantic_model = semantic_model.to(device).eval()
            semantic_labels = {
                int(key): str(value) for key, value in semantic_model.config.id2label.items()
            }
            semantic_mapping = _semantic_target_map(semantic_labels)
            with Image.open(frame_paths[0]) as opened:
                warm_image = opened.convert("RGB")
            warm_inputs = processor(images=warm_image, return_tensors="pt")
            warm_inputs = {name: value.to(device) for name, value in warm_inputs.items()}
            with torch.inference_mode():
                _ = semantic_model(**warm_inputs).logits
                torch.cuda.synchronize()
            del warm_inputs

            target_pixels: Counter[str] = Counter()
            with metadata_path.open("x", encoding="utf-8", newline="\n") as metadata_stream:
                with torch.inference_mode():
                    for order, (frame, path, instances) in enumerate(
                        zip(frames, frame_paths, instances_by_frame, strict=True),
                        start=1,
                    ):
                        with Image.open(path) as opened:
                            image = np.asarray(opened.convert("RGB"), dtype=np.uint8)
                        inputs = processor(images=Image.fromarray(image), return_tensors="pt")
                        inputs = {name: value.to(device) for name, value in inputs.items()}
                        torch.cuda.synchronize()
                        started = time.perf_counter()
                        logits = semantic_model(**inputs).logits
                        resized = functional.interpolate(
                            logits,
                            size=image.shape[:2],
                            mode="bilinear",
                            align_corners=False,
                        )
                        source_semantic = (
                            resized.argmax(dim=1)[0].detach().cpu().numpy().astype(np.uint8)
                        )
                        torch.cuda.synchronize()
                        semantic_latencies.append((time.perf_counter() - started) * 1000.0)
                        semantic = semantic_mapping[source_semantic]
                        semantic = semantic.copy()
                        semantic[~valid_mask] = 0
                        instance_map = np.asarray(
                            Image.open(instance_root / f"image-{order:03d}.png"),
                            dtype=np.uint16,
                        )
                        for instance in instances:
                            semantic[instance_map == int(instance["instance_id"])] = int(
                                instance["draft_category_id"]
                            )
                        labels, counts = np.unique(semantic[valid_mask], return_counts=True)
                        for label_id, count in zip(labels, counts, strict=True):
                            target_pixels[TARGET_NAMES[int(label_id)]] += int(count)
                        _write_png(semantic_root / f"image-{order:03d}.png", semantic)
                        if order - 1 in preview_indices:
                            _write_png(
                                preview_root / f"image-{order:03d}.png",
                                _preview(image, semantic, instance_map, instances),
                            )
                        metadata_stream.write(
                            json.dumps(
                                {
                                    "schema_version": FRAME_SCHEMA,
                                    "image_id": frame["image_id"],
                                    "frame_index": frame["frame_index"],
                                    "session_seconds": frame["session_seconds"],
                                    "role": frame["role"],
                                    "group_id": frame["group_id"],
                                    "instances": instances,
                                    "review_state": "unreviewed-model-draft",
                                },
                                ensure_ascii=False,
                                sort_keys=True,
                                separators=(",", ":"),
                                allow_nan=False,
                            )
                            + "\n"
                        )
                        del inputs, logits, resized
                        if order % 16 == 0 or order == len(frames):
                            print(
                                json.dumps(
                                    {
                                        "phase": "semantic",
                                        "processed": order,
                                        "total": len(frames),
                                    },
                                    sort_keys=True,
                                ),
                                flush=True,
                            )
                metadata_stream.flush()
                os.fsync(metadata_stream.fileno())
            model_files = _model_files(semantic_snapshot, checkpoint_path)
            del semantic_model
            torch.cuda.empty_cache()
        telemetry_summary = telemetry.summary()

    wall_seconds = time.perf_counter() - wall_started
    resource_after = _resource_snapshot()
    identity = {
        "schema_version": RESULT_IDENTITY_SCHEMA,
        "evaluation_pack_id": pack["generation_id"],
        "evaluation_identity_sha256": pack["identity_sha256"],
        "valid_fov_generation_id": pack["identity"]["valid_fov_generation_id"],
        "pipeline": "maskrcnn-beit-e2-prelabels/v1",
        "mapping_profile": MAPPING_PROFILE,
        # String keys make the content identity stable across a JSON round trip.
        "target_categories": {str(category_id): name for category_id, name in TARGET_NAMES.items()},
        "instance_mapping": INSTANCE_MAPPING,
        "semantic_mapping": SEMANTIC_MAPPING,
        "instance_score_threshold": INSTANCE_SCORE_THRESHOLD,
        "instance_mask_threshold": INSTANCE_MASK_THRESHOLD,
        "runner_sha256": _sha256(Path(__file__).resolve(strict=True)),
        "models": {
            "instance": {
                "id": "torchvision/maskrcnn_resnet50_fpn_v2",
                "weights": str(instance_weights),
            },
            "semantic": {"id": SEMANTIC_MODEL_ID, "revision": SEMANTIC_REVISION},
            "files": model_files,
        },
    }
    identity_sha256 = hashlib.sha256(_canonical_json(identity)).hexdigest()
    result_id = f"evaluation-prelabels-{identity_sha256}"
    report = {
        "schema_version": REPORT_SCHEMA,
        "result_id": result_id,
        "created_at_utc": datetime.now(UTC)
        .isoformat(timespec="milliseconds")
        .replace("+00:00", "Z"),
        "state": "completed-unreviewed-model-draft",
        "identity": identity,
        "warning": (
            "These outputs are model-assisted prelabels, not ground truth and not an accuracy "
            "measurement. Every accepted label requires human review."
        ),
        "metrics": {
            "frames": len(frames),
            "instances": sum(instance_counts.values()),
            "instances_by_target": dict(instance_counts.most_common()),
            "ignored_instances_by_source": dict(ignored_instance_counts.most_common()),
            "semantic_pixels_by_target_inside_valid_fov": dict(target_pixels.most_common()),
            "instance_forward_ms": _percentiles(instance_latencies),
            "semantic_forward_ms": _percentiles(semantic_latencies),
            "combined_forward_mean_ms": round(
                _percentiles(instance_latencies)["mean"]
                + _percentiles(semantic_latencies)["mean"],
                6,
            ),
            "wall_seconds": round(wall_seconds, 6),
        },
        "runtime": {
            "hostname": platform.node(),
            "platform": platform.platform(),
            "python": platform.python_version(),
            "torch": torch.__version__,
            "torchvision": torchvision.__version__,
            "transformers": transformers.__version__,
            "cuda_runtime": torch.version.cuda,
            "gpu": torch.cuda.get_device_name(),
            "gpu_compute_capability": list(torch.cuda.get_device_capability()),
            "cuda_peak_memory_allocated_mib": round(
                torch.cuda.max_memory_allocated() / 2**20,
                3,
            ),
            "cuda_peak_memory_reserved_mib": round(
                torch.cuda.max_memory_reserved() / 2**20,
                3,
            ),
            "process_peak_rss_mib": round(
                resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024,
                3,
            ),
            "resource_delta": _resource_delta(resource_before, resource_after),
            "gpu_telemetry": telemetry_summary,
        },
    }
    _write_json(output_root / "run-report.json", report)
    artifacts = [
        _artifact(path, output_root)
        for path in sorted(output_root.rglob("*"))
        if path.is_file()
    ]
    result = {
        "schema_version": RESULT_SCHEMA,
        "result_id": result_id,
        "identity_sha256": identity_sha256,
        "identity": identity,
        "review_state": "unreviewed-model-draft",
        "artifacts": artifacts,
    }
    _write_json(output_root / "result.json", result)
    print(
        json.dumps(
            {
                "state": "completed-unreviewed-model-draft",
                "result_id": result_id,
                "frames": len(frames),
                "wall_seconds": round(wall_seconds, 6),
            },
            sort_keys=True,
        ),
        flush=True,
    )
    return 0


def main() -> int:
    return _run(_arguments())


if __name__ == "__main__":
    raise SystemExit(main())
