import cv2
import numpy as np
import os
from ultralytics import YOLO
from typing import List, Dict, Any
from PIL import Image
from huggingface_hub import hf_hub_download

class PatternDetector:
    def __init__(self, model_path: str = "foduucom/stockmarket-pattern-detection-yolov8"):
        if "/" in model_path and not os.path.exists(model_path):
            print(f"⬇️ Downloading model from Hugging Face: {model_path}")
            try:
                local_model_path = hf_hub_download(
                    repo_id=model_path,
                    filename="model.pt",
                    cache_dir="./models"
                )
                print(f"✅ Model downloaded to: {local_model_path}")
            except Exception as e:
                print(f"❌ Failed to download model: {e}")
                print("⚠️ Falling back to default YOLOv8n model.")
                local_model_path = "yolov8n.pt"
        else:
            local_model_path = model_path

        self.model = YOLO(local_model_path)
        self.supported_patterns = [
            'Head and shoulders bottom', 'Head and shoulders top',
            'M_Head', 'W_Bottom', 'Triangle', 'Flag', 'Pennant',
            'Double Top', 'Double Bottom', 'Cup and Handle'
        ]

    def detect_patterns(self, image: np.ndarray, conf_threshold: float = 0.3) -> List[Dict[str, Any]]:
        if isinstance(image, Image.Image):
            image = np.array(image)
            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

        results = self.model(image, conf=conf_threshold)
        detections = []

        for result in results:
            boxes = result.boxes
            if boxes is not None:
                for box in boxes:
                    x1, y1, x2, y2 = box.xyxy[0].tolist()
                    conf = float(box.conf[0])
                    cls_id = int(box.cls[0])
                    pattern_name = self.model.names[cls_id]
                    detections.append({
                        'pattern': pattern_name,
                        'confidence': conf,
                        'bbox': (int(x1), int(y1), int(x2), int(y2)),
                        'center': ((x1 + x2) / 2, (y1 + y2) / 2)
                    })
        return detections

    def analyze_pattern_sentiment(self, detections: List[Dict]) -> Dict[str, Any]:
        bullish_patterns = ['W_Bottom', 'Double Bottom', 'Cup and Handle', 'Head and shoulders bottom']
        bearish_patterns = ['M_Head', 'Double Top', 'Head and shoulders top']

        bullish_score = 0
        bearish_score = 0

        for det in detections:
            pattern = det['pattern']
            conf = det['confidence']
            if pattern in bullish_patterns:
                bullish_score += conf
            elif pattern in bearish_patterns:
                bearish_score += conf

        total = bullish_score + bearish_score
        sentiment = "NEUTRAL"
        strength = 0.0

        if total > 0:
            if bullish_score > bearish_score:
                sentiment = "BULLISH"
                strength = bullish_score / total
            else:
                sentiment = "BEARISH"
                strength = bearish_score / total

        return {
            'sentiment': sentiment,
            'strength': strength,
            'bullish_score': bullish_score,
            'bearish_score': bearish_score,
            'patterns_found': [d['pattern'] for d in detections],
            'highest_confidence_pattern': max(detections, key=lambda x: x['confidence']) if detections else None
        }

    def draw_detections(self, image: np.ndarray, detections: List[Dict]) -> np.ndarray:
        img_copy = image.copy()
        for det in detections:
            x1, y1, x2, y2 = det['bbox']
            label = f"{det['pattern']}: {det['confidence']:.2f}"
            bullish_patterns = ['W_Bottom', 'Double Bottom', 'Cup and Handle']
            bearish_patterns = ['M_Head', 'Double Top']
            if det['pattern'] in bullish_patterns:
                color = (0, 255, 0)
            elif det['pattern'] in bearish_patterns:
                color = (0, 0, 255)
            else:
                color = (255, 255, 0)
            cv2.rectangle(img_copy, (x1, y1), (x2, y2), color, 2)
            cv2.putText(img_copy, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
        return img_copy