import matplotlib.pyplot as plt
from skimage.measure import label
from sklearn.cluster import DBSCAN
from collections import Counter, defaultdict
import numpy as np

img = plt.imread('figures_and_colors/balls_and_rects.png')

def rect_or_circle(figure):
    if np.all(figure == figure[0][0]):
        return 'rect'
    return 'circle'

gray_img = img.mean(axis=2)
binary_img = gray_img.copy()
binary_img[binary_img > 0] = 1

labels = label(binary_img)

# Строим график распределения кружков и квадратиков по цветам
fig = plt.figure()
ax_before = fig.add_subplot(projection='3d')
ax_before.set_xlabel('Red')
ax_before.set_xlabel('Green')
ax_before.set_xlabel('Blue')

counter = Counter()
figuremap = {'s':{'xs':[], 'ys':[], 'zs':[],'color':[]}, 'o':{'xs':[], 'ys':[], 'zs':[],'color':[]}}
for lbl in np.unique(labels)[1:]:
    x, y = np.where(labels==lbl)
    figure = img[x.min():x.max(), y.min():y.max(), :]
    counter[rect_or_circle(figure)] += 1
    marker = 's' if rect_or_circle(figure) == 'rect' else 'o'
    color = figure[int(x.mean())-x.min()][int(y.mean())-y.min()]
    figuremap[marker]['xs'].append(color[0])
    figuremap[marker]['ys'].append(color[1])
    figuremap[marker]['zs'].append(color[2])
    figuremap[marker]['color'].append(color)

ax_before.scatter(**figuremap['s'], marker='s')
ax_before.scatter(**figuremap['o'], marker='o')

# Кластеризуем по цветам
colors = figuremap['s']['color'] + figuremap['o']['color']
clustering = DBSCAN(eps=0.115, min_samples=1).fit(colors)

clusters_colors = defaultdict(lambda: [])
for color, lbl in zip(colors, clustering.labels_):
    clusters_colors[lbl].append(color)

mean_colors = defaultdict(lambda: [])
for lbl in clusters_colors.keys():
    mean_colors[lbl] = np.mean(clusters_colors[lbl], axis=0)

print(figuremap['s']['color'])

for i,_ in enumerate(figuremap['s']['color']):
    lbl = clustering.labels_[np.where(np.all(colors==figuremap['s']['color'][i], axis=1))[0][0]]
    figuremap['s']['color'][i] = mean_colors[lbl]

for i,_ in enumerate(figuremap['o']['color']):
    lbl = clustering.labels_[np.where(np.all(colors==figuremap['o']['color'][i], axis=1))[0][0]]
    figuremap['o']['color'][i] = mean_colors[lbl]

# Выводим результат кластеризации
ax_after = fig.add_subplot(projection='3d')
ax_after.set_xlabel('Red')
ax_after.set_xlabel('Green')
ax_after.set_xlabel('Blue')

ax_after.scatter(**figuremap['s'], marker='s')
ax_after.scatter(**figuremap['o'], marker='o')

plt.show()
print(counter)
