import os
import warp as wp
import warp.sim
import warp.sim.render
import matplotlib.pyplot as plt
import numpy as np
os.environ['PYGLET_HEADLESS'] = 'true'

def create_dummy_model():
    # Create an empty model builder
    builder = wp.sim.ModelBuilder()
    builder.default_particle_radius = 0.01
    
    cell_dim = 14
    cell_size = (1.0 / cell_dim)
    center = cell_size * cell_dim * 0.5
    density = 1.0
    builder.add_soft_grid(
        pos=wp.vec3(-center, 0.0, -center),
        rot=wp.quat_identity(), 
        vel=wp.vec3(0.0, 0.0, 0.0),
        dim_x=cell_dim,
        dim_y=cell_dim,
        dim_z=cell_dim,
        cell_x=cell_size,
        cell_y=cell_size,
        cell_z=cell_size,
        density=density,
        fix_left=True,
        # doesnt matter what these are because we will assign them later
        k_mu=1e5,
        k_lambda=1e5,
        k_damp=0,
    )
    # Finalize the model
    model = builder.finalize()
    model.gravity[1] = 0.0
    model.ground = False
    return model

# Create renderer with dummy model
model = create_dummy_model()
renderer = wp.sim.render.SimRendererOpenGL(
    model=model,
    path='rendered_frames/opengl',
    scaling=1.0,
    headless=True,  # Set to False if you want to see the window
    up_axis="Y"
)

# model2 = create_dummy_model()
renderer2 = wp.sim.render.SimRendererOpenGL(
    model=model,
    path='rendered_frames/opengl2',
    scaling=1.0,
    headless=True,  # Set to False if you want to see the window
    up_axis="Y"
)

# Create a dummy state (needed for rendering)
state = model.state(requires_grad=True)
split_up_tiles = False
num_tiles = 1
# Create pixel buffer with correct shape
channels = 3  # RGB
if split_up_tiles:
    pixel_shape = (num_tiles, renderer.tile_height, renderer.tile_width, channels)
else:
    pixel_shape = (renderer.screen_height, renderer.screen_width, channels)

pixels = wp.zeros(pixel_shape, dtype=wp.float32)

# Different rendering modes
render_mode = "rgb"  # or "depth"

# Render loop
while renderer.is_running():
    renderer.begin_frame()
    renderer.render(state)
    
    print(renderer.screen_width, renderer.screen_height)
    
    # Get pixels from renderer
    renderer.get_pixels(
        pixels, 
        split_up_tiles=split_up_tiles,  # Set to True if using tiled rendering
        mode=render_mode,
    )
    
    # Save pixels to file
    output_dir = "rendered_frames"
    os.makedirs(output_dir, exist_ok=True)
    frame_num = int(renderer.clock_time * 30)  # 30 fps
    
    # If using tiled rendering
    if split_up_tiles:
        # Handle each tile separately
        for tile in range(num_tiles):
            tile_frame = pixels[tile].numpy()
            output_path = os.path.join(output_dir, f"frame_{frame_num:04d}_tile_{tile}.png")
            plt.imsave(output_path, tile_frame)
    else:
        # Save single frame
        frame = pixels.numpy()
        print(f"Min pixel value: {np.min(frame)}, Max pixel value: {np.max(frame)}")
        output_path = os.path.join(output_dir, f"frame_{frame_num:04d}.png")
        plt.imsave(output_path, frame)
    
    renderer.end_frame()
    
renderer.clear()