"""What a signed distance field represents, compared with the mesh it was built from.

agxOSG.createVisual draws the triangles a signed distance field was built from, which says
nothing about how well the field represents them. agxOSG.createSignedDistanceFieldVisual draws
the field itself:
* The surface of the field, extracted from it: the shape the colliders see.
* Colored by its distance to the source mesh: gray where they agree, red where the field is
  outside the mesh (filled concave corners, grooves and gaps), blue where it is inside (rounded
  convex edges, thinned walls). The colors saturate at a quarter voxel.
* In magenta, the source triangles the field does not resolve: walls too thin and gaps too narrow
  for the voxel size, where the contacts will be wrong.

The same data is available without graphics, from agxCollide.SignedDistanceFieldGrid:
extractSurface, computeExactDistances and findUnresolvedTriangles.

Keys:
  +/-  halve/double the voxel size (rebuilds the field)
  m    show the field or the source mesh
  w    source mesh wireframe on/off
  e    error colors on/off

Run: python signed_distance_field_visualization.agxPy [--mesh file.obj] [--scale s] [--voxelSize v]
     or agxViewer signed_distance_field_visualization.agxPy (default mesh)
Default: models/nut.obj. Default voxel size: 1/100 of the mesh bound diagonal.
"""

import math
import time

import agx
import agxCollide
import agxIO
import agxOSG
import agxRender
import agxSDK

from agxPythonModules.utils.environment import application, init_app, root, simulation

# Python objects the simulation only holds raw pointers to.
_keep = []


def argument(name, default):
    """The value after name on the command line, or default."""
    arguments = application().getArguments()
    index = arguments.find(name)
    if index < 0 or index + 1 >= arguments.getNumArguments():
        return default
    return arguments.getArgumentName(index + 1)


def readMesh(fileName, scale):
    """(vertices, indices) of a mesh file, scaled and centered at the origin."""
    reader = agxIO.MeshReader()
    assert reader.readFile(fileName), f"Can't read {fileName}"
    bound = agxCollide.BoundingAABB()
    for vertex in reader.getVertices():
        bound.expand(vertex * scale)
    center = (bound.min() + bound.max()) * 0.5
    vertices = agx.Vec3Vector()
    for vertex in reader.getVertices():
        vertices.append(vertex * scale - center)
    # Copy: getIndices() refers into the reader, which is deleted on return.
    return vertices, agx.UInt32Vector(reader.getIndices()), (bound.max() - bound.min()).length()


class FieldViewer(agxSDK.GuiEventListener):
    """Shows the field of a mesh, or the mesh, and rebuilds the field when the voxel size changes."""

    def __init__(self, name, vertices, indices, voxelSize):
        super().__init__(agxSDK.GuiEventListener.KEYBOARD)
        self.name = name
        self.vertices = vertices
        self.indices = indices
        self.voxelSize = voxelSize
        self.showField = True
        self.wireframe = False
        self.errorColors = True
        self.node = None
        self.geometry = agxCollide.Geometry()
        self.mesh = agxCollide.Trimesh(vertices, indices, name)
        self.geometry.add(self.mesh)
        self.shape = self.mesh
        simulation().add(self.geometry)
        self.field = None
        self.buildField()

    def buildField(self):
        """Build the field at the current voxel size and measure how well it represents the mesh."""
        options = agxCollide.SignedDistanceFieldGridOptions()
        options.voxelSize = self.voxelSize
        start = time.perf_counter()
        grid = agxCollide.SignedDistanceFieldGrid.createFromMesh(self.vertices, self.indices, options)
        buildTime = time.perf_counter() - start

        # The error of the field surface: the exact distance from each vertex of the surface,
        # extracted from the field, to the source mesh.
        vertices, indices = agx.Vec3Vector(), agx.UInt32Vector()
        grid.extractSurface(vertices, indices)
        errors = agx.RealVector()
        grid.computeExactDistances(vertices, errors)
        errors = [abs(errors[i]) for i in range(len(errors))]
        unresolved = len(grid.findUnresolvedTriangles())

        self.field = agxCollide.SignedDistanceField(grid)
        _keep.append(agxCollide.SignedDistanceFieldRef(self.field))
        self.stats = [
            f"{self.name}: {len(self.indices) // 3} triangles, voxel size {self.voxelSize * 1000:.3g} mm, "
            f"{grid.getNumBricks()} bricks, {grid.getMemoryUsage() / 1e6:.2f} MB, built in {buildTime:.2f} s",
            f"Field surface: error mean {sum(errors) / max(1, len(errors)) * 1000:.3g} mm, max {max(errors, default=0) * 1000:.3g} mm "
            f"({max(errors, default=0) / self.voxelSize:.2f} voxels), volume {grid.getVolume() * 1e6:.4g} cm3",
            f"Unresolved source triangles (magenta): {unresolved} of {len(self.indices) // 3}",
        ]
        print("\n".join(self.stats))
        self.updateVisual()

    def updateVisual(self):
        # The geometry holds the field or the mesh, which is what createVisual draws.
        self.geometry.remove(self.shape)
        self.shape = self.field if self.showField else self.mesh
        self.geometry.add(self.shape)

        if self.node is not None:
            root().removeChild(self.node)
        if self.showField:
            mode = agxOSG.SDF_FIELD_SURFACE | agxOSG.SDF_UNRESOLVED_TRIANGLES
            if self.errorColors:
                mode |= agxOSG.SDF_ERROR_COLORS
            if self.wireframe:
                mode |= agxOSG.SDF_SOURCE_WIREFRAME
            self.node = agxOSG.createSignedDistanceFieldVisual(self.geometry, root(), mode)
        else:
            self.node = agxOSG.createVisual(self.geometry, root())
            agxOSG.setDiffuseColor(self.node, agxRender.Color.Goldenrod())
        self.showText()

    def showText(self):
        shown = "the surface of the field" if self.showField else "the source mesh"
        lines = self.stats + [
            "",
            f"Showing {shown}.",
            "Keys: +/- voxel size, m field/mesh,",
            f"w wireframe ({'on' if self.wireframe else 'off'}), e error colors ({'on' if self.errorColors else 'off'})",
            "Gray: on the mesh, red: outside it, blue: inside it",
        ]
        decorator = application().getSceneDecorator()
        for row, line in enumerate(lines):
            decorator.setText(row, line, agxRender.Color.White())

    def keyboard(self, key, x, y, alt, down):
        if not down:
            return False
        if key in (ord("+"), ord("-")):
            self.voxelSize *= 0.5 if key == ord("+") else 2.0
            self.buildField()
        elif key == ord("m"):
            self.showField = not self.showField
            self.updateVisual()
        elif key == ord("w"):
            self.wireframe = not self.wireframe
            self.updateVisual()
        elif key == ord("e"):
            self.errorColors = not self.errorColors
            self.updateVisual()
        else:
            return False
        return True


def buildScene():
    fileName = argument("--mesh", "models/nut.obj")
    scale = float(argument("--scale", "1"))
    vertices, indices, diagonal = readMesh(fileName, scale)
    voxelSize = float(argument("--voxelSize", str(diagonal / 100)))

    # Nothing moves: the scene is only for looking.
    simulation().setUniformGravity(agx.Vec3())
    viewer = FieldViewer(fileName.replace("\\", "/").split("/")[-1], vertices, indices, voxelSize)
    simulation().add(viewer)
    _keep.append(viewer)

    application().setCameraHome(
        agx.Vec3(0.7, -1.4, 1.1) * diagonal, agx.Vec3(0, 0, 0.15 * diagonal), agx.Vec3.Z_AXIS()
    )


init = init_app(name=__name__, scenes=[(buildScene, "1")], autoStepping=True)