import imageio.v3 as iio
import numpy as np
import wgpu
import pygfx as gfx
import pylinalg as la
from wgpu.gui.offscreen import WgpuCanvas
import matplotlib.pyplot as plt

def snapshot(renderer):

    device = renderer._shared.device
    texture = renderer._blender.color_tex
    size = texture.size
    bytes_per_pixel = 16

    data = device.queue.read_texture(
        {
            "texture": texture,
            "mip_level": 0,
            "origin": (0, 0, 0),
        },
        {
            "offset": 0,
            "bytes_per_row": bytes_per_pixel * size[0],
            "rows_per_image": size[1],
        },
        size,
    )

    return np.frombuffer(data, np.float32).reshape(size[1], size[0], 4)

class FloatFragmentBlender(gfx.renderers.wgpu._blender.BaseFragmentBlender):

    passes = [gfx.renderers.wgpu._blender.OpaquePass()]

    def __init__(self):
        super().__init__()

        usg = wgpu.TextureUsage

        self._texture_info["color"] = (
            wgpu.TextureFormat.rgba32float,
            usg.RENDER_ATTACHMENT | usg.COPY_SRC | usg.TEXTURE_BINDING,
        )


# Create offscreen canvas, renderer and scene
canvas = WgpuCanvas(size=(640, 480), pixel_ratio=1)
renderer = gfx.renderers.WgpuRenderer(canvas)
renderer._blender = FloatFragmentBlender() # Override the default blender method
scene = gfx.Scene()

im = iio.imread("imageio:astronaut.png").astype("float32") / 255
tex = gfx.Texture(im, dim=2)

geometry = gfx.box_geometry(200, 200, 200)
material = gfx.MeshBasicMaterial(map=tex)
cube = gfx.Mesh(geometry, material)
scene.add(cube)

camera = gfx.PerspectiveCamera(70, 16 / 9)
camera.local.z = 400

rot = la.quat_from_euler((0.5, 1.0), order="XY")
cube.local.rotation = la.quat_mul(rot, cube.local.rotation)

canvas.request_draw(lambda: renderer.render(scene, camera))


if __name__ == "__main__":
    # Render the image and extract the memory snapshot
    canvas.draw()
    im1 = snapshot(renderer)
    print("image from snapshot:", im1.shape)  # (960, 1280, 4)
    print(f"array type: {im1.dtype}, min={im1.min()}, max={im1.max()}")

    # Check if the float array still contains the image data
    plt.imshow(im1)
    plt.show()