import cv2
import numpy as np
import math
import argparse
import os
import json
# step4analyzeOrtho.py
# Ball radius in pixels on 2000x2000 ortho image (approx)
MIN_R = 30
MAX_R = 74

ORTHO_W = 2000
ORTHO_H = 2000


# ------------------------------------------------------------
# Helper: convert image pixel (px,py) → world coordinates (mm)
# ------------------------------------------------------------
def image_to_world(px, py):
    x_world = (px / ORTHO_W) * 2000 - 1000
    y_world = 1000 - (py / ORTHO_H) * 2000
    return x_world, y_world


# ------------------------------------------------------------
# Circle detection helper
# ------------------------------------------------------------
def detect_color_balls(hsv, mask, color_name):
    detections = []

    mask = cv2.medianBlur(mask, 5)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area < 200:
            continue

        (x, y), radius = cv2.minEnclosingCircle(cnt)
        radius = int(radius)

        if radius < MIN_R or radius > MAX_R:
            continue

        perimeter = cv2.arcLength(cnt, True)
        if perimeter == 0:
            continue

        circularity = 4 * np.pi * (area / (perimeter * perimeter))
        if circularity < 0.6:
            continue

        detections.append({
            "color": color_name,
            "center_px": (int(x), int(y)),
            "radius": radius
        })

    return detections


# ------------------------------------------------------------
# Load original JSON so we can restore IDs
# ------------------------------------------------------------
def load_original_ids(base_name):
    json_name = f"{base_name}.json"
    if not os.path.exists(json_name):
        print(f"WARNING: Could not find original JSON {json_name}. IDs will be sequential.")
        return None

    with open(json_name, "r") as f:
        return json.load(f)


# ------------------------------------------------------------
# Match measured balls to original IDs
# ------------------------------------------------------------
def assign_original_ids(measured, original):
    if original is None:
        # fallback: sequential IDs
        for i, b in enumerate(measured, start=1):
            b["id"] = i
        return measured

    # Build lists of original world coords
    orig_yellow = [(b["id"], b["x"], b["y"]) for b in original["yellowBalls"]]
    orig_blue   = [(b["id"], b["x"], b["y"]) for b in original["blueBalls"]]

    # Assign IDs by nearest neighbor
    for b in measured:
        xw, yw = b["world"]
        if b["color"] == "yellow":
            candidates = orig_yellow
        else:
            candidates = orig_blue

        # find closest original ball
        best_id = None
        best_dist = 1e12
        for oid, ox, oy in candidates:
            d = (xw - ox)**2 + (yw - oy)**2
            if d < best_dist:
                best_dist = d
                best_id = oid

        b["id"] = best_id

    return measured


# ------------------------------------------------------------
# Main
# ------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser(description="Analyze ortho image for ball positions.")
    parser.add_argument("input_png", help="Path to ortho PNG")
    args = parser.parse_args()

    input_path = args.input_png

    # Determine base name (before first underscore)
    base = os.path.splitext(os.path.basename(input_path))[0]
    if "_" in base:
        file_name_base = base.split("_")[0]
    else:
        file_name_base = base

    output_json = f"{file_name_base}_measured.json"

    # Load original JSON for ID restoration
    original_json = load_original_ids(file_name_base)

    # Load image
    img = cv2.imread(input_path)
    if img is None:
        raise FileNotFoundError(f"Could not load {input_path}")

    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    # Masks
    blue_lower = np.array([90, 80, 40])
    blue_upper = np.array([130, 255, 255])
    mask_blue = cv2.inRange(hsv, blue_lower, blue_upper)

    yellow_lower = np.array([20, 80, 80])
    yellow_upper = np.array([35, 255, 255])
    mask_yellow = cv2.inRange(hsv, yellow_lower, yellow_upper)

    # Detect balls
    blue_balls = detect_color_balls(hsv, mask_blue, "blue")
    yellow_balls = detect_color_balls(hsv, mask_yellow, "yellow")

    balls = blue_balls + yellow_balls

    # Compute world coords + distance
    for b in balls:
        px, py = b["center_px"]
        xw, yw = image_to_world(px, py)
        b["world"] = (xw, yw)
        b["distance"] = math.sqrt(xw*xw + yw*yw)

    # Sort by distance
    balls_sorted = sorted(balls, key=lambda b: b["distance"])

    # Restore original IDs
    balls_with_ids = assign_original_ids(balls_sorted, original_json)

    # Print table
    print("\nBalls sorted by distance from center:\n")
    print(f"{'id':<4}{'color':<8}{'x':>10}{'y':>10}{'dist':>12}")

    for b in balls_with_ids:
        xw, yw = b["world"]
        print(f"{b['id']:<4}{b['color']:<8}{xw:>10.1f}{yw:>10.1f}{b['distance']:>12.1f}")

    print(f"\nTotal balls detected: {len(balls_with_ids)}")

    # --------------------------------------------------------
    # Build simplified JSON output (same structure as input)
    # --------------------------------------------------------
    yellow_list = []
    blue_list = []

    for b in balls_with_ids:
        xw, yw = b["world"]
        entry = {
            "id": b["id"],
            "x": round(xw, 1),
            "y": round(yw, 1)
        }

        if b["color"] == "yellow":
            yellow_list.append(entry)
        else:
            blue_list.append(entry)

    # Sort each list by ID
    yellow_list = sorted(yellow_list, key=lambda b: b["id"])
    blue_list   = sorted(blue_list, key=lambda b: b["id"])

    output_data = {
        "yellowBalls": yellow_list,
        "blueBalls": blue_list
    }

    # Save JSON
    with open(output_json, "w") as f:
        json.dump(output_data, f, indent=4)

    print(f"\nSaved measurements to: {output_json}")
    # Suggest next step
    print(f"next step: python step5reportResults.py {file_name_base}.png")

    # Visualization
    vis = img.copy()
    for b in balls_with_ids:
        cx, cy = b["center_px"]
        r = b["radius"]
        cv2.circle(vis, (cx, cy), r, (0, 0, 0), 3)
        label = f"{b['id']} {b['color']} {b['distance']:.0f}"
        cv2.putText(vis, label, (cx+55, cy),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0,0,0), 3, cv2.LINE_AA)

    vis_small = cv2.resize(vis, (660,660), interpolation=cv2.INTER_AREA)
    cv2.imshow("Detected Balls with Distances (scaled)", vis_small)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()