import numpy as np
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d

from nitransforms.analysis.utils import sample_unit_sphere

values = (1, 2, 8, 10, 12, 20)

fig = plt.figure(figsize=(10, 6))
for i, n_pts in enumerate(values, start=1):
    X = sample_unit_sphere(n_pts)
    ax = fig.add_subplot(2, 3, i, projection="3d")
    ax.scatter(X[:, 0], X[:, 1], X[:, 2], s=30)
    ax.set_title(f"n={n_pts}")
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("z")
    ax.set_box_aspect((1, 1, 1))
fig.tight_layout()