import torch


def create_tau( fval, gradf, d1x, d2x, smoothing_operator=None ):
    """
    tau = create_tau( fval, gradf, d1x, d2x )

    In:
    fval: torch.FloatTensor
        of shape B
    gradf: torch.FloatTensor
        of shape B*C*H*W
    d1x: torch.FloatTensor
        of shape B*C*H*W
    d2x: torch.FloatTensor
        of shape B*C*H*W
    smoothing_operator: function
        A self-adjoint smoothing operator sending torch.FloatTensor
        of shape B*2*H*W to torch.FloatTensor of shape B*2*H*W.

    Out:
    tau: torch.FloatTensor
        of shape B*H*W*2
    """
    
    B,C,H,W = gradf.shape
    
    # Sum over color channels
    alpha1 = torch.sum( gradf*d1x, 1).unsqueeze_(1)
    alpha2 = torch.sum( gradf*d2x, 1).unsqueeze_(1)
    
    # stack vector field components into shape B*2*H*W
    tau = torch.cat([alpha1,alpha2], 1)
    
    # Smoothing
    if smoothing_operator:
        tau = smoothing_operator( tau )
        # torch can't sum over multiple axes.
        norm_squared_alpha = (tau**2).sum(1).sum(1).sum(1)
        # In theory, we need to apply the filter a second time.
        tau = smoothing_operator( tau )
    else:
        # torch can't sum over multiple axes.
        norm_squared_alpha = (tau**2).sum(1).sum(1).sum(1)
        
    scale = -fval/norm_squared_alpha
    tau *= scale.view(B,1,1,1)

    # rearrange for compatibility with compose(), B*2*H*W -> B*H*W*2
    return tau.permute( 0, 2, 3, 1 ).detach()