import sys
import os
import cv2
import numpy as np
import copy
import dense_correspondence_manipulation.utils.utils as utils
import time
from PIL import Image
dc_source_dir = utils.getDenseCorrespondenceSourceDir()
sys.path.append(dc_source_dir)
sys.path.append(os.path.join(dc_source_dir, "dense_correspondence", "correspondence_tools"))
from dense_correspondence.dataset.spartan_dataset_masked import SpartanDataset, ImageType

import dense_correspondence
from dense_correspondence.evaluation.evaluation import *
from dense_correspondence.evaluation.plotting import normalize_descriptor
from dense_correspondence.network.dense_correspondence_network import DenseCorrespondenceNetwork
import dense_correspondence.correspondence_tools.correspondence_finder as correspondence_finder

sys.path.append(os.path.join(os.path.dirname(__file__), "../simple-pixel-correspondence-labeler"))
from annotate_correspondences import label_colors, draw_reticle, pil_image_to_cv2, drawing_scale_config, numpy_to_cv2, label_colors

COLOR_RED = np.array([0, 0, 255])
COLOR_GREEN = np.array([0,255,0])

utils.set_default_cuda_visible_devices()
eval_config_filename = os.path.join(utils.getDenseCorrespondenceSourceDir(), 'config', 'dense_correspondence', 'evaluation', 'evaluation.yaml')
EVAL_CONFIG = utils.getDictFromYamlFilename(eval_config_filename)


class HeatmapVisualization(object):
    """
    Launches a live interactive heatmap visualization.
    Edit config/dense_correspondence/heatmap_vis/heatmap.yaml to specify which networks
    to visualize. Specifically add the network you want to visualize to the "networks" list.
    Make sure that this network appears in the file pointed to by EVAL_CONFIG
    Usage: Launch this file with python after sourcing the environment with
    `use_pytorch_dense_correspondence`
    Then `python live_heatmap_visualization.py`.
    """

    def __init__(self, config):
        self._config = config
        self._dce = DenseCorrespondenceEvaluation(EVAL_CONFIG)
        self._load_networks()
        self._reticle_color = COLOR_GREEN
        self.load_specific_dataset() # uncomment if you want to load a specific dataset

    def _load_networks(self):
        # we will use the dataset for the first network in the series
        self._dcn_dict = dict()

        self._dataset = None
        self._network_reticle_color = dict()

        for idx, network_name in enumerate(self._config["networks"]):
            dcn = self._dce.load_network_from_config(network_name)
            dcn.eval()
            self._dcn_dict[network_name] = dcn

            if len(self._config["networks"]) == 1:
                self._network_reticle_color[network_name] = COLOR_RED
            else:
                self._network_reticle_color[network_name] = label_colors[idx]

            if self._dataset is None:
                self._dataset = dcn.load_training_dataset()

    def load_specific_dataset(self):
        dataset_config_filename = os.path.join(utils.getDenseCorrespondenceSourceDir(), 'config',
                                               'dense_correspondence',
                                               'dataset', 'composite', 'rope_nonrigid_412vert_only.yaml')

        dataset_config = utils.getDictFromYamlFilename(dataset_config_filename)
        self._dataset = SpartanDataset(config=dataset_config)
        #self._dataset.get_knots_info('rope_nonrigid_412vert')

    def get_random_image_pair(self):
        object_id = self._dataset.get_random_object_id()
        scene_name_a = self._dataset.get_random_single_object_scene_name(object_id)
	scene_name_b = self._dataset.get_random_single_object_scene_name(object_id)

        image_a_idx = self._dataset.get_random_image_index(scene_name_a)
        image_b_idx = self._dataset.get_random_image_index(scene_name_b)

        return scene_name_a, scene_name_b, image_a_idx, image_b_idx

    def _get_new_images(self):
        """
        Gets a new pair of images
        :return:
        :rtype:
        """
	if random.random() <= 1.0:
            self._dataset.set_train_mode()
    	else:
	    self._dataset.set_test_mode()

        if self._config["same_object"]:
	    print "getting random image pair"
            scene_name_1, scene_name_2, image_1_idx, image_2_idx = self.get_random_image_pair()
        else:
            raise ValueError("At least one of the image types must be set tot True")

	print "got pair"
        self.img1_pil = self._dataset.get_rgb_image_from_scene_name_and_idx(scene_name_1, image_1_idx)
	print "got img1"
        self.img1_knots = self._dataset._knots_info[scene_name_1][image_1_idx]
	print "got img1 knots"
        self.img2_pil = self._dataset.get_rgb_image_from_scene_name_and_idx(scene_name_2, image_2_idx)
        self.img2_knots = self._dataset._knots_info[scene_name_2][image_2_idx]

    def _get_task_images(self):
	self.img1_pil = self._dataset.get_rgb_image('./images/000025_rgb.png')
        img1_mask = self._dataset.get_mask_image('./image_masks/000025_mask.png')
	self.img2_pil = self._dataset.get_rgb_image('./images/000018_rgb.png')
        pixs = correspondence_finder.random_sample_from_masked_image_torch(np.asarray(img1_mask), 25)
	self.img1_knots = list(zip(pixs[0], pixs[1]))
	#self.img1_knots = utils.getDictFromJSONFilename('./images_start/knots_info.json')["0"][0]
	#self.img2_knots = utils.getDictFromJSONFilename('./images_goal/knots_info.json')["0"][0]

    def _compute_descriptors(self, knot_idx):
        """
        Computes the descriptors for image 1 and image 2 for each network
        :return:
        :rtype:
        """
        self.img1 = pil_image_to_cv2(self.img1_pil)
        self.img2 = pil_image_to_cv2(self.img2_pil)
        self.rgb_1_tensor = self._dataset.rgb_image_to_tensor(self.img1_pil)
        self.rgb_2_tensor = self._dataset.rgb_image_to_tensor(self.img2_pil)
        self.img1_gray = cv2.cvtColor(self.img1, cv2.COLOR_RGB2GRAY) / 255.0
        self.img2_gray = cv2.cvtColor(self.img2, cv2.COLOR_RGB2GRAY) / 255.0

        self._res_a = dict()
        self._res_b = dict()
        for network_name, dcn in self._dcn_dict.iteritems():
            self._res_a[network_name] = dcn.forward_single_image_tensor(self.rgb_1_tensor).data.cpu().numpy()
            self._res_b[network_name] = dcn.forward_single_image_tensor(self.rgb_2_tensor).data.cpu().numpy()
	print(self.img1_knots)
        u, v = self.img1_knots[knot_idx]
        source, blended, target, p = self.find_best_match(None, u, v, None, None)
        return (source, blended, target, p)

    def scale_norm_diffs_to_make_heatmap(self, norm_diffs, threshold):
        """
        Scales the norm diffs to make a heatmap. This will be scaled between 0 and 1.
        0 corresponds to a match, 1 to non-match
        :param norm_diffs: The norm diffs
        :type norm_diffs: numpy.array [H,W]
        :return:
        :rtype:
        """


        heatmap = np.copy(norm_diffs)
        greater_than_threshold = np.where(norm_diffs > threshold)
        heatmap = heatmap / threshold * self._config["heatmap_vis_upper_bound"] # linearly scale [0, threshold] to [0, 0.5]
        heatmap[greater_than_threshold] = 1 # greater than threshold is set to 1
        heatmap = heatmap.astype(self.img1_gray.dtype)
        return heatmap


    def find_best_match(self, event,u,v,flags,param):

        """
        For each network, find the best match in the target image to point highlighted
        with reticle in the source image. Displays the result
        :return:
        :rtype:
        """

        img_1_with_reticle = np.copy(self.img1)
        draw_reticle(img_1_with_reticle, u, v, self._reticle_color)
	source = img_1_with_reticle

        alpha = self._config["blend_weight_original_image"]
        beta = 1 - alpha

        img_2_with_reticle = np.copy(self.img2)

        self._res_uv = dict()

        for network_name in self._dcn_dict:
            res_a = self._res_a[network_name]
            res_b = self._res_b[network_name]
            best_match_uv, best_match_diff, norm_diffs = \
                DenseCorrespondenceNetwork.find_best_match((u, v), res_a, res_b)
            print "network_name:", network_name
            self._res_uv[network_name] = dict()
            self._res_uv[network_name]['source'] = res_a[v, u, :].tolist()
            self._res_uv[network_name]['target'] = res_b[v, u, :].tolist()

            print "%s best match diff: %.3f" %(network_name, best_match_diff)

            threshold = self._config["norm_diff_threshold"]
            if network_name in self._config["norm_diff_threshold_dict"]:
                threshold = self._config["norm_diff_threshold_dict"][network_name]

            heatmap = self.scale_norm_diffs_to_make_heatmap(norm_diffs, threshold)

            reticle_color = self._network_reticle_color[network_name]
            draw_reticle(heatmap, best_match_uv[0], best_match_uv[1], reticle_color)
            draw_reticle(img_2_with_reticle, best_match_uv[0], best_match_uv[1], reticle_color)
            blended = cv2.addWeighted(self.img2_gray, alpha, heatmap, beta, 0)
	    
	target = img_2_with_reticle
        return (source, blended, target, [best_match_uv[0], best_match_uv[1]])

    def run(self):
        self._get_task_images()
	pixels = []
#        for i in range(len(utils.getDictFromJSONFilename('./images_start/knots_info.json')["0"][0])):
        for i in range(25):
	    print "computing descriptors"
            source, blended, target, p = self._compute_descriptors(i)
	    pixels.append(p)
	    print "computed descriptors"
            vis = np.concatenate((source, target), axis=1)
	    print "concatenated, writing image"
	    cv2.imwrite("/home/priya/code/data_volume/annotated/%06d.png" % i, vis)
	np.savetxt('pixels_pred.txt', pixels)

if __name__ == "__main__":
    dc_source_dir = utils.getDenseCorrespondenceSourceDir()
    config_file = os.path.join(dc_source_dir, 'config', 'dense_correspondence', 'heatmap_vis', 'heatmap.yaml')
    config = utils.getDictFromYamlFilename(config_file)

    heatmap_vis = HeatmapVisualization(config)
    print "starting heatmap vis"
    heatmap_vis.run()
