import mss
import cv2
import numpy as np
import time

# ========= Dein Spielausschnitt =========
monitor_area = {"top": 120, "left": 330, "width": 1900, "height": 1263}

# ========= HSV-Grenzen =========
yellow_lower = np.array([15, 40, 200], dtype=np.uint8)
yellow_upper = np.array([25, 120, 255], dtype=np.uint8)
white_lower  = np.array([0, 0, 220], dtype=np.uint8)
white_upper  = np.array([180, 50, 255], dtype=np.uint8)

black_lower  = np.array([0, 0, 0], dtype=np.uint8)
black_upper  = np.array([180, 80, 60], dtype=np.uint8)

green1_lower = np.array([30,  80,  80], dtype=np.uint8)
green1_upper = np.array([45, 255, 255], dtype=np.uint8)
green2_lower = np.array([65, 100,  80], dtype=np.uint8)
green2_upper = np.array([90, 255, 255], dtype=np.uint8)

kernel = np.ones((3,3), np.uint8)

# Radien
EAT_RADIUS  = 95
COLL_RADIUS = 115

# Bomben-Filter
BOMB_MIN_AREA     = 400   # angepasst!
BOMB_CIRC_MIN     = 0.60
BOMB_ASPECT_TOL   = 0.35
BOMB_EXTENT_MIN   = 0.60
BOMB_SOLIDITY_MIN = 0.85

# Fenster-Skalierung (0.7 = 70 % Größe)
WINDOW_SCALE = 0.8

# Anzeige-Modi
MODE_OVERLAY, MODE_FLOWER_MASK, MODE_BOMB_MASK, MODE_TURTLE_MASK = 0,1,2,3
mode = MODE_OVERLAY


def centroid(mask):
    cnt = int(cv2.countNonZero(mask))
    if cnt == 0: return None, None, 0
    M = cv2.moments(mask)
    if M["m00"] == 0: return None, None, cnt
    return int(M["m10"]/M["m00"]), int(M["m01"]/M["m00"]), cnt


def bomb_centroids_filtered(mask):
    contours,_ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    out=[]
    for c in contours:
        area = float(cv2.contourArea(c))
        if area < BOMB_MIN_AREA:
            continue
        x,y,w,h = cv2.boundingRect(c)
        if w == 0 or h == 0:
            continue
        aspect = w/float(h)
        if not (1.0 - BOMB_ASPECT_TOL <= aspect <= 1.0 + BOMB_ASPECT_TOL):
            continue
        per = float(cv2.arcLength(c, True))
        if per <= 0:
            continue
        circularity = 4.0 * np.pi * area / (per * per)
        if circularity < BOMB_CIRC_MIN:
            continue
        hull = cv2.convexHull(c)
        hull_area = float(cv2.contourArea(hull))
        if hull_area <= 0:
            continue
        solidity = area / hull_area
        extent   = area / float(w*h)
        if solidity < BOMB_SOLIDITY_MIN or extent < BOMB_EXTENT_MIN:
            continue
        M = cv2.moments(c)
        if M["m00"] == 0:
            continue
        cx = int(M["m10"]/M["m00"])
        cy = int(M["m01"]/M["m00"])
        out.append((cx, cy, int(area)))
    return out


def detect_all(frame_bgr):
    hsv = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2HSV)
    mw = cv2.inRange(hsv, white_lower, white_upper)
    my = cv2.inRange(hsv, yellow_lower, yellow_upper)
    mw = cv2.morphologyEx(mw, cv2.MORPH_DILATE, kernel, iterations=1)
    my = cv2.morphologyEx(my, cv2.MORPH_DILATE, kernel, iterations=1)
    mf = cv2.bitwise_and(mw, my)
    mf = cv2.morphologyEx(mf, cv2.MORPH_CLOSE, kernel, iterations=1)
    fx, fy, _ = centroid(mf)

    mb = cv2.inRange(hsv, black_lower, black_upper)
    bombs = bomb_centroids_filtered(mb)

    g1 = cv2.inRange(hsv, green1_lower, green1_upper)
    g2 = cv2.inRange(hsv, green2_lower, green2_upper)
    mg = cv2.bitwise_or(g1, g2)
    mg = cv2.morphologyEx(mg, cv2.MORPH_OPEN,  kernel, iterations=1)
    mg = cv2.morphologyEx(mg, cv2.MORPH_DILATE, kernel, iterations=1)
    tx, ty, _ = centroid(mg)

    masks = {"flower": mf, "bomb": mb, "turtle": mg}
    return (fx,fy), bombs, (tx,ty), masks


def draw_overlay(frame, flower, bombs, turtle, fps):
    fx, fy = flower
    tx, ty = turtle
    if fx is not None:
        cv2.circle(frame, (fx,fy), 8, (0,255,255), 2)
        cv2.putText(frame, "Flower", (fx+10, fy-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,255), 2)
    nearest = None
    if bombs and tx is not None:
        nearest = min(bombs, key=lambda b: np.hypot(b[0]-tx, b[1]-ty))
    for (bx,by,_) in bombs:
        color = (60,60,60); thick = 2
        if nearest and (bx,by)==(nearest[0],nearest[1]):
            color = (0,0,255); thick = 3
        cv2.circle(frame, (bx,by), 10, color, thick)
    if tx is not None:
        cv2.circle(frame, (tx,ty), 8, (0,200,0), 2)
        cv2.putText(frame, "Turtle", (tx+10, ty-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,200,0), 2)
        cv2.circle(frame, (tx,ty), EAT_RADIUS,  (0,255,0), 1)
        cv2.circle(frame, (tx,ty), COLL_RADIUS, (0,0,255), 1)
    cv2.putText(frame, f"FPS: {fps:.1f}", (20,40),  cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,255), 2)
    return frame


def colorize(mask):
    return cv2.applyColorMap(
        cv2.normalize(mask, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8),
        cv2.COLORMAP_JET
    )


def main():
    global mode
    sct = mss.mss()
    prev = time.time()
    fps = 0.0
    while True:
        raw = np.array(sct.grab(monitor_area))
        frame = cv2.cvtColor(raw, cv2.COLOR_BGRA2BGR)
        flower, bombs, turtle, masks = detect_all(frame)
        now = time.time()
        dt = now - prev; prev = now
        if dt > 0: fps = 1.0/dt
        if mode == MODE_OVERLAY:
            out = draw_overlay(frame.copy(), flower, bombs, turtle, fps)
        elif mode == MODE_FLOWER_MASK:
            out = colorize(masks["flower"])
        elif mode == MODE_BOMB_MASK:
            out = colorize(masks["bomb"])
        elif mode == MODE_TURTLE_MASK:
            out = colorize(masks["turtle"])
        # --- hier skalieren ---
        if WINDOW_SCALE != 1.0:
            out = cv2.resize(out, (int(out.shape[1]*WINDOW_SCALE), int(out.shape[0]*WINDOW_SCALE)))
        cv2.imshow("Debug Viewer", out)
        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'): break
        elif key == ord('0'): mode = MODE_OVERLAY
        elif key == ord('1'): mode = MODE_FLOWER_MASK
        elif key == ord('2'): mode = MODE_BOMB_MASK
        elif key == ord('3'): mode = MODE_TURTLE_MASK
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()