import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d.art3d as art3d
import numpy as np
from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d import proj3d


class myArrow3D(FancyArrowPatch):
    def __init__(self, xs, ys, zs, *args, **kwargs):
        FancyArrowPatch.__init__(self, (0, 0), (0, 0), *args, **kwargs)
        self._verts3d = xs, ys, zs

    def do_3d_projection(self, renderer=None):
        xs3d, ys3d, zs3d = self._verts3d
        xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, self.axes.M)
        self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
        return np.min(zs)


class otherArrow3D(FancyArrowPatch):
    # https://stackoverflow.com/a/74122407/31611660
    # https://stackoverflow.com/a/22867877/31611660

    def __init__(self, xdir, zdir, *args, **kwargs):
        super().__init__(*args, **kwargs)
        transform_2d_to_3d(self, veca=xdir, vecb=zdir)


def rotation_matrix(x, z):
    x = np.array(x) / np.linalg.norm(x)
    y = np.cross(np.array(z), x)
    z = np.cross(x, y)
    return np.array([x, y / np.linalg.norm(y), z / np.linalg.norm(z)]).T


def transform_2d_to_3d(pathpatch, veca=[-1, 0, 0], vecb=[0, 1, 0], zs=0):
    """
    Transforms a 2D Patch to a 3D patch using the given normal vector.

    The patch is projected into they XY plane, rotated about the origin
    and finally translated by z.
    https://stackoverflow.com/a/18228967/31611660
    """
    path = pathpatch.get_path()  # Get the path and the associated transform
    trans = pathpatch.get_patch_transform()
    path = trans.transform_path(path)  # Apply the transform
    pathpatch.__class__ = art3d.PathPatch3D  # Change the class
    pathpatch._code3d = path.codes  # Copy the codes
    pathpatch._facecolor3d = pathpatch.get_facecolor  # Get the face color
    verts = path.vertices  # Get the vertices in 2D
    pathpatch._segment3d = np.array(
        [np.dot(rotation_matrix(veca, vecb), (x, y, 0)) + (0, 0, zs) for x, y in verts]
    )


fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
ax.set_xlim([0, 4])
ax.set_ylim([4, 0])
ax.set_zlim([0, 4])
ax.plot([0, 0], [0, 0], [0, 4], "r-")
ax.plot([0, 4], [0, 4], [0, 4], "g-")
ax.add_artist(
    myArrow3D(
        [0, 1],
        [0, 1],
        [1, 1],
        mutation_scale=25,
        lw=0.5,
        arrowstyle="simple",
        color="b",
    )
)
ax.add_artist(
    myArrow3D(
        [0, 2],
        [0, 2],
        [2, 2],
        mutation_scale=25,
        lw=1,
        arrowstyle="simple",
        connectionstyle="arc3, rad=-0.5",
        color="r",
    )
)
ax.add_artist(
    otherArrow3D(
        [0, 0, 1],
        np.cross([0, 0, 1], [1, 1, 1]),
        (3, 0),
        (3, 3),
        lw=1,
        mutation_scale=1,
        arrowstyle="simple",
        connectionstyle="arc3, rad=0.75",
        color="g",
    )
)
ax.set(xlabel="X", ylabel="Y", zlabel="Z")
plt.show()