"""Create and log an image."""

import numpy as np
import rerun as rr

# Create an image with numpy
image = np.zeros((200, 300, 3), dtype=np.uint8)
image[:, :, 0] = 255
image[50:150, 50:150] = (0, 255, 0)

rr.init("rerun_example_image", spawn=True)

rr.log("image", rr.Image(image))

"""Create and log an image with various formats."""

import numpy as np
import rerun as rr

rr.init("rerun_example_image_formats", spawn=True)

# Simple gradient image, logged in different formats.
image = np.array([[[x, min(255, x + y), y] for x in range(0, 256)] for y in range(0, 256)], dtype=np.uint8)
rr.log("image_rgb", rr.Image(image))
rr.log("image_green_only", rr.Image(image[:, :, 1], color_model="l"))  # Luminance only
rr.log("image_bgr", rr.Image(image[:, :, ::-1], color_model="bgr"))  # BGR

# New image with Separate Y/U/V planes with 4:2:2 chroma downsampling
y = bytes([128 for y in range(0, 256) for x in range(0, 256)])
u = bytes([x * 2 for y in range(0, 256) for x in range(0, 128)])  # Half horizontal resolution for chroma.
v = bytes([y for y in range(0, 256) for x in range(0, 128)])
rr.log("image_yuv422", rr.Image(bytes=y + u + v, width=256, height=256, pixel_format=rr.PixelFormat.Y_U_V16_FullRange))

"""Log an image."""

import tempfile

import cv2
import numpy as np
import rerun as rr
from PIL import Image as PILImage, ImageDraw

# Save a transparent PNG to a temporary file.
_, file_path = tempfile.mkstemp(suffix=".png")
image = PILImage.new("RGBA", (300, 200), color=(0, 0, 0, 0))
draw = ImageDraw.Draw(image)
draw.rectangle((0, 0, 300, 200), outline=(255, 0, 0), width=6)
draw.rounded_rectangle((50, 50, 150, 150), fill=(0, 255, 0), radius=20)
image.save(file_path)


rr.init("rerun_example_image_advanced", spawn=True)

# Log the image from the file.
rr.log("from_file", rr.EncodedImage(path=file_path))

# Read with Pillow and NumPy, and log the image.
image = np.array(PILImage.open(file_path))
rr.log("from_pillow_rgba", rr.Image(image))

# Drop the alpha channel from the image.
image_rgb = image[..., :3]
rr.log("from_pillow_rgb", rr.Image(image_rgb))

# Read with OpenCV.
image = cv2.imread(file_path)
# OpenCV uses BGR ordering, we need to make this known to Rerun.
rr.log("from_opencv", rr.Image(image, color_model="BGR"))

"""
Update an image over time.

See also the `image_column_updates` example, which achieves the same thing in a single operation.
"""

import numpy as np
import rerun as rr

rr.init("rerun_example_image_row_updates", spawn=True)

for t in range(20):
    rr.set_time_sequence("time", t)

    image = np.zeros((200, 300, 3), dtype=np.uint8)
    image[:, :, 2] = 255
    image[50:150, (t * 10) : (t * 10 + 100)] = (0, 255, 255)

    rr.log("image", rr.Image(image))

"""
Update an image over time, in a single operation.

This is semantically equivalent to the `image_row_updates` example, albeit much faster.
"""

import numpy as np
import rerun as rr

rr.init("rerun_example_image_column_updates", spawn=True)

# Timeline on which the images are distributed.
times = np.arange(0, 20)

# Create a batch of images with a moving rectangle.
width, height = 300, 200
images = np.zeros((len(times), height, width, 3), dtype=np.uint8)
images[:, :, :, 2] = 255
for t in times:
    images[t, 50:150, (t * 10) : (t * 10 + 100), 1] = 255

# Log the ImageFormat and indicator once, as static.
format = rr.components.ImageFormat(width=width, height=height, color_model="RGB", channel_datatype="U8")
rr.log("images", rr.Image.from_fields(format=format), static=True)

# Send all images at once.
rr.send_columns(
    "images",
    indexes=[rr.TimeSequenceColumn("step", times)],
    # Reshape the images so `Image` can tell that this is several blobs.
    #
    # Note that the `Image` consumes arrays of bytes, so we should ensure that we take a
    # uint8 view of it. This way, this also works when working with datatypes other than `U8`.
    columns=rr.Image.columns(buffer=images.view(np.uint8).reshape(len(times), -1)),
)