import numpy as np 
import matplotlib.pyplot as plt 
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
r = 10
Init = [0,0,0] ; Final = [5,5,5]
u, v = np.mgrid[0:2 * np.pi:30j, 0:np.pi:20j]
x = Init[0] + r*np.cos(u) * np.sin(v)
y = Init[1] + r*np.sin(u) * np.sin(v)
z = Init[2] + r*np.cos(v)
ax.plot_surface(x, y, z, color = 'r', alpha = 0.5, linewidth=0)
    
u, v = np.mgrid[0:2 * np.pi:30j, 0:np.pi:20j]
x = Final[0] + r*np.cos(u) * np.sin(v)
y = Final[1] + r*np.sin(u) * np.sin(v)
z = Final[2] + r*np.cos(v)
ax.plot_surface(x, y, z, color = 'r', alpha = 0.5, linewidth=0)
plt.show()

fig = plt.figure()
ax = fig.add_subplot(projection='3d')
# centers of the spheres
c1 = np.array([-1, -2, -3])
c2 = np.array([1, 2, 3])
# draw centers and a segment connecting them
centers = np.stack([c1, c2])
ax.scatter(centers[:, 0], centers[:, 1], centers[:, 2])
ax.plot(centers[:, 0], centers[:, 1], centers[:, 2])
r1 = 5
r2 = 3
# intersection of two sphere is a circle
# compute the radius of this circle.
# https://mathworld.wolfram.com/Sphere-SphereIntersection.html
# distance between centers
d = np.linalg.norm(c1 - c2)
# radius of the circle
a = np.sqrt((-d + r2 - r1) * (-d - r2 + r1) * (-d + r2 + r1) * (d + r2 + r1)) / (2 * d)
# we want to draw sphere caps that are cut at the intersecting circle.
# these are the limiting angles for the parameterization
alpha = np.arcsin(a / r2)
beta = np.arcsin(a / r1)
# parameterization such that the "poles" intersect the x-axis
u, v = np.mgrid[0:2 * np.pi:30j, beta:np.pi:20j]
x1 = r1 * np.cos(v)
y1 = r1 * np.cos(u) * np.sin(v)
z1 = r1 * np.sin(u) * np.sin(v)
 
u, v = np.mgrid[0:2 * np.pi:30j, 0:np.pi-alpha:20j]
x2 = d + r2 * np.cos(v)
y2 = r2 * np.cos(u) * np.sin(v)
z2 = r2 * np.sin(u) * np.sin(v)
# let's compute the two rotation angles that aligns the
# segment connecting the centers to the x-axis
c2new = c2 - c1
theta = np.arctan2(c2new[1], c2new[0])
phi = np.arctan2(c2new[2], np.sqrt(c2new[0]**2 + c2new[1]**2))
# rotations matrices. 4x4 because they are easier to apply
# to a matrix of coordinates
Ry = lambda t: np.array([
    [np.cos(t), 0, np.sin(t), 0],
    [0, 1, 0, 0],
    [-np.sin(t), 0, np.cos(t), 0],
    [0, 0, 0, 1]
])
Rz = lambda t: np.array([
    [np.cos(t), np.sin(t), 0, 0],
    [-np.sin(t), np.cos(t), 0, 0],
    [0, 0, 1, 0],
    [0, 0, 0, 1]
])
# first, rotate by theta around the z-axis,
# then rotate by phi around y-axis
rot = Ry(phi) @ Rz(theta)
def rotate_points(x, y, z):
    shape = x.shape
    # add a column of 1s so that we can multiply the
    # resulting 3xN array with the rotation matrix
    coord = np.stack([t.flatten() for t in [x, y, z, np.ones_like(x)]])
    coord = rot.T @ coord
    x, y, z, _ = [t.reshape(shape) for t in coord]
    return x, y, z
x1, y1, z1 = rotate_points(x1, y1, z1)
x2, y2, z2 = rotate_points(x2, y2, z2)
# apply the necessary offset to the spheres and plot them
ax.plot_surface(x1 + c1[0], y1 + c1[1], z1 + c1[2], color = 'r', alpha = 0.5, linewidth=0)
ax.plot_surface(x2 + c1[0], y2 + c1[1], z2 + c1[2], color = 'g', alpha = 0.5, linewidth=0)
ax.set_aspect("equal")
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.view_init(30, -20)
plt.show()