"""
Example node showing realtime update from gathered data computed on the GPU.
"""


class OgnBouncingCubesGpu:
    @staticmethod
    def compute(db):

        # This is how it should work when gather is fully supported
        # velocities = db.state.velocities
        # translations = db.state.translations
        attr_velocities = db.node.get_attribute("state:velocities")
        attr_translations = db.node.get_attribute("state:translations")

        velocities = db.context_helper.get_attr_value(attr_velocities, isGPU=True, isTensor=True)
        translations = db.context_helper.get_attr_value(attr_translations, isGPU=True, isTensor=True)

        # When there is no data there is nothing to do
        if velocities.size == 0 or translations.size == 0:
            return True

        h = 1.0 / 60.0
        g = -1000.0

        velocities += h * g
        translations[:, 2] += h * velocities

        translations_z = translations.split(1, dim=1)[2].squeeze(-1)
        update_idx = (translations_z < 1.0).nonzero().squeeze(-1)
        if update_idx.shape[0] > 0:
            translations[update_idx, 2] = 1
            velocities[update_idx] = -5 * h * g

        translations[:, 0] += 0.03

        # We no longer need to set values because we do not copy memory of the numpy arrays
        # Mutating the numpy array will also mutate the attribute data
        # Computegraph APIs using torch tensors currently do not have support for setting values.

        # context.set_attr_value(translations, attr_translations)
        # context.set_attr_value(velocities, attr_translations)

        return True
