import cv2
import numpy as np
import json
import os
#10 - фигурирует как константа для максимального числа голов
check_glob = 0
g = np.array([-1 for i in range(10)])


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=-1):
    if number != -1:
        path = make_name(path, number)
    with open(path) as json_data:
        d = json.load(json_data)
        d1 = d['objects']
        h = []
        dig = -1
        for i in d1:
            for j in i['tags']:
                if j.isdigit():
                    dig = int(j)
                if j == 'h':
                    h.append([i['data'], dig])
        h1 = np.array(h)
        return h1


def give_cur_heads(ans):#скорее всего тут умирает, но я пытался...
    seq = []
    for i in ans:
        seq.append(i[1])
    return seq


def compare(roi, ans):
    size = ans.shape[0]
    seq = give_cur_heads(ans)
    ans = fulfil_ans(ans)
    res = []
    print(seq)
    if size == 1:
        print("********")
        #print(roi)
        print(len(roi))
        print(g[seq[0]])
        return compare_head(roi[g[seq[0]]], ans)
    for i in range(size):
        num_head = g[seq[i]]
        print("Number of head", num_head)
        res.append(compare_head(roi[num_head], ans[i]))
    return np.array(res)


def get_res(res):
    return res


def compare_head(track, ans):
    s_track = set()
    s_ans = set()
    print(track.shape)
    for i in track:
        for j in i:
            x, y = j
            s_track.add((int(x), int(y)))
    for i in ans:
        for j in i:
            x, y = j
            s_ans.add((int(x), int(y)))
    s_dif = s_ans & s_track
    return get_res(s_dif)


def upd(roi, flow):
    global check_glob
    if not check_glob:
        return roi
    if len(roi) == 1:  # only one head
        roi[0] =  upd_head(roi[0], flow)
        return roi
    for i in range(len(roi)):
        roi[i] = upd_head(roi[i], flow)
    return roi


def upd_head(roi_head, flow):
    for i in range(roi_head.shape[0]):
        for j in range(roi_head.shape[1]):
            roi_head[i][j] = upd_cell(roi_head[i][j], flow)
    return np.array(roi_head)


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


def fulfil(roi_border): #we think that in json y, x not x, y    in fact, initialize empty roi/ans
    if roi_border.shape[0] == 0:
        return
    if roi_border.shape[0] == 1:
        global check_glob
        check_glob = 1
        g[roi_border[0][1]] = 0
        roi = []
        roi.append(fulfil_ins(roi_border[0]))
        return roi
    ans = []
    for i in roi_border:
        check_glob = 1
        ans.append(fulfil_ins(i))
        g[i[1]] = len(ans) - 1
    return ans


def fulfil_ans(roi_border): #we think that in json y, x not x, y    in fact, initialize empty roi/ans
    if roi_border.shape[0] == 0:
        return
    if roi_border.shape[0] == 1:
        return fulfil_ins(roi_border[0])
    ans = []
    for i in roi_border:
        ans.append(fulfil_ins(i))
    return np.array(ans)


def fulfil_ins(roi_border):  #we think that in json y, x not x
    x1, y1, x2, y2 = roi_border[0]
    ind = 0
    res = []
    for i in range(int(x2), int(x1 + 1), 1): #round to int real values!!!
        row = []
        for j in range(int(y2), int(y1 + 1), 1):
            row.append([j, i])
            ind += 1
        res.append(row)
    return np.array(res)


def add_head(roi, path):
    h = find_head(path)
    global check_glob
    if not check_glob:
        check_glob = 1
        return fulfil(h)
    for i in range(h.shape[0]):# suppose, problems here....
        num = h[i][1]
        print(h.shape)
        print(h)
        if g[num] == -1:
            roi.append(fulfil_ins(h[i]))
            g[h[i][1]] = len(roi) - 1
    return roi


def check_new_head(roi, path_js, ind, ans):
    path = make_name(path_js, ind)
    global check_glob
    if not check_glob:
        roi = add_head(roi, path)
        return roi
    seq = give_cur_heads(ans)
    print(seq)
    for i in seq:
        if g[i] == -1:
            roi = add_head(roi, path)
            break
    return roi


def track_head(path_vid, path_js):
    cap = cv2.VideoCapture(path_vid)
    ret, frame1 = cap.read()
    roi = []
    roi_border = find_head(path_js, 0)
    #roi.append(fulfil(roi_border)) для 1 кадра с головами не работает
    prvs = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
    ind = 1
    test = []
    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)
        roi = upd(roi, flow)
        if find(path_js, ind):
            ans = find_head(path_js, ind)
            roi = check_new_head(roi, path_js, ind, ans)
            test.append(compare(roi, ans))
        print(ind)
        if ind == 401:
            break
        ind += 1
        k = cv2.waitKey(30) & 0xff #несколько странная остановка, но вроде работала...
        if k == 27:
            break
    return 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)
test = track_head(path_to_video, path_to_json)
for i in test:
    print(len(i))