__init__.py

##
# GPL License
##
# Blender Addon | SKkeeper
# Copyright (C) 2020  Johannes Rauch
##
# 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 3 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, see <https://www.gnu.org/licenses/>.

from bpy.props import BoolProperty, CollectionProperty
from bpy.types import Operator, PropertyGroup
import bpy
from enum import Enum
import time
bl_info = {
    "name": "SKkeeper",
    "description": "Applies modifiers and keeps shapekeys",
    "author": "Johannes Rauch",
    "version": (1, 8, 1),
    "blender": (2, 80, 3),
    "location": "View3D > Object",
    "doc_url": "https://github.com/smokejohn/SKkeeper",
    "tracker_url": "https://github.com/smokejohn/SKkeeper/issues",
    "category": "Object",
}


class Mode(Enum):
    ALL = 0
    SUBD = 1
    SELECTED = 2

#####################
# UTILITY FUNCTIONS #
#####################


def log(msg):
    """ prints to console in the following format:
        <SKkeeper Time(HH:MM)> message
    """
    t = time.localtime()
    current_time = time.strftime("%H:%M", t)
    print("<SKkeeper {}> {}".format(current_time, (msg)))


def copy_object(obj, times=1, offset=0):
    """ copies the given object and links it to the main collection"""

    objects = []
    for i in range(0, times):
        copy_obj = obj.copy()
        copy_obj.data = obj.data.copy()
        copy_obj.name = obj.name + "_shapekey_" + str(i+1)
        copy_obj.location.x += offset*(i+1)

        bpy.context.collection.objects.link(copy_obj)
        objects.append(copy_obj)

    return objects


def apply_shapekey(obj, sk_keep):
    """ deletes all shapekeys except the one with the given index """

    shapekeys = obj.data.shape_keys.key_blocks

    # check for valid index
    if sk_keep < 0 or sk_keep > len(shapekeys):
        return

    # remove all other shapekeys
    for i in reversed(range(0, len(shapekeys))):
        if i != sk_keep:
            obj.shape_key_remove(shapekeys[i])

    # remove the chosen one and bake it into the object
    obj.shape_key_remove(shapekeys[0])


def apply_modifiers(obj):
    """ applies all modifiers in order """
    # now uses object.convert to circumvent errors with disabled modifiers

    modifiers = obj.modifiers
    for modifier in modifiers:
        if modifier.type == 'SUBSURF':
            modifier.show_only_control_edges = False

    for o in bpy.context.scene.objects:
        o.select_set(False)

    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj

    bpy.ops.object.convert(target='MESH')


def apply_selected_modifiers(obj, resource_list):
    """ applies only the user selected modifiers to the object"""

    for entry in resource_list:
        if entry.selected:
            log("Applying modifier {} on object {}".format(entry.name, obj.name))
            apply_modifier(obj, entry.name)


def remove_modifiers(obj):
    """ removes all modifiers from the object """

    for i in reversed(range(0, len(obj.modifiers))):
        modifier = obj.modifiers[i]
        obj.modifiers.remove(modifier)


def apply_subdmod(obj):
    """ applies subdivision surface modifier """

    # get subsurface modifier/s
    modifiers = [mod for mod in obj.modifiers if mod.type == 'SUBSURF']

    for o in bpy.context.scene.objects:
        o.select_set(False)
    bpy.context.view_layer.objects.active = obj

    modifiers[0].show_only_control_edges = False
    bpy.ops.object.modifier_apply(modifier=modifiers[0].name)


def apply_modifier(obj, modifier_name):
    """ applies a specific modifier """

    log("Applying chosen modifier")

    modifier = [mod for mod in obj.modifiers if mod.name == modifier_name][0]

    # deselect all
    for o in bpy.context.scene.objects:
        o.select_set(False)

    bpy.context.view_layer.objects.active = obj
    bpy.ops.object.modifier_apply(modifier=modifier.name)


def add_objs_shapekeys(destination, sources):
    """ takes an array of objects and adds them as shapekeys to the destination
    object """
    for o in bpy.context.scene.objects:
        o.select_set(False)

    for src in sources:
        src.select_set(True)

    bpy.context.view_layer.objects.active = destination
    bpy.ops.object.join_shapes()


def common_validation(self):
    """Checks for common user errors for all operators and informs user of
    mistake"""

    # GUARD CLAUSES | USER ERROR

    # check for valid selection
    if not self.obj:
        self.report({'ERROR'}, "No Active object. Please select an object")
        return {'CANCELLED'}

    # check for valid obj-type
    if self.obj.type != 'MESH':
        self.report(
            {'ERROR'}, "Wrong object type. Please select a MESH object")
        return {'CANCELLED'}


def keep_shapekeys(self, mode=Mode.ALL):
    """
    Function which is used by the blender operators to collapse modifier
    stack and keep shapekeys. The given mode parameter will determine which
    execution style will be used. The available modes match those of the
    available blender operators (SUBD, SELECTED, ALL) which collapse only
    subdivision surface, the selected or all modifiers respectively.
    """

    # Check if the object has shapekeys
    has_shapekeys = self.obj.data.shape_keys is not None and len(
        self.obj.data.shape_keys.key_blocks) > 0

    if not has_shapekeys:
        # Apply modifiers directly if no shapekeys to preserve
        log("Object {} has no shapekeys, applying modifiers directly".format(self.obj.name))
        if mode == Mode.ALL:
            apply_modifiers(self.obj)
        elif mode == Mode.SUBD:
            apply_subdmod(self.obj)
        elif mode == Mode.SELECTED:
            apply_selected_modifiers(self.obj, self.resource_list)
        return {'FINISHED'}

    shapekey_names = [
        block.name for block in self.obj.data.shape_keys.key_blocks]

    # create receiving object that will contain all collapsed shapekeys
    receiver = copy_object(self.obj, times=1, offset=0)[0]
    receiver.name = "shapekey_receiver"
    apply_shapekey(receiver, 0)
    if mode == Mode.ALL:
        apply_modifiers(receiver)
    elif mode == Mode.SUBD:
        apply_subdmod(receiver)
    elif mode == Mode.SELECTED:
        apply_selected_modifiers(receiver, self.resource_list)

    num_shapekeys = len(self.obj.data.shape_keys.key_blocks)

    shapekeys_to_process = len(shapekey_names) - 1
    log("Processing {} shapekeys on {} in mode {}".format(
        shapekeys_to_process, self.obj.name, mode))

    # create a copy for each shapekey and transfer it to the receiver one after the other
    # start the loop at 1 so we skip the base shapekey
    for shapekey_index in range(1, num_shapekeys):
        log("Processing shapekey {} with name {}".format(
            shapekey_index, shapekey_names[shapekey_index]))

        # copy of baseobject / shapekey donor
        shapekey_obj = copy_object(self.obj, times=1, offset=0)[0]
        apply_shapekey(shapekey_obj, shapekey_index)
        if mode == Mode.ALL:
            apply_modifiers(shapekey_obj)
        elif mode == Mode.SUBD:
            apply_subdmod(shapekey_obj)
        elif mode == Mode.SELECTED:
            apply_selected_modifiers(shapekey_obj, self.resource_list)

        # add the copy as a shapekey to the receiver
        add_objs_shapekeys(receiver, [shapekey_obj])

        # check if the shapekey could be added
        # due to problematic modifier stack
        help_url = "https://github.com/smokejohn/SKkeeper/blob/master/readme.md#troubleshooting-problems"
        if receiver.data.shape_keys is None:
            error_msg = ("IMPOSSIBLE TO TRANSFER SHAPEKEY BECAUSE OF VERTEX COUNT MISMATCH\n\n"
                         "The processed shapekey {} with name {} cannot be transferred.\n"
                         "The shapekey doesn't have the same vertex count as the base after applying modifiers.\n"
                         "This is most likely due to a problematic modifier in your modifier stack (Decimate, Weld)\n\n"
                         "For help on how to fix problems visit: {}).\n\n"
                         "Press UNDO to return to your previous working state."
                         )
            self.report({'ERROR'}, error_msg.format(
                shapekey_index, shapekey_names[shapekey_index], help_url))
            return {'CANCELLED'}

        # due to problematic shape key
        num_transferred_keys = len(receiver.data.shape_keys.key_blocks) - 1
        if num_transferred_keys != shapekey_index:
            error_msg = ("IMPOSSIBLE TO TRANSFER SHAPEKEY BECAUSE OF VERTEX COUNT MISMATCH\n\n"
                         "The processed shapekey {} with name {} cannot be transferred.\n"
                         "The shapekey doesn't have the same vertex count as the base after applying modifiers.\n"
                         "For help on how to fix problems visit: {}).\n\n"
                         "Press UNDO to return to your previous working state."
                         )
            self.report({'ERROR'}, error_msg.format(
                shapekey_index, shapekey_names[shapekey_index], help_url))
            return {'CANCELLED'}

        # restore the shapekey name
        receiver.data.shape_keys.key_blocks[shapekey_index].name = shapekey_names[shapekey_index]

        # delete the shapekey donor and its mesh datablock (save memory)
        mesh_data = shapekey_obj.data
        bpy.data.objects.remove(shapekey_obj)
        bpy.data.meshes.remove(mesh_data)

    orig_name = self.obj.name
    orig_data = self.obj.data

    # transfer over drivers on shapekeys if they exist
    if orig_data.shape_keys.animation_data is not None:
        receiver.data.shape_keys.animation_data_create()
        for orig_driver in orig_data.shape_keys.animation_data.drivers:
            receiver.data.shape_keys.animation_data.drivers.from_existing(
                src_driver=orig_driver)

        # if the driver has variable targets that refer to the original object
        # we need to retarget them to the new receiver because we delete the
        # original object later
        for fcurve in receiver.data.shape_keys.animation_data.drivers:
            for variable in fcurve.driver.variables:
                for target in variable.targets:
                    if target.id == self.obj:
                        target.id = receiver

    # delete the original and its mesh data
    bpy.data.objects.remove(self.obj)
    bpy.data.meshes.remove(orig_data)

    # rename the receiver
    receiver.name = orig_name

    return {'FINISHED'}


def process_multiple_objects(self, context, mode):
    """
    Process multiple selected objects, applying modifiers based on the
    specified mode
    """
    # Get all selected objects
    selected_objects = [
        obj for obj in context.selected_objects if obj.type == 'MESH']

    if not selected_objects:
        self.report({'ERROR'}, "No mesh objects selected")
        return {'CANCELLED'}

    processed_count = 0
    skipped_count = 0
    error_count = 0

    log("Processing {} objects in mode {}".format(len(selected_objects), mode))

    # Process each selected object
    for obj in selected_objects[:]:
        obj_name = obj.name

        # Make this the active object
        context.view_layer.objects.active = obj
        self.obj = obj

        log("Processing object: {}".format(obj_name))

        # Check for modifiers based on mode
        if mode == Mode.ALL and len(obj.modifiers) == 0:
            log("Skipping {}: No modifiers".format(obj_name))
            skipped_count += 1
            continue

        if mode == Mode.SUBD:
            subd = [mod for mod in obj.modifiers if mod.type == 'SUBSURF']
            if len(subd) == 0:
                log("Skipping {}: No subdivision modifiers".format(obj_name))
                skipped_count += 1
                continue

        # For selected modifiers mode, populate the resource list
        if mode == Mode.SELECTED:
            # Since we're bypassing invoke(), we need to populate
            # the resource_list here
            self.resource_list.clear()
            for mod in obj.modifiers:
                entry = self.resource_list.add()
                entry.name = mod.name
                entry.selected = True  # Select all by default when batch processing

        # Process the object
        try:
            result = keep_shapekeys(self, mode=mode)
            if result == {'CANCELLED'}:
                log("Error processing {}".format(obj_name))
                error_count += 1
            else:
                log("Successfully processed {}".format(obj_name))
                processed_count += 1
        except Exception as e:
            log("Exception processing {}: {}".format(obj_name, str(e)))
            error_count += 1

    # Report results
    if processed_count > 0:
        message = "Applied modifiers on {} objects".format(processed_count)
        if skipped_count > 0:
            message += ", skipped {} objects".format(skipped_count)
        if error_count > 0:
            message += ", encountered errors on {} objects".format(error_count)
        self.report({'INFO'}, message)
    else:
        if error_count > 0:
            self.report(
                {'ERROR'}, "Failed to process any objects, encountered {} errors".format(error_count))
        else:
            self.report({'WARNING'}, "No valid objects processed")

    if processed_count > 0:
        return {'FINISHED'}
    else:
        return {'CANCELLED'}

#####################
# BLENDER OPERATORS #
#####################


class SK_TYPE_Resource(PropertyGroup):
    selected: BoolProperty(name="Selected", default=False)


class SK_OT_apply_mods_SK(Operator):
    """ Applies modifiers and keeps shapekeys """

    bl_idname = "sk.apply_mods_sk"
    bl_label = "Apply All Modifiers (Keep Shapekeys)"
    bl_options = {'REGISTER', 'UNDO'}

    def validate_input(self, obj):
        if common_validation(self) == {'CANCELLED'}:
            return {'CANCELLED'}

        # check for modifiers
        if len(self.obj.modifiers) == 0:
            self.report(
                {'ERROR'}, "The selected object doesn't have any modifiers")
            return {'CANCELLED'}

    def execute(self, context):
        # Check if multiple objects are selected
        if len(context.selected_objects) > 1:
            return process_multiple_objects(self, context, Mode.ALL)

        # Single-object behavior
        self.obj = context.active_object

        # Exit out if the selected object is not valid
        if self.validate_input(self.obj) == {'CANCELLED'}:
            return {'CANCELLED'}

        return keep_shapekeys(self, mode=Mode.ALL)


class SK_OT_apply_subd_SK(Operator):
    """ Applies modifiers and keeps shapekeys """

    bl_idname = "sk.apply_subd_sk"
    bl_label = "Apply All Subdivision (Keep Shapekeys)"
    bl_options = {'REGISTER', 'UNDO'}

    def validate_input(self, obj):
        if common_validation(self) == {'CANCELLED'}:
            return {'CANCELLED'}

        # check for subd modifiers
        subd = [mod for mod in self.obj.modifiers if mod.type == 'SUBSURF']
        if len(subd) == 0:
            self.report(
                {'ERROR'}, "The selected object doesn't have any subdivision surface modifiers")
            return {'CANCELLED'}

    def execute(self, context):
        # Check if multiple objects are selected
        if len(context.selected_objects) > 1:
            return process_multiple_objects(self, context, Mode.SUBD)

        # Single-object behavior
        self.obj = context.active_object

        # Exit out if the selected object is not valid
        if self.validate_input(self.obj) == {'CANCELLED'}:
            return {'CANCELLED'}

        return keep_shapekeys(self, mode=Mode.SUBD)


class SK_OT_apply_mods_choice_SK(Operator):
    """ Applies modifiers and keeps shapekeys """

    bl_idname = "sk.apply_mods_choice_sk"
    bl_label = "Apply Chosen Modifiers (Keep Shapekeys)"
    bl_options = {'REGISTER', 'UNDO'}

    resource_list: CollectionProperty(
        name="Modifier List", type=SK_TYPE_Resource)

    def invoke(self, context, event):
        self.obj = context.active_object

        # Handle multiple selection case first
        if len(context.selected_objects) > 1:
            # Skip common_validation for multiple objects as it will be handled
            # per-object
            # Ask user if they want to process all selected objects
            # with all modifiers selected by default
            return context.window_manager.invoke_props_dialog(self, width=350)

        # Single object validation
        if common_validation(self) == {'CANCELLED'}:
            return {'CANCELLED'}

        # check for modifiers
        if len(self.obj.modifiers) == 0:
            self.report(
                {'ERROR'}, "The selected object doesn't have any modifiers")
            return {'CANCELLED'}

        # For single object, regular behavior
        # populate the resource_list
        self.resource_list.clear()
        for mod in self.obj.modifiers:
            entry = self.resource_list.add()
            entry.name = mod.name

        # display floating gui
        return context.window_manager.invoke_props_dialog(self, width=350)

    def execute(self, context):
        # Check for multiple selected objects
        if len(context.selected_objects) > 1:
            return process_multiple_objects(self, context, Mode.SELECTED)

        # Single-object behavior
        return keep_shapekeys(self, mode=Mode.SELECTED)

    def draw(self, context):
        """ Draws the resource selection GUI """

        layout = self.layout

        # If multiple objects are selected, show a different UI
        if len(context.selected_objects) > 1:
            layout.label(text="Process {} selected objects?".format(
                len(context.selected_objects)))
            layout.label(text="All modifiers will be applied for each object")
            return

        # Regular UI for single object
        obj = context.active_object
        layout.label(text="Select modifiers to apply on {}".format(obj.name))
        col = layout.column(align=True)
        for entry in self.resource_list:
            row = col.row()
            row.prop(entry, 'selected', text=entry.name)


classes = (
    SK_TYPE_Resource,
    SK_OT_apply_mods_SK,
    SK_OT_apply_subd_SK,
    SK_OT_apply_mods_choice_SK
)


def modifier_panel(self, context):
    layout = self.layout
    layout.separator()
    layout.operator("sk.apply_mods_sk")
    layout.operator("sk.apply_subd_sk")
    layout.operator("sk.apply_mods_choice_sk")


def register():
    from bpy.utils import register_class
    for cls in classes:
        register_class(cls)

    log("Registered SKKeeper addon")
    bpy.types.VIEW3D_MT_object.append(modifier_panel)


def unregister():
    from bpy.utils import unregister_class
    for cls in classes:
        unregister_class(cls)

    log("Unregistered SKKeeper addon")
    bpy.types.VIEW3D_MT_object.remove(modifier_panel)