import csv


def load_subset(subset):
  """Loads one csv file.

  Args:
    subset: Either train or val.

  Returns:
    boxes: A dictionary: {(video_id,  frame_no, person_id): bounding_box}.
    categories: A dictionary:
      {combined_label: [(video_id, person_id, [frame_nos])]}.
    cat2: A dictionary: {combined_label: {video_id: {frame_no: [person_ids]}}}.
  """
  instances = {}
  boxes = {}
  categories = {}
  cat2 = {}
  with open('data/ava_%s_v2.1.csv' % subset, 'r') as f:
    reader = csv.reader(f)
    for row in reader:
      index = (row[0], row[7])
      if index in instances:
        instances[index].setdefault(int(row[1]), []).append(int(row[6]))
      else:
        instances[index] = {}
        instances[index].setdefault(int(row[1]), []).append(int(row[6]))
      index = (row[0], row[1], row[7])
      box = map(float, row[2:6])
      if index in boxes:
        assert boxes[index] == box, index
      else:
        boxes[index] = box

  for k in instances.keys():
    v = instances[k]
    val = []
    label_frame = {}
    for f, l in v.items():
      l.sort()
      label_frame.setdefault(tuple(l), []).append(f)
    for l, f in label_frame.items():
      if l not in cat2:
        cat2[l] = {}
      if k[0] not in cat2[l]:
        cat2[l][k[0]] = {}
      target = cat2[l][k[0]]
      for i in f:
        target.setdefault(i, []).append(k[1])
      f.sort()
      f.append(99999)
      prev = 0
      for i in range(len(f) - 1):
        if f[i] + 1 != f[i + 1]:
          val.append((f[prev:i + 1], l))
          categories.setdefault(l, []).append(k + (f[prev:i + 1],))
          prev = i + 1
    instances[k] = val
  return boxes, categories, cat2