import sys
import os
import cv2
import torch
import time
import numpy as np
from ultralytics import YOLO
from shapely.geometry import Polygon, Point

def main():
    print("=== Detailed Jetson Detection Test ===")
    
    # 1. Load Zones
    zones_path = "/home/jetson/karung/zones.json"
    truck_pts = []
    if os.path.exists(zones_path):
        import json
        with open(zones_path, 'r') as f:
            data = json.load(f)
            truck_pts = data.get('truck', [])
            print(f"Zones.json truck points: {truck_pts}")
    
    # Target resolution
    target_w, target_h = 1280, 720
    
    # Scale factors assuming zones were drawn on 1920x1080
    orig_w, orig_h = 1920, 1080
    scale_x = target_w / orig_w
    scale_y = target_h / orig_h
    
    scaled_truck_pts = [[int(p[0] * scale_x), int(p[1] * scale_y)] for p in truck_pts] if truck_pts else [
        [389, 294], [398, 718], [885, 719], [885, 277]
    ]
    truck_polygon = Polygon(scaled_truck_pts)
    print(f"Scaled truck polygon: {scaled_truck_pts}")
    
    # 2. Open Stream
    source = "rtsp://192.168.192.96:8554/cam"
    print(f"Connecting to RTSP stream: {source}...")
    cap = cv2.VideoCapture(source)
    if not cap.isOpened():
        print("Error: Could not open RTSP source.")
        return
        
    # Wait for the stream to warm up and buffer
    print("Warming up stream reader for 3 seconds...")
    time.sleep(3.0)
    
    # 3. Load Model
    model_path = "/home/jetson/karung/model_karung_truk.engine"
    print(f"Loading TensorRT Model: {model_path}...")
    model = YOLO(model_path)
    
    print("Running 10 frames of inference...")
    detections_summary = {}
    
    frame_count = 0
    attempts = 0
    while frame_count < 10 and attempts < 100:
        ret, frame = cap.read()
        attempts += 1
        if not ret or frame is None:
            time.sleep(0.1)
            continue
            
        frame_count += 1
        frame_resized = cv2.resize(frame, (target_w, target_h))
        results = model(frame_resized, conf=0.01, imgsz=640, device="cuda", verbose=False)
        result = results[0]
        
        detected_in_frame = []
        for box in result.boxes:
            cls_id = int(box.cls[0])
            name = model.names[cls_id]
            conf = float(box.conf[0])
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            tcx = (x1 + x2) / 2.0
            tcy = (y1 + y2) / 2.0
            
            # Check if inside truck polygon
            pt = Point(tcx, tcy)
            in_poly = truck_polygon.contains(pt)
            
            detected_in_frame.append(f"{name} ({conf:.3f}) at ({tcx:.1f},{tcy:.1f}) in_poly={in_poly}")
            detections_summary[name] = detections_summary.get(name, 0) + 1
            
        print(f"Frame {frame_count} (attempt {attempts}): {', '.join(detected_in_frame) if detected_in_frame else 'None'}")
            
    cap.release()
    print("\nSummary of detected objects over processed frames:")
    print(detections_summary)

if __name__ == "__main__":
    main()