import torch


def get_transforms_to_optimize(chain, to_optimize_quat, to_optimize_trans,
                               device):
    """Returns pytorch tensors to optimize along the chain as per
    to_optimize_trans and to_optimize_quat."""
    opt_pyt_trans, opt_pyt_quat = [], []
    for i in range(len(chain) - 1):
        t = None
        q = None
        if chain[i + 1] in to_optimize_trans:
            t = torch.zeros(1, 3, device=device, dtype=torch.double,
                            requires_grad=True)
        if chain[i + 1] in to_optimize_quat:
            qxyz = torch.zeros(1, 3, device=device, dtype=torch.double,
                               requires_grad=True)
            qw = torch.ones(1, 1, device=device, dtype=torch.double,
                            requires_grad=True)
            q = [qw, qxyz]
        opt_pyt_trans.append(t)
        opt_pyt_quat.append(q)
    return opt_pyt_quat, opt_pyt_trans