import pyvista as pv
pv.start_xvfb()

# Configuration - Set input_vtp_path to your VTP file or leave None for default
  # e.g., "path/to/your/file.vtp"
dataset = "plane_transonic_48"

#input_mesh_path = f"Data/{dataset[:-3]}/{dataset}.vtp"  # Input geometry mesh
#output_mesh_path = f"metrics/{dataset[:-3]}/Transformer_baseline-variable-velocity-long/vtk_files/{dataset}.vtk"  # Results with pressure fields


# Configuration - Set paths for input mesh and output results
input_mesh_path = f"Data/{dataset[:-3]}/{dataset}.vtp"  # Input geometry mesh
output_mesh_path = f"metrics/{dataset[:-3]}/Transformer_baseline_ep1000/vtk_files/{dataset}.vtk"  # Results with pressure fields

# Load input mesh (geometry only)
print(f"Loading input mesh: {input_mesh_path}")
input_mesh = pv.read(input_mesh_path)
print(f"Input mesh info: {input_mesh.n_points} points, {input_mesh.n_cells} cells")

# Load output mesh (with pressure fields)
print(f"Loading output mesh: {output_mesh_path}")
output_mesh = pv.read(output_mesh_path)
print("Available fields in output mesh:", output_mesh.array_names)
print(f"Output mesh type: {type(output_mesh)}")
print(f"Output mesh info: {output_mesh.n_points} points, {output_mesh.n_cells} cells")

# Check mesh cell types to understand the data structure
if hasattr(output_mesh, 'celltypes'):
    unique_types = set(output_mesh.celltypes)
    print(f"Cell types in mesh: {unique_types}")

# ParaView likely treats this as a surface - force surface rendering
print("Converting to surface representation like ParaView...")
if output_mesh.n_cells > 0:
    # If it has cells but they're vertices, convert to proper surface
    if output_mesh.n_points == output_mesh.n_cells:
        print("Point cloud data - using points as surface")
        # Keep as is but ensure proper rendering
    else:
        print("Extracting outer surface...")
        output_mesh = output_mesh.extract_surface()
        
print(f"Final mesh: {output_mesh.n_points} points, {output_mesh.n_cells} cells")

# Input mesh is already clean (no need to remove arrays)
output_name = dataset

# Get pressure ranges for consistent colorbars from output mesh
gt_pressure = output_mesh["gt_pressure"]
pred_pressure = output_mesh["pred_pressure"] 

pressure_min = min(gt_pressure.min(), pred_pressure.min())
pressure_max = max(gt_pressure.max(), pred_pressure.max())
pressure_range = [pressure_min, pressure_max]

# Create plotter with 3 columns, 1 row
plotter = pv.Plotter(shape=(1, 3), off_screen=True, window_size=(1920, 540))  # 16:9 aspect ratio for presentations

# Set black background and consistent camera settings for all subplots
for i in range(3):
    plotter.subplot(0, i)
    plotter.set_background('black')

# Column 1: Input (Solid STL)
plotter.subplot(0, 0)
actor = plotter.add_mesh(input_mesh, color="white", 
                        show_edges=False, edge_color="black", line_width=1,
                        smooth_shading=False, show_scalar_bar=False)


# Column 2: Ground Truth
plotter.subplot(0, 1)
# Force surface-like rendering similar to ParaView
actor = plotter.add_mesh(output_mesh, scalars="gt_pressure", clim=pressure_range, cmap="jet", 
                        show_edges=False, smooth_shading=False, show_scalar_bar=False,
                        render_points_as_spheres=False,
                        point_size=3.0,  # Larger points if it's point cloud
                        style='surface')  # Force surface style
plotter.add_scalar_bar(title='Pressure', title_font_size=18, label_font_size=18,
                      position_x=0.1, position_y=0.06, width=0.8, height=0.08, color='white')


# Column 3: Prediction (ensure same orientation as GT)
plotter.subplot(0, 2)
# Force surface-like rendering similar to ParaView
actor = plotter.add_mesh(output_mesh, scalars="pred_pressure", clim=pressure_range, cmap="jet",
                        show_edges=False, smooth_shading=False, show_scalar_bar=False,
                        render_points_as_spheres=False,
                        point_size=1.0,  # Larger points if it's point cloud
                        style='surface')  # Force surface style
plotter.add_scalar_bar(title='Pressure', title_font_size=18, label_font_size=18,
                      position_x=0.1, position_y=0.06, width=0.8, height=0.08, color='white')


# Open movie file with optimized settings  
plotter.open_movie(f"{output_name}.mp4", framerate=8, quality=5)  # Lower settings for speed

# Ultra-fast execution options:
n_frames = 60  # Fewer frames for speed (30° steps)

# Set initial camera position ONCE for all subplots
for subplot_idx in range(3):
    plotter.subplot(0, subplot_idx)
    plotter.camera.elevation = -5
    plotter.reset_camera()  # Only call once at the beginning

plotter.show(auto_close=False)

print(f"Creating {n_frames} frames (speed optimized)...")
for i in range(n_frames):
    angle = i * (360 / n_frames)  # Calculate angle for full 360° rotation
    
    # Only change azimuth (fastest camera operation)
    for subplot_idx in range(3):
        plotter.subplot(0, subplot_idx)
        plotter.camera.azimuth = angle
        # NO reset_camera() here - this is the major slowdown!
    plotter.write_frame()
    
    # Progress indicator
    print(f"Frame {i+1}/{n_frames} ({(i+1)/n_frames*100:.0f}%)")

print("Finalizing video...")
plotter.close()
print(f"✅ Fast video created: {output_name}.mp4")