import cv2
import numpy as np
import torch
from PIL import Image, ImageGrab
from transformers import CLIPProcessor, CLIPModel
import time

# --- ASETUKSET ---
TARGET_PROMPT = "a paperclip icon" # MitC$ etsitC$C$n?
DEBUG_OUTPUT = "/tmp/semantic_search_result.png"

def load_vision_model():
    print("[*] Ladataan semanttinen Vision-malli (CLIP)...")
    model_id = "openai/clip-vit-base-patch32"
    model = CLIPModel.from_pretrained(model_id)
    processor = CLIPProcessor.from_pretrained(model_id)
    print("[+] Malli ladattu!")
    return model, processor

def extract_icon_candidates(cv_img):
    """ Etsii ruudulta kaikki ikonikokoiset elementit. """
    gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 50, 150)
    
    # Pieni sulatus yhdistC$C$ ikonin osat (esim. hymynaaman silmC$t ja suun)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    dilated = cv2.dilate(edges, kernel, iterations=1)
    
    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
    
    candidates = []
    for c in contours:
        x, y, w, h = cv2.boundingRect(c)
        # Suodatetaan: Vain ikonikokoiset neliC6mC$iset alueet (15-80
pikseliC$)
        if 15 <= w <= 80 and 15 <= h <= 80:
            # Otetaan hieman marginaalia ympC$rille, jotta malli nC$kee ikonin
kunnolla
            pad = 5
            x1, y1 = max(0, x-pad), max(0, y-pad)
            x2, y2 = min(cv_img.shape[1], x+w+pad), min(cv_img.shape[0],
y+h+pad)
            candidates.append({"x": int(x + w/2), "y": int(y + h/2), "box":
(x1, y1, x2, y2)})
            
    return candidates

def find_best_match(model, processor, pil_img, candidates, target_prompt):
    """ SyC6ttC$C$ kandidaatit CLIPille ja etsii parhaan osuman. """
    print(f"[*] EtsitC$C$n kandidaattien ({len(candidates)} kpl) joukosta:
'{target_prompt}'")
    
    # Leikataan kaikki kandidaatit kuviksi
    crops = []
    for cand in candidates:
        crop = pil_img.crop(cand["box"])
        crops.append(crop)
        
    if not crops:
        return None
        
    # Annetaan mallille negatiivisia vaihtoehtoja kohinan suodattamiseksi
    labels = [target_prompt, "a text word", "random background noise", "a
profile picture"]
    
    # Prosessoidaan yhtenC$ erC$nC$ (Batch processing on todella nopeaa!)
    inputs = processor(text=labels, images=crops, return_tensors="pt",
padding=True)
    
    with torch.no_grad():
        outputs = model(**inputs)
        
    # outputs.logits_per_image: [kuvien_mC$C$rC$, labelien_mC$C$rC$]
    logits = outputs.logits_per_image.softmax(dim=1).numpy()
    
    best_score = 0
    best_candidate = None
    
    for i, probs in enumerate(logits):
        # probs[0] on meidC$n target_promptimme todennC$kC6isyys
        target_score = probs[0] * 100
        
        if target_score > best_score:
            best_score = target_score
            best_candidate = candidates[i]
            best_candidate["score"] = best_score
            
    return best_candidate

def main():
    model, processor = load_vision_model()
    
    start_time = time.time()
    print("
[*] Kaapataan ruutu...")
    pil_img = ImageGrab.grab()
    cv_img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
    
    candidates = extract_icon_candidates(cv_img)
    print(f"[*] LC6ydettiin {len(candidates)} potentiaalista ikonia.")
    
    winner = find_best_match(model, processor, pil_img, candidates,
TARGET_PROMPT)
    
    if winner and winner["score"] > 20.0: # Varmuusraja 20%
        print("
" + "="*50)
        print(f"[+] LCYTYI: '{TARGET_PROMPT}'")
        print(f"[+] Koordinaatit: X:{winner['x']}, Y:{winner['y']}")
        print(f"[+] Varmuus (Confidence): {winner['score']:.2f}%")
        print("="*50 + "
")
        
        # PiirretC$C$n kirkkaanvihreC$ laatikko voittajan ympC$rille!
        bx1, by1, bx2, by2 = winner["box"]
        cv2.rectangle(cv_img, (bx1, by1), (bx2, by2), (0, 255, 0), 4)
        cv2.putText(cv_img, f"{winner['score']:.1f}%", (bx1, max(20, by1-10)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        
        cv2.imwrite(DEBUG_OUTPUT, cv_img)
        print(f"[*] Tallennettiin tuloskuva: {DEBUG_OUTPUT}")
    else:
        print(f"[-] Ei luotettavaa osumaa kohteelle '{TARGET_PROMPT}'.")
        
    print(f"[*] Kokonaisaika: {time.time() - start_time:.2f}s")

if __name__ == "__main__":
    main()