# -*- coding: utf-8 -*-
import maya.cmds as cmds


WINDOW_NAME = "furcraeaSmoothNormalTool"


def get_selected_mesh_components():
    """
    現在の選択を取得します。
    オブジェクト選択の場合は、そのメッシュ全体を返します。
    コンポーネント選択の場合は、そのまま返します。
    """
    selection = cmds.ls(selection=True, long=True, flatten=True) or []

    if not selection:
        cmds.warning("メッシュ、フェース、エッジ、頂点を選択してください。")
        return []

    result = []

    for item in selection:
        # コンポーネント選択
        if "." in item:
            result.append(item)
            continue

        # TransformからShapeを取得
        if cmds.nodeType(item) == "transform":
            shapes = cmds.listRelatives(
                item,
                shapes=True,
                noIntermediate=True,
                fullPath=True
            ) or []

            mesh_shapes = [
                shape for shape in shapes
                if cmds.nodeType(shape) == "mesh"
            ]

            if mesh_shapes:
                result.append(item)

        # Shapeが直接選択されている場合
        elif cmds.nodeType(item) == "mesh":
            parents = cmds.listRelatives(
                item,
                parent=True,
                fullPath=True
            ) or []

            if parents:
                result.append(parents[0])

    if not result:
        cmds.warning("選択内にポリゴンメッシュがありません。")

    return result


def unlock_normals():
    """選択したメッシュまたはコンポーネントの法線ロックを解除します。"""
    targets = get_selected_mesh_components()
    if not targets:
        return

    cmds.undoInfo(openChunk=True)

    try:
        for target in targets:
            try:
                cmds.polyNormalPerVertex(
                    target,
                    unFreezeNormal=True
                )
            except RuntimeError:
                pass

        cmds.inViewMessage(
            amg="<hl>法線ロックを解除しました</hl>",
            pos="midCenter",
            fade=True
        )

    finally:
        cmds.undoInfo(closeChunk=True)


def soften_normals(angle=None):
    """
    指定角度以下のエッジをソフト化します。

    angle:
        180の場合、基本的にすべてのエッジを滑らかにします。
        30～60程度なら、鋭い角を残しやすくなります。
    """
    targets = get_selected_mesh_components()
    if not targets:
        return

    if angle is None:
        angle = cmds.floatSliderGrp(
            "smoothNormalAngleSlider",
            query=True,
            value=True
        )

    cmds.undoInfo(openChunk=True)

    try:
        for target in targets:
            # ロックされた法線がある場合に備えて解除
            try:
                cmds.polyNormalPerVertex(
                    target,
                    unFreezeNormal=True
                )
            except RuntimeError:
                pass

            cmds.polySoftEdge(
                target,
                angle=angle,
                constructionHistory=True
            )

        cmds.inViewMessage(
            amg="<hl>法線を滑らかにしました：</hl> {}°".format(angle),
            pos="midCenter",
            fade=True
        )

    finally:
        cmds.undoInfo(closeChunk=True)


def soften_all():
    """選択したメッシュまたはコンポーネントを完全にソフト化します。"""
    soften_normals(180.0)


def harden_normals():
    """選択した範囲をハードエッジ化します。"""
    targets = get_selected_mesh_components()
    if not targets:
        return

    cmds.undoInfo(openChunk=True)

    try:
        for target in targets:
            try:
                cmds.polyNormalPerVertex(
                    target,
                    unFreezeNormal=True
                )
            except RuntimeError:
                pass

            cmds.polySoftEdge(
                target,
                angle=0,
                constructionHistory=True
            )

        cmds.inViewMessage(
            amg="<hl>ハードエッジ化しました</hl>",
            pos="midCenter",
            fade=True
        )

    finally:
        cmds.undoInfo(closeChunk=True)


def average_vertex_normals():
    """
    選択頂点の法線を平均化します。
    頂点以外を選択した場合は頂点へ変換します。
    """
    selection = cmds.ls(selection=True, long=True, flatten=True) or []

    if not selection:
        cmds.warning("平均化するメッシュまたはコンポーネントを選択してください。")
        return

    vertices = cmds.polyListComponentConversion(
        selection,
        toVertex=True
    )

    vertices = cmds.ls(
        vertices,
        flatten=True,
        long=True
    ) or []

    if not vertices:
        cmds.warning("頂点を取得できませんでした。")
        return

    cmds.undoInfo(openChunk=True)

    try:
        try:
            cmds.polyNormalPerVertex(
                vertices,
                unFreezeNormal=True
            )
        except RuntimeError:
            pass

        cmds.polyAverageNormal(
            vertices,
            distance=0.0,
            replaceNormalXYZ=(0, 0, 0)
        )

        cmds.inViewMessage(
            amg="<hl>頂点法線を平均化しました</hl>",
            pos="midCenter",
            fade=True
        )

    finally:
        cmds.undoInfo(closeChunk=True)


def conform_normals():
    """
    法線の向きを統一します。
    裏返ったフェースが混在している場合に使用します。
    """
    targets = get_selected_mesh_components()
    if not targets:
        return

    cmds.undoInfo(openChunk=True)

    try:
        for target in targets:
            cmds.polyNormal(
                target,
                normalMode=2,
                userNormalMode=0,
                constructionHistory=True
            )

        cmds.inViewMessage(
            amg="<hl>法線方向を統一しました</hl>",
            pos="midCenter",
            fade=True
        )

    finally:
        cmds.undoInfo(closeChunk=True)


def reverse_normals():
    """選択したフェースまたはメッシュの法線を反転します。"""
    targets = get_selected_mesh_components()
    if not targets:
        return

    cmds.undoInfo(openChunk=True)

    try:
        for target in targets:
            cmds.polyNormal(
                target,
                normalMode=0,
                userNormalMode=0,
                constructionHistory=True
            )

        cmds.inViewMessage(
            amg="<hl>法線を反転しました</hl>",
            pos="midCenter",
            fade=True
        )

    finally:
        cmds.undoInfo(closeChunk=True)


def delete_normal_history():
    """
    選択メッシュの履歴を削除します。
    法線処理以外の履歴も削除されるため注意してください。
    """
    targets = get_selected_mesh_components()
    if not targets:
        return

    transforms = set()

    for target in targets:
        object_name = target.split(".")[0]

        if cmds.nodeType(object_name) == "mesh":
            parents = cmds.listRelatives(
                object_name,
                parent=True,
                fullPath=True
            ) or []

            if parents:
                transforms.add(parents[0])
        else:
            transforms.add(object_name)

    if transforms:
        cmds.delete(list(transforms), constructionHistory=True)

        cmds.inViewMessage(
            amg="<hl>ヒストリを削除しました</hl>",
            pos="midCenter",
            fade=True
        )


def show_smooth_normal_tool():
    """UIを表示します。"""
    if cmds.window(WINDOW_NAME, exists=True):
        cmds.deleteUI(WINDOW_NAME)

    window = cmds.window(
        WINDOW_NAME,
        title="Smooth Normal Tool",
        sizeable=False,
        widthHeight=(340, 390)
    )

    cmds.columnLayout(
        adjustableColumn=True,
        rowSpacing=8
    )

    cmds.separator(height=8, style="none")

    cmds.text(
        label="メッシュ・フェース・エッジ・頂点に対応",
        align="center"
    )

    cmds.separator(height=8, style="in")

    cmds.floatSliderGrp(
        "smoothNormalAngleSlider",
        label="Soft Angle",
        field=True,
        minValue=0.0,
        maxValue=180.0,
        fieldMinValue=0.0,
        fieldMaxValue=180.0,
        value=60.0,
        columnWidth3=(80, 60, 180)
    )

    cmds.button(
        label="指定角度で滑らかにする",
        height=36,
        command=lambda *_: soften_normals()
    )

    cmds.button(
        label="すべて滑らかにする  180°",
        height=32,
        command=lambda *_: soften_all()
    )

    cmds.button(
        label="ハードエッジにする  0°",
        height=32,
        command=lambda *_: harden_normals()
    )

    cmds.separator(height=10, style="in")

    cmds.button(
        label="頂点法線を平均化",
        height=32,
        command=lambda *_: average_vertex_normals()
    )

    cmds.button(
        label="法線ロックを解除",
        height=32,
        command=lambda *_: unlock_normals()
    )

    cmds.separator(height=10, style="in")

    cmds.rowLayout(
        numberOfColumns=2,
        adjustableColumn=1,
        columnWidth2=(165, 165)
    )

    cmds.button(
        label="法線方向を統一",
        height=30,
        command=lambda *_: conform_normals()
    )

    cmds.button(
        label="法線を反転",
        height=30,
        command=lambda *_: reverse_normals()
    )

    cmds.setParent("..")

    cmds.separator(height=10, style="in")

    cmds.button(
        label="ヒストリを削除",
        height=30,
        command=lambda *_: delete_normal_history()
    )

    cmds.separator(height=8, style="none")

    cmds.showWindow(window)


show_smooth_normal_tool()

# -*- coding: utf-8 -*-

import maya.cmds as cmds
import maya.api.OpenMaya as om


WINDOW_NAME = "furcraeaSeamNormalTool"


def get_selected_mesh_dag_paths():
    """
    選択されたTransformまたはMesh Shapeから、
    MeshのMDagPathを取得する。
    """
    selection = om.MGlobal.getActiveSelectionList()
    mesh_paths = []

    for index in range(selection.length()):
        try:
            dag_path = selection.getDagPath(index)
        except RuntimeError:
            continue

        if dag_path.node().hasFn(om.MFn.kTransform):
            try:
                dag_path.extendToShape()
            except RuntimeError:
                continue

        if dag_path.node().hasFn(om.MFn.kMesh):
            mesh_paths.append(dag_path)

    return mesh_paths


def unlock_mesh_normals(dag_path):
    """
    メッシュ全体の頂点法線ロックを解除する。
    """
    mesh_name = dag_path.fullPathName()

    try:
        cmds.polyNormalPerVertex(
            mesh_name,
            unFreezeNormal=True
        )
    except RuntimeError:
        pass


def get_vertex_data(dag_path):
    """
    各頂点について以下を取得する。

    - 頂点番号
    - ワールド座標
    - 頂点法線
    """
    mesh_fn = om.MFnMesh(dag_path)

    points = mesh_fn.getPoints(om.MSpace.kWorld)
    normals = mesh_fn.getVertexNormals(
        True,
        om.MSpace.kWorld
    )

    result = []

    for vertex_id, point in enumerate(points):
        normal = normals[vertex_id]

        result.append({
            "dag_path": dag_path,
            "mesh_fn": mesh_fn,
            "vertex_id": vertex_id,
            "position": om.MVector(point.x, point.y, point.z),
            "normal": om.MVector(normal.x, normal.y, normal.z)
        })

    return result


def build_spatial_hash(vertex_data, tolerance):
    """
    頂点座標をグリッド化して、近接頂点を高速に検索する。
    """
    spatial_hash = {}

    for data in vertex_data:
        position = data["position"]

        key = (
            int(round(position.x / tolerance)),
            int(round(position.y / tolerance)),
            int(round(position.z / tolerance))
        )

        spatial_hash.setdefault(key, []).append(data)

    return spatial_hash


def get_neighbor_keys(position, tolerance):
    """
    誤差境界をまたいだ頂点も検索できるように、
    周囲27セルを返す。
    """
    center_x = int(round(position.x / tolerance))
    center_y = int(round(position.y / tolerance))
    center_z = int(round(position.z / tolerance))

    for offset_x in (-1, 0, 1):
        for offset_y in (-1, 0, 1):
            for offset_z in (-1, 0, 1):
                yield (
                    center_x + offset_x,
                    center_y + offset_y,
                    center_z + offset_z
                )


def average_seam_normals(tolerance=0.001, border_only=True):
    """
    選択した複数メッシュ間で、
    同じ位置にある頂点法線を平均化する。

    tolerance:
        同じ位置と判定する距離。

    border_only:
        Trueの場合は境界頂点のみを対象にする。
    """
    dag_paths = get_selected_mesh_dag_paths()

    if len(dag_paths) < 2:
        cmds.warning(
            "接合部を修正する2つ以上のメッシュを選択してください。"
        )
        return

    cmds.undoInfo(openChunk=True)

    try:
        all_vertex_data = []

        for dag_path in dag_paths:
            unlock_mesh_normals(dag_path)

            mesh_vertex_data = get_vertex_data(dag_path)

            if border_only:
                border_ids = get_border_vertex_ids(dag_path)

                mesh_vertex_data = [
                    data
                    for data in mesh_vertex_data
                    if data["vertex_id"] in border_ids
                ]

            all_vertex_data.extend(mesh_vertex_data)

        spatial_hash = build_spatial_hash(
            all_vertex_data,
            tolerance
        )

        processed = set()
        matched_groups = []

        for source_data in all_vertex_data:
            source_key = (
                source_data["dag_path"].fullPathName(),
                source_data["vertex_id"]
            )

            if source_key in processed:
                continue

            matching_vertices = []

            for neighbor_key in get_neighbor_keys(
                source_data["position"],
                tolerance
            ):
                candidates = spatial_hash.get(
                    neighbor_key,
                    []
                )

                for candidate in candidates:
                    candidate_key = (
                        candidate["dag_path"].fullPathName(),
                        candidate["vertex_id"]
                    )

                    if candidate_key in processed:
                        continue

                    # 同一メッシュ内の近接頂点は対象外
                    if (
                        candidate["dag_path"].fullPathName()
                        == source_data["dag_path"].fullPathName()
                    ):
                        continue

                    distance = (
                        candidate["position"]
                        - source_data["position"]
                    ).length()

                    if distance <= tolerance:
                        matching_vertices.append(candidate)

            if not matching_vertices:
                continue

            group = [source_data] + matching_vertices

            # 同一頂点の重複を除去
            unique_group = {}
            for data in group:
                key = (
                    data["dag_path"].fullPathName(),
                    data["vertex_id"]
                )
                unique_group[key] = data

            group = list(unique_group.values())

            if len(group) < 2:
                continue

            matched_groups.append(group)

            for data in group:
                processed.add((
                    data["dag_path"].fullPathName(),
                    data["vertex_id"]
                ))

        if not matched_groups:
            cmds.warning(
                "指定距離内に一致する接合頂点が見つかりませんでした。"
            )
            return

        apply_average_normals(matched_groups)

        cmds.inViewMessage(
            amg=(
                "<hl>接合部法線を平均化しました</hl>"
                "<br/>一致グループ数: {}"
            ).format(len(matched_groups)),
            pos="midCenter",
            fade=True
        )

    finally:
        cmds.undoInfo(closeChunk=True)


def get_border_vertex_ids(dag_path):
    """
    メッシュの境界エッジに含まれる頂点番号を取得する。
    """
    border_vertex_ids = set()

    edge_iterator = om.MItMeshEdge(dag_path)

    while not edge_iterator.isDone():
        if edge_iterator.onBoundary():
            border_vertex_ids.add(
                edge_iterator.vertexId(0)
            )
            border_vertex_ids.add(
                edge_iterator.vertexId(1)
            )

        edge_iterator.next()

    return border_vertex_ids


def apply_average_normals(matched_groups):
    """
    法線グループごとに平均法線を計算し、
    各頂点へ適用する。
    """
    mesh_updates = {}

    for group in matched_groups:
        average_normal = om.MVector()

        for data in group:
            average_normal += data["normal"]

        if average_normal.length() < 0.000001:
            continue

        average_normal.normalize()

        for data in group:
            mesh_name = data["dag_path"].fullPathName()

            if mesh_name not in mesh_updates:
                mesh_updates[mesh_name] = {
                    "mesh_fn": data["mesh_fn"],
                    "vertex_ids": [],
                    "normals": []
                }

            mesh_updates[mesh_name]["vertex_ids"].append(
                data["vertex_id"]
            )

            mesh_updates[mesh_name]["normals"].append(
                om.MVector(
                    average_normal.x,
                    average_normal.y,
                    average_normal.z
                )
            )

    for update_data in mesh_updates.values():
        mesh_fn = update_data["mesh_fn"]

        mesh_fn.setVertexNormals(
            update_data["normals"],
            update_data["vertex_ids"],
            om.MSpace.kWorld
        )


def run_from_ui(*args):
    tolerance = cmds.floatField(
        "seamNormalToleranceField",
        query=True,
        value=True
    )

    border_only = cmds.checkBox(
        "seamNormalBorderOnlyCheckBox",
        query=True,
        value=True
    )

    if tolerance <= 0:
        cmds.warning("許容距離は0より大きくしてください。")
        return

    average_seam_normals(
        tolerance=tolerance,
        border_only=border_only
    )


def show_seam_normal_tool():
    if cmds.window(WINDOW_NAME, exists=True):
        cmds.deleteUI(WINDOW_NAME)

    window = cmds.window(
        WINDOW_NAME,
        title="Seam Normal Smoother",
        sizeable=False,
        widthHeight=(360, 190)
    )

    cmds.columnLayout(
        adjustableColumn=True,
        rowSpacing=8
    )

    cmds.separator(height=8, style="none")

    cmds.text(
        label="接合する2つ以上のメッシュを選択",
        align="center"
    )

    cmds.text(
        label="同じ位置にある境界頂点の法線を平均化します",
        align="center"
    )

    cmds.separator(height=8, style="in")

    cmds.floatFieldGrp(
        "seamNormalToleranceField",
        numberOfFields=1,
        label="許容距離",
        value1=0.001,
        precision=6,
        columnWidth2=(100, 220)
    )

    cmds.checkBox(
        "seamNormalBorderOnlyCheckBox",
        label="境界頂点のみを処理",
        value=True
    )

    cmds.separator(height=8, style="in")

    cmds.button(
        label="接合部の法線を滑らかにする",
        height=38,
        command=run_from_ui
    )

    cmds.separator(height=8, style="none")

    cmds.showWindow(window)


show_seam_normal_tool()