#!/usr/bin/env python3
"""
Quick traffic sign detection smoke test.

Tests the traffic sign detection using the YOLO-TS model at:
  models/traffic-sign-detection/YOLO-TS/weights/best.pt

Environment variables:
  YOLO_TS_MODEL_PATH (optional absolute/relative path to .pt)
  YOLO_TS_MODEL (default: models/traffic-sign-detection/YOLO-TS/weights/best.pt)
  YOLO_DEVICE (e.g., "cpu" or "0")
  YOLO_TS_CONF (default: 0.25)
  YOLO_TS_IMGSZ (default: 640)

Usage:
  python tools/test_traffic_signs.py path/to/image_with_traffic_signs.jpg
  python tools/test_traffic_signs.py path/to/image.jpg --conf 0.20 --imgsz 1024
  python tools/test_traffic_signs.py path/to/image.jpg --device cpu
"""
import argparse
import os
import sys
from pathlib import Path


def main():
    ap = argparse.ArgumentParser(description="Test traffic sign detection with YOLO-TS")
    ap.add_argument("image", type=str, help="path to input image")
    ap.add_argument("--conf", type=float, default=None, help="confidence threshold (overrides YOLO_TS_CONF)")
    ap.add_argument("--imgsz", type=int, default=None, help="inference image size (overrides YOLO_TS_IMGSZ)")
    ap.add_argument("--device", type=str, default=None, help="device (overrides YOLO_DEVICE)")
    ap.add_argument("--show", action="store_true", help="display annotated image")
    ap.add_argument("--save", type=str, default=None, help="save annotated image to path")
    args = ap.parse_args()

    # Validate input
    img_path = Path(args.image)
    if not img_path.exists():
        print(f"error: image not found: {img_path}", file=sys.stderr)
        sys.exit(1)

    # Load environment
    try:
        from dotenv import load_dotenv
        load_dotenv()
    except Exception:
        pass  # .env optional

    # Import detector after env loaded
    try:
        from ultralytics import YOLO  # type: ignore
    except ImportError as e:
        print(f"error: ultralytics import failed: {e}")
        print("install: pip install ultralytics")
        sys.exit(1)

    # Find model
    repo_root = Path(__file__).resolve().parents[1]
    default_model_path = repo_root / "models" / "traffic-sign-detection" / "YOLO-TS" / "weights" / "best.pt"
    
    model_path = os.getenv("YOLO_TS_MODEL_PATH") or os.getenv("YOLO_TS_MODEL") or str(default_model_path)
    device = args.device or os.getenv("YOLO_DEVICE", "cpu")
    conf = args.conf if args.conf is not None else float(os.getenv("YOLO_TS_CONF", "0.25"))
    imgsz = args.imgsz if args.imgsz is not None else int(os.getenv("YOLO_TS_IMGSZ", "640"))

    if not Path(model_path).exists():
        print(f"error: model not found: {model_path}")
        print(f"Please download a YOLO-TS model from https://github.com/Heqiang-Huang/YOLO-TS")
        print(f"Expected location: {default_model_path}")
        sys.exit(1)

    # Load model
    print(f"loading YOLO-TS model: {model_path}")
    print(f"  device: {device}, conf: {conf}, imgsz: {imgsz}")
    try:
        model = YOLO(model_path)
        print("loaded traffic sign detection model")
    except Exception as e:
        print(f"error loading model: {e}", file=sys.stderr)
        sys.exit(1)

    # Run inference
    print(f"detecting traffic signs in: {img_path}")
    try:
        results = model.predict(
            source=str(img_path),
            conf=conf,
            imgsz=imgsz,
            device=device,
            verbose=True,
        )
    except Exception as e:
        print(f"error during inference: {e}", file=sys.stderr)
        sys.exit(1)

    # Process results
    total_detections = 0
    for r in results:
        boxes = getattr(r, "boxes", None)
        if boxes is None:
            continue
        
        num_boxes = len(boxes)
        total_detections += num_boxes
        
        if num_boxes > 0:
            print(f"\nFound {num_boxes} traffic sign(s):")
            
            # Get box data
            try:
                import torch
                xyxy = boxes.xyxy.cpu().numpy() if hasattr(boxes.xyxy, "cpu") else boxes.xyxy
                confs = boxes.conf.cpu().numpy() if hasattr(boxes.conf, "cpu") else boxes.conf
                cls_ids = boxes.cls.cpu().numpy() if hasattr(boxes.cls, "cpu") else boxes.cls
            except Exception:
                xyxy = boxes.xyxy
                confs = getattr(boxes, "conf", None)
                cls_ids = getattr(boxes, "cls", None)
            
            # Print detections
            for i, box in enumerate(xyxy):
                x1, y1, x2, y2 = box[:4]
                conf_val = float(confs[i]) if confs is not None else None
                cls_id = int(cls_ids[i]) if cls_ids is not None else None
                cls_name = model.names.get(cls_id, f"class_{cls_id}") if cls_id is not None and hasattr(model, 'names') else "unknown"
                
                print(f"  [{i+1}] {cls_name} (class {cls_id})")
                print(f"      confidence: {conf_val:.3f}")
                print(f"      bbox: ({int(x1)}, {int(y1)}) -> ({int(x2)}, {int(y2)})")
                print(f"      size: {int(x2-x1)}x{int(y2-y1)}px")
        
        # Visualization
        if args.show or args.save:
            import cv2
            import numpy as np
            
            # Get annotated image
            annotated = r.plot()
            
            if args.save:
                save_path = Path(args.save)
                save_path.parent.mkdir(parents=True, exist_ok=True)
                cv2.imwrite(str(save_path), annotated)
                print(f"\nSaved annotated image to: {save_path}")
            
            if args.show:
                cv2.imshow(f"Traffic Sign Detection - {img_path.name}", annotated)
                print("\nPress any key to close window...")
                cv2.waitKey(0)
                cv2.destroyAllWindows()

    # Summary
    print(f"\n{'='*50}")
    print(f"Detection complete: {total_detections} traffic sign(s) found")
    print(f"{'='*50}")
    
    if total_detections == 0:
        print("\nTips if no signs were detected:")
        print("  - Try lowering confidence: --conf 0.15")
        print("  - Try increasing image size: --imgsz 1024")
        print("  - Verify the model is trained on appropriate dataset")
        print("  - Check if image actually contains traffic signs")


if __name__ == "__main__":
    main()
