from sympy import tan, cos, sin, pi, symbols
from spb import *
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 to discretize vector field (low numb of discret points)
cone_discr = surface_parametric(
    u / tan(t), u * cos(v), u * sin(v), (u, 0, 1), (v, 0 , 2*pi),
    n1=n1, n2=n2)[0]
graphics(
    surface_parametric(
        u / tan(t), u * cos(v), u * sin(v), (u, 0, 1), (v, 0 , 2*pi),
        rendering_kw={"opacity": 1}, wireframe=True,
        wf_n1=n1, wf_n2=n2, wf_rendering_kw={"width": 0.004}),
    vector_field_3d(
        n, range1=(xn, -5, 5), range2=(yn, -5, 5), range3=(zn, -5, 5),
        use_cm=False, slice=cone_discr,
        quiver_kw={"scale": 0.5, "pivot": "tail"}
    ),
    backend=KB)