from functools import partial

import cv2
import numpy as np

# ======================================

def partition(items, predicate):
    N = len(items)

    # // The first O(N) pass: create N single-vertex trees
    parents = [-1] * N
    ranks = [0] * N

    def _find_root(i):
        _root = i
        while parents[_root] >= 0:
            _root = parents[_root]
        return _root

    def _compress_path(i, target):
        _k = i
        while True:
            parent = parents[_k]
            if parent < 0:
                break
            parents[_k] = target
            _k = parent

    # The main O(N^2) pass: merge connected components
    for i in range(N):
        # Find root
        root = _find_root(i)

        for j in range(N):
            if i == j or not predicate(items[i], items[j]):
                continue

            root2 = _find_root(j)

            if root != root2:
                # Unite both trees
                rank, rank2 = ranks[root], ranks[root2]
                if rank > rank2:
                    parents[root2] = root
                else:
                    parents[root] = root2
                    ranks[root2] += 1 if rank == rank2 else 0
                    root = root2
                assert parents[root] < 0

                _compress_path(j, root)
                _compress_path(i, root)

    # Final O(N) pass: enumerate classes
    labels = [0] * N
    nclasses = 0

    for i in range(N):
        root = _find_root(i)
        # re-use the rank as the class label
        if ranks[root] >= 0:
            ranks[root] = ~nclasses
            nclasses += 1
        labels[i] = ~ranks[root]

    return nclasses, labels

# ======================================

def create_mask(image):
    lower_brown = np.array([29, 39, 52], dtype = "uint8")
    upper_brown = np.array([33, 44, 60], dtype = "uint8")
    return cv2.inRange(image, lower_brown, upper_brown)

# ======================================

def get_coordinates(mask):
    coord = cv2.findNonZero(mask)
    return coord.reshape(coord.shape[0], 2)

# ======================================

# Simple metric -- distance along x and y axis has to be below threshold
def points_are_close_by(pt1, pt2, threshold):
    return (abs(pt1[0] - pt2[0]) < threshold) and (abs(pt1[1] - pt2[1]) < threshold)

# ======================================

def cluster_coordinates(coord, threshold):
    predicate = partial(points_are_close_by, threshold=threshold)
    cluster_count, labels = partition(coord, predicate)

    cluster_map = {}
    for i in range(cluster_count):
        cluster_map[i] = []

    for label, pos in zip(labels, coord):
        cluster_map[label].append(tuple(pos))

    return cluster_map

# ======================================

CLOSE_BY_THRESHOLD = 20
MIN_CLUSTER_SIZE = 5

image = cv2.imread("bloodweb.png")
mask = create_mask(image)
cv2.imwrite("bloodweb_mask.png", mask)

coord = get_coordinates(mask)
clusters = cluster_coordinates(coord, CLOSE_BY_THRESHOLD)

print("Found %d clusters." % len(clusters))

result = image.copy()
for label, points in clusters.items():
    if len(points) >= MIN_CLUSTER_SIZE:
        center, radius = cv2.minEnclosingCircle(np.array(points))
        center, radius = np.int32(center), np.int32(radius) # Brute, you may wanna do this better
        print("Cluster #%d: center=(%0.3f, %0.3f) radius=%0.3f size=%d"
            % (label, center[0], center[1], radius, len(points)))
        cv2.circle(result, center, radius, (0,255,0), 2)
        cv2.putText(result, str(label), center+radius, cv2.FONT_HERSHEY_DUPLEX, 0.6, (0,255,0), 2)
    else:
        print("Cluster #%d does not contain enough points (%d)." % (label, len(points)))

cv2.imwrite("bloodweb_result.png", result)