import sys
import numpy
import random as rnd
import pygame
from scipy.spatial import Voronoi
from coopgame.gameTemplate import GameTemplate, tt
from cooptools.colors import Color
from cooptools.geometry_utils import vector_utils as vec
from coopgame.label_drawing import label_drawing_utils as lbl
from coopgame.pointdrawing import point_draw_utils as pnt
from coopgame.linedrawing import line_draw_utils as line
from coopgame.label_drawing.pyLabel import TextAlignmentType
import numpy as np
from collections import defaultdict
from shapely.geometry import Polygon
from cooptools.coopEnum import CoopEnum, auto
from dataclasses import dataclass, field
from cooptools.coopDataclass import BaseIdentifiedDataClass, UniqueIdentifier
from typing import List, Dict, Set, Tuple, Self, Iterator
from cooptools.decor import timer
import logging
logger = logging.getLogger(__name__)
@timer(logger=logger, log_level=0, time_tracking_class=tt)
def voronoi_polygons(voronoi, diameter):
    """Generate shapely.geometry.Polygon objects corresponding to the
    regions of a scipy.spatial.Voronoi object, in the order of the
    input points. The polygons for the infinite regions are large
    enough that all points within a distance 'diameter' of a Voronoi
    vertex are contained in one of the infinite polygons.

    """
    centroid = voronoi.points.mean(axis=0)

    # Mapping from (input point index, Voronoi point index) to list of
    # unit vectors in the directions of the infinite ridges starting
    # at the Voronoi point and neighbouring the input point.
    ridge_direction = defaultdict(list)
    for (p, q), rv in zip(voronoi.ridge_points, voronoi.ridge_vertices):
        u, v = sorted(rv)
        if u == -1:
            # Infinite ridge starting at ridge point with index v,
            # equidistant from input points with indexes p and q.
            t = voronoi.points[q] - voronoi.points[p] # tangent
            n = np.array([-t[1], t[0]]) / np.linalg.norm(t) # normal
            midpoint = voronoi.points[[p, q]].mean(axis=0)
            direction = np.sign(np.dot(midpoint - centroid, n)) * n
            ridge_direction[p, v].append(direction)
            ridge_direction[q, v].append(direction)

    for i, r in enumerate(voronoi.point_region):
        region = voronoi.regions[r]
        if -1 not in region:
            # Finite region.
            yield Polygon(voronoi.vertices[region])
            continue
        # Infinite region.
        inf = region.index(-1)              # Index of vertex at infinity.
        j = region[(inf - 1) % len(region)] # Index of previous vertex.
        k = region[(inf + 1) % len(region)] # Index of next vertex.
        if j == k:
            # Region has one Voronoi vertex with two ridges.
            dir_j, dir_k = ridge_direction[i, j]
        else:
            # Region has two Voronoi vertices, each with one ridge.
            dir_j, = ridge_direction[i, j]
            dir_k, = ridge_direction[i, k]

        # Length of ridges needed for the extra edge to lie at least
        # 'diameter' away from all Voronoi vertices.
        length = 2 * diameter / np.linalg.norm(dir_j + dir_k)

        # Polygon consists of finite part plus an extra edge.
        finite_part = voronoi.vertices[region[inf + 1:] + region[:inf]]
        extra_edge = [voronoi.vertices[j] + dir_j * length,
                      voronoi.vertices[k] + dir_k * length]
        yield Polygon(np.concatenate((finite_part, extra_edge)))

@timer(logger=logger, log_level=0, time_tracking_class=tt)
def build_voroni_neigbors_map(v: Voronoi):
    vert_to_point_map = {}

    # create a map of the point_region and which point it maps to
    region_point_map = {int(x): ii for ii, x in enumerate(v.point_region)}
    point_region_map = {ii: int(x) for ii, x in enumerate(v.point_region)}

    '''Error control if duplicate points were supplied'''
    seen = set()
    dupes = {int(x) for ii, x in enumerate(v.point_region) if x in seen or seen.add(x)}
    deb = {}
    for x in dupes:
        deb[x] = []
        [deb[x].append(ii) for ii, v in enumerate(v.point_region) if x == v]

    if len(dupes) > 0:
        raise ValueError(f"Input points resulted in shared regions: {deb}")

    """We want to know which points define areas that share verticies. Therefore, we first wish to create a
    map of the verticies and which points they are a member of for fast lookup.
    For each vert in each defined region, add the area that it maps to."""
    for ii, verticies in enumerate([x for x in v.regions]):
        if len(verticies) == 0:
            continue
        for vert in verticies:
            vert_to_point_map.setdefault(vert, set()).add(region_point_map[ii])

    neighbor_map = {}

    for vert, points in vert_to_point_map.items():
        if vert == -1: continue

        for point in points:
            neighbor_map.setdefault(point, set())
            for sub in points:
                if sub != point:
                    neighbor_map[point].add(sub)

    border_map = {}
    for point, neighbors in neighbor_map.items():
        for n in neighbors:
            my_vert_indicies = [x for x in v.regions[point_region_map[point]] if x != -1]
            n_vert_indicies = [x for x in v.regions[point_region_map[n]] if x != -1]

            shared = [x for x in my_vert_indicies if x in n_vert_indicies]

            if len(shared) > 2: raise Exception(f"Neighbors share more than 2 verts: {my_vert_indicies}, {n_vert_indicies}")

            shared_verts = [tuple(float(x) for x in v.vertices[x]) for x in shared]

            border_map[frozenset((point, n))] = {x: shared_verts[ii] for ii, x in enumerate(shared)}

    return neighbor_map, border_map


@timer(logger=logger, log_level=0, time_tracking_class=tt)
def generate_voronoi(points: vec.IterVec):
    """
    Randomly chooses various points within the x and y dimensions of the map.
    Then, uses SciPy to generate a voronoi diagram with them, and returns it.

    :return: SciPy voronoi diagram
    """
    point_arr = numpy.zeros([len(points), 2], numpy.float32)

    for i in range(len(points)):
        point_arr[i][0] = points[i][0]
        point_arr[i][1] = points[i][1]

    v = Voronoi(point_arr)

    # print(point_arr[161], point_arr[749])
    vn_map, border_map = build_voroni_neigbors_map(v)
    polys = [x for x in voronoi_polygons(v, diameter=1000)]

    return v, vn_map, point_arr, polys, border_map
