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

# Logging einrichten
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

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

# OCR für Preise
reader = easyocr.Reader(['en'], gpu=False)

def analyze_image(image, prompt):
    logger.info("Starting image analysis with prompt: %s", prompt)
    
    # Konvertiere PIL-Bild zu OpenCV-Format
    image_np = np.array(image)
    image_cv = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
    
    # Bildvorverarbeitung: Kontrast erhöhen
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    gray = cv2.cvtColor(image_cv, cv2.COLOR_BGR2GRAY)
    enhanced = clahe.apply(gray)
    image_cv = cv2.cvtColor(enhanced, cv2.COLOR_GRAY2BGR)
    logger.info("Image preprocessed: shape=%s", image_np.shape)

    # Führe Objekterkennung durch
    results = model.predict(source=image_np, conf=0.3, iou=0.5, save=False)
    logger.info("YOLO predictions: %d boxes detected", len(results[0].boxes))

    # Extrahiere Kerzen
    detections = []
    for result in results:
        for box in result.boxes:
            label = result.names[int(box.cls)]
            confidence = float(box.conf)
            xmin, ymin, xmax, ymax = box.xyxy[0].tolist()
            logger.info("Detected: %s, confidence=%.2f, box=(%.0f, %.0f, %.0f, %.0f)", 
                       label, confidence, xmin, ymin, xmax, ymax)

            # Extrahiere Farbe (Fokus auf Kerzenkörper)
            candle_roi = image_cv[int(ymin):int(ymax), int(xmin):int(xmax)]
            if candle_roi.size == 0:
                logger.warning("Empty ROI for box: (%.0f, %.0f, %.0f, %.0f)", xmin, ymin, xmax, ymax)
                continue
            mean_color = np.mean(candle_roi, axis=(0, 1)).astype(int)
            color_rgb = f"RGB({mean_color[2]},{mean_color[1]},{mean_color[0]})"

            # OCR für Preise (erweitere ROI für Achsen)
            price_roi = image_cv[max(0, int(ymin)-50):min(image_np.shape[0], int(ymax)+50), 
                                max(0, int(xmin)-50):min(image_np.shape[1], int(xmax)+50)]
            price_text = reader.readtext(price_roi, detail=0, allowlist='0123456789.')
            prices = ' '.join(price_text) if price_text else "No price detected"
            logger.info("OCR prices: %s", prices)

            detections.append({
                "pattern": label,
                "confidence": confidence,
                "color": color_rgb,
                "prices": prices,
                "x_center": (xmin + xmax) / 2
            })

    # Sortiere nach x-Position (rechts nach links = neueste Kerzen)
    detections = sorted(detections, key=lambda x: x["x_center"], reverse=True)
    logger.info("Sorted detections: %d", len(detections))

    # Begrenze auf die letzten 8 Kerzen
    if "last 8 candles" in prompt.lower() or "letzte 8 kerzen" in prompt.lower():
        detections = detections[:8]

    # Debugging: Wenn leer, gib Hinweis
    if not detections:
        logger.warning("No detections found. Check image quality or model configuration.")
        return {"prompt": prompt, "description": "No candlesticks detected. Ensure clear image and visible candles."}
    
    return {
        "prompt": prompt,
        "detections": 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 8 candles with their colors'")
    ],
    outputs="json",
    title="Stock Chart Analysis with YOLOv8",
    description="Upload a TradingView screenshot to detect the last 8 candlesticks, their colors, and prices."
)

iface.launch()