import numpy as np
import torch


def get_shifted_data(data, di, dj):
    H, W = data.shape
    shifted_data = data[max(di, 0): H + min(di, 0), max(dj, 0): W + min(dj, 0)]
    if isinstance(data, np.ndarray):
        if di < 0:
            shifted_data = np.concatenate([np.zeros_like(shifted_data[di:]), shifted_data], axis=0)
        if di > 0:
            shifted_data = np.concatenate([shifted_data, np.zeros_like(shifted_data[:di])], axis=0)
        if dj < 0:
            shifted_data = np.concatenate([np.zeros_like(shifted_data[:, dj:]), shifted_data], axis=1)
        if dj > 0:
            shifted_data = np.concatenate([shifted_data, np.zeros_like(shifted_data[:, :dj])], axis=1)
    elif isinstance(data, torch.Tensor):
        shifted_data = data[max(di, 0): H + min(di, 0), max(dj, 0): W + min(dj, 0)]
        if di < 0:
            shifted_data = torch.cat([torch.zeros_like(shifted_data[di:]), shifted_data], dim=0)
        if di > 0:
            shifted_data = torch.cat([shifted_data, torch.zeros_like(shifted_data[:di])], dim=0)
        if dj < 0:
            shifted_data = torch.cat([torch.zeros_like(shifted_data[:, dj:]), shifted_data], dim=1)
        if dj > 0:
            shifted_data = torch.cat([shifted_data, torch.zeros_like(shifted_data[:, :dj])], dim=1)
    return shifted_data