#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
from direct.showbase.ShowBase import ShowBase
from panda3d.core import (
    Point3, Vec3, Vec4, CardMaker, Texture, GraphicsOutput,
    FrameBufferProperties, WindowProperties, RenderState,
    ColorWriteAttrib, DepthTestAttrib, BitMask32, PNMImage,
    Camera, DisplayRegion, GraphicsStateGuardian, RenderState
)

class ColorPickingTest(ShowBase):
    def __init__(self):
        ShowBase.__init__(self)
        
        # 相机设置
        self.cam.setPos(-3.87655, -188.38084, 82.602684)
        self.cam.lookAt(0, 0, 0)
        
        print("=== 颜色编码点击检测测试 ===")
        
        # 创建测试对象
        self.createTestObjects()
        
        # 设置颜色拾取缓冲区
        self.setupColorPickingBuffer()
        
        # 测试颜色拾取
        self.testColorPicking()
    
    def createTestObjects(self):
        """创建测试对象"""
        self.objects = []
        self.id_to_object = {}  # ID到对象的映射
        
        # 创建几个测试立方体
        positions = [
            (0, 10, 3, "中心立方体"),
            (-5, 10, 3, "左侧立方体"), 
            (5, 10, 3, "右侧立方体"),
            (0, 10, 8, "上方立方体"),
            (0, 15, 3, "远处立方体")
        ]
        
        display_colors = [
            (1, 0, 0, 1),  # 红色
            (0, 1, 0, 1),  # 绿色
            (0, 0, 1, 1),  # 蓝色
            (1, 1, 0, 1),  # 黄色
            (1, 0, 1, 1),  # 紫色
        ]
        
        for i, ((x, y, z, name), display_color) in enumerate(zip(positions, display_colors)):
            obj_id = i + 1  # 从1开始编号
            
            # 创建立方体
            cm = CardMaker(f"cube_{i}")
            cm.setFrame(-1, 1, -1, 1)
            cube = self.render.attachNewNode(cm.generate())
            cube.setPos(x, y, z)
            cube.setScale(2, 2, 2)
            cube.setColor(*display_color)
            cube.setBillboardAxis()
            cube.setName(name)
            
            # 为颜色拾取创建唯一颜色
            # 将ID编码为RGB颜色
            pick_color = self.encodeIdToColor(obj_id)
            
            # 存储对象信息
            obj_info = {
                'node': cube,
                'name': name,
                'world_pos': Point3(x, y, z),
                'id': obj_id,
                'pick_color': pick_color,
                'display_color': display_color
            }
            self.objects.append(obj_info)
            self.id_to_object[obj_id] = obj_info
            
            print(f"创建 {name} 在位置 ({x}, {y}, {z}), ID: {obj_id}, 拾取颜色: {pick_color}")
    
    def encodeIdToColor(self, obj_id):
        """将对象ID编码为RGB颜色"""
        # 将ID分解为RGB分量（24位）
        r = (obj_id & 0xFF0000) >> 16
        g = (obj_id & 0x00FF00) >> 8
        b = (obj_id & 0x0000FF)
        
        # 转换为0-1范围的浮点数
        return Vec4(r / 255.0, g / 255.0, b / 255.0, 1.0)
    
    def decodeColorToId(self, color):
        """将RGB颜色解码为对象ID"""
        if color is None:
            return 0
        
        # 转换为0-255范围的整数
        r = int(color[0] * 255 + 0.5)
        g = int(color[1] * 255 + 0.5)
        b = int(color[2] * 255 + 0.5)
        
        # 组合为ID
        obj_id = (r << 16) | (g << 8) | b
        return obj_id
    
    def setupColorPickingBuffer(self):
        """设置颜色拾取渲染缓冲区"""
        # 创建离屏渲染缓冲区
        fbp = FrameBufferProperties()
        fbp.setRgbColor(True)
        fbp.setColorBits(8, 8, 8, 0)  # RGB，不需要Alpha
        fbp.setDepthBits(24)
        
        wp = WindowProperties()
        wp.setSize(self.win.getXSize(), self.win.getYSize())
        
        # 创建缓冲区
        self.pick_buffer = self.graphicsEngine.makeOutput(
            self.pipe, "pick_buffer", -1, fbp, wp, 
            GraphicsOutput.M_offscreen, self.win.getGsg(), self.win
        )
        
        if self.pick_buffer is None:
            print("无法创建颜色拾取缓冲区")
            return
        
        # 创建拾取相机
        self.pick_cam = self.makeCamera(self.pick_buffer)
        self.pick_cam.node().setLens(self.cam.node().getLens())
        self.pick_cam.setPos(self.cam.getPos())
        self.pick_cam.setHpr(self.cam.getHpr())
        
        # 创建拾取场景（带有颜色编码的对象）
        self.pick_scene = self.render.attachNewNode("pick_scene")
        
        # 为每个对象创建拾取版本
        for obj_info in self.objects:
            # 复制几何体
            pick_node = obj_info['node'].copyTo(self.pick_scene)
            pick_node.setColor(obj_info['pick_color'])
            pick_node.setRenderModeWireframe(False)
            pick_node.setLightOff()  # 关闭光照以确保颜色准确
        
        # 设置拾取相机只渲染拾取场景
        self.pick_cam.node().setCameraMask(BitMask32.allOn())
        
        # 创建颜色纹理
        self.pick_tex = Texture()
        self.pick_tex.setFormat(Texture.FRgb8)
        self.pick_buffer.addRenderTexture(self.pick_tex, GraphicsOutput.RTMCopyRam)
        
        print("颜色拾取缓冲区设置完成")
    
    def getColorAtPixel(self, mouseX, mouseY):
        """获取指定像素的颜色值"""
        if self.pick_buffer is None:
            return None
        
        # 获取窗口大小
        winWidth = self.pick_buffer.getXSize()
        winHeight = self.pick_buffer.getYSize()
        
        # 确保坐标在窗口范围内
        if mouseX < 0 or mouseX >= winWidth or mouseY < 0 or mouseY >= winHeight:
            return None
        
        # 强制渲染拾取缓冲区
        self.pick_buffer.setActive(True)
        self.graphicsEngine.renderFrame()
        
        # 读取颜色纹理
        if self.pick_tex.hasRamImage():
            # 获取颜色图像
            pnm = PNMImage()
            self.pick_tex.store(pnm)
            
            # Y坐标需要翻转（OpenGL坐标系）
            flippedY = winHeight - 1 - mouseY
            
            if flippedY < pnm.getYSize() and mouseX < pnm.getXSize():
                # 获取颜色值
                r = pnm.getRed(mouseX, flippedY)
                g = pnm.getGreen(mouseX, flippedY)
                b = pnm.getBlue(mouseX, flippedY)
                return (r, g, b)
        
        return None
    
    def pickObjectByColor(self, mouseX, mouseY):
        """使用颜色编码选择对象"""
        print(f"\n=== 颜色拾取: 鼠标位置 ({mouseX}, {mouseY}) ===")
        
        # 获取点击位置的颜色
        color = self.getColorAtPixel(mouseX, mouseY)
        if color is None:
            print("无法读取颜色值")
            return None
        
        print(f"拾取颜色: RGB({color[0]:.3f}, {color[1]:.3f}, {color[2]:.3f})")
        
        # 解码颜色为对象ID
        obj_id = self.decodeColorToId(color)
        print(f"解码的对象ID: {obj_id}")
        
        # 查找对应的对象
        if obj_id in self.id_to_object:
            obj_info = self.id_to_object[obj_id]
            print(f"选中对象: {obj_info['name']}")
            return obj_info
        else:
            print("没有找到对应的对象")
            return None
    
    def testColorPicking(self):
        """测试颜色拾取"""
        print(f"\n=== 颜色拾取测试 ===")
        
        # 等待几帧以确保场景渲染完成
        for i in range(5):
            self.graphicsEngine.renderFrame()
        
        # 模拟不同的点击位置
        test_clicks = [
            (619, 362, "屏幕中心"),
            (400, 362, "屏幕偏左"),
            (800, 362, "屏幕偏右"),
            (619, 200, "屏幕上方"),
            (619, 500, "屏幕下方"),
        ]
        
        for mouseX, mouseY, description in test_clicks:
            print(f"\n--- 测试 {description}: ({mouseX}, {mouseY}) ---")
            selected = self.pickObjectByColor(mouseX, mouseY)

if __name__ == "__main__":
    app = ColorPickingTest()
    print("颜色拾取测试完成") 