import torch


def one_hot_encode(data, code, numeric=False):
  """Encodes a sequence in one-hot format, given a code.

  Args:
    data (Sequence): sequence of non-encoded data points.
    code (Int | Sequence): sequence of codes for one-hot encoding or a the alphabet length if data is already a list of indices.
  Returns:
    FloatTensor containing a one-hot encoding of the input data.
  """
  try:
    if isinstance(code, int):
      coded = torch.tensor(data, dtype=torch.int64)
      alpha_len = code
    else:
      coded = torch.tensor(list(map(
        code.index, data
      )), dtype=torch.int64)
      alpha_len = len(code)
  except:
    print(data)
    exit()
  if numeric:
    result = coded
  else:
    data_len = len(data)
    result = torch.zeros(alpha_len, data_len, dtype=torch.float)
    result[coded, torch.arange(data_len)] = 1
  return result