from panda3d.core import Mat4
from direct.showbase.ShowBase import ShowBase

import p3dimgui # pip install panda3d-imgui
from imgui_bundle import imgui, imguizmo # pip install imgui-bundle

gizmo = imguizmo.im_guizmo

def convertPandaMatrixToGizmo(matrix):
    # Create a new matrix in case it's a constant variable.
    matrix = Mat4(matrix)

    mat = gizmo.Matrix16([
        matrix[0][0], matrix[0][1], matrix[0][2], matrix[0][3],
        matrix[1][0], matrix[1][1], matrix[1][2], matrix[1][3],
        matrix[2][0], matrix[2][1], matrix[2][2], matrix[2][3],
        matrix[3][0], matrix[3][1], matrix[3][2], matrix[3][3]
    ])

    return mat

def convertGizmoMatrixToPanda(matrix: gizmo.Matrix16):
    pandaMat = Mat4(*matrix.values.tolist())
    return pandaMat

class Demo(ShowBase):
    def __init__(self):
        ShowBase.__init__(self)
        self.disableMouse()

        p3dimgui.init()

        self.smiley = self.loader.loadModel('models/smiley.egg')
        self.smiley.reparentTo(self.render)
        self.smiley.setY(5)
        self.smiley.place()

        self.camera.place()

        self.operation = gizmo.OPERATION.translate

        self.accept('imgui-new-frame', self.__draw)

    def __draw(self):
        gizmo.begin_frame()
        size = imgui.get_io().display_size
        gizmo.set_rect(0, 0, size.x, size.y)
        gizmo.set_drawlist(imgui.get_background_draw_list())

        if not gizmo.is_using():
            if imgui.is_key_pressed(imgui.Key.g):
                self.operation = gizmo.OPERATION.translate
            elif imgui.is_key_pressed(imgui.Key.r):
                self.operation = gizmo.OPERATION.rotate
            elif imgui.is_key_pressed(imgui.Key.s):
                self.operation = gizmo.OPERATION.scale

        objMat = convertPandaMatrixToGizmo(self.smiley.getMat())
        if gizmo.manipulate(convertPandaMatrixToGizmo(self.camLens.getViewMat() * self.cam.getNetTransform().getInverse().getMat()),
                         convertPandaMatrixToGizmo(self.camLens.getProjectionMat()),
                         self.operation, gizmo.MODE.local,
                         objMat):
            self.smiley.setMat(convertGizmoMatrixToPanda(objMat))

app = Demo()
app.run()
