import torch


def spatial_grad( func ):
    """
    Approximate derivatives of the functions func[b,c,:,:].
    
    dfdx, dfdy = spatial_grad( func )

    In:
    func: torch.FloatTensor
        of shape BxCxhxw with B >= 1 (batch size), C = 1 or C = 3 (color channels),
        h,w >= 3, and [type] is 'Float' or 'Double'.
        Contains the values of functions f_b: R^2 -> R^C, b=1,...,B,
        on the grid {0,...,h-1}x{0,...,w-1}.

    Out:
    dfdx: torch.FloatTensor
    dfdy: torch.FloatTensor
        of shape BxCxhxw contain the x and y derivatives of f_1, ..., f_B
        at the points on the grid, approximated by central differences (except on boundaries):
        For b=0,...,B-1, c=0,...,C, i=1,...,h-2, j=1,...,w-2
        dfdx[b,c,i,j] = (func[b,c,i,j+1] - func[b,c,i,j-1])/2
        dfdx[b,c,i,j] = (func[b,c,i+1,j] - func[b,c,i-1,j])/2

    positive x-direction is along rows from left to right.
    positive y-direction is along columns from above to below.
    """

    # Derivative in x direction (rows from left to right)
    dfdx = torch.zeros_like( func )
    # forward difference in first column
    dfdx[:,:,:,0] = func[:,:,:,1] - func[:,:,:,0]
    # backwards difference in last column
    dfdx[:,:,:,-1] = func[:,:,:,-1] - func[:,:,:,-2]
    # central difference elsewhere
    dfdx[:,:,:,1:-1] = 0.5*(func[:,:,:,2:] - func[:,:,:,:-2] )

    # Derivative in y direction (columns from above to below)
    dfdy = torch.zeros_like( func )
    # forward difference in first row
    dfdy[:,:,0,:] = func[:,:,1,:] - func[:,:,0,:]
    # backwards difference in last row
    dfdy[:,:,-1,:] = func[:,:,-1,:] - func[:,:,-2,:]
    # central difference elsewhere
    dfdy[:,:,1:-1,:] = 0.5*(func[:,:,2:,:] - func[:,:,:-2,:] )

    return dfdx.detach(), dfdy.detach()