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

# In[ ]:


import k3d
from k3d.headless import k3d_remote, get_headless_driver
import numpy as np
from IPython.display import Image


# In[ ]:


plot = k3d.plot(screenshot_scale=1.0)
headless = k3d_remote(plot, get_headless_driver())


# In[ ]:


T = 1.618033988749895
r = 4.77
zmin, zmax = -r, r
xmin, xmax = -r, r
ymin, ymax = -r, r
Nx, Ny, Nz = 77, 77, 77

x = np.linspace(xmin, xmax, Nx)
y = np.linspace(ymin, ymax, Ny)
z = np.linspace(zmin, zmax, Nz)
x, y, z = np.meshgrid(x, y, z, indexing='ij')
p = 2 - (np.cos(x + T * y) + np.cos(x - T * y) + 
         np.cos(y + T * z) + np.cos(y - T * z) + 
         np.cos(z - T * x) + np.cos(z + T * x)).astype(np.float32)
iso = k3d.marching_cubes(p, level=0.0)


# In[ ]:


plot += iso


# In[ ]:


headless.sync()


# In[ ]:


Image(headless.get_browser_screenshot()) 


# In[ ]:


Image(headless.get_screenshot(True))


# In[ ]:


with open('screenshot.png', 'wb') as f:
    f.write(headless.get_screenshot())


# In[ ]:


# movie creator

import cv2
import tqdm

fourcc = cv2.VideoWriter_fourcc(*'MP4V')
video = cv2.VideoWriter("headless_movie.mp4", fourcc, 60.0, (1280, 720))
r = 2

for fi in tqdm.tqdm(np.linspace(0, 2 * np.pi, 60 * 4)):
    plot.camera = [
        np.sin(fi) * r, np.cos(fi) * r, 0.5,
        0, 0, 0,
        0, 0, 1
    ]
    
    headless.sync()
    
    file_bytes = np.asarray(bytearray(headless.get_screenshot(True)), dtype=np.uint8)    
    img = cv2.imdecode(file_bytes, cv2.IMREAD_UNCHANGED)[:,:,0:3]
    video.write(img)    
    
video.release()


# In[ ]:


headless.close()


# In[ ]: