#!/usr/bin/env python3
"""
Fixed Mitsuba 3 render + sRGB tonemap script for the Agents' Last Exam
Cornell-glass scene reconstruction task.

Agents must not edit this script.
It enforces variant, resolution, sampler, SPP, seed, film, and tonemap.
It preserves the submitted scene camera. The task description instructs agents
to use the provided scaffold camera; camera accuracy is ultimately enforced by
the image metric. Camera recovery is not intended to be part of the task.
It does not read input/reference.png and executes no agent-supplied Python.
"""
from __future__ import annotations

import argparse
import os
from pathlib import Path
import struct
import tempfile
import time
import zlib
import xml.etree.ElementTree as ET

import numpy as np

FIXED_WIDTH = 1024
FIXED_HEIGHT = 1024
FIXED_SPP = 512
FIXED_SEED = 17419
DEFAULT_VARIANT = "scalar_rgb"  # CPU, deterministic, no GPU dependency.


def png_write_rgb8(path: Path, img: np.ndarray) -> None:
    """Write an 8-bit RGB PNG using only Python's standard library."""
    if img.dtype != np.uint8 or img.ndim != 3 or img.shape[2] != 3:
        raise ValueError("png_write_rgb8 expects HxWx3 uint8")
    h, w, _ = img.shape

    def chunk(tag: bytes, data: bytes) -> bytes:
        return (
            struct.pack(">I", len(data))
            + tag
            + data
            + struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF)
        )

    raw = b"".join(b"\x00" + img[y].tobytes() for y in range(h))
    png = b"\x89PNG\r\n\x1a\n"
    png += chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
    png += chunk(b"IDAT", zlib.compress(raw, level=6))
    png += chunk(b"IEND", b"")
    path.write_bytes(png)


def linear_to_srgb(x: np.ndarray) -> np.ndarray:
    x = np.clip(x, 0.0, 1.0)
    return np.where(x <= 0.0031308, 12.92 * x, 1.055 * np.power(x, 1.0 / 2.4) - 0.055)


def tonemap_to_srgb(linear_rgb: np.ndarray) -> np.ndarray:
    """Fixed display transform: clamp to [0,1] and apply sRGB transfer."""
    x = np.nan_to_num(linear_rgb.astype(np.float32), nan=0.0, posinf=1e6, neginf=0.0)
    x = np.maximum(x, 0.0)
    return linear_to_srgb(x)


def force_render_settings(scene_path: Path, spp: int) -> Path:
    """Create a temporary scene that keeps the scaffold camera but fixes film/sampler."""
    tree = ET.parse(scene_path)
    root = tree.getroot()

    sensors = root.findall("sensor")
    if not sensors:
        raise RuntimeError("Submitted scene must contain a <sensor>; camera matching is part of the task.")
    sensor = sensors[0]

    # Remove existing sampler/film under the first sensor.
    for child in list(sensor):
        if child.tag in {"sampler", "film"}:
            sensor.remove(child)

    sampler = ET.SubElement(sensor, "sampler", {"type": "independent"})
    ET.SubElement(sampler, "integer", {"name": "sample_count", "value": str(spp)})

    film = ET.SubElement(sensor, "film", {"type": "hdrfilm"})
    ET.SubElement(film, "integer", {"name": "width", "value": str(FIXED_WIDTH)})
    ET.SubElement(film, "integer", {"name": "height", "value": str(FIXED_HEIGHT)})
    ET.SubElement(film, "string", {"name": "pixel_format", "value": "rgb"})
    ET.SubElement(film, "string", {"name": "component_format", "value": "float32"})
    rfilter = ET.SubElement(film, "rfilter", {"type": "gaussian"})
    ET.SubElement(rfilter, "float", {"name": "stddev", "value": "0.45"})

    tmp = tempfile.NamedTemporaryFile(prefix="ale_scene_", suffix=".xml", delete=False)
    tmp_path = Path(tmp.name)
    tmp.close()
    tree.write(tmp_path, encoding="unicode")
    return tmp_path


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--scene", default="output/scene.xml")
    ap.add_argument("--out-exr", default="output/render.exr")
    ap.add_argument("--out-png", default="output/render.png")
    ap.add_argument("--variant", default=os.environ.get("MI_VARIANT", DEFAULT_VARIANT))
    ap.add_argument("--spp", type=int, default=FIXED_SPP)
    ap.add_argument("--seed", type=int, default=FIXED_SEED)
    args = ap.parse_args()

    scene_path = Path(args.scene).resolve()
    out_exr = Path(args.out_exr).resolve()
    out_png = Path(args.out_png).resolve()
    out_exr.parent.mkdir(parents=True, exist_ok=True)
    out_png.parent.mkdir(parents=True, exist_ok=True)

    import mitsuba as mi

    mi.set_variant(args.variant)
    print(f"mitsuba_version {getattr(mi, '__version__', 'unknown')}")

    fixed_scene_path = force_render_settings(scene_path, args.spp)
    try:
        scene = mi.load_file(str(fixed_scene_path))
        t0 = time.time()
        image = mi.render(scene, spp=args.spp, seed=args.seed)
        mi.util.write_bitmap(str(out_exr), image)

        linear = np.array(image, dtype=np.float32)
        if linear.ndim == 3 and linear.shape[2] > 3:
            linear = linear[:, :, :3]
        srgb = tonemap_to_srgb(linear)
        png_write_rgb8(out_png, (np.clip(srgb, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8))
    finally:
        try:
            fixed_scene_path.unlink(missing_ok=True)
        except Exception:
            pass

    print(f"rendered {FIXED_WIDTH}x{FIXED_HEIGHT}, spp={args.spp}, seed={args.seed}, variant={args.variant}")
    print(f"wrote {out_exr}")
    print(f"wrote {out_png}")
    print(f"seconds {time.time() - t0:.3f}")
    return 0


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