#!/usr/bin/env python3
"""
Quick debug script to test model locally and inspect detections
Run this in your container or locally to see what the model returns
"""

import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""

from PIL import Image
import numpy as np
from rfdetr import RFDETRSegPreview
import supervision as sv

# Path to your checkpoint
CHECKPOINT_PATH = "/tmp/checkpoint_best_total.pth"

print("Loading model...")
model = RFDETRSegPreview(pretrain_weights=CHECKPOINT_PATH)
print("Model loaded!")

# Create a test image (or load your own)
test_img = Image.new("RGB", (640, 480), color=(73, 109, 137))
print(f"Test image size: {test_img.size}")

# Run prediction
print("\nRunning prediction with threshold=0.25...")
detections = model.predict(test_img, threshold=0.25)

print("\n" + "="*60)
print("DETECTION RESULTS:")
print("="*60)
print(f"Type: {type(detections)}")
print(f"Length: {len(detections)}")

# Check attributes
attrs = dir(detections)
print(f"\nAvailable attributes: {[a for a in attrs if not a.startswith('_')]}")

if hasattr(detections, 'confidence'):
    print(f"\nConfidence scores: {detections.confidence}")
    print(f"Confidence count: {len(detections.confidence)}")
    print(f"Confidence type: {type(detections.confidence)}")

if hasattr(detections, 'class_id'):
    print(f"\nClass IDs: {detections.class_id}")
    print(f"Class count: {len(detections.class_id)}")

if hasattr(detections, 'masks'):
    masks = detections.masks
    print(f"\nMasks present: {masks is not None}")
    if masks is not None:
        print(f"Masks type: {type(masks)}")
        if hasattr(masks, 'shape'):
            print(f"Masks shape: {masks.shape}")
        elif hasattr(masks, '__len__'):
            print(f"Masks length: {len(masks)}")
            if len(masks) > 0:
                print(f"First mask type: {type(masks[0])}")
                if hasattr(masks[0], 'shape'):
                    print(f"First mask shape: {masks[0].shape}")
else:
    print("\nNO MASKS ATTRIBUTE FOUND!")

# Try to annotate
if len(detections) > 0:
    print("\n" + "="*60)
    print("Testing annotation...")
    print("="*60)
    try:
        palette = sv.ColorPalette.from_hex(["#ffff00", "#ff9b00"])
        mask_annotator = sv.MaskAnnotator(color=palette)
        annotated = mask_annotator.annotate(test_img, detections)
        print("✓ Mask annotation successful!")
        print(f"Annotated image size: {annotated.size}")
    except Exception as e:
        print(f"✗ Mask annotation failed: {e}")
        import traceback
        traceback.print_exc()

print("\n" + "="*60)
print("Debug complete!")
print("="*60)