Py.Cafe

import solara
import ipyreact
from pathlib import Path
import numpy as np

# Ipyreact module defn
ipyreact.define_module("glue-xr-viewer", Path("./glue-xr-viewer.esm.js"))


def array_to_binary(ar, obj=None, force_contiguous=True):
    if ar is None:
        return None
    if ar.dtype.kind not in ['u', 'i', 'f']:  # ints and floats
        raise ValueError("unsupported dtype: %s" % (ar.dtype))
    if ar.dtype == np.float64:  # WebGL does not support float64, case it here
        ar = ar.astype(np.float32)
    if ar.dtype == np.int64:  # JS does not support int64
        ar = ar.astype(np.int32)
    if force_contiguous and not ar.flags["C_CONTIGUOUS"]:  # make sure it's contiguous
        ar = np.ascontiguousarray(ar)
    return {'dataView': memoryview(ar), 'dtype': str(ar.dtype), 'shape': ar.shape}

sar = array_to_binary


@solara.component
def Canvas(props={}, events={}, children=[]):
    return ipyreact.Widget.element(props=props, events=events, children=children, _type="Canvas", _module="glue-xr-viewer")


@solara.component
def OrbitControls(props={}, events={}, children=[]):
    return ipyreact.Widget.element(_type="OrbitControls", _module="glue-xr-viewer", props=props, events=events, children=children)


@solara.component
def ConditionalOrbitControls(props={}, events={}, children=[]):
    return ipyreact.Widget.element(_type="ConditionalOrbitControls", _module="glue-xr-viewer", props=props, events=events, children=children)


@solara.component
def DirectionalLight(props={}, events={}, children=[]):
    # starts with a lower case, should be available globally, so we don't need to pass _module
    return ipyreact.Widget.element(_type="directionalLight", props=props, events=events, children=children)


@solara.component
def AmbientLight(props={}, events={}, children=[]):
    return ipyreact.Widget.element(_type="ambientLight", props=props, events=events, children=children)


@solara.component
def Div(style={}, props={}, events={}, children=[]):
    # we use a ipyreact based div to avoid an extra wrapper div which will affect layout
    return ipyreact.Widget.element(_type="div", props={**props, **dict(style=style)}, children=children, events=events)


@solara.component
def InstancedMesh(limit, x, y, z, props={}, events={}, children=[]):
    return ipyreact.Widget.element(_type="InstancedMesh", _module="glue-xr-viewer", props={**props, **dict(limit=limit, x=sar(x), y=sar(y), z=sar(z))}, children=children, events=events)


y0 = 1
boxes = solara.reactive([
    ([-1, y0, 3], "#18a36e"),
    ([1, y0, 3], "#f56f42"),
])


# too boring to 3d selection
x, y, z = (np.random.rand(3, 1000*10) * 0.5 - 0.25).astype(np.float32)


# penguins dataset is a nice example for 3d selection
from palmerpenguins import load_penguins
df = load_penguins()
x = df["bill_length_mm"].values
y = df["bill_depth_mm"].values
z = df["flipper_length_mm"].values
# remove nans
nans = np.isnan(x) | np.isnan(y) | np.isnan(z)
x = x[~nans]
y = y[~nans]
z = z[~nans]

# normalize
s = 1
x = (x - x.min()) / (x.max() - x.min()) * s
y = (y - y.min()) / (y.max() - y.min()) * s
z = (z - z.min()) / (z.max() - z.min()) * s

y += 0.5
z -= 0.5


def clear():
    boxes.value = boxes.value[:2]


inside = solara.reactive(np.array([]))

force_update = solara.reactive(0)

@solara.component
def Page():
    solara.use_memo(lambda: update_inside())
    solara.Markdown("Click to add a new box")
    with Div(style={"height": "1024px", "flex-grow": 1}):
        with ipyreact.Widget.element(_type="XRStoreProvider", _module="glue-xr-viewer"):
            ipyreact.Widget.element(_type="XRButtons", _module="glue-xr-viewer")
            with Canvas(events={"onClick": lambda: print("clickkkkoooo")}):
                with ipyreact.Widget.element(_type="XRStore", _module="glue-xr-viewer"):
                    colors = ["#f56f42", "#18a36e", "#f5d142", "#42f5d1", "#4218a3"]
                    color = colors[force_update.value % len(colors)]
                    color = "red"
                    size = 1.
                    ConditionalOrbitControls()
                    DirectionalLight(props=dict(color="#ffffff", intensity=12, position=[-1, 2, 4]))                    
                    AmbientLight(props=dict(color="#ffffff", intensity=0.5))

                    ipyreact.Widget.element(_type="HandSelection", _module="glue-xr-viewer", events={'onHull': on_hull})

                    with InstancedMesh(limit=len(x), x=x, y=y, z=z):
                        ipyreact.Widget.element(props=dict(args=[size, 0]), _type="dodecahedronGeometry")
                        ipyreact.Widget.element(props=dict(color=color), _type="meshPhongMaterial")
                    if len(inside.value) > 0:
                        with InstancedMesh(limit=len(x), x=x[inside.value], y=y[inside.value], z=z[inside.value]):
                            ipyreact.Widget.element(props=dict(args=[size*1.1, 0]), _type="dodecahedronGeometry")
                            ipyreact.Widget.element(props=dict(color="orange"), _type="meshPhongMaterial")


from scipy.spatial import Delaunay

@solara.lab.task
async def update_inside():
    import asyncio
    await asyncio.sleep(0.1)
    print("update_inside.cancelled", update_inside.is_current())
    if not update_inside.is_current():
        return
    force_update.value += 1
    
    hull = Delaunay(hull_array)

    inside_indices = []
    for i, p in enumerate(zip(x, y, z)):
        s = hull.find_simplex(p)
        # print(p, s)
        if s >= 0:
            inside_indices.append(i)
    inside.value = np.array(inside_indices)
    print("inside", inside_indices)

def on_hull(event_data):
    global hull_array
    ar = np.array(event_data)
    hull_array = ar.reshape((-1, 3))
    update_inside()
    
hull_array = None