"""Signed distance fields floating in water (agxModel.WindAndWaterController).

A SignedDistanceField is also a Trimesh of the triangles it was built from. Hydrodynamics
(buoyancy, drag and lift) works on that mesh, so a field floats exactly as the same triangle
mesh does, while its collisions use the field.

Scenes (keys 1-2):
1. Fields and their meshes: a back row of signed distance fields and a front row of triangle
   meshes of the same triangles, dropped into water. The text shows the largest height
   difference between a field and its mesh twin.
2. A field basin: signed distance fields of different shapes and densities dropped into water in
   a static signed distance field bowl. They float, bump into each other and into the bowl,
   while a dense bolt sinks to the bottom.

Run: agxViewer signed_distance_field_hydrodynamics.agxPy, or python signed_distance_field_hydrodynamics.agxPy
See also data/python/signed_distance_field_shapes.agxPy and data/python/buoyancy.agxPy.
"""

import math
import random

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

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

# -----------------------------------------------------------------------------
# Shapes
# -----------------------------------------------------------------------------

# Closed meshes from data/models (hydrodynamics needs a closed mesh for the right volume).
MESH_FILES = {
    "torus": "models/torusForCollision.obj",
    "round torus": "models/torusRoundForCollision.obj",
    "knight": "models/knight.obj",
    "nut": "models/nut.obj",
    "bolt": "models/bolt.obj",
    "mushroom": "models/mushroom.obj",
    "gear": "models/gear24T.obj",
}

_meshes = {}
_grids = {}


def mesh(name, size):
    """The vertices and indices of a mesh in MESH_FILES, centered, scaled to the largest extent size."""
    key = (name, size)
    if key not in _meshes:
        reader = agxIO.MeshReader()
        assert reader.readFile(MESH_FILES[name]), MESH_FILES[name]
        bound = agxCollide.BoundingAABB()
        for vertex in reader.getVertices():
            bound.expand(vertex)
        extent = bound.max() - bound.min()
        scale = size / max(extent.x(), extent.y(), extent.z())
        center = (bound.min() + bound.max()) * 0.5
        vertices = agx.Vec3Vector()
        for vertex in reader.getVertices():
            vertices.append((vertex - center) * scale)
        # Copy: getIndices() refers into the reader.
        _meshes[key] = (vertices, agx.UInt32Vector(reader.getIndices()))
    return _meshes[key]


def field(name, size):
    """A signed distance field shape of a mesh. Shapes of the same mesh and size share one grid."""
    key = (name, size)
    if key not in _grids:
        options = agxCollide.SignedDistanceFieldGridOptions()
        options.voxelSize = size / 60
        _grids[key] = agxCollide.SignedDistanceFieldGrid.createFromMesh(*mesh(name, size), options)
    return agxCollide.SignedDistanceField(_grids[key])


def trimesh(name, size):
    return agxCollide.Trimesh(*mesh(name, size), name)


def bowlField(scale):
    """models/bowl.obj as a signed distance field, about 7.2 * scale across, base at z = 0."""
    reader = agxIO.MeshReader()
    assert reader.readFile("models/bowl.obj")
    vertices = agx.Vec3Vector()
    for vertex in reader.getVertices():
        vertices.append(vertex * scale)
    options = agxCollide.SignedDistanceFieldGridOptions()
    options.voxelSize = 0.02
    grid = agxCollide.SignedDistanceFieldGrid.createFromMesh(vertices, agx.UInt32Vector(reader.getIndices()), options)
    return agxCollide.SignedDistanceField(grid)


# -----------------------------------------------------------------------------
# Water, bodies and graphics
# -----------------------------------------------------------------------------

COLORS = [
    agxRender.Color.Gold(),
    agxRender.Color.OrangeRed(),
    agxRender.Color.YellowGreen(),
    agxRender.Color.Violet(),
    agxRender.Color.Coral(),
    agxRender.Color.HotPink(),
    agxRender.Color.Khaki(),
]


def addWater(halfSize, depth):
    """A water box with its surface at z = 0, and the controller that makes things float in it."""
    water = agxCollide.Geometry(agxCollide.Box(agx.Vec3(halfSize, halfSize, 0.5 * depth)))
    water.setPosition(0, 0, -0.5 * depth)
    simulation().add(water)
    node = agxOSG.createVisual(water, root())
    agxOSG.setDiffuseColor(node, agx.Vec4f(0.0, 0.75, 1.0, 1))
    agxOSG.setAlpha(node, 0.4)

    controller = agxModel.WindAndWaterController()
    controller.addWater(water)
    simulation().add(controller)
    return controller


def addBody(shape, position, rotation, density, color):
    body = agx.RigidBody(agxCollide.Geometry(shape))
    material = agx.Material(f"density {density}")
    material.getBulkMaterial().setDensity(density)
    body.getGeometries()[0].setMaterial(material)
    body.updateMassProperties()
    body.setPosition(position)
    body.setRotation(rotation)
    simulation().add(body)
    agxOSG.setDiffuseColor(agxOSG.createVisual(body, root()), color)
    return body


def setText(lines):
    decorator = application().getSceneDecorator()
    for i, line in enumerate(lines):
        decorator.setText(i, line, agxRender.Color.White())


def setup(cameraEye, cameraCenter):
    simulation().setTimeStep(1.0 / 120.0)
    # Collision detection and the solver run in parallel; four threads make a large difference here.
    # A thread count given on the command line (--numThreads) is kept.
    if application().getArguments().find("--numThreads") < 0:
        agx.setNumThreads(4)
    application().setCameraHome(cameraEye, cameraCenter, agx.Vec3.Z_AXIS())


class TwinHeightDifference(agxSDK.StepEventListener):
    """Shows the largest height difference between a field and its mesh twin."""

    def __init__(self, pairs, line):
        super().__init__(agxSDK.StepEventListener.POST_STEP)
        self.pairs = pairs
        self.line = line

    def post(self, time):
        difference = max(abs(field.getPosition().z() - mesh.getPosition().z()) for field, mesh in self.pairs)
        application().getSceneDecorator().setText(
            self.line, f"Largest height difference field - mesh: {difference * 1000:.3f} mm", agxRender.Color.Yellow())


# -----------------------------------------------------------------------------
# Scenes
# -----------------------------------------------------------------------------


def buildFieldsAndMeshes():
    """Signed distance fields and triangle meshes of the same triangles float alike."""
    setText(["Back row: signed distance fields. Front row: triangle meshes of the same triangles.",
             "Hydrodynamics uses the mesh of a field, so both rows float the same."])
    setup(agx.Vec3(0, -20, 8), agx.Vec3(0, 0, 0))
    addWater(15, 10)

    names = list(MESH_FILES)
    densities = [450, 500, 600, 550, 700, 400, 650]
    pairs = []
    for k, name in enumerate(names):
        x = (k - 0.5 * (len(names) - 1)) * 2.2
        rotation = agx.Quat(0.6 * k + 0.3, agx.Vec3(1, 1, 0))
        fieldBody = addBody(field(name, 1.5), agx.Vec3(x, 2, 2), rotation, densities[k], COLORS[k])
        meshBody = addBody(trimesh(name, 1.5), agx.Vec3(x, -2, 2), rotation, densities[k], agxRender.Color.SteelBlue())
        pairs.append((fieldBody, meshBody))
    simulation().add(TwinHeightDifference(pairs, 2))


def buildFieldBasin():
    """Signed distance fields floating in a signed distance field bowl."""
    setText(["Signed distance fields floating, and bumping into each other, in a static signed distance",
             "field bowl. Densities 300-800 kg/m3, except the dark bolt (7800 kg/m3), which sinks."])
    setup(agx.Vec3(0, -12, 9), agx.Vec3(0, 0, -0.5))
    addWater(10, 6)

    # The bowl, 11 m across with its rim 1.5 m above the water.
    bowl = agxCollide.Geometry(bowlField(1.5))
    bowl.setPosition(0, 0, -3.1)
    simulation().add(bowl)
    agxOSG.setDiffuseColor(agxOSG.createVisual(bowl, root()), agxRender.Color.LightGray())

    random.seed(3)
    names = [name for name in MESH_FILES if name != "bolt"]
    for layer in range(3):
        for k in range(10):
            name = names[(k + layer) % len(names)]
            angle = 2 * math.pi * (k + 0.5 * layer) / 10
            radius = 1.2 + 1.3 * (k % 2)
            position = agx.Vec3(radius * math.cos(angle), radius * math.sin(angle), 1.0 + 1.2 * layer)
            rotation = agx.Quat(random.uniform(0, math.pi), agx.Vec3(random.uniform(-1, 1), random.uniform(-1, 1), 1))
            density = random.uniform(300, 800)
            addBody(field(name, 0.9), position, rotation, density, COLORS[names.index(name) % len(COLORS)])
    addBody(field("bolt", 1.0), agx.Vec3(0, 0, 5), agx.Quat(), 7800, agxRender.Color.DimGray())


def addRemainingScenes():
    """Adds the other scene so that it can be started with the key 2 from agxViewer."""
    scriptFileName = application().getArguments().getArgumentName(1).replace("agxscene:", "")
    application().addScene(scriptFileName, "buildFieldBasin", ord(str(application().getNumScenes() + 1)), True)


def buildScene():
    """Entry point when this script is run with agxViewer."""
    if application().getNumScenes() == 1:
        addRemainingScenes()
    buildFieldsAndMeshes()


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