import numpy as np
import trimesh
from utils.config import get_dataset, get_args
import os

# Since there are hundreds of objects in the scene, assigning visually distinguishable colors to each object is difficult. You can change the random seed to check if the two objects are actually segmented apart.
np.random.seed(4)

def vis_one_object(point_ids, scene_points):
    points = scene_points[point_ids]
    color = (np.random.rand(3) * 0.7 + 0.3) * 255
    colors = np.tile(color, (points.shape[0], 1))
    return point_ids, points, colors, color, np.mean(points, axis=0)


def main(args):
    dataset = get_dataset(args)
    
    # Загружаем меш/облако точек используя trimesh
    geometry = trimesh.load(dataset.mesh_path)
    scene_points = np.array(geometry.vertices)
    # scene_points = scene_points - np.mean(scene_points, axis=0)
    
    # Исходные цвета
    if hasattr(geometry.visual, 'vertex_colors'):
        scene_colors = geometry.visual.vertex_colors[:, :3].astype(float)
    else:
        scene_colors = np.ones((len(scene_points), 3)) * 200  # Default color if not available
    
    # Тон-маппинг для осветления цветов
    scene_colors = np.power(scene_colors/255, 1/2.2) * 255
    
    # Создаем массивы для хранения меток экземпляров и их цветов
    instance_ids = np.zeros(len(scene_points), dtype=np.int32)
    instance_colors = np.zeros_like(scene_colors)

    pred = np.load(f'data/prediction/scannet_dust3r_posed_class_agnostic/{args.seq_name}.npz')
    # pred = np.load(f'data/prediction/{args.config}/{args.seq_name}.npz')

    masks = pred['pred_masks']
    num_instances = masks.shape[1]
    
    # Задаем фоновый цвет для точек без инстанса (серый)
    background_color = np.array([128, 128, 128])
    instance_colors[:] = background_color  # Устанавливаем фоновый цвет для всех точек
    
    for idx in range(num_instances):
        mask = masks[:, idx]
        point_ids = np.where(mask)[0]
        
        point_ids, points, colors, label_color, center = vis_one_object(point_ids, scene_points)
        instance_colors[point_ids] = label_color
        instance_ids[point_ids] = idx + 1  # Присваиваем ID экземпляра (начиная с 1)

    # Создаем новое облако точек с цветами инстансов
    cloud = trimesh.PointCloud(
        vertices=scene_points,
        colors=instance_colors.astype(np.uint8)  # Используем цвета инстансов вместо исходных
    )
    
    # Добавляем ID инстансов как атрибут вершин
    cloud.metadata['instance_ids'] = instance_ids
    
    # Создаем директорию для сохранения, если она не существует
    os.makedirs(f'data/ply_output', exist_ok=True)
    
    # Сохраняем в PLY формат
    output_path = f'data/ply_output/{args.seq_name}_segmented.ply'
    cloud.export(output_path)
    
    print(f"Сегментированное облако точек сохранено в {output_path}")
    print(f"Найдено {num_instances} объектов")


if __name__ == '__main__':
    args = get_args()
    main(args)