import json
import pathlib
import sys

import bpy


def group_weight(vertex, group_index):
    return next((item.weight for item in vertex.groups if item.group == group_index), 0.0)


def set_weight(obj, vertex_index, group_index, value):
    group = obj.vertex_groups[group_index]
    if value == 0.0:
        try:
            group.remove([vertex_index])
        except RuntimeError:
            pass
    else:
        group.add([vertex_index], value, "REPLACE")


def normalize_vertex(obj, vertex_index):
    vertex = obj.data.vertices[vertex_index]
    total = sum(item.weight for item in vertex.groups)
    if total > 0.0:
        for item in list(vertex.groups):
            set_weight(obj, vertex_index, item.group, item.weight / total)


def limit_vertex(obj, vertex_index, limit):
    vertex = obj.data.vertices[vertex_index]
    while len(vertex.groups) > limit:
        remove = min(vertex.groups, key=lambda item: (item.weight, -item.group))
        obj.vertex_groups[remove.group].remove([vertex_index])


def mirror_map(obj, axis, tolerance):
    result = {}
    for source in obj.data.vertices:
        reflected = source.co.copy()
        reflected[axis] = -reflected[axis]
        candidate = min(obj.data.vertices, key=lambda item: ((item.co - reflected).length, item.index))
        if (candidate.co - reflected).length > tolerance:
            raise RuntimeError("mesh is not mirror symmetric")
        result[source.index] = candidate.index
    if any(result[result[index]] != index for index in result):
        raise RuntimeError("mesh mirror map is not reciprocal")
    return result


def snapshot(obj):
    groups = [group.name for group in obj.vertex_groups]
    return {
        "groups": groups,
        "vertices": [
            {"index": vertex.index, "weights": {
                groups[item.group]: round(item.weight, 9) for item in vertex.groups
            }}
            for vertex in obj.data.vertices
        ],
    }


def main(blend_path, output_path):
    bpy.ops.wm.open_mainfile(filepath=str(pathlib.Path(blend_path).resolve()), load_ui=False)
    obj = bpy.data.objects.get("RiggedShapeObject")
    if obj is None or obj.type != "MESH":
        raise RuntimeError("weight paint fixture object is missing")
    group = obj.vertex_groups.get("WebPaintGroup") or obj.vertex_groups.new(name="WebPaintGroup")

    set_weight(obj, 0, group.index, 0.75)
    set_weight(obj, 1, group.index, 0.25)
    steps = [{"name": "initial", **snapshot(obj)}]

    set_weight(obj, 2, group.index, 0.5)
    normalize_vertex(obj, 2)
    steps.append({"name": "normalize", **snapshot(obj)})

    set_weight(obj, 3, group.index, 0.4)
    limit_vertex(obj, 3, 2)
    normalize_vertex(obj, 3)
    steps.append({"name": "limit-normalize", **snapshot(obj)})

    mapping = mirror_map(obj, 0, 1e-4)
    for source in (0, mapping[0]):
        set_weight(obj, source, group.index, 0.9)
    steps.append({"name": "mirror", **snapshot(obj)})

    result = {
        "schemaVersion": 1,
        "fixture": pathlib.Path(blend_path).name,
        "blenderVersion": bpy.app.version_string,
        "object": obj.name,
        "mesh": obj.data.name,
        "operations": [
            {"name": "normalize", "vertex": 2},
            {"name": "limit-normalize", "vertex": 3, "limit": 2},
            {"name": "mirror", "vertices": [0, 1], "axis": 0, "tolerance": 1e-4},
        ],
        "steps": steps,
        "tolerance": 1e-6,
    }
    pathlib.Path(output_path).write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")


if __name__ == "__main__":
    if "--" not in sys.argv or len(sys.argv) <= sys.argv.index("--") + 2:
        raise SystemExit("usage: blender -b --python generate-weight-paint-golden.py -- input.blend output.json")
    main(*sys.argv[sys.argv.index("--") + 1:sys.argv.index("--") + 3])