import numpy as np
import cv2
import random

def findClosestCentroids3(img, centroids):
    closestCentroi = np.zeros((img.shape[0],img.shape[1]), np.uint8)
    klusters = centroids.shape[0]
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            dist = np.zeros(klusters, np.float32)
            for l in range(klusters):
                # calculate the distance between the pixel and the centroid
                dist[l] = np.sqrt(
                    ((img[i][j][0] - centroids[l][0]) ** 2)+
                    ((img[i][j][1] - centroids[l][1]) ** 2)+
                    ((img[i][j][2] - centroids[l][2]) ** 2))
               

            #assign the pixel to the cluster with the nearest centroid human readable
            #closestCentroi[i, j] = np.argmin(dist) * 255 / (klusters - 1)
            closestCentroi[i, j] = np.argmin(dist)
            
    return closestCentroi

def calculateNewCenter3(img, centroids, closestCentroi):
    newCentroids = np.zeros(centroids.shape, np.float32)
    klusters = centroids.shape[0]
    centroidsCount = np.zeros(klusters, np.float32)
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            newCentroids[closestCentroi[i][j]] += img[i][j]
            centroidsCount[closestCentroi[i][j]] += 1

    for i in range(klusters):
        if centroidsCount[i] == 0:
            centroidsCount[i] = 1
        newCentroids[i] /= centroidsCount[i]

    return newCentroids

def kmeans3(img, centroids, iterations):
    for i in range(iterations):
        print("Iteration: ", i)
        closestCentroi = findClosestCentroids3(img, centroids)
        # cv2.imshow('closestCentroi', closestCentroi)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()
        centroids = calculateNewCenter3(img, centroids, closestCentroi)

    commpressedImg = np.zeros(img.shape, np.uint8)
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            commpressedImg[i][j] = centroids[closestCentroi[i][j]]

    return commpressedImg

def findClosestCentroids5(img, centroids, locations):
    closestCentroi = np.zeros((img.shape[0],img.shape[1]), np.uint8)
    klusters = centroids.shape[0]
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            dist = np.zeros(klusters, np.float32)
            for l in range(klusters):
                # calculate the distance between the pixel and the centroid
                dist[l] = np.sqrt(
                        ((img[i][j][0] - centroids[l][0]) ** 2)+
                        ((img[i][j][1] - centroids[l][1]) ** 2)+
                        ((img[i][j][2] - centroids[l][2]) ** 2)+
                        ((i - locations[l][0]) ** 2)+
                        ((j - locations[l][1]) ** 2))
                

            # assign the pixel to the cluster with the nearest centroid human readable
            #closestCentroi[i, j] = np.argmin(dist) * 255 / (klusters - 1)
            closestCentroi[i, j] = np.argmin(dist)

    return closestCentroi

def calculateNewCenter5(img, centroids, locations, closestCentroi):
    newCentroids = np.zeros(centroids.shape, np.float32)
    newLocations = np.zeros(locations.shape, np.float32)
    klusters = centroids.shape[0]
    centroidsCount = np.zeros(klusters, np.float32)
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            newCentroids[closestCentroi[i][j]] += img[i][j]
            newLocations[closestCentroi[i][j]] += [i,j]
            centroidsCount[closestCentroi[i][j]] += 1

    for i in range(klusters):
        if centroidsCount[i] == 0:
            centroidsCount[i] = 1
        newCentroids[i] /= centroidsCount[i]
        newLocations[i] /= centroidsCount[i]

    return newCentroids, newLocations

def kmeans5(img, centroids, locations, iterations):
    for i in range(iterations):
        print("Iteration: ", i)
        closestCentroi = findClosestCentroids5(img, centroids, locations)
        centroids, locations = calculateNewCenter5(img, centroids, locations, closestCentroi)

    commpressedImg = np.zeros(img.shape, np.uint8)
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            commpressedImg[i][j] = centroids[closestCentroi[i][j]]

    return commpressedImg

def kmeans(img, centroids, locations, iterations, dim):
    
    if dim == 5:
        commpressedImg = kmeans5(img, centroids, locations, iterations)
    else:
        commpressedImg = kmeans3(img, centroids, iterations)    
    
    return commpressedImg

def click_event(event, x, y, flags, params):
    global points
    global locations
    if event == cv2.EVENT_LBUTTONDOWN:
        print(f'({y},{x})')
        points = np.append(points, [img[y,x]] , axis=0)
        locations = np.append(locations, [[y,x]] , axis=0)

def selectPointsM(img):
    cv2.namedWindow('Select color')
    cv2.setMouseCallback('Select color', click_event)

    while True:
        cv2.imshow('Select color', img)
        if cv2.waitKey(1) & 0xFF == 27:
            break

    cv2.destroyAllWindows()

T = 50
def selectPointsR(img, stDimenzij):
    global points
    global locations
    global T
    numOfPoints = input('Vnesi stevilo tock:')

    i = 0
    while i < int(numOfPoints):
        x = random.randint(0, img.shape[0]-1)
        y = random.randint(0, img.shape[1]-1)
        point = img[x,y]
        location = [x,y]

        add = True

        if stDimenzij == 3:
           for j in range(points.shape[0]):
                dist = np.sqrt(
                    ((points[j][0] - point[0]) ** 2)+
                    ((points[j][1] - point[1]) ** 2)+
                    ((points[j][2] - point[2]) ** 2))
                if dist < T:
                    add = False
                    break
        else:
            for j in range(points.shape[0]):
                dist = np.sqrt(
                    ((points[j][0] - point[0]) ** 2)+
                    ((points[j][1] - point[1]) ** 2)+
                    ((points[j][2] - point[2]) ** 2)+
                    ((locations[j][0] - location[0]) ** 2)+
                    ((locations[j][1] - location[1]) ** 2))
                if dist < T:
                    add = False
                    break

        if add:
            points = np.append(points, [point] , axis=0)
            locations = np.append(locations, [location] , axis=0)
            i += 1
    

# read the input image
img = cv2.imread('veg.jpg' , 1)
points = np.empty((0, 3), dtype=np.float32)
locations = np.empty((0, 2), dtype=np.float32)

stDimenzij = input('Vnesi stevilo dimenzij 3 ali 5:')
stIteracij = input('Vnesi stevilo iteracij:')
selectType = input('Vnesi 1 za rocno izbiro, 2 za random izbiro:')

if selectType == '1':
    selectPointsM(img)
else:
    selectPointsR(img, int(stDimenzij))

commpressedImg = kmeans(img, points, locations, int(stIteracij), int(stDimenzij))

cv2.imshow('Original image',img)
cv2.imshow('Compressed image',commpressedImg)
cv2.waitKey(0)
cv2.destroyAllWindows()