# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

###########################################################################
# Example Sim Granular Collision SDF
#
# Shows how to set up a particle-based granular material model using the
# wp.sim.ModelBuilder(). This version shows how to create collision geometry
# objects from SDFs.
#
###########################################################################

import math
import os

import numpy as np

import warp as wp
import warp.sim
import warp.sim.render

wp.init()


class Example:
    def __init__(self, stage):
        self.frame_dt = 1.0 / 60
        self.frame_count = 400

        self.sim_substeps = 64
        self.sim_dt = self.frame_dt / self.sim_substeps
        self.sim_steps = self.frame_count * self.sim_substeps
        self.sim_time = 0.0

        self.radius = 0.1

        builder = wp.sim.ModelBuilder()
        builder.default_particle_radius = self.radius

        builder.add_particle_grid(
            dim_x=16,
            dim_y=32,
            dim_z=16,
            cell_x=self.radius * 2.0,
            cell_y=self.radius * 2.0,
            cell_z=self.radius * 2.0,
            pos=wp.vec3(0.0, 20.0, 0.0),
            rot=wp.quat_identity(),
            vel=wp.vec3(2.0, 0.0, 0.0),
            mass=0.1,
            jitter=self.radius * 0.1,
        )
        rock_file = open(os.path.join(os.path.dirname(__file__), "assets/rocks.nvdb"), "rb")
        rock_vdb = wp.Volume.load_from_nvdb(rock_file.read())
        rock_file.close()

        rock_sdf = wp.sim.SDF(rock_vdb)

        builder.add_shape_sdf(
            ke=1.0e4,
            kd=1000.0,
            kf=1000.0,
            mu=0.5,
            sdf=rock_sdf,
            body=-1,
            pos=wp.vec3(0.0, 0.0, 0.0),
            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -0.5 * math.pi),
            scale=wp.vec3(0.01, 0.01, 0.01),
        )

        mins = np.array([-3.0, -3.0, -3.0])
        voxel_size = 0.2
        maxs = np.array([3.0, 3.0, 3.0])
        nums = np.ceil((maxs - mins) / (voxel_size)).astype(dtype=int)
        center = np.array([0.0, 0.0, 0.0])
        rad = 2.5
        sphere_sdf_np = np.zeros(tuple(nums))
        for x in range(nums[0]):
            for y in range(nums[1]):
                for z in range(nums[2]):
                    pos = mins + voxel_size * np.array([x, y, z])
                    dis = np.linalg.norm(pos - center)
                    sphere_sdf_np[x, y, z] = dis - rad

        sphere_vdb = wp.Volume.load_from_numpy(sphere_sdf_np, mins, voxel_size, rad + 3.0 * voxel_size)
        sphere_sdf = wp.sim.SDF(sphere_vdb)

        self.sphere_pos = wp.vec3(3.0, 15.0, 0.0)
        self.sphere_scale = 1.0
        self.sphere_radius = rad
        builder.add_shape_sdf(
            ke=1.0e4,
            kd=1000.0,
            kf=1000.0,
            mu=0.5,
            sdf=sphere_sdf,
            body=-1,
            pos=self.sphere_pos,
            scale=wp.vec3(self.sphere_scale, self.sphere_scale, self.sphere_scale),
        )

        self.model = builder.finalize()
        self.model.particle_kf = 25.0

        self.model.soft_contact_kd = 100.0
        self.model.soft_contact_kf *= 2.0

        self.state_0 = self.model.state()
        self.state_1 = self.model.state()

        self.integrator = wp.sim.SemiImplicitIntegrator()

        self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=20.0)

    def update(self):
        with wp.ScopedTimer("simulate", active=True):
            self.model.particle_grid.build(self.state_0.particle_q, self.radius * 2.0)

            for _ in range(self.sim_substeps):
                self.state_0.clear_forces()
                wp.sim.collide(self.model, self.state_0)
                self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)

                # swap states
                (self.state_0, self.state_1) = (self.state_1, self.state_0)

            self.sim_time += self.frame_dt

    def render(self, is_live=False):
        with wp.ScopedTimer("render", active=True):
            time = 0.0 if is_live else self.sim_time

            self.renderer.begin_frame(time)

            # Note the extra wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi) is because .usd is oriented differently from .nvdb
            self.renderer.render_ref(
                name="collision",
                path=os.path.join(os.path.dirname(__file__), "assets/rocks.usd"),
                pos=wp.vec3(0.0, 0.0, 0.0),
                rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -0.5 * math.pi)
                * wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi),
                scale=wp.vec3(0.01, 0.01, 0.01),
            )

            self.renderer.render_sphere(
                name="sphere",
                pos=self.sphere_pos,
                radius=self.sphere_scale * self.sphere_radius,
                rot=wp.quat(0.0, 0.0, 0.0, 1.0),
            )

            self.renderer.render(self.state_0)
            self.renderer.end_frame()


if __name__ == "__main__":
    stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_sim_sdf_shape.usd")

    example = Example(stage_path)

    for _ in range(example.frame_count):
        example.update()
        example.render()

    example.renderer.save()