Py.Cafe

import solara
import ipyreact
from pathlib import Path
ipyreact.define_module("@react-three/fiber", Path("./threejs-fiber.bundle.js"))
ipyreact.define_module("@react-three/xr", Path("./react-three-xr.bundle.js"))
ipyreact.define_module("vizxr", Path("./vizxr.bundle.js"))



@solara.component
def Canvas(props={}, events={}, children=[]):
    return ipyreact.Widget.element(props=props, events=events, children=children, _type="Canvas", _module="@react-three/fiber")


@solara.component
def OrbitControls(props={}, events={}, children=[]):
    return ipyreact.Widget.element(_type="OrbitControls", _module="@react-three/fiber", props=props, events=events, children=children)

@solara.component
def ConditionalOrbitControls(props={}, events={}, children=[]):
    return ipyreact.Widget.element(_type="ConditionalOrbitControls", _module="vizxr", 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="@react-three/fiber"
    return ipyreact.Widget.element(_type="directionalLight", 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="@react-three/fiber"
    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)

y0 = 1
boxes = solara.reactive([
    ([-1, y0, 3], "#18a36e"),
    ([1, y0, 3], "#f56f42"),
])
    
def add(event_data=None):
    x = random.random() * 4 - 2
    z = random.random() * 4 - 1
    color = random_color()  # Call the random_color function to get a random color
    boxes.value = [*boxes.value, ([x, y0, z], color)]


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


def add_10():
    for i in range(10):
        add()


class InstancedMeshWidget(ipyreact.Widget):
    _esm = Path("instanced_mesh.tsx")


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 InstancedMesh(limit, x, y, z, props={}, events={}, children=[]):
    return InstancedMeshWidget.element(props={**props, **dict(limit=limit, x=sar(x), y=sar(y), z=sar(z))}, children=children, events=events)


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


# looks nice, but too busy for 3d selection
# def zeldovich( dim=2, N=256, n=-2.5, t=None, seed=None, scale=1):

#     if seed is not None:
#         np.random.seed(seed)
#     #sys.exit(0)
#     shape = (N,) * dim
#     A = np.random.normal(0.0, 1.0, shape)
#     F = np.fft.fftn(A)
#     K = np.fft.fftfreq(N, 1./(2*np.pi))[np.indices(shape)]
#     k = (K**2).sum(axis=0)
#     k_max = np.pi
#     F *= np.where(np.sqrt(k) > k_max, 0, np.sqrt(k**n) * np.exp(-k*4.0))
#     F.flat[0] = 0
#     #plt.imshow(np.where(sqrt(k) > k_max, 0, np.sqrt(k**-2)), interpolation='nearest')
#     grf = np.fft.ifftn(F).real
#     Q = np.indices(shape) / float(N-1) - 0.5
#     s = np.array(np.gradient(grf)) / float(N)
#     #plt.imshow(s[1], interpolation='nearest')
#     #plt.show()
#     s /= s.max() * 100.
#     #X = np.zeros((4, 3, N, N, N))
#     #for i in range(4):
#     #if t is None:
#     #	s = s/s.max()
#     t = t or 1.
#     X = Q + s * t

#     # # for d, name in zip(list(range(dim)), "xyzw"):
#     #     self.add_column(name, X[d].reshape(-1) * scale)
#     # for d, name in zip(list(range(dim)), "xyzw"):
#     #     self.add_column("v"+name, s[d].reshape(-1) * scale)
#     # for d, name in zip(list(range(dim)), "xyzw"):
#     #     self.add_column(name+"0", Q[d].reshape(-1) * scale)
#     return [X[d].reshape(-1) * scale for d in range(dim)]

# #x, y, z = zeldovich(3, 30, t=130.5, seed=42, scale=0.5)


# 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

print(x, y, z)
y += 0.5
z -= 0.5
print("z", z, z.dtype)



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

force_update = solara.reactive(0)

@solara.component
def Page():
    print("force_update", force_update.value)
    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="vizxr"):
            ipyreact.Widget.element(_type="XRButtons", _module="vizxr")
            with Canvas(events={"onClick": add}):
                with ipyreact.Widget.element(_type="XRStore", _module="vizxr"):
                    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="vizxr", 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
    #return
    
    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