import numpy as np

def sample_sphere(n_points=500, radius=1.0, seed=0):
    rng = np.random.default_rng(seed)
    phi = rng.uniform(0, 2*np.pi, n_points)
    cos_theta = rng.uniform(-1, 1, n_points)
    theta = np.arccos(cos_theta)

    x = radius * np.sin(theta) * np.cos(phi)
    y = radius * np.sin(theta) * np.sin(phi)
    z = radius * np.cos(theta)
    return np.stack([x, y, z], axis=1)

def sample_torus(n_points=500, R=2.0, r=0.7, seed=0):
    """
    R: 大きい半径（ドーナツの中心からチューブ中心まで）
    r: 小さい半径（チューブの太さ）
    """
    rng = np.random.default_rng(seed)
    u = rng.uniform(0, 2*np.pi, n_points)
    v = rng.uniform(0, 2*np.pi, n_points)

    x = (R + r * np.cos(v)) * np.cos(u)
    y = (R + r * np.cos(v)) * np.sin(u)
    z = r * np.sin(v)
    return np.stack([x, y, z], axis=1)


from ripser import ripser
from persim import plot_diagrams
import matplotlib.pyplot as plt

# 点群生成
X_sphere = sample_sphere(n_points=800, radius=1.0, seed=0)
X_torus  = sample_torus(n_points=800, R=2.0, r=0.7, seed=0)

# 計算（0次〜2次くらいまで見る）
result_sphere = ripser(X_sphere, maxdim=2)
result_torus  = ripser(X_torus,  maxdim=2)

dgms_sphere = result_sphere["dgms"]
dgms_torus  = result_torus["dgms"]

# 図を並べて表示
fig, axes = plt.subplots(2, 3, figsize=(12, 8))

# 元データの散布図（簡単に 3D 風）
axes[0,0].scatter(X_sphere[:,0], X_sphere[:,1], s=5)
axes[0,0].set_title("Sphere (projection)")
axes[0,0].set_aspect("equal")

axes[1,0].scatter(X_torus[:,0], X_torus[:,1], s=5)
axes[1,0].set_title("Torus (projection)")
axes[1,0].set_aspect("equal")

# 球のバーコード
plot_diagrams(dgms_sphere, show=False, ax=axes[0,1])
axes[0,1].set_title("Sphere: persistence diagrams")

# トーラスのバーコード
plot_diagrams(dgms_torus, show=False, ax=axes[1,1])
axes[1,1].set_title("Torus: persistence diagrams")

# 次元1だけを拡大して比較
plot_diagrams([dgms_sphere[1]], show=False, ax=axes[0,2])
axes[0,2].set_title("Sphere: H1 only")

plot_diagrams([dgms_torus[1]], show=False, ax=axes[1,2])
axes[1,2].set_title("Torus: H1 only")

plt.tight_layout()
plt.show()