import cv2
import numpy as np

locations = np.empty((0,2), dtype=np.int32)

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

def selectPointsM(img):
    global locations

    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()
    return locations

def MeanShift(src, window, criteria):

    for _ in range(criteria[1]):
        ret, window = Iteration(src, window, criteria[2])
        if ret:
            return (True, window)

    return (False, window)

def Iteration(src, window, criteria):
    if window[0] < 0:
        window = (0, window[1], window[2], window[3])
    if window[1] < 0:
        window = (window[0], 0, window[2], window[3])

    roi = src[window[1]:window[1]+window[3], window[0]:window[0]+window[2]]

    #iteracija 1
    # x = 0
    # y = 0
    # num = 0
    # for i in range(roi.shape[0]):
    #     for j in range(roi.shape[1]):
    #         if roi[i][j] > 0:
    #             x += j
    #             y += i
    #             num += 1

    #iteracija 2
    # nonzero_indices = np.nonzero(roi)
    # nonzero_values = roi[nonzero_indices]

    # x = np.sum(nonzero_indices[1])
    # y = np.sum(nonzero_indices[0])
    # num = np.count_nonzero(nonzero_values)

    #iteracija 3
    mom = cv2.moments(roi)
    x = 0
    y = 0
    if mom['m00'] != 0:
        x = int(mom['m10'] / mom['m00'])
        y = int(mom['m01'] / mom['m00'])
        x = int(x - window[2] / 2)
        y = int(y - window[3] / 2)
        window = (window[0] + x, window[1] + y, window[2], window[3])

    if x  < criteria and y < criteria:
        return (True, window)

    return  (False, window)

def CamShift(src, window, criteria):
    ret = False
    for _ in range(criteria[1]):
        ret, window = Iteration(src, window, criteria[2])
        window = changeSize(src, window)
           

    return (ret, window)

def changeSize(src, window):
    roi = src[window[1]:window[1]+window[3], window[0]:window[0]+window[2]]
    mom = cv2.moments(roi)

    ratio = window[2] / window[3]
    if mom['m00'] != 0:
        w = 2 * np.sqrt(mom['m00'] / 256)
        h = w / ratio
        window = (window[0], window[1], int(w), int(h))

    if window[0] < 0:
        window = (0, window[1], window[2], window[3])
    if window[1] < 0:
        window = (window[0], 0, window[2], window[3])
    if window[0] + window[2] > src.shape[1]:
        window = (src.shape[1] - window[2], window[1], window[2], window[3])
    if window[1] + window[3] > src.shape[0]:
        window = (window[0], src.shape[0] - window[3], window[2], window[3])
    if window[2] < 10:
        window = (window[0], window[1], 10, window[3])
    if window[3] < 10:
        window = (window[0], window[1], window[2], 10)

    return window

def selectObjectH(frame, target_hist):

    target_hist = target_hist.astype(np.float32)
    cv2.normalize(target_hist, target_hist, 0, 1, cv2.NORM_MINMAX)

    hsvImg = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    h_channel = hsvImg[:, :, 0]
    #s_channel = hsvImg[:, :, 1]

    #probabilityImg = np.zeros(shape=(hsvImg.shape[0], hsvImg.shape[1], 1) , dtype=np.float32)

    # Stack the H and S channels
    # hs_img = np.stack((h_channel, s_channel), axis=-1)

    # for i in range(hs_img.shape[0]):
    #     for j in range(hs_img.shape[1]):
    #         probabilityImg[i][j] = target_hist[int(h_channel[i][j])]
    #         satProp = target_hist[int(hs_img[i][j][1])]
    #         probabilityImg[i][j] = hueProp * satProp

    probabilityImg = target_hist[h_channel.astype(np.uint8)]
    

    cv2.normalize(probabilityImg, probabilityImg, 0, 255, cv2.NORM_MINMAX)
    probabilityImg = probabilityImg.astype(np.uint8)

    mask = cv2.inRange(probabilityImg, 100, 255)

    contour, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    maxContour = max(contour, key=cv2.contourArea) 
    x, y, w, h = cv2.boundingRect(maxContour)
    return (x, y, w, h)  

def selectObjectHS(frame, target_hist):

    target_hist = target_hist.astype(np.float32)
    cv2.normalize(target_hist, target_hist, 0, 1, cv2.NORM_MINMAX)
    hsvImg = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

    # probabilityImg = np.zeros(shape=(hsvImg.shape[0], hsvImg.shape[1], 1) , dtype=np.float32)
    # for i in range(hsvImg.shape[0]):
    #     for j in range(hsvImg.shape[1]):
    #         hueProp = target_hist[int(hsvImg[i][j][0])][0]
    #         satProp = target_hist[int(hsvImg[i][j][1])][1]
    #         probabilityImg[i][j] = hueProp * satProp

 
    hue_values = hsvImg[..., 0].astype(np.uint8)
    sat_values = hsvImg[..., 1].astype(np.uint8)
    hue_props = target_hist[hue_values.flatten(), 0].reshape(hsvImg.shape[:2])
    sat_props = target_hist[sat_values.flatten(), 1].reshape(hsvImg.shape[:2])
    probabilityImg = (hue_props * sat_props).astype(np.float32)
    

    cv2.normalize(probabilityImg, probabilityImg, 0, 255, cv2.NORM_MINMAX)
    probabilityImg = probabilityImg.astype(np.uint8)

    mask = cv2.inRange(probabilityImg, 100, 255)

    contour, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    maxContour = max(contour, key=cv2.contourArea) 
    x, y, w, h = cv2.boundingRect(maxContour)
    return (x, y, w, h)  

def selectObject(frame, target_hist):

    if target_hist.shape[0] == 180:
        return selectObjectH(frame, target_hist)
    else:
        return selectObjectHS(frame, target_hist)

def calcHistH(roi):
    histogram = np.zeros(180, dtype=int)
    for i in range(roi.shape[0]):
        for j in range(roi.shape[1]):
            histogram[roi[i][j][0]] += 1

    return histogram

def calcHistHS(roi):
    histogramH = np.zeros(256, dtype=int)
    histogramS = np.zeros(256, dtype=int)
    for i in range(roi.shape[0]):
        for j in range(roi.shape[1]):
            histogramH[roi[i][j][0]] += 1
            histogramS[roi[i][j][1]] += 1

    histogram = np.stack(arrays=(histogramH, histogramS), axis=-1)
    return histogram

def calcHist(roi, numOfchannels=1):
    if numOfchannels == 1:
        return calcHistH(roi)
    else:
        return calcHistHS(roi)

def calcBackProjectH(frame, target_hist):
    target_hist = target_hist.astype(np.float32)
    cv2.normalize(target_hist, target_hist, 0, 1, cv2.NORM_MINMAX)

    hsvImg = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    h_channel = hsvImg[:, :, 0]
    probabilityImg = target_hist[h_channel.astype(np.uint8)]
    cv2.normalize(probabilityImg, probabilityImg, 0, 255, cv2.NORM_MINMAX)
    return probabilityImg.astype(np.uint8)

def calcBackProjectHS(frame, target_hist):
    target_hist = target_hist.astype(np.float32)
    cv2.normalize(target_hist, target_hist, 0, 1, cv2.NORM_MINMAX)
    hsvImg = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
 
    hue_values = hsvImg[..., 0].astype(np.uint8)
    sat_values = hsvImg[..., 1].astype(np.uint8)
    hue_props = target_hist[hue_values.flatten(), 0].reshape(hsvImg.shape[:2])
    sat_props = target_hist[sat_values.flatten(), 1].reshape(hsvImg.shape[:2])
    probabilityImg = (hue_props * sat_props).astype(np.float32)
    

    cv2.normalize(probabilityImg, probabilityImg, 0, 255, cv2.NORM_MINMAX)
    return probabilityImg.astype(np.uint8)

def calcBackProject(frame, target_hist):
    if target_hist.shape[0] == 180:
        return calcBackProjectH(frame, target_hist)
    else:
        return calcBackProjectHS(frame, target_hist)