# Setup
from __future__ import print_function
import os
import numpy as np
import json
import math
import sys
from scipy.spatial import distance
from scipy import spatial
import cv2
import argparse
import utils
from bounding_box import BoundingBox
import pyximport
pyximport.install()
import cv_wrap_module
datadir, imgdir, workdir, outdir = os.getcwd(), os.getcwd(), os.getcwd(), os.getcwd()

def get_image_top_left(ts, tile_index):
    xloc = ts[tile_index]["bbox"][0]
    yloc = ts[tile_index]["bbox"][2]
    return [xloc, yloc]

def get_mfov_centers_from_json(indexed_ts):
    mfov_centers = {}
    for mfov in indexed_ts.keys():
        tile_bboxes = []
        mfov_tiles = indexed_ts[mfov].values()
        tile_bboxes = zip(*[tile["bbox"] for tile in mfov_tiles])
        min_x = min(tile_bboxes[0])
        max_x = max(tile_bboxes[1])
        min_y = min(tile_bboxes[2])
        max_y = max(tile_bboxes[3])
        # center = [(min_x + min_y) / 2.0, (min_y + max_y) / 2.0], but w/o overflow
        mfov_centers[mfov] = np.array([(min_x / 2.0 + max_x / 2.0, min_y / 2.0 + max_y / 2.0)])
    return mfov_centers

    

def get_best_transformations(pre_mfov_matches):
    """Returns a dictionary that maps an mfov number to a matrix that best describes the transformation to the other section.
       As not all mfov's may be matched, some mfovs will be missing from the dictionary"""
    transforms = {}
    for m in pre_mfov_matches["matches"]:
        transforms[m["mfov1"]] = m["transformation"]["matrix"]
    return transforms

def find_best_mfov_transformation(mfov, best_transformations, mfov_centers):
    """Returns a matrix that represnets the best transformation for a given mfov to the other section"""
    if mfov in best_transformations.keys():
        return best_transformations[mfov]
    # Need to find a more intelligent way to do this, but this suffices for now
    # Uses transformation of another mfov (should be changed to the closest, unvisited mfov)
    mfov_center = mfov_centers[mfov]
    trans_keys = best_transformations.keys()
    closest_mfov_index = np.argmin([distance.euclidean(mfov_center, mfov_centers[mfovk]) for mfovk in trans_keys])
    # ** Should we add this transformation? maybe not, because we don't want to get a different result when a few
    # ** missing mfov matches occur and the "best transformation" can change when the centers are updated
    return best_transformations[trans_keys[closest_mfov_index]]
    
    

def get_tile_centers_from_json(ts):
    tiles_centers = []
    for tile in ts:
        center_x = (tile["bbox"][0] + tile["bbox"][1]) / 2.0
        center_y = (tile["bbox"][2] + tile["bbox"][3]) / 2.0
        tiles_centers.append(np.array([center_x, center_y]))
    return tiles_centers


def get_closest_index_to_point(point, centerstree):
    distanc, closest_index = centerstree.query(point)
    return closest_index
    

def get_img_matches(ts1, tile_centers1, ts2, tile_centers2, best_transformations, mfov_centers1):
    """For each tile in section1 find the closest 10 tiles in the second image (after applying the preliminary transformation)"""
    img_matches = []

    # TODO - Build a kd-tree of section2 centers

    # Iterate over the first section tiles, and get the approximated location on section 2 (after transformation)
    # and then get the closest 10 tiles (according to their centers) to that location
    for ind1, tile1 in enumerate(ts1):
        center1 = tile_centers1[ind1]
        trans_matrix = find_best_mfov_transformation(tile1["mfov"], best_transformations, mfov_centers1)
        expected_new_center = np.array(np.dot(trans_matrix, np.append(center1, [1]))[0:2])
        distances_to_sec2_mfovs = [np.linalg.norm(expected_new_center - tile_center) for tile_center in tile_centers2]
        closest_indices = np.array(distances_to_sec2_mfovs).argsort()[0:10]
        img_matches.append(closest_indices)
        
    return img_matches

def is_point_in_img(tile_ts, point):
    """Returns True if the given point lies inside the image as denoted by the given tile_tilespec"""
    # TODO - instead of checking inside the bbox, need to check inside the polygon after transformation
    img_bbox = tile_ts["bbox"]

    if point[0] > img_bbox[0] and point[1] > img_bbox[2] and \
       point[0] < img_bbox[1] and point[1] < img_bbox[3]:
        return True
    return False


def get_images_from_indices_and_point(ts, img_indices, point):
    """Returns all the images at the given img_indices that are overlapping with the given point"""
    img_arr = []
    for img_ind in img_indices:
        if is_point_in_img(ts[img_ind], point):