import cv2
import numpy as np
import os
import torch


from torch.nn.functional import upsample
from dataloaders import utils
import networks.deeplab_resnet as resnet

from glob import glob
from copy import deepcopy


drawing = False
start = False

# TODO: import test image path
image_list = glob(os.path.join('ims/', '*.'+'jpg'))
image = cv2.imread(image_list[0])
img_shape = (450, 450)
image = utils.fixed_resize(image, img_shape).astype(np.uint8)

w = img_shape[0]
h = img_shape[1]
output = np.zeros((w, h, 3), np.uint8)
thres = 0.8

left = 0xFFF
right = 0
up = 0xFFF
down = 0

gpu_id = 0
device = torch.device("cuda:"+str(gpu_id) if torch.cuda.is_available() else "cpu")

#  Create the network and load the weights
net = resnet.resnet101(1, nInputChannels=4, classifier='psp')
print("Initializing weights from: {}".format(os.path.join('models/', 'deepgc_pascal_epoch-99.pth')))
state_dict_checkpoint = torch.load(os.path.join('models/', 'deepgc_pascal_epoch-99.pth'),
                                   map_location=lambda storage, loc: storage)

net.load_state_dict(state_dict_checkpoint)
net.eval()
net.to(device)


def interactive_drawing(event, x, y, flag, params):
    global xs, ys, ix, iy, drawing, image, output, left, right, up, down

    if event == cv2.EVENT_LBUTTONDOWN:
        drawing = True
        ix, iy = x, y
        xs, ys = x, y
        left = min(left, x)
        right = max(right, x)
        up = min(up, y)
        down = max(down, y)

    elif event == cv2.EVENT_MOUSEMOVE:
        if drawing is True:
            cv2.line(image, (ix, iy), (x, y), (0, 0, 255), 2)
            cv2.line(output, (ix, iy), (x, y), (255, 255, 255), 1)
            ix = x
            iy = y
            left = min(left, x)
            right = max(right, x)
            up = min(up, y)
            down = max(down, y)

    elif event == cv2.EVENT_LBUTTONUP:
        drawing = False
        cv2.line(image, (ix, iy), (x, y), (0, 0, 255), 2)
        cv2.line(output, (ix, iy), (x, y), (255, 255, 255), 1)
        ix = x
        iy = y
        cv2.line(image, (ix, iy), (xs, ys), (0, 0, 255), 2)
        cv2.line(output, (ix, iy), (xs, ys), (255, 255, 255), 1)
    return x, y

def main():
    global image, output
    cv2.namedWindow('draw', flags=cv2.WINDOW_NORMAL)
    cv2.setMouseCallback('draw', interactive_drawing)

    image_idx = 0
    while(1):
        cv2.imshow('draw', image)
        k = cv2.waitKey(1) & 0xFF
        if k != 255:
            # print(k)
            pass
        if k == 100: # D
            cv2.imwrite('./' + str(image_idx) + 'out.png', image)
        if k == 115:
            global left, right, up, down
            left = 0xFFF
            right = 0
            up = 0xFFF
            down = 0
            drawing = False  # true if mouse is pressed
            image = cv2.imread(image_list[image_idx])
            image = utils.fixed_resize(image, (450, 450)).astype(np.uint8)
            sp = image.shape
            w = sp[0]
            h = sp[1]
            output = np.zeros((w, h, 3), np.uint8)
            while (1):
                cv2.imshow('draw', image)
                k = cv2.waitKey(1) & 0xFF
                if k == 32:
                    break
                if k == 27:
                    image = cv2.imread(image_list[image_idx])
                    image = utils.fixed_resize(image, (450, 450)).astype(np.uint8)
                    output = np.zeros((w, h, 3), np.uint8)

            tmp = (output[:, :, 0] > 0).astype(np.uint8)
            tmp_ = deepcopy(tmp)
            fill_mask = np.ones((tmp.shape[0] + 2, tmp.shape[1] + 2))
            fill_mask[1:-1, 1:-1] = tmp_
            fill_mask = fill_mask.astype(np.uint8)
            cv2.floodFill(tmp_, fill_mask, (int((left + right) / 2), int((up + down) / 2)), 5)
            tmp_ = tmp_.astype(np.int8)

            output = cv2.resize(output, img_shape)

            tmp_ = tmp_.astype(np.int8)
            tmp_[tmp_ == 5] = -1  # pixel inside bounding box
            tmp_[tmp_ == 0] = 1  # pixel on and outside bounding box

            tmp = (tmp == 0).astype(np.uint8)

            dismap = cv2.distanceTransform(tmp, cv2.DIST_L2, cv2.DIST_MASK_PRECISE)  # compute distance inside and outside bounding box
            dismap = tmp_ * dismap + 128

            dismap[dismap > 255] = 255
            dismap[dismap < 0] = 0
            dismap = dismap

            dismap = utils.fixed_resize(dismap, (450, 450)).astype(np.uint8)

            cv2.imwrite('dismap.png', dismap)
            cv2.imwrite('tmp.png', tmp*255)

            dismap = np.expand_dims(dismap, axis=-1)

            image = image[:, :, ::-1] # change to rgb
            merge_input = np.concatenate((image, dismap), axis=2).astype(np.float32)
            inputs = torch.from_numpy(merge_input.transpose((2, 0, 1))[np.newaxis, ...])

            # Run a forward pass
            inputs = inputs.to(device)
            outputs = net.forward(inputs)
            outputs = upsample(outputs, size=(450, 450), mode='bilinear', align_corners=True)
            outputs = outputs.to(torch.device('cpu'))

            prediction = np.transpose(outputs.data.numpy()[0, ...], (1, 2, 0))
            prediction = 1 / (1 + np.exp(-prediction))
            prediction = np.squeeze(prediction)
            prediction[prediction>thres] = 255
            prediction[prediction<=thres] = 0

            prediction = np.expand_dims(prediction, axis=-1).astype(np.uint8)
            image = image[:, :, ::-1] # change to bgr
            display_mask = np.concatenate([np.zeros_like(prediction), np.zeros_like(prediction), prediction], axis=-1)
            image = cv2.addWeighted(image, 0.9, display_mask, 0.5, 0.1)

        if k == 99:
            break

        if k == 110:
            image_idx += 1
            if image_idx >= len(image_list):
                print('Already the last image. Starting from the beginning.')
                image_idx = 0
            image = cv2.imread(image_list[image_idx])
            image = utils.fixed_resize(image, (450, 450)).astype(np.uint8)
            sp = image.shape
            w = sp[0]
            h = sp[1]
            output = np.zeros((w, h, 3), np.uint8)

        if k == 112:
            image_idx -= 1
            if image_idx < 0:
                print('Reached the first image. Starting from the end.')
                image_idx = len(image_list)-1
            image = cv2.imread(image_list[image_idx])
            image = utils.fixed_resize(image, (450, 450)).astype(np.uint8)
            sp = image.shape
            w = sp[0]
            h = sp[1]
            output = np.zeros((w, h, 3), np.uint8)


    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
