import cv2
import numpy as np
import json
import os


def set_green(path):
    cap = cv2.VideoCapture(path)
    ret, frame1 = cap.read()
    cv2.circle(frame1, (40, 500), 20, (0, 255, 0), -1)
    cv2.imwrite('frame_change.png', frame1)
    return frame1


def draw_circle(x, y, flow, img):
    height, width, _ = img.shape
    x = x + flow[int(y)][int(x)][0]
    y = y + flow[int(y)][int(x)][1]
    cv2.circle(img, (int(x), int(y)), 20, (0, 255, 0), -1)
    cv2.imshow('green', img)
    return x, y


def print_track(path):
    cap = cv2.VideoCapture(path)
    ret, frame1 = cap.read()
    prvs = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
    hsv = np.zeros_like(frame1)
    hsv[..., 1] = 255
    while (1):
        ret, frame2 = cap.read()
        cv2.imwrite("framefirst.png", frame1)
        cv2.imshow('frame_smth', frame2)
        next = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)
        flow = cv2.calcOpticalFlowFarneback(prvs, next, None, 0.5, 3, 15, 3, 5, 1.2, 0)
        mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])
        hsv[..., 0] = ang * 180 / np.pi / 2
        hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
        bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
        cv2.imshow('frame2', bgr)
        k = cv2.waitKey(30) & 0xff
        if k == 27:
            break
        elif k == ord('s'):
            cv2.imwrite('opticalfb.png', frame2)
            cv2.imwrite('opticalhsv.png', bgr)
        prvs = next
    cap.release()
    cv2.destroyAllWindows()


def track_green_circle(path):  # track specified region in the image
    cap = cv2.VideoCapture(path)
    frame1 = set_green(path)
    prvs = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
    hsv = np.zeros_like(frame1)
    hsv[..., 1] = 255
    x, y = 40, 500
    while (1):
        ret, frame2 = cap.read()
        cv2.imwrite("framefirst.png", frame1)
        cv2.imshow('frame_smth', frame2)
        next = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)
        flow = cv2.calcOpticalFlowFarneback(prvs, next, None, 0.5, 3, 15, 3, 5, 1.2, 0)
        print(flow.shape)
        x, y = draw_circle(x, y, flow, frame2)
        mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])
        hsv[..., 0] = ang * 180 / np.pi / 2
        hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
        bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
        cv2.imshow('frame2', bgr)
        k = cv2.waitKey(30) & 0xff
        if k == 27:
            break
        elif k == ord('s'):
            cv2.imwrite('opticalfb.png', frame2)
            cv2.imwrite('opticalhsv.png', bgr)
        prvs = next
    cap.release()
    cv2.destroyAllWindows()


def inside(point, coord, img):
    x, y = img.shape
    if (not coord):
        return point >= 0 and point <= x
    return point >= 0 and point <= y
    pass


def create_test_image(path,
                      change):  # we give out image and draw circles with some certain optical flow in vector change(x, y) when we reach the end *= -1 return correct_optical_flow
    cap = cv2.VideoCapture(path)
    ret, frame = cap.read()
    float = []
    x, y, r = 40, 500, 20
    fourcc = cv2.VideoWriter_fourcc(*'mp4')
    fps = cap.get(cv2.CV_CAP_PROP_FPS)
    out = cv2.VideoWriter('green_circles_video.mp4', fourcc, fps, (frame.shape))
    while (not frame.empty()):
        cv2.circle(frame, (x, y), r, (0, 255, 0), -1)
        if (not inside((x + change[0]), 0, frame)):
            change[0] *= -1
        if (not inside(y + change[1], 1, frame)):
            change[1] *= -1
        float.append(change)
        x += change[0]
        y += change[1]
    pass


def drawOptFlowMap(flow, prvs, step, _, color):
    cflowmap = cv2.cvtColor(prvs, cv2.COLOR_GRAY2BGR)
    for y in range(0, cflowmap.shape[0], step):
        for x in range(0, cflowmap.shape[1], step):
            fxy = flow[y][x]
            cv2.line(cflowmap, (x, y), (int(x + fxy[0]), int(y + fxy[1])), color)
            cv2.circle(cflowmap, (x, y), 2, color, -1)
    cv2.imshow("flow", cflowmap)


def draw_hsv(flow, frame1):
    hsv = np.zeros_like(frame1)
    hsv[..., 1] = 255
    mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])
    hsv[..., 0] = ang * 180 / np.pi / 2
    hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
    bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
    cv2.imshow('frame2', bgr)
    pass


def track_map(path):  # with drawOptFlowMap
    cap = cv2.VideoCapture(path)
    ret, frame1 = cap.read()
    prvs = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
    while (1):
        ret, frame2 = cap.read()
        next = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)
        flow = cv2.calcOpticalFlowFarneback(prvs, next, None, 0.5, 3, 15, 3, 5, 1.2, 0)
        drawOptFlowMap(flow, prvs, 16, 1.5, (0, 255, 0))
        draw_hsv(flow, frame1)
        k = cv2.waitKey(30) & 0xff
        if k == 27:
            break
        elif k == ord('s'):
            cv2.imwrite('opticalfb.png', frame2)
            #cv2.imwrite('opticalhsv.png', bgr)
        prvs = next
    cap.release()
    cv2.destroyAllWindows()
#####################################################################


def find(path, number):
    return os.path.isfile(make_name(path, number))


def make_name(path, number):
    s = str(number)
    k = len(s)
    zer = ''
    for i in range(0, 6 - k):
        zer += '0'
    end = 'rc.jpg.json'
    res = path + zer + s + end
    return res


def find_head(path, number):
    with open(make_name(path, number)) as json_data:
        d = json.load(json_data)
        d1 = d['objects']
        h = []
        for i in d1:
            for j in i['tags']:
                if j == 'h':
                    h.append(i['data'])
        return h


def compare(roi, ans):
    size = ans.shape.size
    res = []
    if size == 1:
        return compare_head(roi, ans)
    for i in range(size):
        res.append(compare_head(roi[i], ans[i]))
    return res


def get_res(res):
    return res


def compare_head(track,  ans):
    res = 0
    s_track = set()
    s_ans = set()
    for i in track:
        s_track.add(i)
    for i in ans:
        s_ans.add(i)
    for i in s_ans:
        if s_track.find(i):
            res += 1
    return get_res(res)


def upd(roi, flow):
    for i in range(roi.shape[0]):
        for j in range(roi.shape[1]):
            for u in range(roi.shape[2]):
                roi[i][j][u] = upd_cell(roi[i][j][u], flow)
    return roi


def upd_cell(cell, flow): #we think that in json y, x not x, y
    cell[0] += flow[cell][1]
    cell[1] += flow[cell][0]

def fulfil(roi_border): #we think that in json y, x not x, y
    ans = []
    for i in roi_border:
        y1, x1, y2, x2 = i
        res = []
        for j in range(x1, x2 + 1):
            row = []
            for u in range(y1, y2 + 1):
                row.append([u, j])
            res.append(row)
        ans.append(res)
    return ans

def rewrittten_fulfil(roi_border):  #we think that in json y, x not x, y
    y1, x1, y2, x2 = roi_border
    res = []
    for i in range(x1, x2 + 1):
        row = []
        for j in range(y1, y2 + 1):
            row.append([j, i])
            res.append(row)
    return res

def count_head(path):
    ans = 0
    with open(path) as json_data:
        d = json.load(json_data)
        for i in d['objects']:
            for j in i['tags']:
                if (j == 'h'):
                    ans += 1
    return ans


def add_head(roi, path):
    ans = 0
    with open(path) as json_data:
        d = json.load(json_data)
        for i in d['objects']:
            for j in i['tags']:
                if (j == 'h'):
                    ans += 1
            if (ans > roi.shape[0]):
                roi.append(fulfil(i['data'])) # rewrite fulfil for one dimesion
    return roi


def check_new_head(roi, path_js, ind):
    path = make_name(path_js, ind)
    if count_head(path) == roi.shape[0]:
        return roi
    roi = add_head(roi, path)
    return roi


def track_head(path_vid, path_js):
    cap = cv2.VideoCapture(path_vid)
    ret, frame1 = cap.read()
    roi_border = find_head(path_js, 0)
    roi = fulfil(roi_border)
    prvs = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
    ind = 1
    test = []
    while (1):
        ret, frame2 = cap.read()
        if frame2.empty:
            break
        next = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)
        flow = cv2.calcOpticalFlowFarneback(prvs, next, None, 0.5, 3, 15, 3, 5, 1.2, 0)
        print(np.shape(flow))
        roi = upd(roi, flow)
        if find(path_js, ind):
            roi = check_new_head(roi, path_js, ind)
            ans = find_head(path_js, ind)
            test.append(compare(roi, ans))
    print(test)




path = "/Users/denissurin/Downloads/Telegram Desktop/video_2018-07-13_15-56-11.mp4"
path1 = "/Users/denissurin/Downloads/OneDrive-2018-07-18/real2.avi"
path2 = "/Users/denissurin/Downloads/OneDrive-2018-07-18/record-0000-004C-20180505193003-20180505193014.avi"
path3 = "/Users/denissurin/Downloads/OneDrive-2018-07-18/VNG_180621_3730.avi"
path_to_json = "/Users/denissurin/Downloads/GPN_180319_6103/"
path_to_video = "/Users/denissurin/Downloads/GPN_180319_6103.mp4"
track_green_circle(path)
# track_map(path3)