#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
InteriorGS数据转换脚本
将InteriorGS格式的场景数据转换为项目格式
"""

import os
import json
import sys
from PIL import Image

# 添加项目根目录到路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from src.map_middleware import MapMiddleware
from src.data_structures import Object


def get_bbox_min_and_size(bbox):
    """从bounding_box计算最小点和尺寸"""
    xs = [p['x'] for p in bbox]
    ys = [p['y'] for p in bbox]
    zs = [p['z'] for p in bbox]
    x_min, y_min, z_min = min(xs), min(ys), min(zs)
    x_max, y_max, z_max = max(xs), max(ys), max(zs)
    size = (x_max - x_min, y_max - y_min, z_max - z_min)
    return (x_min, y_min, z_min), size


def convert_one_scene(scene_dir, output_root, middleware):
    """转换单个场景"""
    labels_path = os.path.join(scene_dir, "labels.json")
    if not os.path.exists(labels_path):
        print(f"跳过无labels.json的目录: {scene_dir}")
        return

    print(f"读取场景文件: {labels_path}")
    with open(labels_path, "r", encoding="utf-8") as f:
        labels = json.load(f)

    scene_label = os.path.basename(scene_dir)
    print(f"创建场景: {scene_label}")
    scene = middleware.create_scene(scene_label)

    # 为每个类别维护独立的计数器
    label_counters = {}
    for obj in labels:
        if "bounding_box" not in obj:
            continue
        
        bbox = obj["bounding_box"]
        (x_min, y_min, z_min), size = get_bbox_min_and_size(bbox)
        
        # 为每个类别单独编号
        label = obj['label']
        if label not in label_counters:
            label_counters[label] = 0
        
        unique_label = f"{label}_{label_counters[label]}"
        o = Object(unique_label, size)
        middleware.add_object_to_scene(scene, o, x_min, y_min, z_min)
        label_counters[label] += 1

    # 计算总物体数量
    total_objects = sum(label_counters.values())
    print(f"添加了 {total_objects} 个物体")

    # 自动调整场景尺寸
    print("调整场景尺寸...")
    middleware.auto_resize_scene(scene)

    # 保存场景
    output_scene_dir = os.path.join(output_root, scene_label)
    os.makedirs(output_scene_dir, exist_ok=True)
    scene_save_path = os.path.join(output_scene_dir, "scene.json")
    middleware.save_scene(scene, scene_save_path)
    print(f"场景已保存到: {scene_save_path}")

    # 复制 occupancy.json 和 occupancy.png
    for fname in ["occupancy.json", "occupancy.png"]:
        src = os.path.join(scene_dir, fname)
        dst = os.path.join(output_scene_dir, fname)
        if os.path.exists(src):
            if fname.endswith('.png'):
                # 对于PNG文件，进行左右、上下翻转
                with Image.open(src) as img:
                    # 先左右翻转，再上下翻转
                    flipped_img = img.transpose(Image.FLIP_LEFT_RIGHT).transpose(Image.FLIP_TOP_BOTTOM)
                    flipped_img.save(dst)
                print(f"复制并翻转PNG文件: {fname}")
            else:
                # 对于其他文件，直接复制
                with open(src, "rb") as fsrc, open(dst, "wb") as fdst:
                    fdst.write(fsrc.read())
                print(f"复制文件: {fname}")



def main():
    """主函数"""
    input_root = "interiorGS/InteriorGS"
    output_root = "data/scenes/InteriorGS"
    
    print("=== InteriorGS数据转换 ===")
    print(f"输入目录: {input_root}")
    print(f"输出目录: {output_root}")
    
    middleware = MapMiddleware("config.yaml")
    # 检查输入目录是否存在
    if not os.path.exists(input_root):
        print(f"错误: 输入目录不存在: {input_root}")
        return
    
    # 创建输出目录
    os.makedirs(output_root, exist_ok=True)
    
    # 初始化中间件
    print("初始化地图中间件...")
    middleware = MapMiddleware("config.yaml")
    
    # 处理所有场景
    scene_count = 0
    for scene_name in os.listdir(input_root):
        scene_dir = os.path.join(input_root, scene_name)
        if os.path.isdir(scene_dir):
            print(f"\n处理场景: {scene_name}")
            try:
                convert_one_scene(scene_dir, output_root, middleware)
                scene_count += 1
            except Exception as e:
                print(f"处理场景 {scene_name} 时出错: {e}")
    
    print(f"\n转换完成! 共处理了 {scene_count} 个场景")

    



if __name__ == "__main__":
    main()
