import numpy as np
from mpl_toolkits.mplot3d.art3d import Poly3DCollection, Line3DCollection
import mpl_toolkits.mplot3d.art3d as art3d
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from scipy.spatial.transform import Rotation
plt.rcParams['font.family'] = 'IPAmjMincho'

def descartes_3d(ax, ran_x, ran_y, ran_z, ax_title, 
                 x_label = "x軸", y_label = "y軸", z_label="z軸"):
    ax.set_xlabel(x_label, fontsize = 12)
    ax.set_ylabel(y_label, fontsize = 12)
    ax.set_zlabel(z_label, fontsize = 12)
    ax.set_xlim(ran_x[0], ran_x[1])
    ax.set_ylim(ran_y[0], ran_y[1])
    ax.set_zlim(ran_z[0], ran_z[1])
    ax.set_title(ax_title, fontsize = 16)
    ax.grid()

fig = plt.figure(figsize = (6, 6))
ax = fig.add_subplot(111, projection='3d')
title1 = "Tokyo University 2023 Math 第６問 円弧 NO3_2"

descartes_3d(ax, [-2, 2], [-2, 2], [-2, 2], title1)
    
alpha = np.pi/2-np.arcsin(1/np.sqrt(3))
r= np.sqrt(3) # 半径を指定

theta_3_0 = np.linspace(-(np.pi/2-alpha), (np.pi/2-alpha), 100)
theta_4_0 = 0
x_c = np.cos(theta_4_0) * np.sin(theta_3_0) * r 
y_c = np.sin(theta_4_0) * np.sin(theta_3_0) * r 
z_c = np.cos(theta_3_0) * r 

x_d = x_c.reshape(100,1)
y_d = y_c.reshape(100,1)
z_d = z_c.reshape(100,1)
mm = np.r_["2,3,1", x_d, y_d , z_d]

mm2 = mm.reshape(100,3)

rot = Rotation.from_rotvec(np.array([np.pi/4,0,0]))
e = rot.apply(mm2)

for k in range(4):
    x_t=[]
    y_t=[]
    z_t=[]
    rot = Rotation.from_rotvec(np.array([0,0,np.pi/2*k]))
    e = rot.apply(e)
    for j in range(3):
        for i in range(100):
            if j == 0:
                x_t.append(e[i][j])
            if j == 1:
                y_t.append(e[i][j])
            if j == 2:
                z_t.append(e[i][j])
    
    ax.plot(x_t, y_t, z_t, color="r",lw=3)

ax.set_xlabel('x 軸')
ax.set_ylabel('y 軸')
ax.set_zlabel('z 軸')
ax.set_xlim(-2,2)
ax.set_ylim(2,-2)
ax.set_zlim(-2,2)

plt.show()