import numpy as np
import warp as wp

import newton
import newton.examples
import newton.utils

wp.config.enable_backward = False


def create_torus_mesh(major_radius=1.0, minor_radius=0.3, major_segments=32, minor_segments=16):
    """
    Generate a torus mesh for Newton physics simulation.

    Args:
        major_radius (float): The radius from the center of the torus to the center of the tube
        minor_radius (float): The radius of the tube cross-section
        major_segments (int): Number of segments around the major circumference
        minor_segments (int): Number of segments around the minor circumference

    Returns:
        newton.geometry.Mesh: A Newton Mesh object representing the torus
    """
    vertices = []
    indices = []

    # Generate vertices
    for i in range(major_segments + 1):
        major_angle = 2.0 * np.pi * i / major_segments
        cos_major = np.cos(major_angle)
        sin_major = np.sin(major_angle)

        for j in range(minor_segments + 1):
            minor_angle = 2.0 * np.pi * j / minor_segments
            cos_minor = np.cos(minor_angle)
            sin_minor = np.sin(minor_angle)

            # Calculate vertex position
            x = (major_radius + minor_radius * cos_minor) * cos_major
            y = (major_radius + minor_radius * cos_minor) * sin_major
            z = minor_radius * sin_minor

            vertices.append([x, y, z])

    # Generate triangle indices
    for i in range(major_segments):
        for j in range(minor_segments):
            # Calculate vertex indices for current quad
            v0 = i * (minor_segments + 1) + j
            v1 = i * (minor_segments + 1) + (j + 1)
            v2 = (i + 1) * (minor_segments + 1) + (j + 1)
            v3 = (i + 1) * (minor_segments + 1) + j

            # Add two triangles to form the quad (flipped winding)
            # First triangle
            indices.extend([v0, v2, v1])
            # Second triangle
            indices.extend([v0, v3, v2])

    # Create Newton Mesh object
    mesh = newton.geometry.Mesh(vertices, indices)
    return mesh


class Example:
    def __init__(self, stage_path="example_torus.usda", num_envs=1):
        self.num_envs = num_envs

        # set numpy random seed
        self.seed = 123
        self.rng = np.random.default_rng(self.seed)

        # Create torus mesh
        torus_mesh = create_torus_mesh(major_radius=1.0, minor_radius=0.3, major_segments=32, minor_segments=16)

        builder = newton.ModelBuilder()
        builder.default_shape_cfg.ke = 1e4
        builder.default_shape_cfg.kd = 500.0

        # Add a torus body
        b = builder.add_body(xform=wp.transform(wp.vec3(0.0, 0.0, 2.0), wp.quat_identity()))
        builder.add_joint_free(b)
        builder.add_shape_mesh(
            body=b,
            mesh=torus_mesh,
            scale=wp.vec3(1.0, 1.0, 1.0),
        )

        # Add a sphere to fall on the torus
        b = builder.add_body(xform=wp.transform(wp.vec3(0.0, 0.0, 4.0), wp.quat_identity()))
        builder.add_shape_sphere(
            body=b,
            radius=0.2,
        )
        builder.add_joint_free(b)

        # Add a ground plane
        builder.add_ground_plane()

        builder.simplify_meshes(method="coacd", threshold=0.1)
        # builder.simplify_meshes()

        self.sim_time = 0.0
        fps = 60
        self.frame_dt = 1.0 / fps

        self.sim_substeps = 50
        self.sim_dt = self.frame_dt / self.sim_substeps

        # finalize model
        self.model = builder.finalize()

        self.control = self.model.control()

        self.solver = newton.solvers.MuJoCoSolver(
            self.model,
            use_mujoco=True,
            solver="newton",
            integrator="euler",
            iterations=20,
            ls_iterations=15,
            save_to_mjcf="converted_model.xml",
        )

        self.renderer = None
        if stage_path:
            self.renderer = newton.utils.SimRendererOpenGL(
                path=stage_path, model=self.model, scaling=1.0, show_joints=True
            )

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

        self.use_cuda_graph = not getattr(self.solver, "use_mujoco", False) and wp.get_device().is_cuda

        if self.use_cuda_graph:
            with wp.ScopedCapture() as capture:
                self.simulate()
            self.graph = capture.graph

    def simulate(self):
        if not isinstance(self.solver, newton.solvers.MuJoCoSolver):
            contacts = self.model.collide(self.state_0, rigid_contact_margin=100.0)
        else:
            contacts = None
        for _ in range(self.sim_substeps):
            self.state_0.clear_forces()
            self.solver.step(self.state_0, self.state_1, self.control, contacts, self.sim_dt)
            self.state_0, self.state_1 = self.state_1, self.state_0

    def step(self):
        with wp.ScopedTimer("step", active=False):
            if self.use_cuda_graph:
                wp.capture_launch(self.graph)
            else:
                self.simulate()
        self.sim_time += self.frame_dt

    def render(self):
        if self.renderer is None:
            return

        with wp.ScopedTimer("render", active=False):
            self.renderer.begin_frame(self.sim_time)
            self.renderer.render(self.state_0)
            self.renderer.end_frame()


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
    parser.add_argument(
        "--stage_path",
        type=lambda x: None if x == "None" else str(x),
        default="example_torus.usda",
        help="Path to the output USD file.",
    )
    parser.add_argument("--num_frames", type=int, default=1200, help="Total number of frames.")
    parser.add_argument("--num_envs", type=int, default=1, help="Total number of simulated environments.")

    args = parser.parse_known_args()[0]

    with wp.ScopedDevice(args.device):
        example = Example(stage_path=args.stage_path, num_envs=args.num_envs)

        for _ in range(args.num_frames):
            example.step()
            example.render()

        if example.renderer:
            example.renderer.save()