kinetix/editor.py

import os
os.system(f"cd {os.getcwd()}/Kinetix")
import itertools
import time
from timeit import default_timer as tmr

import optax
from PIL import Image
from flax.serialization import to_state_dict
from flax.training.train_state import TrainState

from matplotlib import pyplot as plt

from kinetix.environment.ued.distributions import sample_kinetix_level
from kinetix.models import make_network_from_config
from kinetix.models.actor_critic import ScannedRNN
from kinetix.render.renderer_symbolic_entity import make_render_entities

ss = tmr()

import jax

jax.config.update("jax_compilation_cache_dir", ".cache-location")

import hydra
from omegaconf import OmegaConf

from kinetix.environment.ued.mutators import (
    make_mutate_change_shape_rotation,
    mutate_add_connected_shape,
    mutate_add_shape,
    make_mutate_change_shape_size,
    mutate_change_shape_location,
    mutate_swap_role,
    mutate_remove_shape,
    mutate_remove_joint,
    mutate_toggle_fixture,
    mutate_add_thruster,
    mutate_remove_thruster,
)
from kinetix.environment.ued.ued import make_mutate_env, ALL_MUTATION_FNS
from kinetix.environment.ued.ued_state import UEDParams
from kinetix.environment.ued.util import rectangle_vertices
from kinetix.util.config import generate_params_from_config, normalise_config


import argparse
import os
import sys

sys.path.append("editor")
import tkinter
import tkinter.filedialog
from enum import Enum
from timeit import default_timer as tmr

import jax.numpy as jnp
import numpy as np
import pygame
import pygame_widgets
from pygame_widgets.slider import Slider
from pygame_widgets.textbox import TextBox
from pygame_widgets.toggle import Toggle

from jax2d.engine import (
    calc_inverse_inertia_circle,
    calc_inverse_inertia_polygon,
    calc_inverse_mass_circle,
    calc_inverse_mass_polygon,
    calculate_collision_matrix,
    recalculate_mass_and_inertia,
    recompute_global_joint_positions,
    select_shape,
)
from jax2d.maths import rmat
from jax2d.sim_state import RigidBody
from kinetix.environment.env import (
    create_empty_env,
    make_kinetix_env_from_name,
)
from kinetix.environment.env_state import EnvParams, EnvState, StaticEnvParams
from kinetix.environment.utils import permute_pcg_state
from kinetix.environment.wrappers import AutoResetWrapper
from kinetix.pcg.pcg import env_state_to_pcg_state, sample_pcg_state
from kinetix.pcg.pcg_state import PCGState
from kinetix.render.renderer_pixels import make_render_pixels
from kinetix.render.textures import (
    CIRCLE_TEXTURE_RGBA,
    EDIT_TEXTURE_RGBA,
    PLAY_TEXTURE_RGBA,
    RECT_TEXTURE_RGBA,
    RJOINT_TEXTURE_RGBA,
    SELECT_TEXTURE_RGBA,
    THRUSTER_TEXTURE_RGBA,
    TRIANGLE_TEXTURE_RGBA,
)
from kinetix.util.saving import (
    expand_pcg_state,
    export_env_state_to_json,
    get_pcg_state_from_json,
    load_from_json_file,
    load_pcg_state_pickle,
    load_world_state_pickle,
    save_pickle,
    load_params_from_wandb_artifact_path,
    load_train_state_from_wandb_artifact_path,
)
from kinetix.util.timing import time_function
from tkinter import Tk

#root = Tk()
#root.destroy()

ee = tmr()
print(f"Imported in {ee - ss} seconds")

editor = None
outer_timer = tmr()
EMPTY_ENV = False


class ObjectType(Enum):
    POLYGON = 0
    CIRCLE = 1
    JOINT = 2
    THRUSTER = 3


class EditMode(Enum):
    ADD_CIRCLE = 0
    ADD_RECTANGLE = 1
    ADD_JOINT = 2
    SELECT = 3
    ADD_TRIANGLE = 4
    ADD_THRUSTER = 5


TOTAL_DUMMY_STEPS_TO_SNAP = 0
SNAPPING_DIST = 0.1


def select_object(state: EnvState, type: int, index: int):
    if type is None:
        type = ObjectType.POLYGON
    li = {0: state.polygon, 1: state.circle, 2: state.joint, 3: state.thruster}[type.value]
    return jax.tree.map(lambda x: x[index], li)


def snap_to_center(shape: RigidBody, position: jnp.ndarray):
    if jnp.linalg.norm(shape.position - position) < SNAPPING_DIST:
        return shape.position
    return position


def snap_to_polygon_center_line(polygon: RigidBody, position: jnp.ndarray):
    # Snap to the center line
    r = rmat(polygon.rotation)
    x = jnp.matmul(r, position - polygon.position)
    if jnp.abs(x[0]) < SNAPPING_DIST:
        x = x.at[0].set(0.0)
    if jnp.abs(x[1]) < SNAPPING_DIST:
        x = x.at[1].set(0.0)
    x = jnp.matmul(r.transpose(), x)
    return x + polygon.position


def snap_to_circle_center_line(circle: RigidBody, position: jnp.ndarray):
    # Snap to the center line, i.e. on the edge of the circle, if the position is close enough to directly below the circle, etc., snap the position to that
    x = position - circle.position
    if jnp.linalg.norm(x) < SNAPPING_DIST:
        return circle.position
    angle = (jnp.arctan2(x[1], x[0]) + 2 * jnp.pi) % (2 * jnp.pi)

    for i in range(0, 8):
        if jnp.abs(angle - i * jnp.pi / 4) < jnp.radians(25):  # 25 degrees
            angle = i * jnp.pi / 4
            break
    x = jnp.array([jnp.cos(angle), jnp.sin(angle)]) * circle.radius
    return x + circle.position


def prompt_file(save=False):
    dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "worlds")
    """Create a Tk file dialog and cleanup when finished"""
    top = tkinter.Tk()
    top.withdraw()  # hide window
    if save:
        file_name = tkinter.filedialog.asksaveasfilename(parent=top, initialdir=dir)
    else:
        file_name = tkinter.filedialog.askopenfilename(parent=top, initialdir=dir)
    top.destroy()
    return file_name


def get_numeric_key_pressed(pygame_events, is_mod=False):
    for event in pygame_events:
        if not is_mod:
            if event.type == pygame.KEYDOWN and event.unicode.isdigit():
                return int(event.unicode)
        else:
            if event.type == pygame.KEYDOWN:
                pass
            if event.type == pygame.KEYDOWN and event.key in [pygame.K_0 + i for i in range(10)]:
                return int(event.key - pygame.K_0)
    return None


def new_env(static_env_params):
    return create_empty_env(static_env_params)


myrng = jax.random.PRNGKey(0)


def make_reset_function(static_env_params):
    def reset(rng):
        return env_state_to_pcg_state(create_empty_env(static_env_params))

    return reset


def new_pcg_env(static_env_params):
    global myrng
    if EMPTY_ENV:
        env_state = create_empty_env(static_env_params)
    else:
        return get_pcg_state_from_json("worlds/l/h0_angrybirds.json")
    return env_state_to_pcg_state(env_state)


class Editor:
    def __init__(self, env, env_params, config, upscale=1):
        self.env = env
        self.upscale = upscale
        self.env_params = env_params
        self.static_env_params = env.static_env_params
        self.ued_params = UEDParams()

        self.side_panel_width = env.static_env_params.screen_dim[0] // 3

        self.rng = jax.random.PRNGKey(0)
        self.config = config
        self.pcg_state = new_pcg_env(self.static_env_params)

        self.rng, _rng = jax.random.split(self.rng)
        self.play_state = sample_pcg_state(_rng, self.pcg_state, self.env_params, self.static_env_params)
        self.last_played_level = None

        self.pygame_events = []

        self.mutate_world = make_mutate_env(env.static_env_params, env_params, self.ued_params)

        self.num_triangle_clicks = 0
        self.triangle_order = jnp.array([0, 1, 2])

        # Init rendering
        pygame.init()
        pygame.key.set_repeat(250, 75)

        self.screen_surface = pygame.display.set_mode(
            tuple(
                (t + extra) * self.upscale
                for t, extra in zip(self.static_env_params.screen_dim, (self.side_panel_width, 0))
            ),
            display=0,
        )
        self.all_widgets = {}
        self._setup_side_panel()
        self.has_done_action = False

        self._setup_rendering(self.static_env_params, env_params)

        self._render_edit_overlay_fn = jax.jit(self._render_edit_overlay)
        self._step_fn = jax.jit(env.step)

        # Agent
        if self.config["agent_taking_actions"]:
            self._entity_renderer = jax.jit(make_render_entities(env_params, self.env.static_env_params))
            self.network = make_network_from_config(env, env_params, config)

            rng = jax.random.PRNGKey(0)
            dones = jnp.zeros((config["num_train_envs"]), dtype=jnp.bool_)
            rng, _rng = jax.random.split(rng)
            init_hstate = ScannedRNN.initialize_carry(config["num_train_envs"])
            obsv = self._entity_renderer(self.play_state)
            obsv = jax.tree.map(lambda x: jnp.repeat(x[None, ...], repeats=config["num_train_envs"], axis=0), obsv)
            init_x = jax.tree.map(lambda x: x[None, ...], (obsv, dones))
            network_params = self.network.init(_rng, init_hstate, init_x)

            tx = optax.chain(
                optax.clip_by_global_norm(config["max_grad_norm"]),
                optax.adam(config["lr"], eps=1e-5),
            )
            self.train_state = TrainState.create(
                apply_fn=self.network.apply,
                params=network_params,
                tx=tx,
            )

            self.train_state = load_train_state_from_wandb_artifact_path(self.train_state, config["agent_wandb_path"])

            self.apply_fn = jax.jit(self.network.apply)

        # JIT Compile
        def _jit_step():
            rng = jax.random.PRNGKey(0)
            ans = self._step_fn(
                rng,
                self.env.reset_env_to_level(rng, self.play_state, self.env_params)[1],
                jnp.zeros(
                    env.static_env_params.num_motor_bindings + env.static_env_params.num_thruster_bindings, dtype=int
                ),
                self.env_params,
            )

        def _jit_render():
            self._render_fn(self.play_state)
            self._render_fn_edit(self.play_state)

        time_function(_jit_step, "_jit_step")
        time_function(_jit_render, "_jit_render")

        # self._step_fn(rng, self.play_state, 0, self.env_params)

        # Editing
        self.is_editing = True

        self.edit_shape_mode = EditMode.ADD_CIRCLE
        self.creating_shape = False
        self.create_shape_position = jnp.array([0.0, 0.0])
        self.creating_shape_index = 0
        self.selected_shape_index = -1
        self.selected_shape_type = ObjectType.POLYGON

        self.all_selected_shapes = []

        self.rng = jax.random.PRNGKey(0)
        time_function(self._jit, "self._jit")
        self._put_state_values_into_gui(self.pcg_state)

        self.mutate_change_shape_size = make_mutate_change_shape_size(self.env_params, self.static_env_params)
        self.mutate_change_shape_rotation = make_mutate_change_shape_rotation(self.env_params, self.static_env_params)

    def _setup_rendering(self, static_env_params: StaticEnvParams, env_params: EnvParams):
        def _make_render(should_do_edit_additions=False):
            def _render(env_state):
                side_panel = self._render_side_panel()
                render_pixels = make_render_pixels(params=env_params, static_params=static_env_params)
                pixels = render_pixels(
                    env_state,
                )
                pixels = jnp.concatenate([side_panel, pixels], axis=0)

                pixels = jnp.repeat(pixels, repeats=static_env_params.downscale * self.upscale, axis=0)
                pixels = jnp.repeat(pixels, repeats=static_env_params.downscale * self.upscale, axis=1)

                return pixels[:, ::-1, :]

            return _render

        def _make_screenshot_render():
            def _render(env_state):
                px_upscale = 4

                static_params = static_env_params.replace(screen_dim=(500 * px_upscale, 500 * px_upscale))
                ss_env_params = env_params.replace(
                    pixels_per_unit=100 * px_upscale,
                )

                render_pixels = make_render_pixels(
                    params=ss_env_params,
                    static_params=static_params,
                    render_rjoint_sectors=False,
                    pixel_upscale=2 * px_upscale,
                )
                pixels = render_pixels(
                    env_state,
                )

                return pixels[:, ::-1, :]

            return _render

        self._render_fn_edit = jax.jit(_make_render(True))
        self._render_fn = jax.jit(_make_render(False))
        self._render_fn_screenshot = jax.jit(_make_screenshot_render())

    def _jit(self):
        self._get_circles_on_mouse(self.pcg_state.env_state)
        self._get_polygons_on_mouse(self.pcg_state.env_state)
        self._get_revolute_joints_on_mouse(self.pcg_state.env_state)
        self._get_thrusters_on_mouse(self.pcg_state.env_state)
        self.pygame_events = list(pygame.event.get())
        self._handle_events(do_dummy=True)

        state = self.play_state
        for mutation_fn in ALL_MUTATION_FNS:
            mutation_fn(jax.random.PRNGKey(0), state, self.env_params, self.static_env_params, self.ued_params)

    def update(self, rng):
        # Update pygame events
        self.pygame_events = list(pygame.event.get())

        for event in self.pygame_events:
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:
                    self.has_done_action = False
                    self.is_editing = not self.is_editing
                    if not self.is_editing:
                        self.pcg_state = self._discard_shape_being_created(self.pcg_state)
                        self.pcg_state = self._reset_select_shape(self.pcg_state)
                        self.pcg_state = self.pcg_state.replace(
                            env_state=self.pcg_state.env_state.replace(
                                collision_matrix=calculate_collision_matrix(
                                    self.static_env_params, self.pcg_state.env_state.joint
                                ),
                            ),
                            env_state_pcg_mask=self.pcg_state.env_state_pcg_mask.replace(
                                collision_matrix=jnp.zeros_like(self.pcg_state.env_state_pcg_mask.collision_matrix)
                            ),
                        )
                        self.rng, _rng = jax.random.split(self.rng)
                        self.play_state = sample_pcg_state(
                            _rng, self.pcg_state, self.env_params, self.static_env_params
                        )
                        self.last_played_level = self.play_state
                elif event.key == pygame.K_s and not self.is_editing:
                    self.take_screenshot()

        if self.is_editing:
            self.pcg_state = self.edit()
        else:
            rng, _rng = jax.random.split(rng)
            # action = []
            action = jnp.zeros(
                self.env.static_env_params.num_motor_bindings + self.env.static_env_params.num_thruster_bindings,
                dtype=jnp.int32,
            )

            keys = pygame.key.get_pressed()
            if keys[pygame.K_LEFT]:
                action = action.at[0].set(1)
            if keys[pygame.K_RIGHT]:
                action = action.at[0].set(2)
            if keys[pygame.K_UP]:
                action = action.at[1].set(1)
            if keys[pygame.K_DOWN]:
                action = action.at[1].set(2)

            if keys[pygame.K_1]:
                action = action.at[0 + self.env.static_env_params.num_motor_bindings].set(1)
            if keys[pygame.K_2]:
                action = action.at[1 + self.env.static_env_params.num_motor_bindings].set(1)
            if keys[pygame.K_3]:
                action = action.at[2 + self.env.static_env_params.num_motor_bindings].set(1)
            if keys[pygame.K_4]:
                action = action.at[3 + self.env.static_env_params.num_motor_bindings].set(1)
            # if self.has_done_action: action = action * 0
            self.has_done_action = self.has_done_action | (action != 0).any()

            if self.config["agent_taking_actions"]:
                obs = self._entity_renderer(self.play_state)
                obs = jax.tree.map(lambda x: x[None, ...], obs)

                last_done = jnp.zeros((1, 1), dtype=bool)
                ac_in = (jax.tree.map(lambda x: x[np.newaxis, :], obs), last_done[np.newaxis, :])
                hstate = ScannedRNN.initialize_carry(1)
                hstate, pi, value = self.apply_fn(self.train_state.params, hstate, ac_in)
                rng, _rng = jax.random.split(rng)
                action = pi.sample(seed=_rng)
                action = action[0, 0]

            _rng, __rng = jax.random.split(_rng)
            obs, self.play_state, reward, done, info = self._step_fn(
                _rng, self.env.reset_to_level(__rng, self.play_state, self.env_params)[1], action, self.env_params
            )
            if done:
                self.rng, _rng = jax.random.split(self.rng)
                self.play_state = sample_pcg_state(_rng, self.pcg_state, self.env_params, self.static_env_params)
        state_to_render = self.pcg_state.env_state if self.is_editing else self.play_state

        self.render(state_to_render)
        self._handle_events()

        # Update screen
        pygame.display.flip()

    def take_screenshot(self):
        print("screenshot!")

        pixels = self._render_fn_screenshot(self.play_state)
        mtime = round(time.time() * 1000)

        pixels = pixels.transpose((1, 0, 2))

        # Black border
        border_thickness = 5
        pixels = pixels.at[:, :border_thickness].set(0.0)
        pixels = pixels.at[:, -border_thickness:].set(0.0)
        pixels = pixels.at[:border_thickness, :].set(0.0)
        pixels = pixels.at[-border_thickness:, :].set(0.0)

        im = Image.fromarray(np.array(pixels.astype(jnp.uint8)))
        im.save(f"results/screenshot_{str(mtime)}.png")

    def _get_selected_shape_global_indices(self):
        def _idx(idx, type):
            if type == ObjectType.CIRCLE:
                return idx + self.static_env_params.num_polygons
            return idx

        indices_to_use = jnp.array([_idx(idx, type) for idx, type in self.all_selected_shapes])
        return indices_to_use

    # flag1
    def _handle_events(self, do_dummy=False):
        pygame_widgets.update(self.pygame_events)
        if do_dummy or self.selected_shape_index < 0:
            gravity_main = self.all_widgets[None]["sldGravity"].getValue()
            gravity_max = self.all_widgets[None]["sldMaxGravity"].getValue()
            gravity_main, gravity_max = min(gravity_main, gravity_max), max(gravity_main, gravity_max)
            gravity_pcg_mask = self.all_widgets[None]["pcgTglGravity"].getValue()

            def _set_single_global(state, gravity):
                return state.replace(
                    gravity=state.gravity.at[1].set(gravity),
                )

            env_state = _set_single_global(self.pcg_state.env_state, gravity_main)
            env_state_max = _set_single_global(self.pcg_state.env_state_max, gravity_max)
            env_state_pcg_mask = _set_single_global(self.pcg_state.env_state_pcg_mask, gravity_pcg_mask)
            if not do_dummy:
                self.pcg_state = self.pcg_state.replace(
                    env_state=env_state,
                    env_state_max=env_state_max,
                    env_state_pcg_mask=env_state_pcg_mask,
                )

        if self.edit_shape_mode == EditMode.SELECT or do_dummy:  # is on the hand.
            if do_dummy or len(self.all_selected_shapes) > 1:
                # this processes the tying together.
                indices_to_use = self._get_selected_shape_global_indices()
                if len(indices_to_use) > 1:
                    toggle_val = self.all_widgets["TIE_TOGETHER"]["tglTieTogether"].getValue()
                    idxs = itertools.product(indices_to_use, indices_to_use)
                    idxs_a = []
                    idxs_b = []
                    for i, j in idxs:
                        idxs_a.append(i)
                        idxs_b.append(j)
                    idxs = jnp.array(idxs_a), jnp.array(idxs_b)
                    if toggle_val:
                        self.pcg_state = self.pcg_state.replace(
                            tied_together=self.pcg_state.tied_together.at[idxs[0], idxs[1]].set(True)
                        )
                    else:
                        self.pcg_state = self.pcg_state.replace(
                            tied_together=self.pcg_state.tied_together.at[idxs[0], idxs[1]].set(False)
                        )
            if self.selected_shape_index < 0 and not do_dummy:
                return
            if do_dummy or self.selected_shape_type in [ObjectType.POLYGON, ObjectType.CIRCLE]:  # rigidbody
                shape_main = select_object(
                    self.pcg_state.env_state, self.selected_shape_type, self.selected_shape_index
                )
                shape_max = select_object(
                    self.pcg_state.env_state_max, self.selected_shape_type, self.selected_shape_index
                )

                parent_container_main = (
                    self.pcg_state.env_state.circle
                    if self.selected_shape_type == ObjectType.CIRCLE
                    else self.pcg_state.env_state.polygon
                )
                parent_container_max = (
                    self.pcg_state.env_state_max.circle
                    if self.selected_shape_type == ObjectType.CIRCLE
                    else self.pcg_state.env_state_max.polygon
                )
                parent_container_pcg_mask = (
                    self.pcg_state.env_state_pcg_mask.circle
                    if self.selected_shape_type == ObjectType.CIRCLE
                    else self.pcg_state.env_state_pcg_mask.polygon
                )

                new_density_main = self.all_widgets[self.selected_shape_type]["sldDensity"].getValue()
                new_density_max = self.all_widgets[self.selected_shape_type]["sldMaxDensity"].getValue()
                density_pcg_mask = self.all_widgets[self.selected_shape_type]["pcgTglDensity"].getValue()
                if density_pcg_mask:
                    new_density_main, new_density_max = min(new_density_main, new_density_max), max(
                        new_density_main, new_density_max
                    )
                fixated = self.all_widgets[self.selected_shape_type]["tglFixate"].getValue()

                fix_val = 0.0 if fixated else 1.0

                def _density_calcs(base, new_density):
                    inv_mass = jax.lax.select(
                        self.selected_shape_type == ObjectType.CIRCLE,
                        calc_inverse_mass_circle(base.radius, new_density),
                        calc_inverse_mass_polygon(base.vertices, base.n_vertices, self.static_env_params, new_density)[
                            0
                        ],
                    )
                    inv_inertia = jax.lax.select(
                        self.selected_shape_type == ObjectType.CIRCLE,
                        calc_inverse_inertia_circle(base.radius, new_density),
                        calc_inverse_inertia_polygon(
                            base.vertices, base.n_vertices, self.static_env_params, new_density
                        ),
                    )

                    return inv_mass, inv_inertia

                inv_mass_main, inv_inertia_main = _density_calcs(shape_main, new_density_main)
                inv_mass_max, inv_inertia_max = _density_calcs(shape_max, new_density_max)

                friction_main = self.all_widgets[self.selected_shape_type]["sldFriction"].getValue()
                friction_max = self.all_widgets[self.selected_shape_type]["sldMaxFriction"].getValue()
                friction_pcg_mask = self.all_widgets[self.selected_shape_type]["pcgTglFriction"].getValue()
                if friction_pcg_mask:
                    friction_main, friction_max = min(friction_main, friction_max), max(friction_main, friction_max)

                restitution = self.all_widgets[self.selected_shape_type]["sldRestitution"].getValue()

                position_main = jnp.array(
                    [
                        self.all_widgets[self.selected_shape_type]["sldPosition_X"].getValue(),
                        self.all_widgets[self.selected_shape_type]["sldPosition_Y"].getValue(),
                    ]
                )

                position_max = jnp.array(
                    [
                        self.all_widgets[self.selected_shape_type]["sldMaxPosition_X"].getValue(),
                        self.all_widgets[self.selected_shape_type]["sldMaxPosition_Y"].getValue(),
                    ]
                )
                position_pcg_mask = self.all_widgets[self.selected_shape_type]["pcgTglPosition_X"].getValue()
                if position_pcg_mask:
                    position_main, position_max = jnp.minimum(position_main, position_max), jnp.maximum(
                        position_main, position_max
                    )

                rotation_main = self.all_widgets[self.selected_shape_type]["sldRotation"].getValue()
                rotation_max = self.all_widgets[self.selected_shape_type]["sldMaxRotation"].getValue()
                rotation_pcg_mask = self.all_widgets[self.selected_shape_type]["pcgTglRotation"].getValue()
                if rotation_pcg_mask:
                    rotation_main, rotation_max = min(rotation_main, rotation_max), max(rotation_main, rotation_max)

                velocity_main = jnp.array(
                    [
                        self.all_widgets[self.selected_shape_type]["sldVelocity_X"].getValue(),
                        self.all_widgets[self.selected_shape_type]["sldVelocity_Y"].getValue(),
                    ]
                )

                velocity_max = jnp.array(
                    [
                        self.all_widgets[self.selected_shape_type]["sldMaxVelocity_X"].getValue(),
                        self.all_widgets[self.selected_shape_type]["sldMaxVelocity_Y"].getValue(),
                    ]
                )
                velocity_main, velocity_max = jnp.minimum(velocity_main, velocity_max), jnp.maximum(
                    velocity_main, velocity_max
                )
                velocity_pcg_mask = self.all_widgets[self.selected_shape_type]["pcgTglVelocity_X"].getValue()

                angular_velocity_main = self.all_widgets[self.selected_shape_type]["sldAngular_Velocity"].getValue()
                angular_velocity_max = self.all_widgets[self.selected_shape_type]["sldMaxAngular_Velocity"].getValue()
                angular_velocity_pcg_mask = self.all_widgets[self.selected_shape_type][
                    "pcgTglAngular_Velocity"
                ].getValue()
                if angular_velocity_pcg_mask:
                    angular_velocity_main, angular_velocity_max = min(angular_velocity_main, angular_velocity_max), max(
                        angular_velocity_main, angular_velocity_max
                    )

                # Circle stuff
                radius_main, radius_max, radius_pcg_mask = None, None, None
                if self.selected_shape_type == ObjectType.CIRCLE:
                    radius_main = self.all_widgets[self.selected_shape_type]["sldRadius"].getValue()
                    radius_max = self.all_widgets[self.selected_shape_type]["sldMaxRadius"].getValue()
                    radius_pcg_mask = self.all_widgets[self.selected_shape_type]["pcgTglRadius"].getValue()

                # Poly stuff
                vertices_main, vertices_max, vertices_pcg_mask = None, None, None
                if self.selected_shape_type == ObjectType.POLYGON:
                    # Triangle
                    new_size_main = self.all_widgets[self.selected_shape_type]["sldSize"].getValue()
                    new_size_max = self.all_widgets[self.selected_shape_type]["sldMaxSize"].getValue()

                    current_size = jnp.abs(self.pcg_state.env_state.polygon.vertices[self.selected_shape_index]).max()
                    scale_main = new_size_main / current_size
                    scale_max = new_size_max / current_size

                    vertices_main = self.pcg_state.env_state.polygon.vertices[self.selected_shape_index] * scale_main
                    vertices_max = self.pcg_state.env_state.polygon.vertices[self.selected_shape_index] * scale_max
                    vertices_pcg_mask = (
                        jnp.ones_like(vertices_main, dtype=bool)
                        * self.all_widgets[self.selected_shape_type]["pcgTglSize"].getValue()
                    )

                def _set_single_state_rbody(
                    state,
                    parent_container,
                    density,
                    inv_mass,
                    inv_inertia,
                    friction,
                    position,
                    radius,
                    rotation,
                    velocity,
                    angular_velocity,
                    vertices,
                ):
                    new = {
                        "friction": parent_container.friction.at[self.selected_shape_index].set(friction),
                        "collision_mode": parent_container.collision_mode.at[self.selected_shape_index].set(
                            int(self.all_widgets[self.selected_shape_type]["sldCollidability"].getValue())
                        ),
                        "inverse_mass": parent_container.inverse_mass.at[self.selected_shape_index].set(
                            inv_mass * fix_val
                        ),
                        "inverse_inertia": parent_container.inverse_inertia.at[self.selected_shape_index].set(
                            inv_inertia * fix_val
                        ),
                        "position": parent_container.position.at[self.selected_shape_index].set(position),
                        "rotation": parent_container.rotation.at[self.selected_shape_index].set(rotation),
                        "velocity": parent_container.velocity.at[self.selected_shape_index].set(velocity),
                        "angular_velocity": parent_container.angular_velocity.at[self.selected_shape_index].set(
                            angular_velocity
                        ),
                        "restitution": parent_container.restitution.at[self.selected_shape_index].set(restitution),
                    }

                    if self.selected_shape_type == ObjectType.CIRCLE:
                        state = state.replace(
                            circle=state.circle.replace(
                                **new,
                                radius=parent_container.radius.at[self.selected_shape_index].set(radius),
                            ),
                            circle_shape_roles=state.circle_shape_roles.at[self.selected_shape_index].set(
                                int(self.all_widgets[self.selected_shape_type]["sldRole"].getValue())
                            ),
                            circle_densities=state.circle_densities.at[self.selected_shape_index].set(density),
                        )
                    else:
                        state = state.replace(
                            polygon=state.polygon.replace(
                                **new,
                                vertices=parent_container.vertices.at[self.selected_shape_index].set(vertices),
                            ),
                            polygon_shape_roles=state.polygon_shape_roles.at[self.selected_shape_index].set(
                                int(self.all_widgets[self.selected_shape_type]["sldRole"].getValue())
                            ),
                            polygon_densities=state.polygon_densities.at[self.selected_shape_index].set(density),
                        )

                    return state

                position_delta = position_main - shape_main.position
                env_state = _set_single_state_rbody(
                    self.pcg_state.env_state,
                    parent_container_main,
                    new_density_main,
                    inv_mass_main,
                    inv_inertia_main,
                    friction_main,
                    position_main,
                    radius_main,
                    rotation_main,
                    velocity_main,
                    angular_velocity_main,
                    vertices_main,
                )
                env_state_max = _set_single_state_rbody(
                    self.pcg_state.env_state_max,
                    parent_container_max,
                    new_density_max,
                    inv_mass_max,
                    inv_inertia_max,
                    friction_max,
                    position_max,
                    radius_max,
                    rotation_max,
                    velocity_max,
                    angular_velocity_max,
                    vertices_max,
                )
                env_state_pcg_mask = _set_single_state_rbody(
                    self.pcg_state.env_state_pcg_mask,
                    parent_container_pcg_mask,
                    density_pcg_mask,
                    density_pcg_mask,
                    density_pcg_mask,
                    friction_pcg_mask,
                    position_pcg_mask,
                    radius_pcg_mask,
                    rotation_pcg_mask,
                    velocity_pcg_mask,
                    angular_velocity_pcg_mask,
                    vertices_pcg_mask,
                )

                # Now, we must also set all of the tied shape's positions

                correct_index = self.selected_shape_index + (
                    self.static_env_params.num_polygons if self.selected_shape_type == ObjectType.CIRCLE else 0
                )

                nonzero_indices = set(
                    jnp.nonzero(self.pcg_state.tied_together[correct_index].at[correct_index].set(False))[0].tolist()
                )
                for idx in nonzero_indices:
                    if idx < self.static_env_params.num_polygons:
                        env_state = env_state.replace(
                            polygon=env_state.polygon.replace(
                                position=env_state.polygon.position.at[idx].set(
                                    env_state.polygon.position[idx] + position_delta
                                )
                            )
                        )
                    else:
                        idx = idx - self.static_env_params.num_polygons
                        env_state = env_state.replace(
                            circle=env_state.circle.replace(
                                position=env_state.circle.position.at[idx].set(
                                    env_state.circle.position[idx] + position_delta
                                )
                            )
                        )

                if not do_dummy:
                    self.pcg_state = PCGState(
                        env_state=env_state,
                        env_state_max=env_state_max,
                        env_state_pcg_mask=env_state_pcg_mask,
                        tied_together=self.pcg_state.tied_together,
                    )

            if do_dummy or self.selected_shape_type == ObjectType.JOINT:  # rjoint
                speed_main = self.all_widgets[ObjectType.JOINT]["sldSpeed"].getValue()
                speed_max = self.all_widgets[ObjectType.JOINT]["sldMaxSpeed"].getValue()
                speed_pcg_mask = self.all_widgets[ObjectType.JOINT]["pcgTglSpeed"].getValue()
                if speed_pcg_mask:
                    speed_main, speed_max = min(speed_main, speed_max), max(speed_main, speed_max)

                power_main = self.all_widgets[ObjectType.JOINT]["sldPower"].getValue()
                power_max = self.all_widgets[ObjectType.JOINT]["sldMaxPower"].getValue()
                power_pcg_mask = self.all_widgets[ObjectType.JOINT]["pcgTglPower"].getValue()
                if power_pcg_mask:
                    power_main, power_max = min(power_main, power_max), max(power_main, power_max)

                motor_binding_val_min = int(self.all_widgets[ObjectType.JOINT]["sldColour"].getValue())
                motor_binding_val_max = int(self.all_widgets[ObjectType.JOINT]["sldMaxColour"].getValue())
                colour_pcg_mask = self.all_widgets[ObjectType.JOINT]["pcgTglColour"].getValue()
                if colour_pcg_mask:
                    motor_binding_val_min, motor_binding_val_max = min(
                        motor_binding_val_min, motor_binding_val_max
                    ), max(motor_binding_val_min, motor_binding_val_max)

                auto_motor_val = self.all_widgets[ObjectType.JOINT]["tglAutoMotor"].getValue()
                joint_limits_val = self.all_widgets[ObjectType.JOINT]["tglJointLimits"].getValue()
                is_fixed_val = self.all_widgets[ObjectType.JOINT]["tglIsFixedJoint"].getValue()
                is_motor_on = self.all_widgets[ObjectType.JOINT]["tglIsMotorOn"].getValue()

                min_rot_val = jnp.deg2rad(self.all_widgets[ObjectType.JOINT]["sldMin_Rotation"].getValue())
                max_rot_val = jnp.deg2rad(self.all_widgets[ObjectType.JOINT]["sldMax_Rotation"].getValue())
                # ensure the min is less than the max
                min_rot_val, max_rot_val = min(min_rot_val, max_rot_val), max(min_rot_val, max_rot_val)

                def _set_single_state_joint(state, speed, power, colour):
                    state = state.replace(
                        joint=state.joint.replace(
                            motor_speed=state.joint.motor_speed.at[self.selected_shape_index].set(speed),
                            motor_power=state.joint.motor_power.at[self.selected_shape_index].set(power),
                            motor_has_joint_limits=state.joint.motor_has_joint_limits.at[self.selected_shape_index].set(
                                joint_limits_val
                            ),
                            min_rotation=state.joint.min_rotation.at[self.selected_shape_index].set(min_rot_val),
                            max_rotation=state.joint.max_rotation.at[self.selected_shape_index].set(max_rot_val),
                            is_fixed_joint=state.joint.is_fixed_joint.at[self.selected_shape_index].set(is_fixed_val),
                            motor_on=state.joint.motor_on.at[self.selected_shape_index].set(is_motor_on),
                        ),
                        motor_bindings=state.motor_bindings.at[self.selected_shape_index].set(colour),
                        motor_auto=state.motor_auto.at[self.selected_shape_index].set(auto_motor_val),
                    )
                    return state

                env_state = _set_single_state_joint(
                    self.pcg_state.env_state,
                    speed_main,
                    power_main,
                    motor_binding_val_min,
                )

                env_state_max = _set_single_state_joint(
                    self.pcg_state.env_state_max,
                    speed_max,
                    power_max,
                    motor_binding_val_max,
                )

                env_state_pcg_mask = _set_single_state_joint(
                    self.pcg_state.env_state_pcg_mask, speed_pcg_mask, power_pcg_mask, colour_pcg_mask
                )
                if not do_dummy:
                    self.pcg_state = self.pcg_state.replace(
                        env_state=env_state,
                        env_state_max=env_state_max,
                        env_state_pcg_mask=env_state_pcg_mask,
                    )

            if do_dummy or self.selected_shape_type == ObjectType.THRUSTER:  # thruster
                power_main = self.all_widgets[ObjectType.THRUSTER]["sldPower"].getValue()
                power_max = self.all_widgets[ObjectType.THRUSTER]["sldMaxPower"].getValue()
                power_pcg_mask = self.all_widgets[ObjectType.THRUSTER]["pcgTglPower"].getValue()
                if power_pcg_mask:
                    power_main, power_max = min(power_main, power_max), max(power_main, power_max)

                def _set_single_state_thruster(state, power):
                    return state.replace(
                        thruster=state.thruster.replace(
                            power=state.thruster.power.at[self.selected_shape_index].set(power),
                        ),
                        thruster_bindings=state.thruster_bindings.at[self.selected_shape_index].set(
                            int(self.all_widgets[ObjectType.THRUSTER]["sldColour"].getValue())
                        ),
                    )

                env_state = _set_single_state_thruster(self.pcg_state.env_state, power_main)
                env_state_max = _set_single_state_thruster(self.pcg_state.env_state_max, power_max)
                env_state_pcg_mask = _set_single_state_thruster(self.pcg_state.env_state_pcg_mask, power_pcg_mask)
                if not do_dummy:
                    self.pcg_state = self.pcg_state.replace(
                        env_state=env_state,
                        env_state_max=env_state_max,
                        env_state_pcg_mask=env_state_pcg_mask,
                    )

    # flag2
    def _put_state_values_into_gui(self, pcg_state=None):
        def _set_toggle_val(toggle, val):
            if toggle.getValue() != val:
                toggle.toggle()

        def _enable_slider(slider):
            slider.enable()
            slider.colour = (200, 200, 200)
            slider.handleColour = (0, 0, 0)

        def _disable_slider(slider):
            slider.disable()
            slider.colour = (255, 0, 0)
            slider.handleColour = (255, 0, 0)

        if pcg_state is None:
            # state = self.edit_state
            raise ValueError

        def pcg_text(main_val, max_val, pcg_mask):
            if pcg_mask:
                return f"{main_val:.2f} - {max_val:.2f}"
            else:
                return f"{main_val:.2f}"

        # global ones
        gravity_pcg_mask = pcg_state.env_state_pcg_mask.gravity[1]
        self.all_widgets[None]["lblGravity"].setText(
            f"Gravity: {pcg_text(pcg_state.env_state.gravity[1], pcg_state.env_state_max.gravity[1], pcg_state.env_state_pcg_mask.gravity[1])}"
        )
        self.all_widgets[None]["sldGravity"].setValue(pcg_state.env_state.gravity[1])
        self.all_widgets[None]["sldMaxGravity"].setValue(pcg_state.env_state_max.gravity[1])
        if not gravity_pcg_mask:
            _disable_slider(self.all_widgets[None]["sldMaxGravity"])
        else:
            _enable_slider(self.all_widgets[None]["sldMaxGravity"])

        if self.edit_shape_mode != EditMode.SELECT or self.selected_shape_index < 0:
            return

        obj_main = select_object(pcg_state.env_state, self.selected_shape_type, self.selected_shape_index)
        obj_max = select_object(pcg_state.env_state_max, self.selected_shape_type, self.selected_shape_index)
        obj_pcg_mask = select_object(pcg_state.env_state_pcg_mask, self.selected_shape_type, self.selected_shape_index)

        if len(self.all_selected_shapes) > 1:
            indices_to_use = self._get_selected_shape_global_indices()
            are_all_tied = pcg_state.tied_together[indices_to_use.min(), indices_to_use].all()
            _set_toggle_val(self.all_widgets["TIE_TOGETHER"]["tglTieTogether"], are_all_tied)

        if self.selected_shape_type == ObjectType.JOINT:
            self.all_widgets[ObjectType.JOINT]["lblSpeed"].setText(
                f"Speed: {pcg_text(obj_main.motor_speed, obj_max.motor_speed, obj_pcg_mask.motor_speed)}"
            )
            self.all_widgets[ObjectType.JOINT]["sldSpeed"].setValue(obj_main.motor_speed)
            self.all_widgets[ObjectType.JOINT]["sldMaxSpeed"].setValue(obj_max.motor_speed)
            _set_toggle_val(self.all_widgets[ObjectType.JOINT]["pcgTglSpeed"], obj_pcg_mask.motor_speed)
            if obj_pcg_mask.motor_speed:
                _enable_slider(self.all_widgets[ObjectType.JOINT]["sldMaxSpeed"])
            else:
                _disable_slider(self.all_widgets[ObjectType.JOINT]["sldMaxSpeed"])

            self.all_widgets[ObjectType.JOINT]["lblPower"].setText(
                f"Power: {pcg_text(obj_main.motor_power, obj_max.motor_power, obj_pcg_mask.motor_power)}"
            )
            self.all_widgets[ObjectType.JOINT]["sldPower"].setValue(obj_main.motor_power)
            self.all_widgets[ObjectType.JOINT]["sldMaxPower"].setValue(obj_max.motor_power)
            _set_toggle_val(self.all_widgets[ObjectType.JOINT]["pcgTglPower"], obj_pcg_mask.motor_power)
            if obj_pcg_mask.motor_power:
                _enable_slider(self.all_widgets[ObjectType.JOINT]["sldMaxPower"])
            else:
                _disable_slider(self.all_widgets[ObjectType.JOINT]["sldMaxPower"])

            self.all_widgets[ObjectType.JOINT]["lblColour"].setText(
                f"Colour: {pcg_state.env_state.motor_bindings[self.selected_shape_index]}"
            )
            self.all_widgets[ObjectType.JOINT]["sldColour"].setValue(
                pcg_state.env_state.motor_bindings[self.selected_shape_index]
            )

            self.all_widgets[ObjectType.JOINT]["sldMaxColour"].setValue(
                pcg_state.env_state_max.motor_bindings[self.selected_shape_index]
            )

            colour_pcg_mask = pcg_state.env_state_pcg_mask.motor_bindings[self.selected_shape_index]
            if not colour_pcg_mask:
                _disable_slider(self.all_widgets[ObjectType.JOINT]["sldMaxColour"])
            else:
                _enable_slider(self.all_widgets[ObjectType.JOINT]["sldMaxColour"])

            self.all_widgets[ObjectType.JOINT]["lblJointLimits"].setText(
                f"Joint Limits: {obj_main.motor_has_joint_limits}"
            )
            widget_motor_has_joint_limits = self.all_widgets[ObjectType.JOINT]["tglJointLimits"].getValue()
            if obj_main.motor_has_joint_limits != widget_motor_has_joint_limits:  # Update the toggle
                self.all_widgets[ObjectType.JOINT]["tglJointLimits"].toggle()

            mini, maxi = jnp.rad2deg(obj_main.min_rotation), jnp.rad2deg(obj_main.max_rotation)
            self.all_widgets[ObjectType.JOINT]["lblMin_Rotation"].setText(f"Min Rot: {int(mini)}")
            self.all_widgets[ObjectType.JOINT]["sldMin_Rotation"].setValue(mini)
            self.all_widgets[ObjectType.JOINT]["lblMax_Rotation"].setText(f"Max Rot: {int(maxi)}")
            self.all_widgets[ObjectType.JOINT]["sldMax_Rotation"].setValue(maxi)

            if not obj_main.motor_has_joint_limits:
                for k in ["min_rotation", "max_rotation"]:
                    self.all_widgets[self.selected_shape_type][f"sld{k.title()}"].disable()
                    self.all_widgets[self.selected_shape_type][f"sld{k.title()}"].colour = (255, 0, 0)
                    self.all_widgets[self.selected_shape_type][f"sld{k.title()}"].handleColour = (255, 0, 0)
            else:
                for k in ["min_rotation", "max_rotation"]:
                    self.all_widgets[self.selected_shape_type][f"sld{k.title()}"].enable()
                    self.all_widgets[self.selected_shape_type][f"sld{k.title()}"].colour = (200, 200, 200)
                    self.all_widgets[self.selected_shape_type][f"sld{k.title()}"].handleColour = (0, 0, 0)

            self.all_widgets[ObjectType.JOINT]["lblAutoMotor"].setText(
                f"Auto: {pcg_state.env_state.motor_auto[self.selected_shape_index]}"
            )
            widget_is_auto_motor = self.all_widgets[ObjectType.JOINT]["tglAutoMotor"].getValue()
            if pcg_state.env_state.motor_auto[self.selected_shape_index] != widget_is_auto_motor:  # Update the toggle
                self.all_widgets[ObjectType.JOINT]["tglAutoMotor"].toggle()

            self.all_widgets[ObjectType.JOINT]["lblIsFixedJoint"].setText(f"Fixed: {obj_main.is_fixed_joint}")
            widget_is_motor_on = self.all_widgets[ObjectType.JOINT]["tglIsFixedJoint"].getValue()
            if obj_main.is_fixed_joint != widget_is_motor_on:  # Update the toggle
                self.all_widgets[ObjectType.JOINT]["tglIsFixedJoint"].toggle()

            self.all_widgets[ObjectType.JOINT]["lblIsMotorOn"].setText(f"Motor On: {obj_main.motor_on}")
            widget_is_motor_on = self.all_widgets[ObjectType.JOINT]["tglIsMotorOn"].getValue()
            if obj_main.motor_on != widget_is_motor_on:  # Update the toggle
                self.all_widgets[ObjectType.JOINT]["tglIsMotorOn"].toggle()

        elif self.selected_shape_type == ObjectType.THRUSTER:
            # thruster
            self.all_widgets[ObjectType.THRUSTER]["lblPower"].setText(
                f"Power: {pcg_text(obj_main.power, obj_max.power, obj_pcg_mask.power)}"
            )
            self.all_widgets[ObjectType.THRUSTER]["sldPower"].setValue(obj_main.power)
            self.all_widgets[ObjectType.THRUSTER]["sldMaxPower"].setValue(obj_max.power)
            _set_toggle_val(self.all_widgets[ObjectType.THRUSTER]["pcgTglPower"], obj_pcg_mask.power)
            if obj_pcg_mask.power:
                _enable_slider(self.all_widgets[ObjectType.THRUSTER]["sldMaxPower"])
            else:
                _disable_slider(self.all_widgets[ObjectType.THRUSTER]["sldMaxPower"])
            self.all_widgets[ObjectType.THRUSTER]["sldColour"].setValue(
                pcg_state.env_state.thruster_bindings[self.selected_shape_index]
            )
            self.all_widgets[ObjectType.THRUSTER]["lblColour"].setText(
                f"Colour: {pcg_state.env_state.thruster_bindings[self.selected_shape_index]}"
            )

        elif self.selected_shape_type in [ObjectType.POLYGON, ObjectType.CIRCLE]:
            # rigidbody

            # Position
            # We use the mask for position_x for the entire position vector
            self.all_widgets[self.selected_shape_type]["lblPosition_X"].setText(
                f"Position X: {pcg_text(obj_main.position[0], obj_max.position[0], obj_pcg_mask.position[0])}"
            )
            self.all_widgets[self.selected_shape_type]["sldPosition_X"].setValue(obj_main.position[0])
            self.all_widgets[self.selected_shape_type]["sldMaxPosition_X"].setValue(obj_max.position[0])
            _set_toggle_val(self.all_widgets[self.selected_shape_type]["pcgTglPosition_X"], obj_pcg_mask.position[0])

            self.all_widgets[self.selected_shape_type]["lblPosition_Y"].setText(
                f"Position Y: {pcg_text(obj_main.position[1], obj_max.position[1], obj_pcg_mask.position[0])}"
            )
            self.all_widgets[self.selected_shape_type]["sldPosition_Y"].setValue(obj_main.position[1])
            self.all_widgets[self.selected_shape_type]["sldMaxPosition_Y"].setValue(obj_max.position[1])

            if obj_pcg_mask.position[0]:
                _enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxPosition_X"])
                _enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxPosition_Y"])
            else:
                _disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxPosition_X"])
                _disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxPosition_Y"])

            # Velocity
            # We use the mask for velocity_x for the entire velocity vector
            self.all_widgets[self.selected_shape_type]["lblVelocity_X"].setText(
                f"Velocity X: {pcg_text(obj_main.velocity[0], obj_max.velocity[0], obj_pcg_mask.velocity[0])}"
            )
            self.all_widgets[self.selected_shape_type]["sldVelocity_X"].setValue(obj_main.velocity[0])
            self.all_widgets[self.selected_shape_type]["sldMaxVelocity_X"].setValue(obj_max.velocity[0])
            _set_toggle_val(self.all_widgets[self.selected_shape_type]["pcgTglVelocity_X"], obj_pcg_mask.velocity[0])

            self.all_widgets[self.selected_shape_type]["lblVelocity_Y"].setText(
                f"Velocity Y: {pcg_text(obj_main.velocity[1], obj_max.velocity[1], obj_pcg_mask.velocity[0])}"
            )
            self.all_widgets[self.selected_shape_type]["sldVelocity_Y"].setValue(obj_main.velocity[1])
            self.all_widgets[self.selected_shape_type]["sldMaxVelocity_Y"].setValue(obj_max.velocity[1])

            if obj_pcg_mask.velocity[0]:
                _enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxVelocity_X"])
                _enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxVelocity_Y"])
            else:
                _disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxVelocity_X"])
                _disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxVelocity_Y"])

            # Density
            is_fixated = obj_main.inverse_mass == 0

            def _calc_density(state):
                if self.selected_shape_type == ObjectType.POLYGON:
                    return state.polygon_densities[self.selected_shape_index]
                elif self.selected_shape_type == ObjectType.CIRCLE:
                    return state.circle_densities[self.selected_shape_index]
                else:
                    raise ValueError

            density_main = _calc_density(pcg_state.env_state)
            density_max = _calc_density(pcg_state.env_state_max)
            density_pcg_mask = obj_pcg_mask.inverse_mass

            self.all_widgets[self.selected_shape_type]["lblDensity"].setText(
                f"Density: {pcg_text(density_main, density_max, density_pcg_mask)}"
            )
            self.all_widgets[self.selected_shape_type]["sldDensity"].setValue(density_main)
            self.all_widgets[self.selected_shape_type]["sldMaxDensity"].setValue(density_max)
            _set_toggle_val(self.all_widgets[self.selected_shape_type]["pcgTglDensity"], density_pcg_mask)
            if is_fixated:
                _disable_slider(self.all_widgets[self.selected_shape_type]["sldDensity"])
            else:
                _enable_slider(self.all_widgets[self.selected_shape_type]["sldDensity"])

            if is_fixated or (not density_pcg_mask):
                _disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxDensity"])
            else:
                _enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxDensity"])

            # Friction
            self.all_widgets[self.selected_shape_type]["lblFriction"].setText(
                f"Friction: {pcg_text(obj_main.friction, obj_max.friction, obj_pcg_mask.friction)}"
            )
            self.all_widgets[self.selected_shape_type]["sldFriction"].setValue(obj_main.friction)
            self.all_widgets[self.selected_shape_type]["sldMaxFriction"].setValue(obj_max.friction)
            _set_toggle_val(self.all_widgets[self.selected_shape_type]["pcgTglFriction"], obj_pcg_mask.friction)

            if not obj_pcg_mask.friction:
                _disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxFriction"])
            else:
                _enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxFriction"])

            # Restitution
            self.all_widgets[self.selected_shape_type]["lblRestitution"].setText(
                f"Restitution: {obj_main.restitution:.2f}"
            )
            self.all_widgets[self.selected_shape_type]["sldRestitution"].setValue(obj_main.restitution)

            # Rotation
            self.all_widgets[self.selected_shape_type]["lblRotation"].setText(
                f"Rotation: {pcg_text(obj_main.rotation, obj_max.rotation, obj_pcg_mask.rotation)}"
            )
            self.all_widgets[self.selected_shape_type]["sldRotation"].setValue(obj_main.rotation)
            self.all_widgets[self.selected_shape_type]["sldMaxRotation"].setValue(obj_max.rotation)
            _set_toggle_val(self.all_widgets[self.selected_shape_type]["pcgTglRotation"], obj_pcg_mask.rotation)

            if not obj_pcg_mask.rotation:
                _disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxRotation"])
            else:
                _enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxRotation"])

            # Angular_Velocity
            self.all_widgets[self.selected_shape_type]["lblAngular_Velocity"].setText(
                f"Angular_Velocity: {pcg_text(obj_main.angular_velocity, obj_max.angular_velocity, obj_pcg_mask.angular_velocity)}"
            )
            self.all_widgets[self.selected_shape_type]["sldAngular_Velocity"].setValue(obj_main.angular_velocity)
            self.all_widgets[self.selected_shape_type]["sldMaxAngular_Velocity"].setValue(obj_max.angular_velocity)
            _set_toggle_val(
                self.all_widgets[self.selected_shape_type]["pcgTglAngular_Velocity"], obj_pcg_mask.angular_velocity
            )

            if not obj_pcg_mask.angular_velocity:
                _disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxAngular_Velocity"])
            else:
                _enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxAngular_Velocity"])

            # Collision mode
            self.all_widgets[self.selected_shape_type]["lblCollidability"].setText(
                f"Collidability: {obj_main.collision_mode}"
            )
            self.all_widgets[self.selected_shape_type]["sldCollidability"].setValue(obj_main.collision_mode)

            # Shape role
            if self.selected_shape_type == ObjectType.POLYGON:
                shape_role = pcg_state.env_state.polygon_shape_roles[self.selected_shape_index]
            else:
                shape_role = pcg_state.env_state.circle_shape_roles[self.selected_shape_index]

            self.all_widgets[self.selected_shape_type]["sldRole"].setValue(shape_role)
            self.all_widgets[self.selected_shape_type]["lblRole"].setText(f"Role: {shape_role}")

            # Fixate
            self.all_widgets[self.selected_shape_type]["lblFixate"].setText(f"Fixate: {is_fixated}")
            widget_is_fixed = self.all_widgets[self.selected_shape_type]["tglFixate"].getValue()
            if is_fixated != widget_is_fixed:  # Update the toggle
                self.all_widgets[self.selected_shape_type]["tglFixate"].toggle()

            # Radius
            if self.selected_shape_type == ObjectType.CIRCLE:
                self.all_widgets[self.selected_shape_type]["lblRadius"].setText(
                    f"Radius: {pcg_text(obj_main.radius, obj_max.radius, obj_pcg_mask.radius)}"
                )
                self.all_widgets[self.selected_shape_type]["sldRadius"].setValue(obj_main.radius)
                self.all_widgets[self.selected_shape_type]["sldMaxRadius"].setValue(obj_max.radius)
                _set_toggle_val(self.all_widgets[self.selected_shape_type]["pcgTglRadius"], obj_pcg_mask.radius)

                if not obj_pcg_mask.radius:
                    _disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxRadius"])
                else:
                    _enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxRadius"])

            elif self.selected_shape_type == ObjectType.POLYGON:
                size_main = jnp.abs(obj_main.vertices).max()
                size_max = jnp.abs(obj_max.vertices).max()
                size_pcg_mask = obj_pcg_mask.vertices.any()

                self.all_widgets[self.selected_shape_type]["lblSize"].setText(
                    f"Size: {pcg_text(size_main, size_max, size_pcg_mask)}"
                )
                self.all_widgets[self.selected_shape_type]["sldSize"].setValue(size_main)
                self.all_widgets[self.selected_shape_type]["sldMaxSize"].setValue(size_max)
                _set_toggle_val(self.all_widgets[self.selected_shape_type]["pcgTglSize"], size_pcg_mask)

                if not size_pcg_mask:
                    _disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxSize"])
                else:
                    _enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxSize"])

    def _render_side_panel(self):
        arr = jnp.ones((self.side_panel_width, self.static_env_params.screen_dim[1], 3)) * (
            jnp.array([135.0, 206.0, 235.0])[None, None] + 20
        )
        return arr

    def make_label_and_slider(
        self,
        start_y,
        label_text,
        slider_min=0.0,
        slider_max=1.0,
        slider_step=0.05,
        font_size=18,
        is_toggle=False,
        is_pcg=False,
        add_pcg_toggle=True,
    ):
        Wl = round(self.W * 0.7)
        pcg_lbl = None
        pcg_toggle = None
        widget_max = None

        label = TextBox(
            self.screen_surface,
            self.MARGIN,
            start_y,
            Wl,
            20,
            fontSize=font_size,
            margin=0,
            placeholderText=label_text,
            font=pygame.font.SysFont("sans-serif", font_size),
        )
        label.disable()  # Act as label instead of textbox
        if is_toggle:
            widget = Toggle(self.screen_surface, self.W // 2 - 20, start_y + 23, 20, 13)
        else:
            widget = Slider(
                self.screen_surface,
                self.MARGIN,
                start_y + 23,
                Wl,
                13,
                min=slider_min,
                max=slider_max,
                step=slider_step,
            )

            if is_pcg:
                widget_max = Slider(
                    self.screen_surface,
                    self.MARGIN,
                    start_y + 50,
                    Wl,
                    13,
                    min=slider_min,
                    max=slider_max,
                    step=slider_step,
                )

                pcg_lbl = TextBox(
                    self.screen_surface,
                    self.MARGIN + Wl + 30,
                    start_y,
                    60,
                    20,
                    fontSize=font_size,
                    margin=0,
                    placeholderText="PCG",
                    font=pygame.font.SysFont("sans-serif", font_size),
                )
                if add_pcg_toggle:
                    pcg_toggle = Toggle(self.screen_surface, Wl + 80, start_y + 23, 20, 13)

        return label, widget, (pcg_toggle, pcg_lbl, widget_max)

    def _setup_side_panel(self):

        W = self.W = int(self.side_panel_width * self.upscale * 0.8)
        MARGIN = self.MARGIN = (self.side_panel_width * self.upscale - W) // 2

        # global values
        G = {}
        gravity_label, gravity_slider, (pcg_toggle, pcg_lbl, slider_max) = self.make_label_and_slider(
            150, "Gravity", -20.0, 0.0, 0.2, is_pcg=True
        )
        G["lblGravity"] = gravity_label
        G["sldGravity"] = gravity_slider
        G["sldMaxGravity"] = slider_max
        G["pcgLblGravity"] = pcg_lbl
        G["pcgTglGravity"] = pcg_toggle

        # thruster values
        T = {}
        thruster_power_label, thruster_power_slider, (pcg_toggle, pcg_lbl, slider_max) = self.make_label_and_slider(
            150, "Power", slider_max=3.0, is_pcg=True
        )
        T["lblPower"] = thruster_power_label
        T["sldPower"] = thruster_power_slider
        T["sldMaxPower"] = slider_max
        T["pcgLblPower"] = pcg_lbl
        T["pcgTglPower"] = pcg_toggle
        thruster_colour_label, thruster_colour_slider, _ = self.make_label_and_slider(
            250, "Colour", 0, self.static_env_params.num_thruster_bindings - 1, 1
        )
        T["lblColour"] = thruster_colour_label
        T["sldColour"] = thruster_colour_slider

        # joints
        D = {}
        for i, (name, (mini, maxi, step), is_pcg) in enumerate(
            zip(
                ["speed", "power", "colour", "min_rotation", "max_rotation"],
                [
                    (-3, 3, 0.05),
                    (0, 3, 0.05),
                    (0, self.static_env_params.num_motor_bindings - 1, 1),
                    (-180, 180, 5),
                    (-180, 180, 5),
                ],
                [True, True, True, False, False],
            )
        ):
            label, slider, (pcg_toggle, pcg_lbl, slider_max) = self.make_label_and_slider(
                150 + 80 * i,
                f"Motor {name.title()}",
                slider_min=mini,
                slider_max=maxi,
                slider_step=step,
                is_pcg=is_pcg,
            )
            D["lbl" + name.title()] = label
            D["sld" + name.title()] = slider
            if is_pcg:
                D["sldMax" + name.title()] = slider_max
                D["pcgLbl" + name.title()] = pcg_lbl
                D["pcgTgl" + name.title()] = pcg_toggle

        label, toggle, _ = self.make_label_and_slider(150 + 80 * 5, "Joint Limits", is_toggle=True)
        D["lblJointLimits"] = label
        D["tglJointLimits"] = toggle

        label, toggle, _ = self.make_label_and_slider(150 + 80 * 6, "Auto", is_toggle=True)
        D["lblAutoMotor"] = label
        D["tglAutoMotor"] = toggle

        label, toggle, _ = self.make_label_and_slider(150 + 80 * 7, "Fixed", is_toggle=True)
        D["lblIsFixedJoint"] = label
        D["tglIsFixedJoint"] = toggle

        label, toggle, _ = self.make_label_and_slider(150 + 80 * 8, "Motor On", is_toggle=True)
        D["lblIsMotorOn"] = label
        D["tglIsMotorOn"] = toggle

        def _create_rigid_body_base_gui():
            D_rigid = {}
            # rigidbodies
            total_toggles = 0
            total_non_toggles = 0
            for i, (name, bounds, is_pcg, add_pcg_toggle) in enumerate(
                zip(
                    [
                        "position_x",
                        "position_y",
                        "rotation",
                        "velocity_x",
                        "velocity_y",
                        "angular_velocity",
                        "density",
                        "friction",
                        "restitution",
                        "collidability",
                        "role",
                    ],
                    [
                        (0, 5.0, 0.01),
                        (0, 5.0, 0.01),
                        (-2 * jnp.pi, 2 * jnp.pi, 0.01),
                        (-10.0, 10.0, 0.1),
                        (-10.0, 10.0, 0.1),
                        (-6, 6.0, 0.01),
                        (0.1, 5.0, 0.1),
                        (0.02, 1.0, 0.02),
                        (0.0, 0.8, 0.02),
                        (0, 2, 1),
                        (0, 3, 1),
                    ],
                    [True, True, True, True, True, True, True, True, False, False, False],
                    [True, False, True, True, False, True, True, True, False, False, False],
                )
            ):
                location = 50 + 80 * total_toggles + 40 * total_non_toggles
                label, slider, (pcg_toggle, pcg_lbl, slider_max) = self.make_label_and_slider(
                    location, name.title(), *bounds, is_pcg=is_pcg, add_pcg_toggle=add_pcg_toggle
                )
                total_toggles += is_pcg
                total_non_toggles += not is_pcg
                D_rigid["lbl" + name.title()] = label
                D_rigid["sld" + name.title()] = slider

                if is_pcg:
                    D_rigid["sldMax" + name.title()] = slider_max
                    if add_pcg_toggle:
                        D_rigid["pcgLbl" + name.title()] = pcg_lbl
                        D_rigid["pcgTgl" + name.title()] = pcg_toggle

            location = 50 + 80 * total_toggles + 40 * total_non_toggles
            # toggles:
            label, toggle, _ = self.make_label_and_slider(location, "Fixate", is_toggle=True)
            D_rigid["lblFixate"] = label
            D_rigid["tglFixate"] = toggle

            return D_rigid, location

        # Circle extras
        D_circle, location = _create_rigid_body_base_gui()
        label, slider, (pcg_toggle, pcg_lbl, slider_max) = self.make_label_and_slider(
            location + 40, "Radius", slider_min=0.1, slider_max=1.0, slider_step=0.02, is_pcg=True, add_pcg_toggle=True
        )
        D_circle["lblRadius"] = label
        D_circle["sldRadius"] = slider
        D_circle["sldMaxRadius"] = slider_max
        D_circle["pcgLblRadius"] = pcg_lbl
        D_circle["pcgTglRadius"] = pcg_toggle

        # Polygon extras
        D_poly, location = _create_rigid_body_base_gui()
        label, slider, (pcg_toggle, pcg_lbl, slider_max) = self.make_label_and_slider(
            location + 40, "Size", slider_min=0.1, slider_max=2.0, slider_step=0.02, is_pcg=True, add_pcg_toggle=True
        )
        D_poly["lblSize"] = label
        D_poly["sldSize"] = slider
        D_poly["sldMaxSize"] = slider_max
        D_poly["pcgLblSize"] = pcg_lbl
        D_poly["pcgTglSize"] = pcg_toggle

        label, toggle, _ = self.make_label_and_slider(150 + 80 * 5, "Tie Positions Together", is_toggle=True)
        TIE_TOGETHER = {
            "lblTieTogether": label,
            "tglTieTogether": toggle,
        }

        self.all_widgets = {
            ObjectType.THRUSTER: T,
            ObjectType.JOINT: D,
            "GENERAL": {
                "lblGeneral": TextBox(
                    self.screen_surface,
                    MARGIN,
                    10,
                    W,
                    30,
                    fontSize=20,
                    margin=0,
                    placeholderText="General",
                    font=pygame.font.SysFont("sans-serif", 35),
                ),
            },
            ObjectType.POLYGON: D_poly,
            ObjectType.CIRCLE: D_circle,
            None: G,
            "TIE_TOGETHER": TIE_TOGETHER,
        }

        self._hide_all_widgets()

    def _render_edit_overlay(self, pixels, is_editing, edit_shape_mode):
        is_editing_texture = jax.lax.select(is_editing, EDIT_TEXTURE_RGBA, PLAY_TEXTURE_RGBA)
        is_editing_texture = jnp.repeat(jnp.repeat(is_editing_texture, self.upscale, axis=0), self.upscale, axis=1)

        offset = self.side_panel_width * self.upscale
        w = 64 * self.upscale
        offset2 = int(w * 1.25)
        offset_y = 16 * self.upscale

        play_tex_with_background = (1 - is_editing_texture[:, :, 3:]) * pixels[
            offset + 0 : offset + w, 0:w
        ] + is_editing_texture[:, :, 3:] * is_editing_texture[:, :, :3]
        pixels = pixels.at[offset : offset + w, 0:w].set(play_tex_with_background)

        edit_shape_texture = jax.lax.switch(
            edit_shape_mode,
            [
                lambda: CIRCLE_TEXTURE_RGBA,
                lambda: RECT_TEXTURE_RGBA,
                lambda: RJOINT_TEXTURE_RGBA,
                lambda: SELECT_TEXTURE_RGBA,
                lambda: TRIANGLE_TEXTURE_RGBA,
                lambda: THRUSTER_TEXTURE_RGBA,
            ],
        )

        edit_shape_texture = jnp.repeat(jnp.repeat(edit_shape_texture, self.upscale, axis=0), self.upscale, axis=1)

        edit_shape_texture_alpha = edit_shape_texture[:, :, 3:] * is_editing

        w = 32 * self.upscale
        edit_shape_texture_with_background = (1 - edit_shape_texture_alpha) * pixels[
            offset + offset2 : offset + offset2 + w, offset_y : offset_y + w
        ] + edit_shape_texture_alpha * edit_shape_texture[:, :, :3]
        pixels = pixels.at[offset + offset2 : offset + offset2 + w, offset_y : offset_y + w].set(
            edit_shape_texture_with_background
        )

        return pixels

    def edit(self):
        self.rng, _rng = jax.random.split(self.rng)
        pcg_state = self.pcg_state

        left_click = False
        right_click = False

        keys = []
        keys_up_this_frame = set()
        for event in self.pygame_events:
            if event.type == pygame.KEYDOWN:
                if (
                    event.key == pygame.K_s
                    and (pygame.key.get_mods() & pygame.KMOD_CTRL)
                    and (pygame.key.get_mods() & pygame.KMOD_SHIFT)
                ):
                    filename = prompt_file(save=True)
                    if filename:
                        filename += ".level.pkl"
                        save_pickle(filename, self.last_played_level)
                        print(f"Saved last sampled level to {filename}")
                elif event.key == pygame.K_s and (pygame.key.get_mods() & pygame.KMOD_CTRL):
                    pcg_state = self._reset_select_shape(pcg_state)
                    filename = prompt_file(save=True)
                    if filename:
                        if filename.endswith(".json"):
                            export_env_state_to_json(
                                filename, pcg_state.env_state, self.static_env_params, self.env_params
                            )
                        elif not filename.endswith(".pcg.pkl"):
                            filename += ".pcg.pkl"
                            save_pickle(filename, pcg_state)
                        print(f"Saved PCG state to {filename}")
                elif event.key == pygame.K_o and (pygame.key.get_mods() & pygame.KMOD_CTRL):
                    filename = prompt_file(save=False)
                    if filename:
                        self._reset_select_shape(pcg_state)
                        if filename.endswith(".pcg.pkl"):
                            pcg_state = load_pcg_state_pickle(filename)
                            pcg_state = expand_pcg_state(pcg_state, self.static_env_params)
                            print(f"Loaded PCG state from {filename}")
                        elif filename.endswith(".level.pkl"):
                            env_state = load_world_state_pickle(filename)
                            pcg_state = env_state_to_pcg_state(env_state)
                            print(f"Converted level state to PCG state from {filename}")
                        elif filename.endswith(".json"):
                            env_state, new_static_env_params, new_env_params = load_from_json_file(filename)
                            self._update_params(new_static_env_params, new_env_params)
                            pcg_state = env_state_to_pcg_state(env_state)
                    self._reset_triangles()
                elif event.key == pygame.K_n and (pygame.key.get_mods() & pygame.KMOD_CTRL):
                    self._reset_select_shape(pcg_state)
                    pcg_state = new_pcg_env(self.static_env_params)
                    self._reset_triangles()
                else:
                    keys.append(event.key)
            elif event.type == pygame.KEYUP:
                keys_up_this_frame.add(event.key)
            if event.type == pygame.MOUSEBUTTONDOWN and self._get_mouse_position_world_space()[0] >= 0:
                if event.button == 1:
                    left_click = True
                if event.button == 3:
                    right_click = True

            if event.type == pygame.MOUSEWHEEL:
                pcg_state = self._handle_scroll_wheel(pcg_state, event.y)

        if self.selected_shape_index == -1:
            num = get_numeric_key_pressed(self.pygame_events)
            if num is not None:
                self.edit_shape_mode = EditMode(num % len(EditMode))

        # We have to do these checks outside the loop, otherwise they get triggered multiple times per key press.
        if pygame.key.get_mods() & pygame.KMOD_SHIFT:
            state = pcg_state.env_state
            if pygame.K_m in keys_up_this_frame:
                state = self.mutate_world(_rng, state, 1)
            if pygame.K_c in keys_up_this_frame:
                state, _ = mutate_add_connected_shape(
                    _rng, state, self.env_params, self.static_env_params, self.ued_params
                )
            elif pygame.K_s in keys_up_this_frame:
                state = mutate_add_shape(_rng, state, self.env_params, self.static_env_params, self.ued_params)
            elif pygame.K_p in keys_up_this_frame:
                state = mutate_swap_role(_rng, state, self.env_params, self.static_env_params, self.ued_params)
            elif pygame.K_r in keys_up_this_frame:
                state = mutate_remove_shape(_rng, state, self.env_params, self.static_env_params, self.ued_params)
            elif pygame.K_j in keys_up_this_frame:
                state = mutate_remove_joint(_rng, state, self.env_params, self.static_env_params, self.ued_params)
            elif pygame.K_t in keys_up_this_frame:
                state = mutate_toggle_fixture(_rng, state, self.env_params, self.static_env_params, self.ued_params)
            elif pygame.K_g in keys_up_this_frame:
                state = mutate_add_thruster(_rng, state, self.env_params, self.static_env_params, self.ued_params)
            elif pygame.K_l in keys_up_this_frame:
                state = mutate_remove_thruster(_rng, state, self.env_params, self.static_env_params, self.ued_params)
            elif pygame.K_b in keys_up_this_frame:
                state = self.mutate_change_shape_size(
                    _rng, state, self.env_params, self.static_env_params, self.ued_params
                )
            elif pygame.K_x in keys_up_this_frame:
                state = mutate_change_shape_location(
                    _rng, state, self.env_params, self.static_env_params, self.ued_params
                )
            elif pygame.K_k in keys_up_this_frame:
                state = self.mutate_change_shape_rotation(
                    _rng, state, self.env_params, self.static_env_params, self.ued_params
                )
            pcg_state = pcg_state.replace(env_state=state)

        if pygame.K_p in keys_up_this_frame:
            global myrng
            myrng, _rng = jax.random.split(myrng)
            # use the same rng
            pcg_state = permute_pcg_state(_rng, pcg_state, self.static_env_params)
        if self.edit_shape_mode == EditMode.SELECT:  # select a shape
            pcg_state = self._edit_select_shape(pcg_state, left_click, right_click, keys)
            self._put_state_values_into_gui(pcg_state)
            pcg_state = self._select_shape_keyboard_shortcuts(pcg_state, left_click, keys)
        else:
            pcg_state = self._reset_select_shape(pcg_state)  # don't highlight
            self._put_state_values_into_gui(pcg_state)
            self._show_correct_widgets(None)

        if self.edit_shape_mode != EditMode.ADD_TRIANGLE or not self.creating_shape:
            self.num_triangle_clicks = 0

        if self.edit_shape_mode == EditMode.ADD_CIRCLE:
            pcg_state = self._edit_circle(pcg_state, left_click, right_click)
        elif self.edit_shape_mode == EditMode.ADD_RECTANGLE:
            pcg_state = self._edit_rect(pcg_state, left_click, right_click)
        elif self.edit_shape_mode == EditMode.ADD_JOINT:
            pcg_state = self._edit_joint(pcg_state, left_click, right_click)
        elif self.edit_shape_mode == EditMode.ADD_TRIANGLE:
            pcg_state = self._edit_triangle(pcg_state, left_click, right_click)
        elif self.edit_shape_mode == EditMode.ADD_THRUSTER:
            pcg_state = self._edit_thruster(pcg_state, left_click, right_click)

        pcg_state = pcg_state.replace(
            env_state=recompute_global_joint_positions(
                pcg_state.env_state.replace(
                    collision_matrix=calculate_collision_matrix(self.static_env_params, pcg_state.env_state.joint),
                ),
                self.static_env_params,
            ),
            env_state_pcg_mask=pcg_state.env_state_pcg_mask.replace(
                collision_matrix=jnp.zeros_like(pcg_state.env_state_pcg_mask.collision_matrix)
            ),
        )
        return pcg_state

    def _update_params(self, new_static_env_params: StaticEnvParams, new_env_params: EnvParams):
        self.static_env_params = new_static_env_params.replace(
            frame_skip=self.config["frame_skip"], downscale=self.config["downscale"]
        )
        self.env_params = new_env_params
        env = make_kinetix_env_from_name("Kinetix-Entity-MultiDiscrete-v1", static_env_params=self.static_env_params)
        self.env = AutoResetWrapper(env, make_reset_function(self.static_env_params))
        self._setup_rendering(self.static_env_params, self.env_params)

    def _discard_shape_being_created(self, pcg_state):
        env_state = pcg_state.env_state
        if self.creating_shape:
            if self.edit_shape_mode == EditMode.ADD_CIRCLE:
                env_state = env_state.replace(
                    circle=env_state.circle.replace(
                        active=env_state.circle.active.at[self.creating_shape_index].set(False)
                    )
                )

            elif self.edit_shape_mode == EditMode.ADD_RECTANGLE:
                env_state = env_state.replace(
                    polygon=env_state.polygon.replace(
                        active=env_state.polygon.active.at[self.creating_shape_index].set(False)
                    )
                )

        self.creating_shape = False

        return pcg_state.replace(env_state=env_state)

    def _handle_scroll_wheel(self, pcg_state, y):
        if y == 0:
            return pcg_state

        state = self._discard_shape_being_created(pcg_state)

        self.edit_shape_mode = EditMode((self.edit_shape_mode.value + y) % len(EditMode))

        return state

    def _get_mouse_position_world_space(self):
        mouse_pos = pygame.mouse.get_pos()
        return (
            jnp.array(
                [
                    mouse_pos[0] / self.upscale - self.side_panel_width,
                    self.static_env_params.screen_dim[1] - mouse_pos[1] / self.upscale,
                ]
            )
            / self.env_params.pixels_per_unit
        )

    def _get_circles_on_mouse(self, state):
        mouse_pos = self._get_mouse_position_world_space()

        cis = []

        for ci in jnp.arange(self.static_env_params.num_circles)[::-1]:
            circle = jax.tree.map(lambda x: x[ci], state.circle)

            if not circle.active:
                continue

            dist = jnp.linalg.norm(mouse_pos - circle.position)
            if dist <= circle.radius:
                cis.append(ci)

        return cis

    def _get_revolute_joints_on_mouse(self, state: EnvState):
        mouse_pos = self._get_mouse_position_world_space()

        ris = []

        for ri in jnp.arange(self.static_env_params.num_joints)[::-1]:
            joint = jax.tree.map(lambda x: x[ri], state.joint)

            if not joint.active:
                continue

            dist = jnp.linalg.norm(mouse_pos - joint.global_position)
            if dist <= 10 / 100:  # arbitrary
                ris.append(ri)

        return ris

    def _get_thrusters_on_mouse(self, state: EnvState):
        mouse_pos = self._get_mouse_position_world_space()

        ris = []

        for ri in jnp.arange(self.static_env_params.num_thrusters)[::-1]:
            thruster = jax.tree.map(lambda x: x[ri], state.thruster)

            if not thruster.active:
                continue

            dist = jnp.linalg.norm(mouse_pos - thruster.global_position)
            if dist <= 16 / 100:  # arbitrary
                ris.append(ri)

        return ris

    def _get_joints_attached_to_shape(self, state, shape_index):
        r_a = jnp.arange(self.static_env_params.num_joints)[state.joint.a_index == shape_index]
        r_b = jnp.arange(self.static_env_params.num_joints)[state.joint.b_index == shape_index]

        t = jnp.arange(self.static_env_params.num_thrusters)[state.thruster.object_index == shape_index]
        return jnp.concatenate([r_a, r_b], axis=0), t

    def _edit_thruster(self, pcg_state: PCGState, left_click: bool, right_click: bool):
        if not self.creating_shape and (1 - pcg_state.env_state.thruster.active.astype(int)).sum() == 0:
            if not right_click:
                return pcg_state

        thruster_pos = self._get_mouse_position_world_space()
        idx = -1
        for ri in self._get_polygons_on_mouse(pcg_state.env_state):
            r = jax.tree.map(lambda x: x[ri], pcg_state.env_state.polygon)
            thruster_pos = snap_to_polygon_center_line(r, thruster_pos)
            relative_pos = jnp.matmul(rmat(r.rotation).transpose((1, 0)), thruster_pos - r.position)
            idx = ri
            break
        if idx == -1:
            for ci in self._get_circles_on_mouse(pcg_state.env_state):
                c = jax.tree.map(lambda x: x[ci], pcg_state.env_state.circle)
                thruster_pos = snap_to_center(c, thruster_pos)
                thruster_pos = snap_to_circle_center_line(c, thruster_pos)
                relative_pos = thruster_pos - c.position
                idx = ci + self.static_env_params.num_polygons
                break
        if left_click:
            if self.creating_shape:
                self.creating_shape = False
            else:
                if idx >= 0:
                    self.creating_shape = True
                    self.creating_shape_position = thruster_pos
                    self.creating_shape_index = jnp.argmin(pcg_state.env_state.thruster.active)
                    shape = select_shape(pcg_state.env_state, idx, self.static_env_params)

                    def _add_thruster_to_state(state):
                        state = state.replace(
                            thruster=state.thruster.replace(
                                object_index=state.thruster.object_index.at[self.creating_shape_index].set(idx),
                                relative_position=state.thruster.relative_position.at[self.creating_shape_index].set(
                                    relative_pos
                                ),
                                power=state.thruster.power.at[self.creating_shape_index].set(
                                    1.0 / jax.lax.select(shape.inverse_mass == 0, 1.0, shape.inverse_mass)
                                ),
                                active=state.thruster.active.at[self.creating_shape_index].set(True),
                                global_position=state.thruster.global_position.at[self.creating_shape_index].set(
                                    thruster_pos
                                ),
                                rotation=state.thruster.rotation.at[self.creating_shape_index].set(0.0),
                            ),
                            thruster_bindings=state.thruster_bindings.at[self.creating_shape_index].set(0),
                        )
                        return state

                    pcg_state = pcg_state.replace(
                        env_state=_add_thruster_to_state(pcg_state.env_state),
                        env_state_max=_add_thruster_to_state(pcg_state.env_state_max),
                    )
        elif right_click:
            for ti in self._get_thrusters_on_mouse(pcg_state.env_state):

                def _remove_thruster_from_state(state):
                    return state.replace(
                        thruster=state.thruster.replace(active=state.thruster.active.at[ti].set(False))
                    )

                return pcg_state.replace(
                    env_state=_remove_thruster_from_state(pcg_state.env_state),
                    env_state_max=_remove_thruster_from_state(pcg_state.env_state_max),
                )
        else:
            if self.creating_shape:
                curr_pos = self._get_mouse_position_world_space()

                normal = pcg_state.env_state.thruster.relative_position[self.creating_shape_index]
                angle = jnp.arctan2(normal[1], normal[0])

                relative_pos = curr_pos - self.creating_shape_position
                # rotation = jnp.arctan2(relative_pos[1], relative_pos[0])
                rotation = jnp.pi + jnp.arctan2(relative_pos[1], relative_pos[0]) + angle
                angle_round = jnp.round(rotation / (jnp.pi / 2))
                angle_norm = rotation / (jnp.pi / 2)
                if jnp.abs(angle_round - angle_norm) < 0.3:
                    rotation = angle_round * (jnp.pi / 2)

                def _update_thruster_rotation(state):
                    return state.replace(
                        thruster=state.thruster.replace(
                            rotation=state.thruster.rotation.at[self.creating_shape_index].set(rotation - angle),
                        )
                    )

                pcg_state = pcg_state.replace(
                    env_state=_update_thruster_rotation(pcg_state.env_state),
                    env_state_max=_update_thruster_rotation(pcg_state.env_state_max),
                )
            else:
                pass

        return pcg_state

    def _edit_circle(self, pcg_state: PCGState, left_click: bool, right_click: bool):
        if right_click:
            for ci in self._get_circles_on_mouse(pcg_state.env_state):
                attached_j, attached_t = self._get_joints_attached_to_shape(
                    pcg_state.env_state, ci + self.static_env_params.num_polygons
                )

                def _remove_circle_from_state(state):
                    return state.replace(
                        circle=state.circle.replace(active=state.circle.active.at[ci].set(False)),
                        joint=state.joint.replace(active=state.joint.active.at[attached_j].set(False)),
                        thruster=state.thruster.replace(active=state.thruster.active.at[attached_t].set(False)),
                    )

                env_state = _remove_circle_from_state(pcg_state.env_state)
                env_state_pcg_mask = _remove_circle_from_state(pcg_state.env_state_pcg_mask)
                env_state_max = _remove_circle_from_state(pcg_state.env_state_max)

                env_state = env_state.replace(
                    collision_matrix=calculate_collision_matrix(self.static_env_params, env_state.joint)
                )

                return PCGState(
                    env_state=env_state,
                    env_state_pcg_mask=env_state_pcg_mask,
                    env_state_max=env_state_max,
                    tied_together=pcg_state.tied_together,
                )

        if not self.creating_shape and (1 - pcg_state.env_state.circle.active.astype(int)).sum() == 0:
            return pcg_state

        radius = jnp.linalg.norm(self._get_mouse_position_world_space() - self.create_shape_position)
        radius = jnp.clip(radius, 5.0 / self.env_params.pixels_per_unit, self.static_env_params.max_shape_size / 2)

        def _add_circle(state, highlight):
            state = state.replace(
                circle=state.circle.replace(
                    position=state.circle.position.at[self.creating_shape_index].set(self.create_shape_position),
                    velocity=state.circle.velocity.at[self.creating_shape_index].set(jnp.array([0.0, 0.0])),
                    radius=state.circle.radius.at[self.creating_shape_index].set(radius),
                    inverse_mass=state.circle.inverse_mass.at[self.creating_shape_index].set(1.0),
                    inverse_inertia=state.circle.inverse_inertia.at[self.creating_shape_index].set(1.0),
                    active=state.circle.active.at[self.creating_shape_index].set(True),
                    collision_mode=state.circle.collision_mode.at[self.creating_shape_index].set(1),
                ),
                circle_shape_roles=state.circle_shape_roles.at[self.creating_shape_index].set(0),
                circle_highlighted=state.circle_highlighted.at[self.creating_shape_index].set(highlight),
                circle_densities=state.circle_densities.at[self.creating_shape_index].set(1.0),
            )
            return state

        if left_click:
            if self.creating_shape:
                env_state = _add_circle(pcg_state.env_state, False)
                env_state_max = _add_circle(pcg_state.env_state_max, False)

                env_state = recalculate_mass_and_inertia(
                    env_state,
                    self.static_env_params,
                    env_state.polygon_densities,
                    env_state.circle_densities,
                )
                env_state_max = recalculate_mass_and_inertia(
                    env_state_max,
                    self.static_env_params,
                    env_state.polygon_densities,
                    env_state.circle_densities,
                )

                pcg_state = pcg_state.replace(env_state=env_state, env_state_max=env_state_max)

                self.creating_shape = False
            else:
                self.creating_shape_index = jnp.argmin(pcg_state.env_state.circle.active)
                self.create_shape_position = self._get_mouse_position_world_space()
                self.creating_shape = True

        else:
            if self.creating_shape:
                env_state = _add_circle(pcg_state.env_state, True)
                env_state_max = _add_circle(pcg_state.env_state_max, True)

                pcg_state = pcg_state.replace(env_state=env_state, env_state_max=env_state_max)
            else:
                pass

        return pcg_state

    def _get_polygons_on_mouse(self, state, n_vertices=None):
        # n_vertices=None selects both triangles and quads
        ris = []

        mouse_pos = self._get_mouse_position_world_space()

        for ri in jnp.arange(self.static_env_params.num_polygons)[::-1]:
            polygon = jax.tree.map(lambda x: x[ri], state.polygon)

            if (not polygon.active) or ((n_vertices is not None) and polygon.n_vertices != n_vertices):
                continue

            mpos = rmat(-polygon.rotation) @ (mouse_pos - polygon.position)

            def _signed_line_distance(a, b, c):
                return (b[0] - a[0]) * (c[1] - a[1]) - (b[1] - a[1]) * (c[0] - a[0])

            inside = True
            for fi in range(polygon.n_vertices):
                v1 = polygon.vertices[fi]
                v2 = polygon.vertices[(fi + 1) % polygon.n_vertices]
                if _signed_line_distance(mpos, v1, v2) > 0:
                    inside = False

            if inside:
                ris.append(ri)

        return ris

    def _edit_rect(self, pcg_state: PCGState, left_click: bool, right_click: bool):
        if right_click:
            for ri in self._get_polygons_on_mouse(pcg_state.env_state, n_vertices=4):
                attached_j, attached_t = self._get_joints_attached_to_shape(pcg_state.env_state, ri)

                def _remove_rect_from_state(state):
                    state = state.replace(
                        polygon=state.polygon.replace(
                            active=state.polygon.active.at[ri].set(False),
                            rotation=state.polygon.rotation.at[ri].set(0.0),
                        ),
                        joint=state.joint.replace(active=state.joint.active.at[attached_j].set(False)),
                        thruster=state.thruster.replace(active=state.thruster.active.at[attached_t].set(False)),
                    )

                    return state

                env_state = _remove_rect_from_state(pcg_state.env_state)
                env_state_max = _remove_rect_from_state(pcg_state.env_state_max)
                env_state_pcg_mask = _remove_rect_from_state(pcg_state.env_state_pcg_mask)

                env_state = env_state.replace(
                    collision_matrix=calculate_collision_matrix(self.static_env_params, env_state.joint)
                )

                return PCGState(
                    env_state=env_state,
                    env_state_max=env_state_max,
                    env_state_pcg_mask=env_state_pcg_mask,
                    tied_together=pcg_state.tied_together,
                )

        if not self.creating_shape and (1 - pcg_state.env_state.polygon.active.astype(int)).sum() == 0:
            return pcg_state

        diff = (self._get_mouse_position_world_space() - self.create_shape_position) / 2
        diff = jnp.clip(
            diff,
            -(self.static_env_params.max_shape_size / 2) / jnp.sqrt(2),
            (self.static_env_params.max_shape_size / 2) / jnp.sqrt(2),
        )
        half_dim = jnp.abs(diff)
        half_dim = jnp.clip(half_dim, a_min=5.0 / self.env_params.pixels_per_unit)
        vertices = rectangle_vertices(half_dim)

        def _add_rect_to_state(state, highlight):
            state = state.replace(
                polygon=state.polygon.replace(
                    position=state.polygon.position.at[self.creating_shape_index].set(
                        self.create_shape_position + diff
                    ),
                    velocity=state.polygon.velocity.at[self.creating_shape_index].set(jnp.array([0.0, 0.0])),
                    vertices=state.polygon.vertices.at[self.creating_shape_index].set(vertices),
                    inverse_mass=state.polygon.inverse_mass.at[self.creating_shape_index].set(1.0),
                    inverse_inertia=state.polygon.inverse_inertia.at[self.creating_shape_index].set(1.0),
                    active=state.polygon.active.at[self.creating_shape_index].set(True),
                    collision_mode=state.polygon.collision_mode.at[self.creating_shape_index].set(1),
                    n_vertices=state.polygon.n_vertices.at[self.creating_shape_index].set(4),
                ),
                polygon_shape_roles=state.polygon_shape_roles.at[self.creating_shape_index].set(0),
                polygon_highlighted=state.polygon_highlighted.at[self.creating_shape_index].set(highlight),
                polygon_densities=state.polygon_densities.at[self.creating_shape_index].set(1.0),
            )

            return state

        if left_click:
            if self.creating_shape:
                env_state = _add_rect_to_state(pcg_state.env_state, False)
                env_state_max = _add_rect_to_state(pcg_state.env_state_max, False)

                env_state = recalculate_mass_and_inertia(
                    env_state,
                    self.static_env_params,
                    env_state.polygon_densities,
                    env_state.circle_densities,
                )
                env_state_max = recalculate_mass_and_inertia(
                    env_state_max,
                    self.static_env_params,
                    env_state.polygon_densities,
                    env_state.circle_densities,
                )

                pcg_state = pcg_state.replace(env_state=env_state, env_state_max=env_state_max)

                self.creating_shape = False
            else:
                self.creating_shape_index = jnp.argmin(pcg_state.env_state.polygon.active)
                self.create_shape_position = self._get_mouse_position_world_space()
                self.creating_shape = True
        else:
            if self.creating_shape:
                env_state = _add_rect_to_state(pcg_state.env_state, True)
                env_state_max = _add_rect_to_state(pcg_state.env_state_max, True)

                pcg_state = pcg_state.replace(env_state=env_state, env_state_max=env_state_max)
            else:
                pass

        return pcg_state

    def _reset_triangles(self):
        self.triangle_order = jnp.array([0, 1, 2])
        self.num_triangle_clicks = 0
        self.creating_shape = False

    def _edit_triangle(self, pcg_state: PCGState, left_click: bool, right_click: bool):
        if right_click:
            self.num_triangle_clicks = 0
            for ri in self._get_polygons_on_mouse(pcg_state.env_state, n_vertices=3):
                attached_r, attached_f = self._get_joints_attached_to_shape(pcg_state.env_state, ri)

                def _remove_triangle_from_state(state):
                    state = state.replace(
                        polygon=state.polygon.replace(
                            active=state.polygon.active.at[ri].set(False),
                            rotation=state.polygon.rotation.at[ri].set(0.0),
                        ),
                        joint=state.joint.replace(active=state.joint.active.at[attached_r].set(False)),
                    )

                    return state

                env_state = _remove_triangle_from_state(pcg_state.env_state)
                env_state_max = _remove_triangle_from_state(pcg_state.env_state_max)
                env_state_pcg_mask = _remove_triangle_from_state(pcg_state.env_state_pcg_mask)

                env_state = env_state.replace(
                    collision_matrix=calculate_collision_matrix(self.static_env_params, env_state.joint)
                )

                pcg_state = PCGState(
                    env_state=env_state,
                    env_state_max=env_state_max,
                    env_state_pcg_mask=env_state_pcg_mask,
                    tied_together=pcg_state.tied_together,
                )

                return pcg_state

        if not self.creating_shape and (1 - pcg_state.env_state.polygon.active.astype(int)).sum() == 0:
            return pcg_state

        def get_correct_center_two_verts(verts):
            return (jnp.max(verts, axis=0) + jnp.min(verts, axis=0)) / 2

        def order_clockwise(verts, loose_ordering=False):
            # verts has shape (3, 2), order them clockwise.
            # https://stackoverflow.com/questions/51074984/sort-vertices-in-clockwise-order
            # Calculate centroid
            centroid = jnp.mean(verts, axis=0)
            # Calculate angles
            angles = jnp.round(jnp.arctan2(verts[:, 1] - centroid[1], verts[:, 0] - centroid[0]), 2)
            # Order vertices
            order = jnp.argsort(-angles, stable=True)
            if loose_ordering:
                order = jnp.arange(len(order))
            ans = verts[order]
            # order is of shape (2, ) or (3, ). I want it to always be of shape 3
            if len(order) < 3:
                order = jnp.concatenate([order, jnp.array([2])])
            return ans, order

        def do_triangle_n_click(pcg_state, how_many_clicks, is_on_a_click=False):
            n = how_many_clicks
            # if we must keep them clockwise all the time, then the one we edit / move around may have varying indices.
            current_index_to_change = self.triangle_order[n]
            sign = 1
            idxs = jnp.arange(n + 1)
            idxs_to_allow = idxs[~(idxs == current_index_to_change)]

            # Get the new vertex and clip its position
            new_tentative_vert = (
                self._get_mouse_position_world_space() - pcg_state.env_state.polygon.position[self.creating_shape_index]
            )
            new_tentative_vert = jnp.clip(
                new_tentative_vert,
                jnp.max(pcg_state.env_state.polygon.vertices[self.creating_shape_index, idxs_to_allow], axis=0)
                - self.static_env_params.max_shape_size * 0.8,
                jnp.min(pcg_state.env_state.polygon.vertices[self.creating_shape_index, idxs_to_allow], axis=0)
                + self.static_env_params.max_shape_size * 0.8,
            )
            new_verts = pcg_state.env_state.polygon.vertices.at[self.creating_shape_index, current_index_to_change].set(
                new_tentative_vert
            )
            new_center_two = get_correct_center_two_verts(new_verts[self.creating_shape_index, : n + 1])

            _, new_center_three = calc_inverse_mass_polygon(
                new_verts[self.creating_shape_index],
                3,
                self.static_env_params,
                1.0,
            )
            new_center = jax.lax.select(n == 1, new_center_two, new_center_three)

            new_verts = new_verts.at[self.creating_shape_index].add(-sign * new_center)
            vvs = new_verts[self.creating_shape_index, : n + 1]

            ordered_vertices, new_permutation = order_clockwise(vvs, loose_ordering=not is_on_a_click)
            self.triangle_order = self.triangle_order[new_permutation]
            new_verts = new_verts.at[self.creating_shape_index, : n + 1].set(ordered_vertices)

            env_state = pcg_state.env_state.replace(
                polygon=pcg_state.env_state.polygon.replace(
                    vertices=new_verts,
                    position=pcg_state.env_state.polygon.position.at[self.creating_shape_index].add(sign * new_center),
                    n_vertices=pcg_state.env_state.polygon.n_vertices.at[self.creating_shape_index].set(n + 1),
                ),
            )
            env_state_max = pcg_state.env_state_max.replace(
                polygon=pcg_state.env_state_max.polygon.replace(
                    vertices=new_verts,
                    position=pcg_state.env_state_max.polygon.position.at[self.creating_shape_index].add(
                        sign * new_center
                    ),
                    n_vertices=pcg_state.env_state_max.polygon.n_vertices.at[self.creating_shape_index].set(n + 1),
                ),
            )

            pcg_state = pcg_state.replace(env_state=env_state, env_state_max=env_state_max)

            return pcg_state

        if left_click:
            if self.creating_shape:
                assert 3 > self.num_triangle_clicks > 0
                if self.num_triangle_clicks == 1:
                    pcg_state = do_triangle_n_click(pcg_state, 1, is_on_a_click=True)
                    self.num_triangle_clicks += 1
                else:  # this finishes it
                    pcg_state = do_triangle_n_click(pcg_state, 2, is_on_a_click=True)
                    self.creating_shape = False
                    self.num_triangle_clicks = 0
                    pcg_state = pcg_state.replace(
                        env_state=recalculate_mass_and_inertia(
                            pcg_state.env_state,
                            self.static_env_params,
                            pcg_state.env_state.polygon_densities,
                            pcg_state.env_state.circle_densities,
                        )
                    )

            else:
                self.triangle_order = jnp.array([0, 1, 2])
                self.creating_shape_index = jnp.argmin(pcg_state.env_state.polygon.active)
                self.create_shape_position = self._get_mouse_position_world_space()
                self.creating_shape = True
                self.num_triangle_clicks = 1
                vertices = jnp.zeros((self.static_env_params.max_polygon_vertices, 2), dtype=jnp.float32)

                def _add_triangle_to_state(state):
                    state = state.replace(
                        polygon=state.polygon.replace(
                            position=state.polygon.position.at[self.creating_shape_index].set(
                                self.create_shape_position
                            ),
                            velocity=state.polygon.velocity.at[self.creating_shape_index].set(jnp.array([0.0, 0.0])),
                            vertices=state.polygon.vertices.at[self.creating_shape_index].set(vertices),
                            inverse_mass=state.polygon.inverse_mass.at[self.creating_shape_index].set(1.0),
                            inverse_inertia=state.polygon.inverse_inertia.at[self.creating_shape_index].set(1.0),
                            active=state.polygon.active.at[self.creating_shape_index].set(True),
                            n_vertices=state.polygon.n_vertices.at[self.creating_shape_index].set(1),
                        ),
                        polygon_shape_roles=state.polygon_shape_roles.at[self.creating_shape_index].set(0),
                        polygon_highlighted=state.polygon_highlighted.at[self.creating_shape_index].set(False),
                        polygon_densities=state.polygon_densities.at[self.creating_shape_index].set(1.0),
                    )

                    return state

                pcg_state = pcg_state.replace(
                    env_state=_add_triangle_to_state(pcg_state.env_state),
                    env_state_max=_add_triangle_to_state(pcg_state.env_state_max),
                )

        elif self.creating_shape:
            assert 1 <= self.num_triangle_clicks <= 2
            pcg_state = do_triangle_n_click(
                pcg_state, self.num_triangle_clicks, is_on_a_click=self.num_triangle_clicks == 1
            )

        return pcg_state

    def _edit_joint(self, pcg_state: PCGState, left_click: bool, right_click: bool):
        if left_click and pcg_state.env_state.joint.active.all():
            return pcg_state

        if left_click:
            joint_index = jnp.argmin(pcg_state.env_state.joint.active)
            joint_position = self._get_mouse_position_world_space()
            # reverse them so that the joint order and rendering order remains the same.
            # We want the first shape to have a lower index than the second shape, with circles always having higher indices compared to rectangles.
            circles = self._get_circles_on_mouse(pcg_state.env_state)[::-1]
            rects = self._get_polygons_on_mouse(pcg_state.env_state)[::-1]

            if len(rects) + len(circles) >= 2:
                r1 = len(rects) >= 1
                r2 = len(rects) >= 2

                a_index = rects[0] if r1 else circles[0]  # + self.static_env_params.num_polygons
                b_index = rects[r1 * 1] if r2 else circles[1 - 1 * r1]  # + self.static_env_params.num_polygons

                a_shape = pcg_state.env_state.polygon if r1 else pcg_state.env_state.circle
                b_shape = pcg_state.env_state.polygon if r2 else pcg_state.env_state.circle

                a = jax.tree.map(lambda x: x[a_index], a_shape)
                b = jax.tree.map(lambda x: x[b_index], b_shape)

                a_index += (not r1) * self.static_env_params.num_polygons
                b_index += (not r2) * self.static_env_params.num_polygons

                joint_position = snap_to_center(a, joint_position)
                joint_position = snap_to_center(b, joint_position)

                a_relative_pos = jnp.matmul(rmat(a.rotation).transpose((1, 0)), joint_position - a.position)
                b_relative_pos = jnp.matmul(rmat(b.rotation).transpose((1, 0)), joint_position - b.position)

                def _add_joint_to_state(state):
                    state = state.replace(
                        joint=state.joint.replace(
                            a_index=state.joint.a_index.at[joint_index].set(a_index),
                            b_index=state.joint.b_index.at[joint_index].set(b_index),
                            a_relative_pos=state.joint.a_relative_pos.at[joint_index].set(a_relative_pos),
                            b_relative_pos=state.joint.b_relative_pos.at[joint_index].set(b_relative_pos),
                            active=state.joint.active.at[joint_index].set(True),
                            global_position=state.joint.global_position.at[joint_index].set(joint_position),
                            motor_on=state.joint.motor_on.at[joint_index].set(True),
                            motor_speed=state.joint.motor_speed.at[joint_index].set(1.0),
                            motor_power=state.joint.motor_power.at[joint_index].set(1.0),
                            rotation=state.joint.rotation.at[joint_index].set(b.rotation - a.rotation),
                        )
                    )

                    return state

                env_state = _add_joint_to_state(pcg_state.env_state)
                env_state_max = _add_joint_to_state(pcg_state.env_state_max)

                env_state = env_state.replace(
                    collision_matrix=calculate_collision_matrix(self.static_env_params, env_state.joint)
                )

                pcg_state = pcg_state.replace(env_state=env_state, env_state_max=env_state_max)

        return pcg_state

    def _reset_select_shape(self, pcg_state):
        pcg_state = pcg_state.replace(
            env_state=pcg_state.env_state.replace(
                polygon_highlighted=jnp.zeros_like(pcg_state.env_state.polygon_highlighted),
                circle_highlighted=jnp.zeros_like(pcg_state.env_state.circle_highlighted),
            )
        )
        self.selected_shape_index = -1
        self.selected_shape_type = ObjectType.POLYGON
        self._hide_all_widgets()
        return pcg_state

    def _hide_all_widgets(self):
        for widget in self.all_widgets.values():
            for w in widget.values():
                w.hide()

    def _show_correct_widgets(self, type: ObjectType | None, do_tie_ui: bool = False):
        for widget in self.all_widgets["GENERAL"].values():
            widget.show()
        if do_tie_ui:
            for widget in self.all_widgets["TIE_TOGETHER"].values():
                widget.show()

            n = len(self.all_selected_shapes)
            #  {[int(i) for (i, t) in self.all_selected_shapes]}
            self.all_widgets["GENERAL"]["lblGeneral"].setText(f"Selected {n} Objects")
            return

        for widget in self.all_widgets[type].values():
            widget.show()
        if type is None:
            self.all_widgets["GENERAL"]["lblGeneral"].setText(f"Global")
        else:
            self.all_widgets["GENERAL"]["lblGeneral"].setText(f"{type.name} (idx {self.selected_shape_index})")

    def _select_shape_keyboard_shortcuts(self, pcg_state: PCGState, left_click: bool, keys: list[int]):
        if left_click:
            return pcg_state
        if len(keys) != 0 and self.selected_shape_index != -1:
            s = 1.0
            ang_s = 0.1
            vel = jnp.array([0.0, 0.0])
            angular_vel = 0.0
            should_toggle_fixed = False
            should_toggle_collidable = False
            change_angle = 0

            def add_step(widget_name, direction, speed=10, overwrite_amount=None):
                widget = self.all_widgets[self.selected_shape_type][widget_name]
                val = widget.getValue()
                step = widget.step
                amount_to_add = overwrite_amount or step * direction * speed
                widget.setValue(jnp.clip(val + amount_to_add, widget.min, widget.max))

            if pygame.K_w in keys:
                add_step("sldPosition_Y", 1)
            if pygame.K_s in keys:
                add_step("sldPosition_Y", -1)
            if pygame.K_a in keys:
                add_step("sldPosition_X", -1)
            if pygame.K_d in keys:
                add_step("sldPosition_X", 1)

            if pygame.K_q in keys:
                add_step("sldRotation", 1)
            if pygame.K_e in keys:
                add_step("sldRotation", -1)

            if pygame.K_f in keys:
                self.all_widgets[self.selected_shape_type]["tglFixate"].toggle()

            if pygame.K_c in keys and not (pygame.key.get_mods() & pygame.KMOD_CTRL):
                widget = self.all_widgets[self.selected_shape_type]["sldCollidability"]
                curr_val = int(widget.getValue())
                widget.setValue((curr_val + 1) % (widget.max + 1))

            if pygame.K_r in keys and not (pygame.key.get_mods() & pygame.KMOD_CTRL):
                widget = self.all_widgets[self.selected_shape_type]["sldRole"]
                curr_val = int(widget.getValue())
                widget.setValue((curr_val + 1) % (widget.max + 1))

            if pygame.K_LEFTBRACKET in keys:
                add_step("sldRotation", 1, 10, jnp.pi / 4)

            if pygame.K_RIGHTBRACKET in keys:
                add_step("sldRotation", -1, 10, -jnp.pi / 4)

            if pygame.K_c in keys and (pygame.key.get_mods() & pygame.KMOD_CTRL):
                # copy
                if self.selected_shape_type == ObjectType.POLYGON:  # rect
                    if not self.pcg_state.env_state.polygon.active.all():
                        where_to_add = jnp.argmin(pcg_state.env_state.polygon.active)
                        if where_to_add < self.static_env_params.num_polygons:

                            def _copy_polygon(state, shift):
                                state = state.replace(
                                    polygon=jax.tree.map(
                                        lambda x: x.at[where_to_add].set(x[self.selected_shape_index]), state.polygon
                                    )
                                )
                                if shift:
                                    state = state.replace(
                                        polygon=state.polygon.replace(
                                            position=state.polygon.position.at[where_to_add].add(0.1),
                                        ),
                                        polygon_highlighted=state.polygon_highlighted.at[where_to_add].set(False),
                                    )
                                return state

                            pcg_state = pcg_state.replace(
                                env_state=_copy_polygon(pcg_state.env_state, shift=True),
                                env_state_max=_copy_polygon(pcg_state.env_state_max, shift=True),
                                env_state_pcg_mask=_copy_polygon(pcg_state.env_state_pcg_mask, shift=False),
                            )

                elif self.selected_shape_type == ObjectType.CIRCLE:  # circle
                    if not self.pcg_state.env_state.circle.active.all():
                        where_to_add = jnp.argmin(pcg_state.env_state.circle.active)
                        if where_to_add < self.static_env_params.num_circles:

                            def _copy_circle(state, shift=True):
                                state = state.replace(
                                    circle=jax.tree.map(
                                        lambda x: x.at[where_to_add].set(x[self.selected_shape_index]), state.circle
                                    )
                                )
                                if shift:
                                    state = state.replace(
                                        circle=state.circle.replace(
                                            position=state.circle.position.at[where_to_add].add(0.1),
                                        ),
                                        circle_highlighted=state.circle_highlighted.at[where_to_add].set(False),
                                    )
                                return state

                            pcg_state = pcg_state.replace(
                                env_state=_copy_circle(pcg_state.env_state),
                                env_state_max=_copy_circle(pcg_state.env_state_max),
                                env_state_pcg_mask=_copy_circle(pcg_state.env_state_pcg_mask, shift=False),
                            )

            if self.selected_shape_index >= 0:
                num = get_numeric_key_pressed(self.pygame_events)
                if num is not None:
                    if self.selected_shape_type in [ObjectType.CIRCLE, ObjectType.POLYGON]:
                        self.all_widgets[self.selected_shape_type]["sldRole"].setValue(num % 4)
                    elif self.selected_shape_type == ObjectType.JOINT:
                        self.all_widgets[self.selected_shape_type]["sldColour"].setValue(
                            num % self.static_env_params.num_motor_bindings
                        )
                    elif self.selected_shape_type == ObjectType.THRUSTER:
                        self.all_widgets[self.selected_shape_type]["sldColour"].setValue(
                            num % self.static_env_params.num_thruster_bindings
                        )

        return pcg_state

    def _edit_select_shape(self, pcg_state: PCGState, left_click: bool, right_click: bool, keys: list[int]):
        def _find_shape(pcg_state):
            found_shape = False
            selected_shape_index, selected_shape_type = -1, ObjectType.POLYGON
            for ri in self._get_revolute_joints_on_mouse(pcg_state.env_state):
                selected_shape_index = ri
                selected_shape_type = ObjectType.JOINT
                found_shape = True
                break
            if not found_shape:
                for ti in self._get_thrusters_on_mouse(pcg_state.env_state):
                    selected_shape_index = ti
                    selected_shape_type = ObjectType.THRUSTER
                    found_shape = True
                    break
            if not found_shape:
                for ri in self._get_polygons_on_mouse(pcg_state.env_state):
                    pcg_state = pcg_state.replace(
                        env_state=pcg_state.env_state.replace(
                            polygon_highlighted=pcg_state.env_state.polygon_highlighted.at[ri].set(True),
                        )
                    )
                    selected_shape_index = ri
                    selected_shape_type = ObjectType.POLYGON
                    found_shape = True
                    break
            if not found_shape:
                for ci in self._get_circles_on_mouse(pcg_state.env_state):
                    pcg_state = pcg_state.replace(
                        env_state=pcg_state.env_state.replace(
                            circle_highlighted=pcg_state.env_state.circle_highlighted.at[ci].set(True),
                        )
                    )
                    selected_shape_index = ci
                    selected_shape_type = ObjectType.CIRCLE
                    found_shape = True
                    break
            return selected_shape_index, selected_shape_type, found_shape, pcg_state
            if found_shape and self.selected_shape_type in self.all_widgets:
                self._show_correct_widgets(self.selected_shape_type)

        # if left and shift
        if left_click and (pygame.key.get_mods() & pygame.KMOD_SHIFT):
            # This is trying to select multiple things.
            idx, type, found, pcg_state = _find_shape(pcg_state)
            if found:
                t = (idx, type)
                if t in self.all_selected_shapes:
                    self.all_selected_shapes.remove(t)
                else:
                    self.all_selected_shapes.append(t)

            self._hide_all_widgets()
            self._show_correct_widgets(None, do_tie_ui=True)
        elif left_click:
            self.all_selected_shapes = []
            self._hide_all_widgets()
            pcg_state = self._reset_select_shape(pcg_state)
            self.selected_shape_index, self.selected_shape_type, found_shape, pcg_state = _find_shape(pcg_state)
            if found_shape:
                self.all_selected_shapes = [(self.selected_shape_index, self.selected_shape_type)]
                if self.selected_shape_type in self.all_widgets:
                    self._show_correct_widgets(self.selected_shape_type)
        if self.selected_shape_index < 0:
            self._show_correct_widgets(None)

        return pcg_state

    def render(self, env_state):
        # Clear
        self.screen_surface.fill((0, 0, 0))

        if self.is_editing:
            pixels = self._render_fn_edit(env_state)
        else:
            pixels = self._render_fn(env_state)
        pixels = self._render_edit_overlay_fn(pixels, self.is_editing, self.edit_shape_mode.value)

        surface = pygame.surfarray.make_surface(np.array(pixels))
        self.screen_surface.blit(surface, (0, 0))

    def is_quit_requested(self):
        for event in self.pygame_events:
            if event.type == pygame.QUIT:
                return True
        return False


@hydra.main(version_base=None, config_path="../configs", config_name="editor")
def main(config):
    config = normalise_config(OmegaConf.to_container(config), "EDITOR", editor_config=True)
    env_params, static_env_params = generate_params_from_config(config)
    static_env_params = static_env_params.replace(frame_skip=config["frame_skip"], downscale=config["downscale"])
    config["env_params"] = to_state_dict(env_params)
    config["static_env_params"] = to_state_dict(static_env_params)

    env = make_kinetix_env_from_name("Kinetix-Entity-MultiDiscrete-v1", static_env_params=static_env_params)
    env = AutoResetWrapper(env, make_reset_function(static_env_params))

    seed = config["seed"]
    print("seed", seed)
    rng = jax.random.PRNGKey(seed)
    outer_timer = tmr()
    editor = Editor(env, env_params, config, upscale=config["upscale"])
    time_e = tmr()
    print("Took {:2f}s to create editor".format(time_e - outer_timer))

    clock = pygame.time.Clock()
    while not editor.is_quit_requested():
        rng, _rng = jax.random.split(rng)
        editor.update(_rng)
        clock.tick(config["fps"])


if __name__ == "__main__":
    main()