import subprocess, sys
# Self-install dependencies — runs before any import
subprocess.check_call([
    sys.executable, '-m', 'pip', 'install',
    '--quiet', '--no-cache-dir',
    'ultralytics>=8.3.0',
    'pillow>=10.0.0',
    'numpy>=1.24.0'
])
import sys
import threading
import gradio as gr
import spaces
from PIL import Image
import numpy as np

def _suppress_fd_warnings(unraisable):
    if isinstance(unraisable.exc_value, ValueError) and 'Invalid file descriptor' in str(unraisable.exc_value):
        return
    sys.__unraisablehook__(unraisable)
sys.unraisablehook = _suppress_fd_warnings

from ultralytics import YOLO

_models = {}

def _load_triage():
    if 'triage' not in _models:
        print('Loading triage model...')
        _models['triage'] = YOLO('natural_seabed.pt', task='classify')
    return _models['triage']

def _load_detector(key):
    if key not in _models:
        print(f'Loading {key} model...')
        _models[key] = YOLO(f'{key}.pt')
    return _models[key]

DETECTOR_KEYS = ['shipwreck', 'subpipes', 'mines', 'cylinder', 'ghostvision']

@spaces.GPU
def predict_sonar(image, detector_choice, conf_cutoff):
    if image is None:
        return None, {'error': 'No image provided'}

    pil_img = image if isinstance(image, Image.Image) else Image.fromarray(image)
    w, h = pil_img.size

    # Stage 1: Triage
    triage = _load_triage()
    t_res  = triage.predict(pil_img, imgsz=224, verbose=False)[0]
    probs  = t_res.probs
    prob_clean  = float(probs.data[0])
    prob_debris = float(probs.data[1])
    is_anomaly  = prob_debris >= 0.35

    # Stage 2: Detection
    detections = []
    models_to_run = (
        DETECTOR_KEYS if detector_choice in ('all', 'auto')
        else [detector_choice.lower()]
    )

    annotated_img = np.array(pil_img)
    for m_key in models_to_run:
        if m_key not in DETECTOR_KEYS:
            continue
        model    = _load_detector(m_key)
        imgsz    = 1024 if m_key == 'shipwreck' else (1536 if m_key == 'cylinder' else 640)
        det_conf = max(conf_cutoff, 0.28 if m_key == 'shipwreck' else (0.32 if m_key == 'subpipes' else conf_cutoff))

        res = model.predict(pil_img, conf=det_conf, imgsz=imgsz, verbose=False)[0]
        if m_key == models_to_run[0]:
            annotated_img = res.plot()

        for box in res.boxes:
            cid   = int(box.cls[0])
            score = float(box.conf[0])
            x1, y1, x2, y2 = [float(c) for c in box.xyxy[0].tolist()]
            bw = max(0.0, x2 - x1)
            bh = max(0.0, y2 - y1)
            detections.append({
                'model': m_key, 'class_id': cid,
                'label': model.names[cid],
                'confidence': round(score, 4),
                'bbox': {
                    'x1': round(x1,1), 'y1': round(y1,1),
                    'x2': round(x2,1), 'y2': round(y2,1),
                    'w':  round(bw,1),  'h':  round(bh,1),
                    'pct_x': round((x1/w)*100,2), 'pct_y': round((y1/h)*100,2),
                    'pct_w': round((bw/w)*100,2), 'pct_h': round((bh/h)*100,2),
                }
            })

    return annotated_img, {
        'triage': {
            'status':      'ABNORMALITY FLAGGED' if is_anomaly else 'CLEAN SEABED',
            'is_anomaly':  is_anomaly,
            'prob_debris': round(prob_debris, 4),
            'prob_clean':  round(prob_clean,  4),
        },
        'detections_count': len(detections),
        'detections':       detections,
    }

demo = gr.Interface(
    fn=predict_sonar,
    inputs=[
        gr.Image(type='pil', label='Side-Scan Sonar Swath'),
        gr.Dropdown(choices=['all'] + DETECTOR_KEYS, value='shipwreck', label='Detector'),
        gr.Slider(0.05, 0.95, value=0.28, step=0.01, label='Confidence Cutoff'),
    ],
    outputs=[
        gr.Image(type='numpy', label='Annotated Swath'),
        gr.JSON(label='Pipeline Intelligence'),
    ],
    title='SonarOps Acoustic Intelligence Engine',
    description='Dual-Stage AI: Natural Seabed Triage + High-Resolution Target Detectors.',
    api_name='predict',
)

def _preload_all():
    print('Pre-loading all 6 models on CPU...')
    _load_triage()
    for key in DETECTOR_KEYS:
        _load_detector(key)
    print('All 6 models ready!')

threading.Thread(target=_preload_all, daemon=True).start()

demo.queue()
demo.launch(ssr_mode=False)