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

def descartes_3d(ax, ran_x, ran_y, ran_z, ax_title, 
                 x_label = "x-axis", y_label = "y-axis", z_label="z-axis"):
    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()

def m_circle(ax,r=1,start_arg_z=0,end_z_arg=2*np.pi,x_arg=0,color="blue",alpha=0.5,rot_vector=[np.pi/4,0,0]):

    xx_arg = np.linspace(0, x_arg, 100) 
    zz_arg = np.linspace(start_arg_z, end_z_arg, 100) 

    
    x1 = np.cos(xx_arg) * np.sin(zz_arg) * r
    y1 = np.sin(xx_arg) * np.sin(zz_arg) * r
    z1 = np.cos(zz_arg) * r

    x2 = x1.reshape(100,1)
    y2 = y1.reshape(100,1)
    z2 = z1.reshape(100,1)

    xyz = np.r_["2,3,1", x2 , y2 , z2]

    xyz2 = xyz.reshape(100,3)

    rot = Rotation.from_rotvec(np.array(rot_vector))
    xyz3 = rot.apply(xyz2)

    x = []
    y = []
    z = []

    x=xyz3[:,0]
    y=xyz3[:,1]
    z=xyz3[:,2]
    
    ax.plot(x, y, z, color="r",lw=2)


fig = plt.figure(figsize = (6, 6))
ax = fig.add_subplot(111, projection='3d')
title1 = "Tokyo University 2023 Math qu.6 arc-shaped"

descartes_3d(ax, [-2, 2], [-2, 2], [-2, 2], title1)

alpha = np.pi/2-np.arcsin(1/np.sqrt(3))
r= np.sqrt(3) 

for i in range(4):
    m_circle(ax,r,-alpha,alpha,rot_vector=[np.pi/4,0,np.pi/2*i])


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

plt.show()