#!/usr/bin/env python3
"""
YOLO-TS Inference Wrapper Script

This script runs in the 'yolots' conda environment and performs traffic sign
detection using YOLO-TS models with custom architecture modules.

Usage:
    conda run -n yolots python scripts/yolo_ts_inference.py <image_path> [--conf 0.25] [--imgsz 640]
    
Output:
    JSON to stdout with detection results in standard format

Author: Auto-generated for YOLO-TS separate environment setup
"""

import argparse
import json
import os
import sys
from pathlib import Path


def main():
    parser = argparse.ArgumentParser(description="YOLO-TS inference in separate environment")
    parser.add_argument("image", type=str, help="Path to input image")
    parser.add_argument("--model", type=str, default=None, help="Path to YOLO-TS model (.pt)")
    parser.add_argument("--conf", type=float, default=None, help="Confidence threshold")
    parser.add_argument("--imgsz", type=int, default=None, help="Inference image size")
    parser.add_argument("--iou", type=float, default=None, help="IOU threshold for NMS")
    parser.add_argument("--device", type=str, default=None, help="Device (cpu, 0, 1, etc.)")
    parser.add_argument("--verbose", action="store_true", help="Verbose output")
    args = parser.parse_args()
    
    # Read from environment variables with fallbacks
    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"))
    iou = args.iou if args.iou is not None else float(os.getenv("YOLO_TS_IOU", "0.45"))
    device = args.device if args.device is not None else os.getenv("YOLO_TS_DEVICE", os.getenv("YOLO_DEVICE", "cpu"))

    # Validate input
    image_path = Path(args.image)
    if not image_path.exists():
        result = {
            "instances": [],
            "meta": {"count": 0, "error": f"Image not found: {image_path}"}
        }
        print(json.dumps(result), file=sys.stdout)
        sys.exit(1)

    # Determine model path
    if args.model:
        model_path = Path(args.model)
    else:
        # Try to find model in main project
        project_root = Path(__file__).parent.parent
        model_path = project_root / "models" / "traffic-sign-detection" / "YOLO-TS" / "weights" / "GTSDB_best.pt"
        
        if not model_path.exists():
            # Try other common names
            for name in ["best.pt", "TT100K_best.pt", "CCTSDB_best.pt"]:
                alt_path = model_path.parent / name
                if alt_path.exists():
                    model_path = alt_path
                    break

    if not model_path.exists():
        result = {
            "instances": [],
            "meta": {
                "count": 0,
                "error": f"Model not found: {model_path}",
                "hint": "Set --model or place model in models/traffic-sign-detection/YOLO-TS/weights/"
            }
        }
        print(json.dumps(result), file=sys.stdout)
        sys.exit(1)

    # Import YOLO (in yolots environment with custom modules)
    try:
        from ultralytics import YOLO
        from PIL import Image
    except ImportError as e:
        result = {
            "instances": [],
            "meta": {"count": 0, "error": f"Import failed: {e}. Make sure you're in yolots environment."}
        }
        print(json.dumps(result), file=sys.stdout)
        sys.exit(1)

    # Load model
    try:
        if args.verbose:
            print(f"Loading model: {model_path}", file=sys.stderr)
        model = YOLO(str(model_path))
    except Exception as e:
        result = {
            "instances": [],
            "meta": {"count": 0, "error": f"Model load failed: {e}"}
        }
        print(json.dumps(result), file=sys.stdout)
        sys.exit(1)

    # Get image dimensions
    try:
        with Image.open(image_path) as img:
            w, h = img.size
    except Exception as e:
        result = {
            "instances": [],
            "meta": {"count": 0, "error": f"Failed to read image: {e}"}
        }
        print(json.dumps(result), file=sys.stdout)
        sys.exit(1)

    # Run inference
    try:
        if args.verbose:
            print(f"Running inference: conf={conf}, imgsz={imgsz}, iou={iou}, device={device}", file=sys.stderr)
        
        results = model.predict(
            source=str(image_path),
            conf=conf,
            imgsz=imgsz,
            iou=iou,
            device=device,
            verbose=args.verbose,
        )
    except Exception as e:
        result = {
            "instances": [],
            "meta": {"count": 0, "error": f"Inference failed: {e}"}
        }
        print(json.dumps(result), file=sys.stdout)
        sys.exit(1)

    # Process results
    instances = []
    try:
        for r in results:
            boxes = getattr(r, "boxes", None)
            if boxes is None:
                continue

            # Extract 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)

            if xyxy is None:
                continue

            # Create instances
            for i, box in enumerate(xyxy):
                x1, y1, x2, y2 = [int(max(0, v)) for v in box[:4]]
                bw = max(1, x2 - x1)
                bh = max(1, y2 - y1)
                
                score = 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"

                instance = {
                    "bbox": [x1, y1, bw, bh],
                    "score": score,
                    "class_id": cls_id,
                    "class_name": cls_name,
                    "description": f"traffic sign ({cls_name})" if cls_name != "unknown" else "traffic sign"
                }
                instances.append(instance)

    except Exception as e:
        result = {
            "instances": [],
            "meta": {"count": 0, "error": f"Result processing failed: {e}"}
        }
        print(json.dumps(result), file=sys.stdout)
        sys.exit(1)

    # Format output
    result = {
        "instances": instances,
        "meta": {
            "count": len(instances),
            "image": str(image_path),
            "model": str(model_path),
            "conf_threshold": conf,
            "iou_threshold": iou,
            "imgsz": imgsz,
            "device": device,
            "source": "yolo-ts-separate-env"
        }
    }

    # Output JSON to stdout
    print(json.dumps(result, indent=2), file=sys.stdout)
    
    if args.verbose:
        print(f"Detection complete: {len(instances)} traffic signs found", file=sys.stderr)


if __name__ == "__main__":
    main()
