visualization/visualizer.py

import tempfile
import os
import spaces
from typing import Optional
from data.loader import load_simulation_data
from visualization.logger import SimulationLogger
import rerun as rr


@spaces.GPU
def visualize_simulation(file, simulation_index: Optional[int]) -> Optional[str]:
    if file is None or simulation_index is None:
        return None

    try:
        simulations, _ = load_simulation_data(file)
        if not simulations or not isinstance(simulation_index, int) or simulation_index < 0 or simulation_index >= len(simulations):
            print(f"Invalid simulation data or index: {simulation_index}")
            return None

        temp_dir = tempfile.mkdtemp()
        rrd_path = os.path.join(temp_dir, "simulation.rrd")

        simulation = simulations[simulation_index]

        logger = SimulationLogger()
        logger.log_metadata(simulation['instructions'])
        logger.log_subjects(simulation['subjects'])

        if 'helper_keyframes' in simulation:
            logger.log_helper_keyframes(simulation['helper_keyframes'])
        
        logger.log_camera_trajectory(simulation['cameraFrames'])
        logger.log_camera_frames(simulation['cameraFrames'])

        rr.save(rrd_path)

        return rrd_path

    except Exception as e:
        print(f"Error processing simulation: {str(e)}")
        return None