"""Chart Pattern Detection API — YOLOv8"""
import torch
import torch.serialization
import gradio as gr
import json

# Allowlist ALL ultralytics model classes for safe loading
try:
    import ultralytics.nn.tasks as tasks
    safe_classes = [getattr(tasks, c) for c in dir(tasks) if isinstance(getattr(tasks, c, None), type)]
    torch.serialization.add_safe_globals(safe_classes)
except: pass

# Also force weights_only=False globally
_orig = torch.load
def _force(*a, **kw):
    kw["weights_only"] = False
    return _orig(*a, **kw)
torch.load = _force

from huggingface_hub import hf_hub_download
from ultralytics import YOLO

model_path = hf_hub_download(repo_id="foduucom/stockmarket-pattern-detection-yolov8", filename="model.pt")
model = YOLO(model_path)
print(f"Classes: {model.names}")

# Quick self-test with a blank image
from PIL import Image
import numpy as np
test_img = Image.fromarray(np.zeros((640,640,3), dtype=np.uint8))
test_res = model.predict(source=test_img, conf=0.01, verbose=False)
print(f"Self-test: {len(test_res[0].boxes)} detections on blank (should be 0)")

def detect_patterns(image):
    try:
        if image is None:
            return json.dumps({"patterns": []})
        results = model.predict(source=image, conf=0.20, iou=0.45, imgsz=640, verbose=False)
        patterns = []
        for r in results:
            if r.boxes is None or len(r.boxes) == 0: continue
            for i in range(len(r.boxes)):
                box = r.boxes[i]
                patterns.append({
                    "label": r.names.get(int(box.cls[0]), "unknown"),
                    "confidence": round(float(box.conf[0]), 3),
                    "bbox": [round(float(x), 1) for x in box.xyxy[0].tolist()],
                })
        patterns.sort(key=lambda p: p["confidence"], reverse=True)
        return json.dumps({"patterns": patterns, "count": len(patterns)})
    except Exception as e:
        return json.dumps({"patterns": [], "error": str(e)})

demo = gr.Interface(fn=detect_patterns, inputs=gr.Image(type="pil"), outputs=gr.Textbox(),
                    title="Chart Pattern Detection — YOLOv8")
demo.launch(show_error=True)