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

# ## Iterated function system  with k3d visualization
# 
# Starting from an array of [0, 0, 0] points, we run a number of iterations that transform each point with one of four functions, with given probabilities (equal by default).

# In[ ]:


import  numpy as np

N = 30000
x = np.zeros(N)
y = np.zeros(N)
z = np.zeros(N)

x1 = np.empty_like(x)
y1 = np.empty_like(y)
z1 = np.empty_like(z)

# Sierpinski triangle iterative functions

def f1(x,y,z,x1,y1,z1,c):
    x1[c] = 1.0/2.0*x[c]
    y1[c] = 1.0/2.0*y[c]
    z1[c] = 1.0/2.0*z[c]

def f2(x,y,z,x1,y1,z1,c):
    x1[c] = 1.0/2.0*x[c] + 1/2.0
    y1[c] = 1.0/2.0*y[c]
    z1[c] = 1.0/2.0*z[c]


def f3(x,y,z,x1,y1,z1,c):
    x1[c] = 1.0/2.0*x[c] + 1/4.
    y1[c] = 1.0/2.0*y[c] + np.sqrt(3)/4
    z1[c] = 1.0/2.0*z[c]

def f4(x,y,z,x1,y1,z1,c):
    x1[c] = 1.0/2.0*x[c] + 1/4.
    y1[c] = 1.0/2.0*y[c] + 1./4
    z1[c] = 1.0/2.0*z[c] + np.sqrt(3)/4

functions = [f1, f2, f3, f4] 
probabilities = [1/4.]*4
assert(len(functions) == len(probabilities))

X,Y,Z = x,y,z
for i in range(20):
    # pick indices for each function to be applied
    r = np.random.choice(len(probabilities), size=N, p=probabilities)
    for i, f in enumerate(functions):
        f(x, y, z, x1, y1, z1, r==i)
    x,x1 = x1,x
    y,y1 = y1,y
    z,z1 = z1,z
    if i > 0:
        X, Y, Z = np.hstack([X,x]), np.hstack([Y,y]), np.hstack([Z,z])

# how much memory are we using, how many points there are
print(3*X.nbytes//1024**2,"MB",X.shape[0])


# Now we turn separate coordinate array into triplets.
# Using `numpy` for what could also be written as `list(zip(X, Y, Z))`.

# In[ ]:


positions = np.vstack([X, Y, Z]).T.copy()


# In[ ]:


import k3d

plot = k3d.plot()
point_plot = k3d.points(positions.astype(np.float32), color=0xff0000, point_size=0.003, shader='3d')
plot += point_plot
plot.display()


# Zoom in a little:

# In[ ]:


plot.camera = [-0.59265772150826,
 1.0966590944867525,
 0.15381683182413644,
 0.35173312413637553,
 0.35752558265043016,
 0.3151305910837551,
 -0.5602813338387698,
 -0.7160643522547137,
 -0.41633720753942915]


# In[ ]: