from sympy import tan, cos, sin, pi, symbols
from spb import plot3d_parametric_surface, plot_vector, KB
from sympy.vector import CoordSys3D, gradient

u, v = symbols("u, v")
N = CoordSys3D("N")
i, j, k = N.base_vectors()
xn, yn, zn = N.base_scalars()

t = 0.35    # half-cone angle in radians
expr = -xn**2 * tan(t)**2 + yn**2 + zn**2    # cone surface equation
g = gradient(expr)
n = g / g.magnitude()    # unit normal vector
n1, n2 = 10, 20 # number of discretization points for the vector field

# cone surface for visualization (high number of discretization points)
p1 = plot3d_parametric_surface(
    u / tan(t), u * cos(v), u * sin(v), (u, 0, 1), (v, 0 , 2*pi),
    {"opacity": 1}, backend=KB, show=False, wireframe=True,
    wf_n1=n1, wf_n2=n2, wf_rendering_kw={"width": 0.004})
# cone surface to discretize vector field (low numb of discret points)
p2 = plot3d_parametric_surface(
    u / tan(t), u * cos(v), u * sin(v), (u, 0, 1), (v, 0 , 2*pi),
    n1=n1, n2=n2, show=False)
# plot vector field on over the surface of the cone
p3 = plot_vector(
    n, (xn, -5, 5), (yn, -5, 5), (zn, -5, 5), slice=p2[0],
    backend=KB, use_cm=False, show=False,
    quiver_kw={"scale": 0.5, "pivot": "tail"})
(p1 + p3).show()