# Copyright (C) 2022 Thomas Hoppe (h0bB1T). All rights reserved.
#
# Unauthorized copying of this file via any medium is strictly prohibited.
# Proprietary and confidential.

import bpy, bmesh

from time import time
from typing import Dict, List, Set
from random import random, seed 

from bpy.types import Operator
from bpy.props import BoolProperty, FloatVectorProperty, IntProperty, StringProperty

from ..properties import Properties
from ..utils.dev import dmp

class UI_OT_support_set_vertex_colors(Operator):
    """
    Depending on parameters, set vertex color of faces using selected strategy.
    """
    bl_idname = 'snw.support_set_vertex_colors'
    bl_label = 'Set vertex colors'
    bl_description = 'Set vertex colors in various ways'
    bl_options = {'REGISTER'}


    mode: StringProperty() # type: ignore
    selection: BoolProperty() # type: ignore
    cols: StringProperty() # type: ignore
    seed: IntProperty() # type: ignore
    col: FloatVectorProperty(subtype='COLOR', size=4) # type: ignore
    final_mode: StringProperty(default='OBJECT') # type: ignore


    def execute(self, context: bpy.types.Context):
        bpy.ops.object.mode_set(mode='OBJECT') 
        obj = context.active_object # type: bpy.types.Object
        mesh = obj.data # type: bpy.types.Mesh

        if self.cols:
            cols = mesh.vertex_colors[self.cols].data
        else:
            cols = mesh.vertex_colors.active.data

        if cols:
            polies = mesh.polygons
            if self.selection:
                polies = [ p for p in polies if p.select ]

            seed(self.seed)
            if self.mode == 'RPV':
                for p in polies:
                    col = (random(), random(), random(), random())
                    for i in range(p.loop_start, p.loop_start + p.loop_total):
                        cols[i].color = col
            elif self.mode == 'CC':
                for p in polies:
                    for i in range(p.loop_start, p.loop_start + p.loop_total):
                        cols[i].color = self.col
            elif self.mode == 'PI':
                # Fast, non-recursive method to find islands ..

                # Create a list of vertices with it's direct connected neighbor verts (by index).
                dmp('Build net list ..')
                bm = bmesh.new()
                bm.from_mesh(mesh)
                nets = {} # type: Dict[int, Set[int]]
                for v in bm.verts:
                    current = set() # type: Set[int] 
                    current.add(v.index)
                    for edge in v.link_edges:
                        ov = edge.other_vert(v)
                        if ov:
                            current.add(ov.index)
                    nets[v.index] = current
                bm.free()

                # Now combust this list into islands.
                dmp('Merge nets ..')
                t = time() + .5
                islands = [] # type: List[Set[int]]
                while nets:
                    # Take an item ..
                    _, current = nets.popitem()
                    # .. and continously merge nets from connected
                    # vertices. Remove them from the net list.
                    while True:
                        if time() > t:
                            dmp(f'Remain {len(nets)}')
                            t = time() + .5

                        # Find at least 2 connected nets. If one is found,
                        # merge to current and remove from available. 
                        new_items = set() # type: Set[int]
                        for i in current:
                            if i in nets:
                                new_items.update(nets.pop(i))
                        
                        # If none has been found, the island is complete.
                        if new_items:
                            current.update(new_items)
                        else:
                            islands.append(current)
                            break

                # Create a vertex index <-> color list, with each island has a randon color.
                dmp('Create colors ..')
                vcolors = {}
                for island in islands:
                    col = (random(), random(), random(), random())
                    for i in island:
                        vcolors[i] = col

                # Apply colors.
                dmp('Apply colors ..')
                for p in polies:
                    for i in range(p.loop_start, p.loop_start + p.loop_total):
                        cols[i].color = vcolors[mesh.loops[i].vertex_index] 

                dmp('Done.')


            bpy.ops.object.mode_set(mode='VERTEX_PAINT')
            bpy.ops.object.mode_set(mode='EDIT')
            bpy.ops.object.mode_set(mode='VERTEX_PAINT')                 

        bpy.ops.object.mode_set(mode=self.final_mode)                         

        return{'FINISHED'}


class UI_OT_support_set_custom_vertex_color(Operator):
    bl_idname = 'snw.support_set_custom_vertex_color'
    bl_label = 'Set custom vertex color'
    bl_description = 'Set custom vertex color to preset'
    bl_options = {'REGISTER'}


    col: FloatVectorProperty(subtype='COLOR', size=4) # type: ignore


    def execute(self, context: bpy.types.Context):
        props = Properties.get()
        props.vertex_color = self.col
        return{'FINISHED'}        
