# Author: Eric Alcaide

import torch

def get_max_dims(*tensors): 
    """ Gets the maximum dimensions of a set of tensors.
        Inputs: 
        * tensors: (n, d) long tensors
    """
    max_s = []
    for t in tensors: 
        # safe-check type
        if t.dtype != torch.long: 
            t = t.long()
        # get size
        s = 1 + (torch.max(t, dim=-1).values)
        max_s.append( s.cpu() if s.is_cuda else s )
        
    return torch.stack(max_s, dim=0).max(dim=0).values.tolist()


def to_sparse(x, s=None, coalesce=True): 
    """ Returns a tensor in sparse format. 
        Inputs: 
        * x: (n, d) tensor. 
        * s: tuple of ints. size of the returned sparse tensor. 
             inferred if None. 
        * coalesce: bool. whether to coalesce before return.
    """
    # ensures type
    if x.dtype != torch.long: 
        x = x.long()
    # infers max size
    if s is None:
        s = get_max_dims(x)
    # creates tensor
    vals = torch.tensor([1]*x.shape[-1], device=x.device)
    x = torch.sparse_coo_tensor(size=s, values=vals, indices=x)
    if coalesce: 
        return x.coalesce()
    return x
    

def unique(x, dim=None):
    """ Returns unique entries of a tensor. 
        Inputs: 
        * x: (n, d). long tensor()
        Outputs: (n_, d) the number of unique entries. 
    """
    return to_sparse(x, coalesce=True).indices()
    

def intersect(*tensors):
    """ Defines the N-D intersection between different arrays. 
        Inputs: 
        * tensors: any number of tensors to compute the intersection
    """
    max_s = get_max_dims(*tensors)
    
    acc = None
    t_uniques = []
    for t in tensors: 
        t_unique = to_sparse(t, s=max_s)
        if acc is None: 
            acc = t_unique.clone() if return_tensors else t_unique
        else: 
            acc *= t_unique
        # record and return if arg is passes
        t_uniques.append(t_unique)
        
    acc = acc.coalesce()
    return acc.indices() if not return_tensors else \
           (acc.indices(), t_uniques)

def union(*tensors, return_tensors=False):
    """ Defines the N-D union between different arrays. 
        Inputs: 
        * tensors: any number of tensors to compute the intersection
    """
    max_s = get_max_dims(*tensors)
    
    acc = None
    t_uniques = []
    for t in tensors: 
        t_unique = to_sparse(t, s=max_s)
        if acc is None: 
            acc = t_unique.clone() if return_tensors else t_unique
        else: 
            acc += t_unique
        # record and return if arg is passes
        t_uniques.append(t_unique)

    acc = acc.coalesce()
    return acc.indices() if not return_tensors else \
           (acc.indices(), t_uniques)

def xor(*tensors, return_tensors=False):
    """ Defines the N-D union between different arrays. 
        Inputs: 
        * tensors: any number of tensors to compute the intersection
    """
    max_s = get_max_dims(*tensors)
    
    acc = None
    for t in tensors: 
        t_unique = to_sparse(t, s=max_s)
        if acc is None: 
            acc = t_unique.clone() if return_tensors else t_unique
        else: 
            acc += t_unique
            
    # get indices present only once
    acc = acc.coalesce()
    mask = acc.values() == 1
    return acc.indices()[:, mask] if not return_tensors else \
           (acc.indices()[:, mask], t_uniques)

