#!/usr/bin/env python
# coding: utf-8

# # Draw a NumPy array directly on the Canvas with `put_image_data`

# In[ ]:


import numpy as np

from ipywidgets import Play, IntProgress, HBox, VBox, link

from ipycanvas import Canvas, hold_canvas


# In[ ]:


n = 1

dx, dy = 0.01, 0.01

y, x = np.mgrid[slice(1, 5 + dy, dy), slice(1, 5 + dx, dx)]

z = np.sin(x) ** n + np.cos(n + y * x) * np.cos(x)


# In[ ]:


min = np.min(z)
max = np.max(z)


def scale(value):
    scaled_value = (value - min) / (max - min)
    return 255 if value > max else scaled_value * 255


vecscale = np.vectorize(scale)


# In[ ]:


data = np.stack((np.zeros_like(z), vecscale(z), vecscale(z)), axis=2)


# In[ ]:


data.shape


# In[ ]:


scale = 1.5

canvas = Canvas(width=scale * data.shape[0], height=scale * data.shape[1])
canvas


# In[ ]:


canvas.scale(scale)
canvas.put_image_data(data, 0, 0)


# ## Make an animation with it!

# In[ ]:


play = Play(interval=500, min=1, max=20, step=1)
progress = IntProgress(min=1, max=20, step=1)

link((play, "value"), (progress, "value"))


def on_update(*args):
    global z

    z = np.sin(x) ** play.value + np.cos(play.value + y * x) * np.cos(x)
    data = np.stack((np.zeros_like(z), vecscale(z), vecscale(z)), axis=2)

    with hold_canvas():
        canvas.put_image_data(data, 0, 0)


play.observe(on_update, "value")

# This is to prevent the Canvas to take the entire available space in the VBox
canvas.layout.width = str(canvas.width) + "px"
canvas.layout.height = str(canvas.height) + "px"

VBox((canvas, HBox((play, progress))))