import torch


def get_uniq_topk(rank, history=None):
  """Based on rank and history, select the top ranked that is not in history 

  Args:
    rank: (N, seq_len) torch.LongTensor, from torch.sort(x)[1]
    history: either None, or a torch.LongTensor of size (N, history_len); history_len <= seq_len

  Returns:
    res: torch.LongTensor of size (N,1)
    history: torch.LongTensor of size (N, history_len); if initially history=None, then history_len = 1;
      otherwise the new history will have res appended as the last column
  """
  res = []
  if history is None:
    res = rank[:, :1]
    history = res
  else:
    for r, h in zip(rank, history):
      for i in r:
        if i in h:
          continue
        else:
          res.append(i)
          break
    res =  torch.stack(res, dim=0).unsqueeze(-1)
    history = torch.cat([history, res], dim=-1) # in fact, dim=1
  return res, history