import os
import cv2
import imageio
import numpy as np
import open3d as o3d
import os.path as osp
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

def render_frames(o3d_geometry, camera_all, output_dir, save_video=True, save_camera=True):
    # Create off-screen renderer
    render = o3d.visualization.rendering.OffscreenRenderer(
        width=camera_all[0].intrinsic.width,
        height=camera_all[0].intrinsic.height
    )
    
    render_frame_path = os.path.join(output_dir, 'render_frames')
    render_camera_path = os.path.join(output_dir, 'render_cameras')
    os.makedirs(render_frame_path, exist_ok=True)
    os.makedirs(render_camera_path, exist_ok=True)

    video_path = os.path.join(output_dir, 'render_frame.mp4')
    if save_video:
        writer = imageio.get_writer(video_path, fps=10)

    material = o3d.visualization.rendering.MaterialRecord()
    material.shader = 'defaultUnlit'  # Use unlit shader for point clouds
    material.point_size = 1.0  # Match original point size
    material.base_color = [1.0, 1.0, 1.0, 1.0] 
    
    for i, camera_params in enumerate(camera_all):
        if camera_params is None:
            continue
            
        # Set camera view
        render.setup_camera(
            camera_params.intrinsic.intrinsic_matrix,
            camera_params.extrinsic,
            camera_params.intrinsic.width,
            camera_params.intrinsic.height
        )

        if save_camera:
            o3d.io.write_pinhole_camera_parameters(
                os.path.join(render_camera_path, f'camera_{i:03d}.json'),
                camera_params
            )

        # Render
        render.scene.add_geometry("points", o3d_geometry, material)
        img = render.render_to_image()
        render.scene.remove_geometry("points")

        # Save frame
        image_uint8 = (np.asarray(img) * 255).astype(np.uint8)
        frame_filename = f'frame_{i:03d}.png'
        imageio.imwrite(osp.join(render_frame_path, frame_filename), image_uint8)

        if save_video:
            writer.append_data(image_uint8)

    if save_video:
        writer.close()

    return video_path

def find_render_cam(pcd, width=1920, height=1080):
    # For headless servers, we'll need to pre-define camera parameters
    # This creates a default viewing angle looking at the center of the point cloud
    
    # Calculate point cloud center and scale
    center = pcd.get_center()
    scale = np.max(pcd.get_max_bound() - pcd.get_min_bound())
    
    # Create default camera parameters
    camera_params = o3d.camera.PinholeCameraParameters()
    
    # Set intrinsic parameters
    intrinsic = o3d.camera.PinholeCameraIntrinsic()
    intrinsic.set_intrinsics(
        width=width,
        height=height,
        fx=width,
        fy=width,
        cx=width/2,
        cy=height/2
    )
    camera_params.intrinsic = intrinsic
    
    # Set extrinsic parameters (looking at center from a 45-degree angle)
    camera_params.extrinsic = np.array([
        [1, 0, 0, 0],
        [0, np.cos(np.pi/4), -np.sin(np.pi/4), 0],
        [0, np.sin(np.pi/4), np.cos(np.pi/4), 2*scale],
        [0, 0, 0, 1]
    ])
    
    return camera_params

def vis_pred_and_imgs(pts_all, save_path, images_all=None, conf_all=None, save_video=True):
    # Set matplotlib backend to non-interactive
    plt.switch_backend('Agg')
    
    # Normalization
    min_val = pts_all.min(axis=(0, 1, 2), keepdims=True)
    max_val = pts_all.max(axis=(0, 1, 2), keepdims=True)
    pts_all = (pts_all - min_val) / (max_val - min_val)

    pts_save_path = osp.join(save_path, 'pts')
    os.makedirs(pts_save_path, exist_ok=True)

    if images_all is not None:
        images_save_path = osp.join(save_path, 'imgs')
        os.makedirs(images_save_path, exist_ok=True)
    
    if conf_all is not None:
        conf_save_path = osp.join(save_path, 'confs')
        os.makedirs(conf_save_path, exist_ok=True)

    if save_video:
        pts_video_path = osp.join(save_path, 'pts.mp4')
        pts_writer = imageio.get_writer(pts_video_path, fps=10)

        if images_all is not None:
            imgs_video_path = osp.join(save_path, 'imgs.mp4')
            imgs_writer = imageio.get_writer(imgs_video_path, fps=10)

        if conf_all is not None:
            conf_video_path = osp.join(save_path, 'confs.mp4')
            conf_writer = imageio.get_writer(conf_video_path, fps=10)
    
    for frame_id in range(pts_all.shape[0]):
        # Points visualization
        pt_vis = pts_all[frame_id].astype(np.float32)
        pt_vis_rgb = mcolors.hsv_to_rgb(1-pt_vis)
        pt_vis_rgb_uint8 = (pt_vis_rgb * 255).astype(np.uint8)

        # Use matplotlib in non-interactive mode
        fig, ax = plt.subplots()
        ax.imshow(pt_vis_rgb_uint8)
        plt.savefig(osp.join(pts_save_path, f'pts_{frame_id:04d}.png'))
        plt.close(fig)

        if save_video:
            pts_writer.append_data(pt_vis_rgb_uint8)

        if images_all is not None:
            image = images_all[frame_id]
            image_uint8 = (image * 255).astype(np.uint8)
            imageio.imwrite(osp.join(images_save_path, f'img_{frame_id:04d}.png'), image_uint8)

            if save_video:
                imgs_writer.append_data(image_uint8)
        
        if conf_all is not None:
            fig, ax = plt.subplots()
            conf_image = plt.cm.jet(conf_all[frame_id])
            ax.imshow(conf_image)
            plt.savefig(osp.join(conf_save_path, f'conf_{frame_id:04d}.png'))
            plt.close(fig)

            conf_image_uint8 = (conf_image * 255).astype(np.uint8)
            if save_video:
                conf_writer.append_data(conf_image_uint8)
    
    if save_video:
        pts_writer.close()
        if images_all is not None:
            imgs_writer.close()
        if conf_all is not None:
            conf_writer.close()