import torch


def compute_times_wrt_weighted_position(C, A, B, t1):
    """
    According to a point C between AB segment and the time t1,
    compute the times at A and B (tA and tB)
    """
    distance = torch.linalg.norm(B - A, axis=1)
    # timing (t1 is the sum tA+tB)
    t1 = torch.squeeze(t1)
    f = t1 / distance
    # prevent nan -> set both time to zero
    mask = distance == 0
    f[mask] = 0
    print(f'Number of distance==0 {mask.sum()}/{len(distance)}')
    # distances
    distA = torch.linalg.norm(C - A, axis=1)
    distB = torch.linalg.norm(C - B, axis=1)
    tA = distA * f
    tB = distB * f
    return tA, tB