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()

def m_heart(ax,r=1,start_z_arg=0,end_z_arg=2*np.pi,
             x_arg=2*np.pi,color="red",alpha=1):

    xx_arg = np.linspace(0, x_arg, 100) 
    zz_arg = -1/(2*np.pi)*(xx_arg-np.pi)**2+3/4*np.pi 

    x = np.cos(xx_arg) * np.sin(zz_arg) * r
    y = np.sin(xx_arg) * np.sin(zz_arg) * r
    z = np.cos(zz_arg) * r

    return x,y,z

def m_rotation(x,y,z,rot_vector=[np.pi/4,0,0]):

    x1 = x.reshape(100,1)
    y1 = y.reshape(100,1)
    z1 = z.reshape(100,1)

    xyz = np.r_["1", x1 , y1 , z1]
    rot = Rotation.from_rotvec(np.array(rot_vector))
    xyz2 = rot.apply(xyz)

    x2=xyz2[:,0]
    y2=xyz2[:,1]
    z2=xyz2[:,2]

    return x2,y2,z2

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

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

r= np.sqrt(3)

x,y,z=m_heart(ax,r,start_z_arg=0,end_z_arg=2*np.pi)


x,y,z=m_rotation(x,y,z,rot_vector=[np.pi/4,0,0])
ax.plot(x, y, z, color="r",lw=2)


plt.show()