import numpy as np
import tqdm
# from scipy.spatial import Voronoi
# import cv2
from cache import persist_to_file, timeit
from spin_image import FFT
from tools import clean_bounds_offset
from abc import abstractmethod, ABC
import logging
logger = logging.getLogger("fft")


class Evaluator(ABC):
    eval_points: list[np.ndarray]

    def extract(self, img: FFT):
        all_val = []
        all_idx = []
        for ev_points in self.eval_points:
            logger.info(f"Extracting points: {ev_points}")
            temp_val = []
            temp_idx = []
            for num in ev_points:
                if np.sum(self.mask == num) == 0:
                    continue
                temp_val.append(np.sum(img.intens[self.mask == num]))
                temp_idx.append(num)
            all_val.append(temp_val)
            all_idx.append(temp_idx)
        all_val.append([np.sum(img.intens[self.mask == -1])])
        all_idx.append([-1])
        return all_idx, all_val

    def debug(self, img: FFT):
        for count, ev_points in enumerate(self.eval_points, start=0):
            for num in ev_points:
                img.intens[self.mask == num] += count
        return img

    def get_mask(self):
        return self.mask

    @timeit
    def generate_mask(self, img: FFT, merge=False):
        hash = str(img.intens.shape)
        self.mask = self.gen_mask_helper(img)
        if merge:
            self.mask = self.merge_mask_helper()
            self.eval_points = [[a] for a in np.arange(len(self.eval_points))]

    @abstractmethod
    def gen_mask_helper(self, img: FFT, hash=str):
        pass

    @abstractmethod
    def merge_mask_helper(self, hash: str):
        pass

    def purge(self, img: FFT):
        img.intens[self.mask != -1] = 0


class Rect_Evaluator(Evaluator):
    # def __init__(self, points, eval_points):
    #    self.eval_points = eval_points
    #    self.points = points
    #    self.length = 4
    def __init__(self, spots: list[tuple[np.ndarray, np.ndarray]], length: int = 6):
        self.spots = spots
        self.length = length
        self.eval_points = []
        start = 0
        for sp in spots:
            self.eval_points.append(np.arange(start, start+sp[0].size))
            start += sp[0].size

    def merge_mask_helper(self):
        new_eval_points = np.arange(len(self.eval_points))
        mask = self.mask
        for nc, ev_points in zip(new_eval_points, self.eval_points):
            maske_low = np.min(ev_points) >= self.mask
            maske_high = np.max(ev_points) <= self.mask
            mask[np.logical_and(maske_high, maske_low)] = nc
        return mask

    def gen_mask_helper(self, img: FFT):
        mask = np.full_like(img.intens, -1)
        count = 0
        for spot_group in self.spots:
            logger.debug(f"Spot: {spot_group}")
            (x, y) = spot_group
            x, y = img.val2pos(x, y)
            for x_p, y_p in zip(x, y):
                xl, yl, xu, yu = clean_bounds_offset(
                    x_p, y_p, self.length, img.intens.shape)
                mask[xl:xu, yl:yu] = count
                count += 1
        return mask

#
# 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()
# 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)
#
#    @persist_to_file("cache_merge_voro")
#    def merge_mask_helper(self):
#        new_eval_points = np.arange(len(self.eval_points))
#        mask = self.mask
#        for nc, ev_points in zip(new_eval_points, self.eval_points):
#            for num in ev_points:
#                mask[self.mask == num] = nc
#        return mask
#
#    @persist_to_file("cache_voro")
#    def gen_mask_helper(self, img: Image_Wrapper):
#        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_2 = np.zeros_like(img.img)
#            cv2.fillConvexPoly(mask_2, pts, 1)
#            mask_2 = mask_2 > 0  # To convert to Boolean
#            mask[mask_2] = counter
#        return mask