import torch


def xform_transpose(xform):
    s = list(range(len(xform.shape)))
    s[-1], s[-2] = s[-2], s[-1]
    return xform.permute(*s)


def xform_fk_vel(lxform, lpos, lvrt, lvel, parents):
    gr, gp, gt, gv = [lxform[..., :1, :, :]], [lpos[..., :1, :]], [lvrt[..., :1, :]], [lvel[..., :1, :]]
    for i in range(1, len(parents)):
        p = parents[i]
        gp.append(gp[p] + torch.matmul(gr[p], lpos[..., i:i + 1, :][..., None])[..., 0])
        gr.append(torch.matmul(gr[p], lxform[..., i:i + 1, :, :]))
        gt.append(gt[p] + torch.matmul(gr[p], lvrt[..., i:i + 1, :][..., None])[..., 0])
        gv.append(gv[p] + torch.matmul(gr[p], lvel[..., i:i + 1, :][..., None])[..., 0] +
                  torch.cross(gt[p], torch.matmul(gr[p], lpos[..., i:i + 1, :][..., None])[..., 0], dim=-1))

    return torch.cat(gr, dim=-3), torch.cat(gp, dim=-2), torch.cat(gt, dim=-2), torch.cat(gv, dim=-2)


def xform_orthogonalize_from_xy(xy, eps=1e-10):
    xaxis = xy[..., 0:1, :]
    zaxis = torch.cross(xaxis, xy[..., 1:2, :])
    yaxis = torch.cross(zaxis, xaxis)

    output = torch.cat([
        xaxis / (torch.norm(xaxis, 2, dim=-1)[..., None] + eps),
        yaxis / (torch.norm(yaxis, 2, dim=-1)[..., None] + eps),
        zaxis / (torch.norm(zaxis, 2, dim=-1)[..., None] + eps)
    ], dim=-2)

    return xform_transpose(output)