import torch


def update_affine_param( cur_af, last_af): # A2(A1*x+b1) + b2 = A2A1*x + A2*b1+b2
    """
       update the current affine parameter A2 based on last affine parameter A1
        A2(A1*x+b1) + b2 = A2A1*x + A2*b1+b2, results in the composed affine parameter A3=(A2A1, A2*b1+b2)
       :param cur_af: current affine parameter
       :param last_af: last affine parameter
       :return: composed affine parameter A3
    """
    cur_af = cur_af.view(cur_af.shape[0], 4, 3)
    last_af = last_af.view(last_af.shape[0],4,3)
    updated_af = torch.zeros_like(cur_af.data).to(cur_af.device)
    dim =3
    updated_af[:,:3,:] = torch.matmul(cur_af[:,:3,:],last_af[:,:3,:])
    updated_af[:,3,:] = cur_af[:,3,:] + torch.squeeze(torch.matmul(cur_af[:,:3,:], torch.transpose(last_af[:,3:,:],1,2)),2)
    updated_af = updated_af.contiguous().view(cur_af.shape[0],-1)
    return updated_af