import numpy as np
from vis.xyz_utils import make_3d_vector
from misc.matrix_building import Rt_3d


def make_thrust_vectors(model, scales=(), Rt=()):
    """
    Draw thrust vectors with direction opposite to the force being
    applied, like the direction of exhaust.
    (since force direction may end up *inside* the 3d model)
    :param model:
    :param scales: multipliers for vector line and head length
    :param Rt: 4, 4 | homogeneous coordinate rigid body transform
        - body frame to given frame
    :return:
    """
    object_list = []
    n_thrusters = model.r.shape[1]
    if len(scales) == 0:
        scales = np.ones((n_thrusters,))
    if len(Rt) == 0:
        Rt = Rt_3d()
    r = Rt[:3, :3].dot(model.r) + Rt[:3, -1][:, np.newaxis]
    f = Rt[:3, :3].dot(model.f)
    for i in range(n_thrusters):
        name = 'thruster {}'.format(i)
        arrow_objs = make_3d_vector(
            r[:, i], -f[:, i], scale=scales[i], name=name)
        object_list.extend(arrow_objs)
    return object_list