import numpy as np
from scipy.spatial import Voronoi
import cv2


class Image_Wrapper:
    # def __init__(self, img, x_lower, x_res, y_lower, y_res):
    #    self.img = img
    #    self.x_lower = x_lower
    #    self.y_lower = y_lower
    #    self.x_res = x_res
    #    self.y_res = y_res

    #    self.x_upper = self.x_lower + self.img.shape[0]*self.x_res - self.x_res
    #    self.y_upper = self.y_lower + self.img.shape[1]*self.y_res - self.x_res

    def __init__(self, img, fx, fy):
        self.img = img
        self.x_lower = np.min(fx)
        self.y_lower = np.min(fy)
        self.x_upper = np.max(fx)
        self.y_upper = np.max(fy)

        self.x_res = (self.x_upper - self.x_lower) / self.img.shape[0]
        self.y_res = (self.y_upper - self.y_lower) / self.img.shape[0]

    def val2pos(self, x, y):
        x = (x - self.x_lower) / self.x_res
        y = (y - self.y_lower) / self.y_res
        return x, y

    def check_bounds(self, xl, yl, xu, yu):
        if xl > self.img.shape[0]:
            print("xl lim")
            return False
        if yl > self.img.shape[1]:
            print("yl lim")
            return False
        if xu < 0:
            print("xu lim")
            return False
        if yu < 0:
            print("yu lim")
            return False
        return True

    def clean_bounds(self, xl, yl, xu, yu):
        if xl < 0:
            xl = 0
        if yl < 0:
            yl = 0

        if xu > self.img.shape[0]:
            xu = self.img.shape[0]
        if yu > self.img.shape[1]:
            yu = self.img.shape[1]

        return xl, yl, xu, yu

    def ext(self):
        return [self.x_lower, self.x_upper, self.y_lower, self.y_upper]


class Evaluator:
    def extract(self, img: Image_Wrapper):
        all_val = []
        all_idx = []
        for ev_points in self.eval_points:
            temp_val = []
            temp_idx = []
            for num in ev_points:
                if np.sum(self.mask == num) == 0:
                    continue
                temp_val.append(np.mean(img.img[self.mask == num]))
                temp_idx.append(num)
            all_val.append(temp_val)
            all_idx.append(temp_idx)
        all_val.append([np.mean(img.img[self.mask == -1])])
        all_idx.append([-1])
        return all_idx, all_val

    def debug(self, img: Image_Wrapper):
        count = 1
        for ev_points in self.eval_points:
            if count != 1 and False:
                count += 1
                continue
            for num in ev_points:
                img.img[self.mask == num] += 10000 * num
            count += 1
        return img.img

    def get_mask(self):
        return self.mask

    def merge_mask(self):
        new_eval_points = np.arange(len(self.eval_points))
        new_eval = []
        for nc, ev_points in zip(new_eval_points, self.eval_points):
            new_eval.append([nc])
            for num in ev_points:
                self.mask[self.mask == num] = nc
        self.eval_points = new_eval


class Voronoi_Evaluator(Evaluator):
    def __init__(self, list_points):
        points = np.concatenate(list_points, axis=0)
        self.eval_points = []
        start = 0
        for l in list_points:
            stop = l.shape[0]
            self.eval_points.append(np.arange(start, start + stop))
            start += stop
        self.vor = Voronoi(points)

    def generate_mask(self, img: Image_Wrapper):
        self.mask = np.full_like(img.img, -1)

        counter = -1
        region_mask = self.vor.point_region
        for i in np.array(self.vor.regions, dtype=list)[region_mask]:
            counter += 1
            if -1 in i:
                continue
            if len(i) == 0:
                continue
            pts = self.vor.vertices[i]
            pts = np.stack(img.val2pos(
                pts[:, 0], pts[:, 1])).astype(np.int32).T
            if np.any(pts < 0):
                continue
            mask = np.zeros_like(img.img)
            cv2.fillConvexPoly(mask, pts, 1)
            mask = mask > 0  # To convert to Boolean
            self.mask[mask] = counter


class Rect_Evaluator(Evaluator):
    # def __init__(self, points, eval_points):
    #    self.eval_points = eval_points
    #    self.points = points
    #    self.length = 4

    def __init__(self, list_points, length=4):
        self.points = np.concatenate(list_points, axis=0)
        self.eval_points = []
        start = 0
        for l in list_points:
            stop = l.shape[0]
            self.eval_points.append(np.arange(start, start + stop))
            start += stop
            print(start)
        print(self.points.shape)
        self.length = length

    def generate_mask(self, img: Image_Wrapper):
        self.mask = np.full_like(img.img, -1)
        count = 0
        for x, y in self.points:
            x, y = img.val2pos(x, y)
            x_lower = int(x - self.length)
            y_lower = int(y - self.length)
            x_upper = int(x + self.length + 1)
            y_upper = int(y + self.length + 1)
            if img.check_bounds(x_lower, y_lower, x_upper, y_upper):
                x_lower, y_lower, x_upper, y_upper = img.clean_bounds(
                    x_lower, y_lower, x_upper, y_upper)
                self.mask[y_lower:y_upper, x_lower:x_upper] = count
                count += 1
        return all

#
# def main():
#     np.random.seed(10)
#     points = (np.random.rand(100, 2)-0.5) * 2
#     voro = Voronoi_Evaluator(points, [[1],[2]])
#     rect = Rect_Evaluator(points, [[1], [2]])
#     Z = np.ones((1000, 1000))
#     img = Image_Wrapper(Z, -5, .01, -5, .01)
#     voro.extract(img)
#     rect.extract(img)
#
#     plt.scatter(points[[1], 0], points[[1], 1])
#     plt.scatter(points[[2], 0], points[[2], 1])
#     plt.imshow(img.img, extent=img.ext(), origin="lower")
#     #plt.imshow(img.img, origin="lower")
#     plt.show()
#
#
# if __name__ == "__main__":
#     main()