# ##### BEGIN GPL LICENSE BLOCK #####
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####

import collections
import numpy as np

import bpy
from bpy.props import StringProperty, EnumProperty, BoolProperty, FloatVectorProperty

from sverchok.node_tree import SverchCustomTreeNode
from sverchok.data_structure import (updateNode, repeat_last, numpy_full_list)

# pylint: disable=E1101
# pylint: disable=W0613

# modified (kosvor, zeffii, vicdoval) 2021

vcol_options = [(k, k, '', i) for i, k in enumerate(["RGB", "RGBA"])]


def set_vertices(loop_count, obj, index_socket, indices, input_colors, colors):
    vertex_index = np.zeros(loop_count, dtype=int)
    loops = obj.data.loops
    loops.foreach_get("vertex_index", vertex_index)
    if index_socket.is_linked:
        idx_lookup = collections.defaultdict(list)
        for idx, v_idx in enumerate(vertex_index):
            idx_lookup[v_idx].append(idx)
        for idx, col in zip(indices, input_colors):
            colors[idx_lookup[idx]] = col
    else:
        if len(obj.data.vertices) > len(input_colors):
            colors[:] = numpy_full_list(input_colors, len(obj.data.vertices))[vertex_index]
        else:
            if isinstance(input_colors, np.ndarray):
                colors[:] = input_colors[vertex_index]
            else:
                colors[:] = np.array(input_colors)[vertex_index]

def set_polygons(polygon_count, obj, index_socket, indices, input_colors, colors):
    p_start = np.empty(polygon_count, dtype=int)
    p_total = np.empty(polygon_count, dtype=int)
    obj.data.polygons.foreach_get("loop_start", p_start)
    obj.data.polygons.foreach_get("loop_total", p_total)
    if index_socket.is_linked:
        for idx, color in zip(indices, input_colors):
            start_slice = p_start[idx]
            stop_slice = start_slice + p_total[idx]
            colors[start_slice:stop_slice] = color
    else:
        for start_slice, total_slice, color in zip(p_start, p_total, repeat_last(input_colors)):
            colors[start_slice: start_slice + total_slice] = color


def set_loops(loop_count, obj, index_socket, indices, input_colors, colors):
    if index_socket.is_linked:
        for idx, color in zip(indices, input_colors):
            colors[idx] = color
    else:
        if len(input_colors) < loop_count:
            colors[:] = numpy_full_list(input_colors, loop_count)
        elif len(input_colors) >= loop_count:
            colors[:] = input_colors[:loop_count]


class SvVertexColorNodeMK3(bpy.types.Node, SverchCustomTreeNode):
    '''
    Triggers: vcol vertex colors
    Tooltip: Set the vertex colors of the named Layer
    '''

    bl_idname = 'SvVertexColorNodeMK3'
    bl_label = 'Vertex color mk3'
    bl_icon = 'COLOR'

    modes = [
        ("vertices", "Vert", "Vcol - color per vertex", 1),
        ("polygons", "Face", "Pcol - color per face", 2),
        ("loops", "Loop", "Color per loop", 3)]

    mode: EnumProperty(items=modes, default='vertices', update=updateNode)
    object_ref: StringProperty(default='', update=updateNode)
    vertex_color: StringProperty(default='SvCol', update=updateNode)
    clear: BoolProperty(name='clear c', default=True, update=updateNode)

    clear_c: FloatVectorProperty(
        name='cl_color', subtype='COLOR', min=0, max=1, size=4,
        default=(0, 0, 0, 1), update=updateNode)

    use_active: BoolProperty(
        default=False, name="Use active layer",
        description="Use active vertex layer")

    unit_color: FloatVectorProperty(
        name='', default=(.3, .3, .2, 1.0),
        size=4, min=0.0, max=1.0, subtype='COLOR', update=updateNode)

    vcol_size: EnumProperty(
        items=vcol_options,
        name='Num Color Components',
        description="3 = rgb, 4 = rgba.  older versions of Blender only support 3 components",
        default="RGBA", update=updateNode)


    def draw_buttons(self, context, layout):
        layout.prop(self, 'use_active')
        layout.prop(self, 'vertex_color')
        layout.prop(self, "mode", expand=True)

    def draw_buttons_ext(self, context, layout):
        row = layout.row(align=True)
        row.prop(self, "clear", text="clear unindexed")
        row.prop(self, "clear_c", text="")
        layout.prop(self, "vcol_size", expand=True)

    def sv_init(self, context):
        inew = self.inputs.new
        inew('SvObjectSocket', 'Object')
        inew('SvStringsSocket', "Index")
        color_socket = inew('SvColorSocket', "Color")
        color_socket.prop_name = 'unit_color'


    def get_vertex_color_layer(self, obj):
        vcols = obj.data.vertex_colors
        vertex_color = None
        if vcols:
            active = self.use_active and vcols.active
            vertex_color = vcols.active if active else vcols.get(self.vertex_color)
        return vertex_color or vcols.new(name=self.vertex_color)

    def process(self):

        color_socket = self.inputs["Color"]
        index_socket = self.inputs["Index"]

        # self upgrade, shall only be called once if encountered.
        if color_socket.bl_idname == 'SvStringsSocket':
            color_socket.replace_socket('SvColorSocket')

        objects = self.inputs["Object"].sv_get(deepcopy=False)
        color_data = color_socket.sv_get(deepcopy=False, default=[None])
        index_data = index_socket.sv_get(deepcopy=False, default=[None])

        num_components = int(len(self.vcol_size))

        for obj, input_colors, indices in zip(objects, repeat_last(color_data), repeat_last(index_data)):
            if not input_colors:
                continue

            loops = obj.data.loops
            loop_count = len(loops)

            vertex_color = self.get_vertex_color_layer(obj)
            colors = np.empty(loop_count * num_components, dtype=np.float32)

            # we have index and colors, set colors of incoming index
            # first get all colors so we can write to them
            if self.clear:
                colors.shape = (loop_count, num_components)
                colors[:] = self.clear_c[:num_components]
            elif index_socket.is_linked:
                vertex_color.data.foreach_get("color", colors)

            colors.shape = (loop_count, num_components)
            standard_params = obj, index_socket, indices, input_colors, colors

            if self.mode == "vertices":
                set_vertices(loop_count, *standard_params)
            elif self.mode == "polygons":
                set_polygons(len(obj.data.polygons), *standard_params)
            elif self.mode == "loops":
                set_loops(loop_count, *standard_params)

            # write out data
            colors.shape = (loop_count * num_components,)
            vertex_color.data.foreach_set("color", colors)
            obj.data.update()


def register():
    bpy.utils.register_class(SvVertexColorNodeMK3)


def unregister():
    bpy.utils.unregister_class(SvVertexColorNodeMK3)
