from absl import app, flags
import cv2
from copy import deepcopy
import glob
import numpy as np
import os
from scipy import ndimage
import tensorflow as tf
from our_func_cvpr18 import build
import time

FLAGS = flags.FLAGS
flags.DEFINE_string("model_path", "", "Path to the sfm model.")


def main(argv):
    # model_path = "/tmp/mickey/sfm/"
    model_path = FLAGS.model_path
    image_ext = ".jpg"
    num_fg_samples = 400
    num_bg_samples = 400

    sess = tf.Session()

    input = tf.placeholder(tf.float32, shape=[None, None, None, 7])
    output = tf.placeholder(tf.float32, shape=[None, None, None, 1])
    sz = tf.placeholder(tf.int32, shape=[2])
    network = build(input, sz)

    saver = tf.train.Saver(var_list=[var for var in tf.trainable_variables() if var.name.startswith('g_')])
    sess.run(tf.initialize_all_variables())

    ckpt = tf.train.get_checkpoint_state("Models/ours_cvpr18")
    if ckpt:
        # print('loaded '+ckpt.model_checkpoint_path)
        saver.restore(sess, ckpt.model_checkpoint_path)

    frame_rgb_images = glob.glob(model_path + "/*" + image_ext)
    frame_rgb_images = sorted(frame_rgb_images)

    for (frame_id, image_path) in enumerate(frame_rgb_images):
        # image_path = os.path.join(model_path, "frame_%06d_rgb_image" % frame_id + image_ext)
        input_image = cv2.imread(image_path, cv2.IMREAD_COLOR)

        # Note: I found that the pre-trained network has very good results on this resolution.
        # input_image = cv2.resize(input_image, (480, 360))

        # Uncomment for debugging.
        # log_output_path = "/tmp/frame_%06d_#fg_%03d_#bg_%03d_rgb_image.jpg" % (frame_id, num_fg_samples, num_bg_samples)
        # cv2.imwrite(log_output_path, input_image)
        # print("Write image: " + log_output_path)

        image_height, image_width, _ = input_image.shape

        fg_mask_path = os.path.join(model_path, "frame_%06d_object_mask.png" % frame_id)
        fg_mask = cv2.imread(fg_mask_path, cv2.IMREAD_GRAYSCALE)
        fg_mask = cv2.resize(fg_mask, (image_width, image_height))

        bg_mask_path = os.path.join(model_path, "frame_%06d_background_mask.png" % frame_id)
        bg_mask = cv2.imread(bg_mask_path, cv2.IMREAD_GRAYSCALE)
        bg_mask = cv2.resize(bg_mask, (image_width, image_height))

        fg_pixels = np.array(np.where(fg_mask >= 255))
        bg_pixels = np.array(np.where(bg_mask >= 255))

        fg_indices = np.random.choice(fg_pixels.shape[1], num_fg_samples, replace=False)
        fg_sample_ys = fg_pixels[0, :][fg_indices]
        fg_sample_xs = fg_pixels[1, :][fg_indices]

        bg_indices = np.random.choice(bg_pixels.shape[1], num_bg_samples, replace=False)
        bg_sample_ys = bg_pixels[0, :][bg_indices]
        bg_sample_xs = bg_pixels[1, :][bg_indices]

        positives = np.uint8(255 * np.ones([image_height, image_width]))
        positive_samples = (positives == 0)
        positive_samples[fg_sample_ys, fg_sample_xs] = 1

        positive_samples = ndimage.distance_transform_edt(1 - positive_samples)
        positive_samples = np.uint8(np.minimum(np.maximum(positive_samples, 0.0), 255.0))

        # Uncomment for debugging.
        # log_positive_samples_path = "/tmp/frame_%06d_#fg_%03d_#bg_%03d_positive_edt.png" % \
        #                             (frame_id, num_fg_samples, num_bg_samples)
        # cv2.imwrite(log_positive_samples_path, positive_samples)
        # print("Write image: " + log_positive_samples_path)

        negatives = np.uint8(255 * np.ones([image_height, image_width]))
        negative_samples = (negatives == 0)
        negative_samples[bg_sample_ys, bg_sample_xs] = 1

        negative_samples = ndimage.distance_transform_edt(1 - negative_samples)
        negative_samples = np.uint8(np.minimum(np.maximum(negative_samples, 0.0), 255.0))

        # Uncomment for debugging.
        # log_negative_samples_path = "/tmp/frame_%06d_#fg_%03d_#bg_%03d_negative_edt.png" % \
        #                             (frame_id, num_fg_samples, num_bg_samples)
        # cv2.imwrite(log_negative_samples_path, negative_samples)
        # print("Write image: " + log_negative_samples_path)

        # sess = tf.Session()

        # input = tf.placeholder(tf.float32, shape=[None, None, None, 7])
        # output = tf.placeholder(tf.float32, shape=[None, None, None, 1])
        # sz = tf.placeholder(tf.int32, shape=[2])
        # network = build(input, sz)
        #
        # saver = tf.train.Saver(var_list=[var for var in tf.trainable_variables() if var.name.startswith('g_')])
        # sess.run(tf.initialize_all_variables())
        #
        # ckpt = tf.train.get_checkpoint_state("Models/ours_cvpr18")
        # if ckpt:
        #     # print('loaded '+ckpt.model_checkpoint_path)
        #     saver.restore(sess, ckpt.model_checkpoint_path)

        start_time = time.time()

        input_positive_samples = deepcopy(positive_samples)
        input_negative_samples = deepcopy(negative_samples)
        input_positive_samples[positive_samples != 0] = 255
        input_negative_samples[negative_samples != 0] = 255

        input_ = np.expand_dims(
            np.float32(np.concatenate(
                [input_image, np.expand_dims(positive_samples, axis=2), np.expand_dims(negative_samples, axis=2),
                 np.expand_dims(input_positive_samples, axis=2), np.expand_dims(input_negative_samples, axis=2)],
                axis=2)), axis=0)
        output_image = sess.run([network], feed_dict={input: input_, sz: [image_height, image_width]})
        output_image = np.minimum(np.maximum(output_image, 0.0), 1.0)
        output_image[np.where(output_image > 0.5)] = 1
        output_image[np.where(output_image <= 0.5)] = 0

        output_mask_path = os.path.join(model_path, "frame_%06d_object_mask_#fg_%03d_#bg_%03d.png" % (
            frame_id, num_fg_samples, num_bg_samples))
        segmask = np.uint8(output_image[0, 0, :, :, 0] * 255.0)
        cv2.imwrite(output_mask_path, segmask)
        print("Write image: " + output_mask_path)

        time_diff = time.time() - start_time
        print(f"Inference timer: {time_diff:0.6f} s")
        start_time = time.time()

        fg_model = np.zeros((1, 65), dtype="float")
        bg_model = np.zeros((1, 65), dtype="float")

        mask = np.uint8(output_image[0, 0, :, :, 0])
        # bbox = cv2.boundingRect(mask)
        # input_mask = np.ones_like(mask) * cv2.GC_PR_BGD
        # input_mask[bbox[0]:bbox[2], bbox[1]:bbox[3]] = cv2.GC_PR_FGD
        input_mask = np.uint8(np.where(mask == 1, cv2.GC_FGD, cv2.GC_PR_BGD))

        (mask, bg_model, fg_model) = cv2.grabCut(input_image, input_mask, None, bg_model, fg_model, iterCount=5,
                                                 mode=cv2.GC_INIT_WITH_MASK)

        output_mask = np.where((mask == cv2.GC_BGD) | (mask == cv2.GC_PR_BGD), 0, 1)
        output_mask = (output_mask * 255).astype("uint8")

        output_mask_grabcut_path = os.path.join(model_path, "frame_%06d_object_mask_#fg_%03d_#bg_%03d_grabcut.png" % (
            frame_id, num_fg_samples, num_bg_samples))
        cv2.imwrite(output_mask_grabcut_path, output_mask)
        print("Write image: " + output_mask_grabcut_path)

        # output = cv2.bitwise_and(input_image, input_image, mask=outputMask)

        time_diff = time.time() - start_time
        print(f"Graphcut timer: {time_diff:0.6f} s")

        # if frame_id > 2:
        #     break


if __name__ == "__main__":
    app.run(main)
