import random as rnd
from coopgame.gameTemplate import GameTemplate, tt
from cooptools.decor import timer
import logging
from examples.voroni_world.models import *
from typing import List, Dict, Set, Tuple, Self, Iterator

logger = logging.getLogger(__name__)

@timer(logger=logger, log_level=0, time_tracking_class=tt)
def __wave_collapse_terrain(vn_map: Dict[int, Set[int]]) -> Dict[int, TerrainType]:
    ret = {r_id: None for r_id, neighbors in vn_map.items()}
    options = {r_id: [x for x in TerrainType] for r_id in vn_map.keys()}

    region_to_set = None
    while len({k: v for k, v in ret.items() if v is None}) > 0:
        if region_to_set is None:
            # any have only one option
            region_to_set = next((r_id for r_id, opts in options.items()
                           if opts is not None and len(opts) == 1), None)

            # choose random if none only have one option
            if region_to_set is None:
                region_to_set = rnd.choice([r_id for r_id, terr in ret.items() if terr is None])

        terrain_to_set = rnd.randint(0, len(options[region_to_set]) - 1)

        # set the choice
        ret[region_to_set] = options[region_to_set][terrain_to_set]

        # update options
        for neighbor in vn_map[region_to_set]:
            if options[neighbor] is None:
                continue

            # which options are no longer valid?
            incompatibles = [x for x in options[neighbor]
                             if x not in TERRAIN_COMPATIBLITY_MAP[ret[region_to_set]]]

            # remove them
            [options[neighbor].remove(x) for x in incompatibles]

            if len(options[neighbor]) == 0:
                blocks = {k: v for k, v in ret.items() if k in vn_map[neighbor]}
                msg = f"The options have been reduced to zero for region {neighbor}." \
                      + f"\n: Neighbors: {blocks}"
                # raise ValueError(msg)
                logger.warning(msg)
                ret[neighbor] = TerrainType.UNDEFINED
                options[neighbor] = None


        options[region_to_set] = None

        # find the neighbor with the min options remaining
        neighbor_len_options = [len(opts) for rid, opts in options.items()
                           if rid in vn_map[region_to_set] and opts is not None]

        min_options = None
        if len(neighbor_len_options) > 0:
            min_options = min(neighbor_len_options)

        if min_options is not None:
            region_to_set = next(x for x in vn_map[region_to_set]
                                 if options[x] is not None and len(options[x]) == min_options)
        else:
            region_to_set = None

    return ret

@timer(logger=logger, log_level=0, time_tracking_class=tt)
def __count_conflicts(r, vn_map, t_map, t):
    conflicts = 0
    for n in vn_map[r]:
        if t not in TERRAIN_COMPATIBLITY_MAP[t_map[n]]:
            conflicts += 1
    return conflicts

@timer(logger=logger, log_level=0, time_tracking_class=tt)
def __min_conflicts_terrain(vn_map: Dict[int, Set[int]], t_map: Dict[int, TerrainType] = None):
    regions = vn_map.keys()
    rnd.shuffle(list(regions))

    if t_map is None:
        t_map = {r: rnd.choice(list(TerrainType)) for r in vn_map.keys()}

    totConflicts = None
    iteration = 0
    while totConflicts is None or totConflicts > 0:
        totConflicts = 0
        iteration += 1
        replacements = []

        for r in regions:
            conflicts = __count_conflicts(r, vn_map, t_map, t_map[r])
            if conflicts > 0:
                totConflicts += 1
                replacements.clear()
                for t in TerrainType:
                    con = __count_conflicts(r, vn_map, t_map, t)
                    if (con <= conflicts) | (rnd.randint(0, 100) < 5):
                        replacements.append(t)

                t_map[r] = rnd.choice(replacements)
        logger.debug(f"Iteration: {iteration}, Total Conflicts: {totConflicts}")
    return t_map

@timer(logger=logger, log_level=0, time_tracking_class=tt)
def __post_process(vn_map: Dict[int, Set[int]], t_map: Dict[int, TerrainType]):
    regions = vn_map.keys()
    rnd.shuffle(list(regions))

    for r in regions:
        # if terrain is diff from all surrounding, change to a surrounding
        if not any(t_map[r]==t_map[x] for x in vn_map[r]):
            t_map[r] = rnd.choice([t_map[x] for x in vn_map[r]])
    return t_map

def wave_collapse_terrains(vn_map):
    terrains = __wave_collapse_terrain(vn_map)
    terrains = __min_conflicts_terrain(vn_map, terrains)
    terrains = __post_process(vn_map, terrains)

    return terrains