Avatar-Toolkit/core/mmd/operators/material.py

# -*- coding: utf-8 -*-
# Copyright 2014 MMD Tools authors
# This file was originally part of the MMD Tools add-on for Blender
# You can find MMD Tools here: https://github.com/MMD-Blender/blender_mmd_tools
# Neoneko has modified this file to work with Avatar Toolkit and may of made changes or improvements.
# MMD Tools is licensed under the terms of the GNU General Public License version 3 (GPLv3) same as Avatar Toolkit.

import bpy
from bpy.props import BoolProperty, StringProperty, FloatProperty
from bpy.types import Operator, Context, Object, Material

from typing import Set, Dict, Any, List, Tuple, Optional, Union, cast

from .. import cycles_converter
from ..core.exceptions import MaterialNotFoundError
from ..core.material import FnMaterial
from ..core.shader import _NodeGroupUtils
from ....core.logging_setup import logger
import traceback


class ConvertMaterialsForCycles(Operator):
    bl_idname = "mmd_tools.convert_materials_for_cycles"
    bl_label = "Convert Materials For Cycles"
    bl_description = "Convert materials of selected objects for Cycles."
    bl_options = {"REGISTER", "UNDO"}

    use_principled: BoolProperty(
        name="Convert to Principled BSDF",
        description="Convert MMD shader nodes to Principled BSDF as well if enabled",
        default=False,
        options={"SKIP_SAVE"},
    )

    clean_nodes: BoolProperty(
        name="Clean Nodes",
        description="Remove redundant nodes as well if enabled. Disable it to keep node data.",
        default=False,
        options={"SKIP_SAVE"},
    )

    @classmethod
    def poll(cls, context: Context) -> bool:
        return next((x for x in context.selected_objects if x.type == "MESH"), None) is not None

    def draw(self, context: Context) -> None:
        layout = self.layout
        layout.prop(self, "use_principled")
        layout.prop(self, "clean_nodes")

    def execute(self, context: Context) -> Set[str]:
        try:
            context.scene.render.engine = "CYCLES"
        except Exception:
            logger.error(f"Failed to change to Cycles render engine: {traceback.format_exc()}")
            self.report({"ERROR"}, " * Failed to change to Cycles render engine.")
            return {"CANCELLED"}
            
        logger.info(f"Converting materials for Cycles with principled={self.use_principled}, clean_nodes={self.clean_nodes}")
        for obj in (x for x in context.selected_objects if x.type == "MESH"):
            logger.debug(f"Converting materials for object: {obj.name}")
            cycles_converter.convertToCyclesShader(obj, use_principled=self.use_principled, clean_nodes=self.clean_nodes)
            
        return {"FINISHED"}


class ConvertMaterials(Operator):
    bl_idname = "mmd_tools.convert_materials"
    bl_label = "Convert Materials"
    bl_description = "Convert materials of selected objects."
    bl_options = {"REGISTER", "UNDO"}

    use_principled: BoolProperty(
        name="Convert to Principled BSDF",
        description="Convert MMD shader nodes to Principled BSDF as well if enabled",
        default=True,
        options={"SKIP_SAVE"},
    )

    clean_nodes: BoolProperty(
        name="Clean Nodes",
        description="Remove redundant nodes as well if enabled. Disable it to keep node data.",
        default=True,
        options={"SKIP_SAVE"},
    )

    subsurface: FloatProperty(
        name="Subsurface",
        default=0.001,
        soft_min=0.000,
        soft_max=1.000,
        precision=3,
        options={"SKIP_SAVE"},
    )

    @classmethod
    def poll(cls, context: Context) -> bool:
        return next((x for x in context.selected_objects if x.type == "MESH"), None) is not None

    def execute(self, context: Context) -> Set[str]:
        logger.info(f"Converting materials with principled={self.use_principled}, clean_nodes={self.clean_nodes}, subsurface={self.subsurface}")
        for obj in context.selected_objects:
            if obj.type != "MESH":
                continue
            logger.debug(f"Converting materials for object: {obj.name}")
            cycles_converter.convertToBlenderShader(obj, use_principled=self.use_principled, clean_nodes=self.clean_nodes, subsurface=self.subsurface)
        return {"FINISHED"}

class ConvertBSDFMaterials(Operator):
    bl_idname = 'mmd_tools.convert_bsdf_materials'
    bl_label = 'Convert Blender Materials'
    bl_description = 'Convert materials of selected objects.'
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context: Context) -> bool:
        return next((x for x in context.selected_objects if x.type == 'MESH'), None) is not None

    def execute(self, context: Context) -> Set[str]:
        logger.info("Converting BSDF materials to MMD shader")
        for obj in context.selected_objects:
            if obj.type != 'MESH':
                continue
            logger.debug(f"Converting BSDF materials for object: {obj.name}")
            cycles_converter.convertToMMDShader(obj)
        return {'FINISHED'}

class _OpenTextureBase:
    """Create a texture for mmd model material."""

    bl_options: Set[str] = {"REGISTER", "UNDO", "INTERNAL"}

    filepath: StringProperty(
        name="File Path",
        description="Filepath used for importing the file",
        maxlen=1024,
        subtype="FILE_PATH",
    )

    use_filter_image: BoolProperty(
        default=True,
        options={"HIDDEN"},
    )

    def invoke(self, context: Context, event: Any) -> Set[str]:
        context.window_manager.fileselect_add(self)
        return {"RUNNING_MODAL"}


class OpenTexture(Operator, _OpenTextureBase):
    bl_idname = "mmd_tools.material_open_texture"
    bl_label = "Open Texture"
    bl_description = "Create main texture of active material"

    def execute(self, context: Context) -> Set[str]:
        mat = context.active_object.active_material
        if not mat:
            logger.error("No active material found")
            return {"CANCELLED"}
            
        logger.info(f"Creating texture for material: {mat.name} from {self.filepath}")
        fnMat = FnMaterial(mat)
        fnMat.create_texture(self.filepath)
        return {"FINISHED"}


class RemoveTexture(Operator):
    """Create a texture for mmd model material."""

    bl_idname = "mmd_tools.material_remove_texture"
    bl_label = "Remove Texture"
    bl_description = "Remove main texture of active material"
    bl_options = {"REGISTER", "UNDO", "INTERNAL"}

    def execute(self, context: Context) -> Set[str]:
        mat = context.active_object.active_material
        if not mat:
            logger.error("No active material found")
            return {"CANCELLED"}
            
        logger.info(f"Removing texture from material: {mat.name}")
        fnMat = FnMaterial(mat)
        fnMat.remove_texture()
        return {"FINISHED"}


class OpenSphereTextureSlot(Operator, _OpenTextureBase):
    """Create a texture for mmd model material."""

    bl_idname = "mmd_tools.material_open_sphere_texture"
    bl_label = "Open Sphere Texture"
    bl_description = "Create sphere texture of active material"

    def execute(self, context: Context) -> Set[str]:
        mat = context.active_object.active_material
        if not mat:
            logger.error("No active material found")
            return {"CANCELLED"}
            
        logger.info(f"Creating sphere texture for material: {mat.name} from {self.filepath}")
        fnMat = FnMaterial(mat)
        fnMat.create_sphere_texture(self.filepath, context.active_object)
        return {"FINISHED"}


class RemoveSphereTexture(Operator):
    """Create a texture for mmd model material."""

    bl_idname = "mmd_tools.material_remove_sphere_texture"
    bl_label = "Remove Sphere Texture"
    bl_description = "Remove sphere texture of active material"
    bl_options = {"REGISTER", "UNDO", "INTERNAL"}

    def execute(self, context: Context) -> Set[str]:
        mat = context.active_object.active_material
        if not mat:
            logger.error("No active material found")
            return {"CANCELLED"}
            
        logger.info(f"Removing sphere texture from material: {mat.name}")
        fnMat = FnMaterial(mat)
        fnMat.remove_sphere_texture()
        return {"FINISHED"}


class MoveMaterialUp(Operator):
    bl_idname = "mmd_tools.move_material_up"
    bl_label = "Move Material Up"
    bl_description = "Moves selected material one slot up"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context: Context) -> bool:
        obj = context.active_object
        valid_mesh = obj and obj.type == "MESH" and obj.mmd_type == "NONE"
        return bool(valid_mesh and obj.active_material_index > 0)

    def execute(self, context: Context) -> Set[str]:
        obj = context.active_object
        current_idx = obj.active_material_index
        prev_index = current_idx - 1
        
        logger.debug(f"Moving material {current_idx} up to position {prev_index} for object {obj.name}")
        try:
            FnMaterial.swap_materials(obj, current_idx, prev_index, reverse=True, swap_slots=True)
        except MaterialNotFoundError:
            logger.error(f"Materials not found for indices {current_idx} and {prev_index}")
            self.report({"ERROR"}, "Materials not found")
            return {"CANCELLED"}
        obj.active_material_index = prev_index

        return {"FINISHED"}


class MoveMaterialDown(Operator):
    bl_idname = "mmd_tools.move_material_down"
    bl_label = "Move Material Down"
    bl_description = "Moves the selected material one slot down"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context: Context) -> bool:
        obj = context.active_object
        valid_mesh = obj and obj.type == "MESH" and obj.mmd_type == "NONE"
        return bool(valid_mesh and obj.active_material_index < len(obj.material_slots) - 1)

    def execute(self, context: Context) -> Set[str]:
        obj = context.active_object
        current_idx = obj.active_material_index
        next_index = current_idx + 1
        
        logger.debug(f"Moving material {current_idx} down to position {next_index} for object {obj.name}")
        try:
            FnMaterial.swap_materials(obj, current_idx, next_index, reverse=True, swap_slots=True)
        except MaterialNotFoundError:
            logger.error(f"Materials not found for indices {current_idx} and {next_index}")
            self.report({"ERROR"}, "Materials not found")
            return {"CANCELLED"}
        obj.active_material_index = next_index
        return {"FINISHED"}


class EdgePreviewSetup(Operator):
    bl_idname = "mmd_tools.edge_preview_setup"
    bl_label = "Edge Preview Setup"
    bl_description = 'Preview toon edge settings of active model using "Solidify" modifier'
    bl_options = {"REGISTER", "UNDO", "INTERNAL"}

    action: bpy.props.EnumProperty(
        name="Action",
        description="Select action",
        items=[
            ("CREATE", "Create", "Create toon edge", 0),
            ("CLEAN", "Clean", "Clear toon edge", 1),
        ],
        default="CREATE",
    )

    def execute(self, context: Context) -> Set[str]:
        from ..core.model import FnModel

        root = FnModel.find_root_object(context.active_object)
        if root is None:
            logger.error("No MMD model root found")
            self.report({"ERROR"}, "Select a MMD model")
            return {"CANCELLED"}

        if self.action == "CLEAN":
            logger.info(f"Cleaning toon edge for model: {root.name}")
            for obj in FnModel.iterate_mesh_objects(root):
                self.__clean_toon_edge(obj)
        else:
            from ..bpyutils import Props

            logger.info(f"Creating toon edge for model: {root.name}")
            scale = 0.2 * getattr(root, Props.empty_display_size)
            counts = sum(self.__create_toon_edge(obj, scale) for obj in FnModel.iterate_mesh_objects(root))
            logger.info(f"Created {counts} toon edge(s)")
            self.report({"INFO"}, "Created %d toon edge(s)" % counts)
        return {"FINISHED"}

    def __clean_toon_edge(self, obj: Object) -> None:
        logger.debug(f"Cleaning toon edge for object: {obj.name}")
        if "mmd_edge_preview" in obj.modifiers:
            obj.modifiers.remove(obj.modifiers["mmd_edge_preview"])

        if "mmd_edge_preview" in obj.vertex_groups:
            obj.vertex_groups.remove(obj.vertex_groups["mmd_edge_preview"])

        FnMaterial.clean_materials(obj, can_remove=lambda m: m and m.name.startswith("mmd_edge."))

    def __create_toon_edge(self, obj: Object, scale: float = 1.0) -> int:
        logger.debug(f"Creating toon edge for object: {obj.name} with scale {scale}")
        self.__clean_toon_edge(obj)
        materials = obj.data.materials
        material_offset = len(materials)
        for m in tuple(materials):
            if m and m.mmd_material.enabled_toon_edge:
                mat_edge = self.__get_edge_material("mmd_edge." + m.name, m.mmd_material.edge_color, materials)
                materials.append(mat_edge)
            elif material_offset > 1:
                mat_edge = self.__get_edge_material("mmd_edge.disabled", (0, 0, 0, 0), materials)
                materials.append(mat_edge)
        if len(materials) > material_offset:
            mod = obj.modifiers.get("mmd_edge_preview", None)
            if mod is None:
                mod = obj.modifiers.new("mmd_edge_preview", "SOLIDIFY")
            mod.material_offset = material_offset
            mod.thickness_vertex_group = 1e-3  # avoid overlapped faces
            mod.use_flip_normals = True
            mod.use_rim = False
            mod.offset = 1
            self.__create_edge_preview_group(obj)
            mod.thickness = scale
            mod.vertex_group = "mmd_edge_preview"
        return len(materials) - material_offset

    def __create_edge_preview_group(self, obj: Object) -> None:
        vertices, materials = obj.data.vertices, obj.data.materials
        weight_map = {i: m.mmd_material.edge_weight for i, m in enumerate(materials) if m}
        scale_map: Dict[int, float] = {}
        vg_scale_index = obj.vertex_groups.find("mmd_edge_scale")
        if vg_scale_index >= 0:
            scale_map = {v.index: g.weight for v in vertices for g in v.groups if g.group == vg_scale_index}
        vg_edge_preview = obj.vertex_groups.new(name="mmd_edge_preview")
        for i, mi in {v: f.material_index for f in reversed(obj.data.polygons) for v in f.vertices}.items():
            weight = scale_map.get(i, 1.0) * weight_map.get(mi, 1.0) * 0.02
            vg_edge_preview.add(index=[i], weight=weight, type="REPLACE")

    def __get_edge_material(self, mat_name: str, edge_color: Tuple[float, float, float, float], materials: List[Material]) -> Material:
        if mat_name in materials:
            return materials[mat_name]
        mat = bpy.data.materials.get(mat_name, None)
        if mat is None:
            mat = bpy.data.materials.new(mat_name)
        mmd_mat = mat.mmd_material
        # note: edge affects ground shadow
        mmd_mat.is_double_sided = mmd_mat.enabled_drop_shadow = False
        mmd_mat.enabled_self_shadow_map = mmd_mat.enabled_self_shadow = False
        # mmd_mat.enabled_self_shadow_map = True # for blender 2.78+ BI viewport only
        mmd_mat.diffuse_color = mmd_mat.specular_color = (0, 0, 0)
        mmd_mat.ambient_color = edge_color[:3]
        mmd_mat.alpha = edge_color[3]
        mmd_mat.edge_color = edge_color
        self.__make_shader(mat)
        return mat

    def __make_shader(self, m: Material) -> None:
        # Note: material.use_nodes is deprecated in Blender 5.0 - materials always use nodes
        nodes, links = m.node_tree.nodes, m.node_tree.links

        node_shader = nodes.get("mmd_edge_preview", None)
        if node_shader is None or not any(s.is_linked for s in node_shader.outputs):
            XPOS, YPOS = 210, 110
            nodes.clear()
            node_shader = nodes.new("ShaderNodeGroup")
            node_shader.name = "mmd_edge_preview"
            node_shader.location = (0, 0)
            node_shader.width = 200
            node_shader.node_tree = self.__get_edge_preview_shader()

            node_out = nodes.new("ShaderNodeOutputMaterial")
            node_out.location = (XPOS * 2, YPOS * 0)
            links.new(node_shader.outputs["Shader"], node_out.inputs["Surface"])

        node_shader.inputs["Color"].default_value = m.mmd_material.edge_color
        node_shader.inputs["Alpha"].default_value = m.mmd_material.edge_color[3]

    def __get_edge_preview_shader(self) -> bpy.types.NodeTree:
        group_name = "MMDEdgePreview"
        shader = bpy.data.node_groups.get(group_name, None) or bpy.data.node_groups.new(name=group_name, type="ShaderNodeTree")
        if len(shader.nodes):
            return shader

        ng = _NodeGroupUtils(shader)

        node_input = ng.new_node("NodeGroupInput", (-5, 0))
        node_output = ng.new_node("NodeGroupOutput", (3, 0))

        ############################################################################
        node_color = ng.new_node("ShaderNodeMixRGB", (-1, -1.5))
        node_color.mute = True

        ng.new_input_socket("Color", node_color.inputs["Color1"])

        ############################################################################
        node_ray = ng.new_node("ShaderNodeLightPath", (-3, 1.5))
        node_geo = ng.new_node("ShaderNodeNewGeometry", (-3, 0))
        node_max = ng.new_math_node("MAXIMUM", (-2, 1.5))
        node_max.mute = True
        node_gt = ng.new_math_node("GREATER_THAN", (-1, 1))
        node_alpha = ng.new_math_node("MULTIPLY", (0, 1))
        node_trans = ng.new_node("ShaderNodeBsdfTransparent", (0, 0))
        node_rgb = ng.new_node("ShaderNodeBackground", (0, -0.5))
        node_mix = ng.new_node("ShaderNodeMixShader", (1, 0.5))

        links = ng.links
        links.new(node_ray.outputs["Is Camera Ray"], node_max.inputs[0])
        links.new(node_ray.outputs["Is Glossy Ray"], node_max.inputs[1])
        links.new(node_max.outputs["Value"], node_gt.inputs[0])
        links.new(node_geo.outputs["Backfacing"], node_gt.inputs[1])
        links.new(node_gt.outputs["Value"], node_alpha.inputs[0])
        links.new(node_alpha.outputs["Value"], node_mix.inputs["Fac"])
        links.new(node_trans.outputs["BSDF"], node_mix.inputs[1])
        links.new(node_rgb.outputs[0], node_mix.inputs[2])
        links.new(node_color.outputs["Color"], node_rgb.inputs["Color"])

        ng.new_input_socket("Alpha", node_alpha.inputs[1])
        ng.new_output_socket("Shader", node_mix.outputs["Shader"])

        return shader