frontend/threads/simulation_thread.py

import sys
import os
import time
import math
import traceback
import threading
from PySide6.QtCore import QThread, Signal, Slot, QTimer, QMutex, QMutexLocker

# Ensure project source root is in path
PROJECT_SRC = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
if PROJECT_SRC not in sys.path:
    sys.path.insert(0, PROJECT_SRC)

import pybullet as pb
import numpy as np
import cv2

from main import PandaSimulation
from llm_panda.plan_verifier import PlanVerifier

def set_pybullet_window_visibility(visible: bool):
    """Finds the PyBullet GUI window on Windows and sets its visibility by moving it on/off-screen."""
    if sys.platform != "win32":
        return False
    import ctypes
    
    EnumWindowsProc = ctypes.WINFUNCTYPE(ctypes.c_bool, ctypes.c_void_p, ctypes.c_void_p)
    
    for _ in range(30):
        found_hwnd = []
        
        def foreach_window(hwnd, lParam):
            length = ctypes.windll.user32.GetWindowTextLengthW(hwnd)
            buff = ctypes.create_unicode_buffer(length + 1)
            ctypes.windll.user32.GetWindowTextW(hwnd, buff, length + 1)
            title = buff.value
            if "bullet physics examplebrowser" in title.lower():
                found_hwnd.append(hwnd)
                return False
            return True
            
        ctypes.windll.user32.EnumWindows(EnumWindowsProc(foreach_window), 0)
        if found_hwnd:
            hwnd = found_hwnd[0]
            # Ensure the window is shown/restored from minimized state first (SW_SHOW = 5)
            ctypes.windll.user32.ShowWindow(hwnd, 5)
            
            # SWP_NOSIZE = 0x0001, SWP_NOZORDER = 0x0004
            if visible:
                # Move window to visible region
                ctypes.windll.user32.SetWindowPos(hwnd, 0, 100, 100, 0, 0, 0x0001 | 0x0004)
            else:
                # Move window far off-screen to hide it without destroying OpenGL context
                ctypes.windll.user32.SetWindowPos(hwnd, 0, -10000, -10000, 0, 0, 0x0001 | 0x0004)
            print(f"[Window Visibility] Set PyBullet window (HWND: {hwnd}) visibility to: {visible}", flush=True)
            return True
        time.sleep(0.1)
    return False

class SimulationThread(QThread):
    """
    Manages the PyBullet simulator. Handles stepping the physics, running plans,
    camera movements, resetting, and emitting captured frames to the UI.
    """
    frame_ready = Signal(object)
    log_message = Signal(str, str)
    fps_updated = Signal(float)
    init_completed = Signal()
    step_completed = Signal(int)  # step number (1-based)
    finished = Signal(bool, str, float)  # executed_successfully, score_md, score_val

    def __init__(self):
        super().__init__()
        self.sim = None
        self.pybullet_lock = threading.RLock()
        self.camera_state = {
            "target": [0.5, 0.0, 0.3],
            "dist": 1.3,
            "yaw": 45.0,
            "pitch": -35.0
        }
        self.show_gui_window = True
        self.running = True
        self.last_view_matrix = None
        self.last_proj_matrix = None
        
        # Command Queue variables
        self.pending_command = None
        self.pending_data = None
        self.cmd_lock = QMutex()
        
        # Frame rate tracking
        self.last_frame_time = 0.0
        self.frame_count = 0
        self.fps_timer_time = 0.0

    def run(self):
        # 1. Initialize PyBullet Simulation in DIRECT mode
        self.log_message.emit("🔌 Khởi tạo môi trường mô phỏng PyBullet...", "info")
        with self.pybullet_lock:
            self.sim = PandaSimulation(gui=False, seed=0)
            
        self.init_completed.emit()
        self.log_message.emit("✅ Môi trường mô phỏng PyBullet đã sẵn sàng!", "success")
        
        # Stream the initial frame
        self.capture_and_emit_frame()
        
        # 2. Start Simulation Loop
        self.last_frame_time = time.time()
        self.fps_timer_time = time.time()
        
        while self.running:
            # Check for commands from UI thread
            cmd = None
            data = None
            
            self.cmd_lock.lock()
            if self.pending_command is not None:
                cmd = self.pending_command
                data = self.pending_data
                self.pending_command = None
                self.pending_data = None
            self.cmd_lock.unlock()
            
            if cmd is not None:
                self.process_command(cmd, data)
                
            # Perform a continuous idle step & render at ~30 FPS if not executing a plan
            # (PyBullet real-time physics is active, so we just capture and emit frames)
            curr_time = time.time()
            elapsed = curr_time - self.last_frame_time
            if elapsed >= 0.033:  # limit to ~30 FPS to save CPU
                self.capture_and_emit_frame()
                self.track_fps(curr_time)
                self.last_frame_time = curr_time
                
            time.sleep(0.005)

    def track_fps(self, current_time):
        self.frame_count += 1
        interval = current_time - self.fps_timer_time
        if interval >= 1.0:
            fps = self.frame_count / interval
            self.fps_updated.emit(fps)
            self.frame_count = 0
            self.fps_timer_time = current_time

    def capture_and_emit_frame(self):
        with self.pybullet_lock:
            if self.sim is None or self.sim.client is None or not pb.isConnected(self.sim.client):
                return
            
            try:
                target = self.camera_state["target"]
                dist = self.camera_state["dist"]
                yaw = self.camera_state["yaw"]
                pitch = self.camera_state["pitch"]
                
                width = 640
                height = 360
                renderer = pb.ER_TINY_RENDERER
                
                view_matrix = pb.computeViewMatrixFromYawPitchRoll(
                    cameraTargetPosition=target,
                    distance=dist,
                    yaw=yaw,
                    pitch=pitch,
                    roll=0.0,
                    upAxisIndex=2,
                    physicsClientId=self.sim.client
                )
                proj_matrix = pb.computeProjectionMatrixFOV(
                    fov=55, aspect=width/height, nearVal=0.1, farVal=5.0,
                    physicsClientId=self.sim.client
                )
                
                # Store view & projection matrices for picking
                self.last_view_matrix = view_matrix
                self.last_proj_matrix = proj_matrix
                
                _, _, rgba, _, _ = pb.getCameraImage(
                    width=width, height=height,
                    viewMatrix=view_matrix,
                    projectionMatrix=proj_matrix,
                    renderer=renderer,
                    physicsClientId=self.sim.client
                )
                
                img_np = np.array(rgba, dtype=np.uint8).reshape(height, width, 4)
                self.frame_ready.emit(img_np)
            except Exception as e:
                print(f"[SimulationThread Frame Capture Error] {e}")

    def queue_command(self, cmd: str, data=None):
        QMutexLocker(self.cmd_lock)
        self.pending_command = cmd
        self.pending_data = data

    def process_command(self, cmd: str, data):
        if cmd == "reset":
            self.log_message.emit("🔄 Đang reset môi trường mô phỏng...", "info")
            with self.pybullet_lock:
                self.sim.reset_objects(seed=0)
                self.sim.controller.reset_to_home()
            self.log_message.emit("🔄 Đã reset môi trường mô phỏng.", "success")
            self.capture_and_emit_frame()
            
        elif cmd == "gui_toggle":
            self.show_gui_window = bool(data)
            self.log_message.emit(f"👁️ Hiển thị cửa sổ PyBullet (Chỉ khả dụng ở chế độ GUI): {self.show_gui_window}", "info")
            
        elif cmd == "camera_rotate":
            dx, dy = data
            self.camera_state["yaw"] += dx * 0.4
            self.camera_state["pitch"] = max(-89, min(-5, self.camera_state["pitch"] - dy * 0.4))
            
        elif cmd == "camera_pan":
            dx, dy = data
            rad_yaw = math.radians(self.camera_state["yaw"])
            rx = -math.sin(rad_yaw)
            ry = math.cos(rad_yaw)
            ux = -math.cos(rad_yaw)
            uy = -math.sin(rad_yaw)
            self.camera_state["target"][0] += (rx * dx + ux * dy) * 0.0015
            self.camera_state["target"][1] += (ry * dx + uy * dy) * 0.0015
            
        elif cmd == "camera_zoom":
            delta = data
            self.camera_state["dist"] = max(0.4, min(3.0, self.camera_state["dist"] + delta * 0.0015))
            
        elif cmd == "viewport_click":
            x, y, w, h = data
            self.perform_picking(x, y, w, h)
            
        elif cmd == "execute_plan":
            plan = data
            self.run_execution(plan)

    def run_execution(self, plan):
        self.log_message.emit("⏳ [3/5] Đang thực thi hành động của robot...", "info")
        
        # Hook pb.stepSimulation to capture frames during movements
        original_step = pb.stepSimulation
        last_stream_time = [0.0]
        
        def wrapped_step():
            original_step()
            curr_time = time.time()
            if curr_time - last_stream_time[0] >= 0.033:
                last_stream_time[0] = curr_time
                self.capture_and_emit_frame()
        
        with self.pybullet_lock:
            pb.stepSimulation = wrapped_step
            
        executed_successfully = True
        try:
            # Send initial frame
            self.capture_and_emit_frame()
            
            for idx, step in enumerate(plan, 1):
                self.log_message.emit(f"🚀 [3/5] Đang chạy bước {idx}/{len(plan)}: {step['skill']}", "info")
                # Build executor format step
                exec_step = {
                    "function": step["skill"],
                    "args": step["args"]
                }
                with self.pybullet_lock:
                    self.sim._execute(exec_step, verbose=True)
                
                self.step_completed.emit(idx)
                
            self.log_message.emit("✅ [3] Thực thi thành công toàn bộ kế hoạch.", "success")
        except Exception as e:
            error_trace = traceback.format_exc()
            print(f"[SimulationThread Execution Error] {error_trace}")
            self.log_message.emit(f"❌ Lỗi trong lúc thực thi mô phỏng: {e}", "error")
            executed_successfully = False
        finally:
            with self.pybullet_lock:
                pb.stepSimulation = original_step
        
        # Step 5: Scoring
        self.log_message.emit("⏳ [5/5] Đang tính toán điểm sắp xếp...", "info")
        score_val, score_md = self.compute_and_format_score(plan, executed_successfully)
        
        self.finished.emit(executed_successfully, score_md, score_val)

    def compute_and_format_score(self, plan: list, executed_successfully: bool):
        """Calculates rearrangement score based on final object locations and collisions."""
        with self.pybullet_lock:
            score = 100
            details = []
            
            # Check target placement if we had a plan
            target_location = None
            target_object = None
            for step in plan:
                func = step.get("skill")
                args = step.get("args", {})
                if func in ("place_at", "place_down", "move_to"):
                    if "location_id" in args:
                        target_location = args["location_id"]
                if func == "pick_up":
                    target_object = args.get("object_id")
                    
            if target_object and target_location:
                obj_id = self.sim.scene.name_to_id(target_object)
                if obj_id is not None:
                    pos, _ = pb.getBasePositionAndOrientation(obj_id)
                    current_loc = self.sim.scene.classify_location(pos)
                    if current_loc == target_location:
                        details.append(f"✅ **Đã đặt {target_object} đúng vị trí {target_location}**: +50 điểm")
                    else:
                        details.append(f"❌ **{target_object} chưa đến đúng vị trí {target_location} (đang ở {current_loc})**: -50 điểm")
                        score -= 50
                else:
                    details.append(f"❌ **Không tìm thấy vật thể {target_object}**: -50 điểm")
                    score -= 50
            else:
                details.append("ℹ️ Không có yêu cầu di chuyển vật thể cụ thể nào được phát hiện.")
                
            # Knock detection
            knocked_count = 0
            for oid in self.sim.scene.object_ids:
                if oid == self.sim.scene.name_to_id(target_object):
                    continue
                pos, _ = pb.getBasePositionAndOrientation(oid)
                loc = self.sim.scene.classify_location(pos)
                if loc == "elsewhere":
                    name = self.sim.scene.get_name(oid)
                    details.append(f"⚠️ **Vật thể '{name}' bị lệch khỏi vị trí quy định**: -15 điểm")
                    score -= 15
                    knocked_count += 1
                    
            if knocked_count == 0:
                details.append("✅ **Không có vật thể nào khác bị va chạm hay xê dịch**: +30 điểm")
            else:
                score = max(0, score)
                
            # Execution success
            if executed_successfully:
                details.append("✅ **Thực thi toàn bộ kế hoạch thành công**: +20 điểm")
            else:
                details.append("❌ **Thực thi gặp lỗi gián đoạn**: -30 điểm")
                score -= 30
                
            score = max(0, min(100, score))
            
            md = f"### 🏆 Rearrangement Evaluation Report\n\n**Tổng điểm: {score}/100**\n\n**Chi tiết:**\n" + "\n".join([f"- {d}" for d in details])
            return score, md

    def perform_picking(self, x, y, w, h):
        """Calculates 3D ray from screen coordinates and tests for intersection in PyBullet."""
        with self.pybullet_lock:
            if self.sim is None or self.last_view_matrix is None or self.last_proj_matrix is None:
                return
            
            try:
                # NDC (Normalized Device Coordinates)
                ndc_x = (2.0 * x) / w - 1.0
                ndc_y = 1.0 - (2.0 * y) / h  # Invert Y for OpenGL NDC
                
                # Near and far vectors in NDC
                ndc_near = np.array([ndc_x, ndc_y, -1.0, 1.0])
                ndc_far = np.array([ndc_x, ndc_y, 1.0, 1.0])
                
                # Reshape matrix list (16 elements) to 4x4 matrix and transpose (column-major)
                view_m = np.array(self.last_view_matrix).reshape(4, 4).T
                proj_m = np.array(self.last_proj_matrix).reshape(4, 4).T
                
                vp_m = proj_m @ view_m
                inv_vp_m = np.linalg.inv(vp_m)
                
                # Transform NDC to World coordinates
                world_near = inv_vp_m @ ndc_near
                world_near /= world_near[3]
                
                world_far = inv_vp_m @ ndc_far
                world_far /= world_far[3]
                
                start_point = world_near[:3]
                end_point = world_far[:3]
                
                # Perform PyBullet raycast test
                ray_test_result = pb.rayTest(start_point, end_point, physicsClientId=self.sim.client)
                if ray_test_result:
                    hit_id, hit_link, hit_fraction, hit_position, hit_normal = ray_test_result[0]
                    if hit_id >= 0:
                        # Safety check: static bodies (plane, table, tray, robot) are not in scene manager registry
                        if hit_id in self.sim.scene._registry:
                            name = self.sim.scene.get_name(hit_id)
                            pos_str = ", ".join([f"{c:.2f}" for c in hit_position])
                            self.log_message.emit(f"🖱️ Đã click chọn vật thể '{name}' tại vị trí 3D [{pos_str}]", "success")
                        else:
                            # Map static bodies to names
                            body_name = "Môi trường"
                            if hasattr(self.sim, 'table_id') and hit_id == self.sim.table_id:
                                body_name = "Bàn làm việc"
                            elif hasattr(self.sim, 'tray_id') and hit_id == self.sim.tray_id:
                                body_name = "Khay chứa đồ"
                            elif hasattr(self.sim, 'robot_id') and hit_id == self.sim.robot_id:
                                body_name = "Robot Franka"
                            elif hit_id == 0:  # Plane is always loaded first
                                body_name = "Sàn nhà"
                                
                            pos_str = ", ".join([f"{c:.2f}" for c in hit_position])
                            self.log_message.emit(f"🖱️ Đã click chọn: {body_name} tại vị trí 3D [{pos_str}]", "info")
                    else:
                        self.log_message.emit("🖱️ Đã click vào không gian trống.", "info")
            except Exception as e:
                print(f"[Picking Error] {e}")

    def shutdown(self):
        self.running = False
        # Wait a moment for loop to exit before disconnecting
        self.wait(1000)
        with self.pybullet_lock:
            if self.sim is not None:
                try:
                    self.sim.disconnect()
                except Exception:
                    pass