import matplotlib.pyplot as plt
import numpy as np
from scipy.spatial import Voronoi, voronoi_plot_2d
import cv2



class Voronoi_Evaulator:
    def __init__(self, points, eval_points):
        self.eval_points = eval_points
        self.vor = Voronoi(points)

    def extract(self, Z):
        for debug_num, ev_points in zip([-10, -5], self.eval_points):
            region_mask = self.vor.point_region[ev_points]
            print(region_mask)
            for i in np.array(self.vor.regions)[region_mask]:
                if -1 in i:
                    print("Containse outside points")
                    continue
                if len(i) == 0:
                    print("Containse outside points")
                    continue
                print(i)
                pts = self.vor.vertices[i]
                pts = (pts * 100).astype(np.int32)
                print(pts)
                mask = np.zeros((Z.shape[0], Z.shape[1]))
                cv2.fillConvexPoly(mask, pts, 1)
                mask = mask > 0  # To convert to Boolean
                Z[mask] = debug_num
        return Z

if __name__ == "__main__":
    np.random.seed(20)
    points = (np.random.rand(100, 2)-0.1) * 2
    voro = Voronoi_Evaulator(points, [[1, 4, 5], [2, 3, 6]])

    x = np.linspace(0, 1, 100)
    y = np.linspace(0, 1, 200)
    X, Y = np.meshgrid(x, y)
    Z = X*2 + Y

    Z = voro.extract(Z)
    plt.imshow(Z)
    plt.show()