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

center = np.array([0, 0, 0])
a = 4
b = 4
c = 4

r = np.arange(0, 4)
x, y = np.meshgrid(r, r)

phi = np.arange(1,10,2)*np.pi/4
Phi, Theta = np.meshgrid(phi, phi)

x = a * np.cos(Phi)*np.sin(Theta) + center[0]
y = b * np.sin(Phi)*np.sin(Theta) + center[1]
z = c * np.cos(Theta)/np.sqrt(2) + center[2]

# r = [2, 0]
# x1, y1 = np.meshgrid(r, r)
# x1, y1

x1 = np.array([
    [2, 2],
    [-2, -2]
])
y1 = np.array([
    [-2, -2+2*np.cos(np.pi/3)],
    [-2, -2+2*np.cos(np.pi/3)]
])

z1 = np.array([
    [2, 2+2*np.sin(np.pi/3)],
    [2, 2+2*np.sin(np.pi/3)]
])

x2 = np.array([
    [2, 2],
    [-2, -2]
])
y2 = np.array([
    [2, 2+2*np.cos(np.pi - np.pi/3)],
    [2, 2+2*np.cos(np.pi - np.pi/3)]
])

z2 = np.array([
    [2, 2+2*np.sin(np.pi - np.pi/3)],
    [2, 2+2*np.sin(np.pi - np.pi/3)]
])

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

ax.plot_surface(x, y, z, color = 'blue', alpha = 0.5)

ax.plot_surface(x1, y1, z1, color = 'red', alpha = 0.5)
ax.plot_surface(x2, y2, z2, color = 'green', alpha = 0.5)

ax.set_xlim(-5,5)
ax.set_ylim(-5,5)
ax.set_zlim(-5,5)
plt.show()