import torch
import torch.nn.functional as F
import unfoldNd

def extract_patches(x, patch_size, stride, loop=False):
    """Extract patches from a motion sequence"""
    b, c, _t = x.shape

    # manually padding to loop
    if loop:
        half = patch_size // 2
        front, tail = x[:,:,:half], x[:,:,-half:]
        x = torch.concat([tail, x, front], dim=-1)

    x_patches = unfoldNd.unfoldNd(x, kernel_size=patch_size, stride=stride).transpose(1, 2).reshape(b, -1, c, patch_size)
    
    return x_patches.view(b, -1, c * patch_size)

def combine_patches(x_shape, ys, patch_size, stride, loop=False):
    """Combine motion patches"""
    # manually handle to loop
    out_shape = [*x_shape]
    if loop:
        padding = patch_size // 2
        out_shape[-1] = out_shape[-1] + padding * 2

    combined = unfoldNd.foldNd(ys.permute(0, 2, 1), output_size=tuple(out_shape[-1:]), kernel_size=patch_size, stride=stride)

    # normal fold matrix
    input_ones = torch.ones(tuple(out_shape), dtype=ys.dtype, device=ys.device)
    divisor = unfoldNd.unfoldNd(input_ones, kernel_size=patch_size, stride=stride)
    divisor = unfoldNd.foldNd(divisor, output_size=tuple(out_shape[-1:]), kernel_size=patch_size, stride=stride)
    combined = (combined / divisor).squeeze(dim=0).unsqueeze(0)
    
    if loop:
        half = patch_size // 2
        front, tail = combined[:,:,:half], combined[:,:,-half:]
        combined[:, :, half:2 * half] = (combined[:, :, half:2 * half] + tail) / 2
        combined[:, :, - 2 * half:-half] = (front + combined[:, :, - 2 * half:-half]) / 2
        combined = combined[:, :, half:-half]

    return combined


def efficient_cdist(X, Y):
    """
    Pytorch efficient way of computing distances between all vectors in X and Y, i.e (X[:, None] - Y[None, :])**2
    Get the nearest neighbor index from Y for each X
    :param X:  (n1, d) tensor
    :param Y:  (n2, d) tensor
    Returns a n2 n1 of indices
    """
    dist = (X * X).sum(1)[:, None] + (Y * Y).sum(1)[None, :] - 2.0 * torch.mm(X, torch.transpose(Y, 0, 1))
    d = X.shape[1]
    dist /= d # normalize by size of vector to make dists independent of the size of d ( use same alpha for all patche-sizes)
    return dist # DO NOT use torch.sqrt


def get_col_mins_efficient(dist_fn, X, Y, b=1024):
    """
    Computes the l2 distance to the closest x or each y.
    :param X:  (n1, d) tensor
    :param Y:  (n2, d) tensor
    Returns n1 long array of L2 distances
    """
    n_batches = len(Y) // b
    mins = torch.zeros(Y.shape[0], dtype=X.dtype, device=X.device)
    for i in range(n_batches):
        mins[i * b:(i + 1) * b] = dist_fn(X, Y[i * b:(i + 1) * b]).min(0)[0]
    if len(Y) % b != 0:
        mins[n_batches * b:] = dist_fn(X, Y[n_batches * b:]).min(0)[0]

    return mins


def get_NNs_Dists(dist_fn, X, Y, alpha=None, b=1024):
    """
    Get the nearest neighbor index from Y for each X.
    Avoids holding a (n1 * n2) amtrix in order to reducing memory footprint to (b * max(n1,n2)).
    :param X:  (n1, d) tensor
    :param Y:  (n2, d) tensor
    Returns a n2 n1 of indices amd distances
    """
    if alpha is not None:
        normalizing_row = get_col_mins_efficient(dist_fn, X, Y, b=b)
        normalizing_row = alpha + normalizing_row[None, :]
    else:
        normalizing_row = 1

    NNs = torch.zeros(X.shape[0], dtype=torch.long, device=X.device)
    Dists = torch.zeros(X.shape[0], dtype=torch.float, device=X.device)

    n_batches = len(X) // b
    for i in range(n_batches):
        dists = dist_fn(X[i * b:(i + 1) * b], Y) / normalizing_row
        NNs[i * b:(i + 1) * b] = dists.min(1)[1]
        Dists[i * b:(i + 1) * b] = dists.min(1)[0]
    if len(X) % b != 0:
        dists = dist_fn(X[n_batches * b:], Y) / normalizing_row
        NNs[n_batches * b:] = dists.min(1)[1]
        Dists[n_batches * b: ] = dists.min(1)[0]

    return NNs, Dists