import matplotlib.pylab as plt
import numpy as np
# %matplotlib widget

def f(s, t):
    laplace = np.exp(s*t)
    return laplace.real, laplace.imag

steps = 1000
theta = np.linspace(0, 10*np.pi, steps)
s = 0.2 + 5j
t = np.linspace(0, 10, steps)
x, y = f(s, t)

origin = np.array([0, 0, 0])

def colormap(angle):
    color = np.array([0, 0, 0])
    if 0 < angle < np.pi/2:
        ratio = np.abs(angle/(np.pi/2))
        red = np.array([1, 0, 0])
        yellow = np.array([1, 1, 0]) * (ratio)
        blue = np.array([0, 0, 1]) * (1-ratio)
        color = np.mean([red, yellow, blue], axis=0)

    elif np.pi/2 < angle < np.pi:
        ratio = np.abs((angle - np.pi/2)/(np.pi/2))
        red = np.array([1, 0, 0]) * (1-ratio)
        yellow = np.array([1, 1, 0]) 
        green = np.array([0, 1, 0]) * (ratio)
        color = np.mean([red, yellow, green], axis=0)

    elif -np.pi < angle < -np.pi/2:
        ratio = np.abs((angle + np.pi/2)/(np.pi/2))
        yellow = np.array([1, 1, 0]) * (ratio)
        blue = np.array([0, 0, 1]) * (1-ratio)
        green = np.array([0, 1, 0])
        color = np.mean([yellow, blue, green], axis=0)
    
    elif -np.pi/2 < angle < 0:
        ratio = np.abs((angle)/(np.pi/2))
        red = np.array([1, 0, 0]) * (1-ratio)
        green = np.array([0, 1, 0]) * (ratio)
        blue = np.array([0, 0, 1])
        color = np.mean([red, green, blue], axis=0)
      
    else:
        color = 'k'
    return color

fig = plt.figure(figsize = (6, 6))
ax = plt.axes(projection='3d')

ax.plot(*origin, color = 'blue', marker = 'o')
ax.text(origin[0] + 0.25, origin[1] + 0.25, origin[2] + 0.25, s = 'O', fontsize = 12)

# ax.plot_surface(X, Y, Z, color= 'red', alpha = 0.5)

for i in range(len(x)):
    angle = np.angle(complex(x[i], y[i]))
    color = colormap(angle)
    ax.plot(x[i], y[i], t[i], color = color, marker = '.')

ax.set_xlim(-10, 10)
ax.set_ylim(-10, 10)
ax.set_zlim(0, 10)
plt.show()