import gradio as gr
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
import cv2
import numpy as np
import re

# Lade das Modell
model_path = hf_hub_download(repo_id="foduucom/stockmarket-pattern-detection-yolov8", filename="model.pt")
model = YOLO(model_path)

def analyze_image(image, prompt):
    # Konvertiere PIL-Bild zu OpenCV-Format
    image_np = np.array(image)
    image_cv = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)

    # Führe Objekterkennung durch
    results = model.predict(source=image_np, save=False)
    
    # Extrahiere Kerzen
    detections = []
    for result in results:
        for box in result.boxes:
            label = result.names[int(box.cls)]
            confidence = float(box.conf)
            # Extrahiere Bounding-Box-Koordinaten
            xmin, ymin, xmax, ymax = box.xyxy[0].tolist()
            
            # Schneide die Kerze aus dem Bild für Farbanalyse
            candle_roi = image_cv[int(ymin):int(ymax), int(xmin):int(xmax)]
            # Berechne dominante Farbe (RGB)
            mean_color = np.mean(candle_roi, axis=(0, 1)).astype(int)
            color_rgb = f"RGB({mean_color[2]},{mean_color[1]},{mean_color[0]})"  # RGB-Format
            
            detections.append({
                "pattern": label,
                "confidence": confidence,
                "color": color_rgb,
                "x_center": (xmin + xmax) / 2,  # Für Sortierung nach Position
                "prompt_used": prompt
            })

    # Sortiere Kerzen nach x-Position (von rechts nach links = neueste zuerst)
    detections = sorted(detections, key=lambda x: x["x_center"], reverse=True)
    
    # Begrenze auf die letzten 10 Kerzen, wenn im Prompt gefordert
    if "last 10 candles" in prompt.lower():
        detections = detections[:10]
    
    return detections

# Erstelle Gradio-Schnittstelle
iface = gr.Interface(
    fn=analyze_image,
    inputs=[
        gr.Image(type="pil", label="Upload TradingView Screenshot"),
        gr.Textbox(label="Prompt", placeholder="Enter your prompt, e.g., 'List last 10 candles with their colors'")
    ],
    outputs="json",
    title="Candlestick Pattern Detection",
    description="Upload a TradingView screenshot and provide a prompt to detect candlestick patterns and their colors."
)

# Starte die App
iface.launch()