# --------------------------------------------------------
# VIRAT Pickel for R-C3D
# --------------------------------------------------------

import os
from six.moves import cPickle as pkl
import numpy as np
import json
import copy
import sys
import argparse
from diva.utils import activity

def generate_testing_video_db(json_index_path, chunk_file, chunk_id):
  # *******************************************************************************
  # Extract Activity times for test set
  video_db = {}
  use_chunks = False
  with open(json_index_path, 'r') as f:
    all_annotations = json.load(f)

  if chunk_id > 0:
    use_chunks = True
    with open(chunk_file, 'r') as f:
      all_chunks = json.load(f)
      chunk_key = "Chunk" + str(chunk_id)
      if chunk_key in all_chunks.keys():
        chunk_videos = all_chunks[chunk_key]["files"]
      else:
        raise Exception("Invalid chunk id " + str(chunk_id) + " provided")
  for video_name, annotation in all_annotations.items():
    if use_chunks:
        if video_name not in chunk_videos:
            continue
    selected_dict = annotation["selected"]
    start_indices = []
    end_indices = []
    for frame_str, selected_flag in selected_dict.items():
        if selected_flag == 1:
            start_indices.append(int(frame_str))
        else:
            end_indices.append(int(frame_str))
    start_indices = sorted(start_indices)
    end_indices = sorted(end_indices)
    assert len(start_indices) == len(end_indices), \
            "Number of start and end durations don't match"
    for index  in range(len(start_indices)):
        start_index = start_indices[index]
        end_index = end_indices[index]
        video_prefix = os.path.splitext(video_name)[0]
        if video_prefix not in video_db.keys():
            video_db[video_prefix] = [[start_index, end_index, -1]]
        else:
            video_db[video_prefix].append([start_index, end_index, -1])
  return video_db
  # *******************************************************************************

def generate_kpf_training_video_db(kpf_dirs, classes):
  # *******************************************************************************
  # Extract Activity times and lables

  video_db = {}
  kpf_activity = activity()
  print(kpf_dirs)
  for kpf_dir in kpf_dirs:
    file_list = os.listdir(kpf_dir)
    print('Analyzing ' + kpf_dir)
    for kpf_file in file_list:
      if kpf_file.endswith(".activities.yml"):  # assumes that all test videos have an activity in the name
        prefix = kpf_file.split('.')[0]
        video_db[prefix] = []
        with open(os.path.join(kpf_dir, kpf_file)) as f:
          print('Analyzing file ' + os.path.join(kpf_dir,kpf_file))
          for kpf in f:
            if kpf.startswith("- { act:"):
              kpf_activity.clear()
              kpf_activity.from_string(kpf)
              label = kpf_activity.get_max_activity_name()

              if label in classes:  # Should need to know this truth information, how is this is used in testing?
                span = kpf_activity.get_frame_id_span()

                for pair in span:
                  start_time = pair[0]
                  end_time = pair[1]
                  video_db[prefix].append([start_time, end_time, classes[label]])

  for video in video_db:
    video_db[video].sort(key=lambda x: x[0])

  return video_db
  # *******************************************************************************

def generate_density_training_video_db(density_dirs, classes):
  # ********************************************************
  # Extract Activity times and lables from the density files
  video_db = {}
  for density_dir in density_dirs:
    file_list = os.listdir(density_dir)

    for density_file in file_list:
      with open(os.path.join(density_dir, density_file)) as f:

        for line in f:
          act = line.split()
          if act[2] in classes and act[5] == '1' and act[6] == '0':
            prefix = act[0]

            if prefix not in video_db.keys():
              video_db[prefix] = []
            video_db[prefix].append([ int(act[3]), int(act[4]), classes[act[2]] ])

  for video in video_db:
    video_db[video].sort(key=lambda x: x[0])

  return video_db
  # ********************************************************

def generate_json_training_video_db(json_files, classes):
  video_db = {}
  for json_file in json_files:
    if os.path.exists(json_file):
      activity_dict = json.load(open(json_file, 'r'))
      activities = activity_dict["activities"]
      for activity in activities:
        class_label = activity["activity"]
        localization_info = activity["localization"]
        for video_name, temporal_span in localization_info.items():
          video_prefix = os.path.splitext(video_name)[0]
          start_frame = -1
          end_frame = -1
          for frame in temporal_span:
            if temporal_span[frame] == 1:
                start_frame = int(frame)
            else:
                end_frame = int(frame)
          if video_prefix in video_db.keys():
            try:
                video_db[video_prefix].append([start_frame, end_frame, classes[class_label]])
            except KeyError:
                pass
          else:
            video_db[video_prefix] = [[start_frame, end_frame, classes[class_label]]]
    else:
      raise Exception("Annotations are not available for " + json_file)
  return video_db

def generate_testing_roi(video_folder, start, end, stride, is_flipped):
  tmp = {}
  tmp['flipped'] = is_flipped
  tmp['frames'] = np.array([[0, start, end, stride]])
  tmp['bg_name'] = video_folder
  tmp['fg_name'] = video_folder
  return tmp

def generate_training_roi(rois, video_folder, start, end, stride, is_flipped):
  tmp = {}
  tmp['gt_classes'] = rois[:, 2]
  tmp['max_classes'] = rois[:, 2]
  tmp['max_overlaps'] = np.ones(len(rois))
  tmp['flipped'] = is_flipped
  tmp['frames'] = np.array([[0, start, end, stride]])
  tmp['wins'] = (rois[:,:2] - start) / stride
  tmp['durations'] = tmp['wins'][:, 1] - tmp['wins'][:, 0] + 1
  tmp['bg_name'] = video_folder
  tmp['fg_name'] = video_folder
  return tmp

def generate_classes(class_index):
    idx = 0
    classes = {'Background': 0}
    with open(class_index) as f:
      for line in f:
        idx = idx + 1
        line = line.strip()
        classes[line] = idx
    return classes

def generate_classes_from_json(activity_index):
    with open(activity_index, 'r') as f:
        activities = json.load(f)
    idx = 0
    classes = {'Background': 0}
    for activity_label in activities:
        idx += 1
        classes[activity_label] = idx
    return classes

def find_leaf_directories(current_directory):
    leaf_directories = []
    is_leaf_directory = True
    for file_name in os.listdir(current_directory):
        file_path = os.path.join(current_directory, file_name)
        if os.path.isdir(file_path):
            is_leaf_directory = False
            ld = find_leaf_directories(file_path)
            leaf_directories.extend(ld)

    if is_leaf_directory:
        leaf_directories.append(current_directory)
    return leaf_directories

def get_video_path(all_video_paths, video):
    plausible_path = [video_path for video_path in all_video_paths if video in video_path]
    assert len(plausible_path) <= 1, "Ambigious paths found for {0}".format(video)
    if len(plausible_path) == 1:
        return plausible_path[0]
    else:
        return None

def generate_roidb_pkl( test_mode, kpf_mode,
                        json_mode,
                        window_length,
                        window_divisor,
                        min_length,
                        sampling_frequency,
                        use_flipped,
                        frame_roots,
                        class_index,
                        pickle_file,
                        file_index=None,
                        annotation_dirs=None,
                        chunk_file=None,
                        chunk_id=-1):
    roidb = []
    remove = 0
    overall = 0  # Total number of activities across all videos
    duration = []  # Total number of frames across all activities
    STEP = window_length//window_divisor
    WINS = [window_length*sampling_frequency]

    if test_mode:
      assert file_index is not None, "Specify the index file path for the test set"
      video_db = generate_testing_video_db(file_index, chunk_file, chunk_id)
    else:
      if json_mode:
          classes = generate_classes_from_json(class_index)
      else:
          classes  = generate_classes(class_index)
      if kpf_mode:
        video_db = generate_kpf_training_video_db(annotation_dirs, classes)
      elif json_mode:
        video_db = generate_json_training_video_db(annotation_dirs, classes)
      else:
        video_db = generate_density_training_video_db(annotation_dirs, classes)

    all_video_paths = []
    for frame_root in frame_roots:
        all_video_paths.extend(find_leaf_directories(frame_root))

    for video in video_db:
      video_folder = get_video_path(all_video_paths, video)
      if video_folder is None:
        print ("Skipping video, could not find its folder : " + video)
        sys.exit()

      num_fr = len(os.listdir(video_folder))

      if test_mode:
        # get rid of one indexed frame numbers
        for duration in video_db[video]:
            start_frame = duration[0] - 1
            end_frame = duration[1] - 1
            for win in WINS:
              stride = win / window_length
              step = stride * STEP
              for start in range(start_frame, max(1, end_frame - win + 1), step):
                end = min(start + win, end_frame)
                tmp = generate_testing_roi(video_folder, start, end, stride,
                                            use_flipped)
                roidb.append(tmp)

                if use_flipped:
                  flipped_tmp = copy.deepcopy(tmp)
                  flipped_tmp['flipped'] = True
                  roidb.append(flipped_tmp)

              for end in range(end_frame, win - 1, -step):
                start = end - win
                assert start >= 0
                # Add data
                tmp = generate_testing_roi(video_folder, start, end, stride,
                                            use_flipped)
                roidb.append(tmp)

                if use_flipped:
                  flipped_tmp = copy.deepcopy(tmp)
                  flipped_tmp['flipped'] = True
                  roidb.append(flipped_tmp)

      else: # training mode
        db = np.array(video_db[video])  # All truth classes for this video and their temporal intervals
        overall += len(db)

        if len(db) == 0: continue

        debug = []

        for win in WINS:
          stride = win // window_length
          step = stride * STEP
          # Create temporal segments going forward
          for start in range(0, max(1, num_fr - win + 1), step):
            end = min(start + win, num_fr)
            rois = db[np.logical_and(db[:, 0] >= start, db[:, 1] < end)]  # all activities within this temporal interval

            if len(rois) > 0:
              duration = rois[:, 1] - rois[:, 0]
              rois[:, 0] = np.maximum(start, rois[:, 0])
              rois[:, 1] = np.minimum(end, rois[:, 1])
              tmp = generate_training_roi(rois, video_folder, start, end,
                                            stride, use_flipped)
              roidb.append(tmp)
              duration = duration + [d for d in tmp['durations']]
              if use_flipped:
                flipped_tmp = copy.deepcopy(tmp)
                flipped_tmp['flipped'] = True
                roidb.append(flipped_tmp)

              for d in rois:
                debug.append(d)
        # Create temporal segments going backward
        for end in range(num_fr - 1, window_length - 1, -step):
          start = end - window_length
          rois = db[np.logical_and(db[:, 0] >= start, db[:, 1] < end)]  # all activities within this temporal interval


          if len(rois) > 0:
            duration = rois[:, 1] - rois[:, 0]
            rois[:, 0] = np.maximum(start, rois[:, 0])
            rois[:, 1] = np.minimum(end, rois[:, 1])
            tmp = generate_training_roi(rois, video_folder, start, end, 
                                        stride, use_flipped)
            roidb.append(tmp)

            duration = duration + [d for d in tmp['durations']]
            if use_flipped:
              flipped_tmp = copy.deepcopy(tmp)
              flipped_tmp['flipped'] = True
              roidb.append(flipped_tmp)

            for d in rois:
              debug.append(d)
        debug_res = [list(x) for x in set(tuple(x) for x in debug)]
        if len(debug_res) < len(db):
          remove += len(db) - len(debug_res)
        print("Removing: ", remove, ' / ', overall)  # Eliminate rois that are smaller than window_length
    print("Save dictionary")
    pkl.dump(roidb, open(pickle_file, 'wb'), pkl.HIGHEST_PROTOCOL)

if __name__ == "__main__":
    argparse = argparse.ArgumentParser("Generate roidb for data")
    argparse.add_argument("--test_mode", action="store_true", 
                        help="Flag to generate test roidb")
    argparse.add_argument("--kpf_mode", action="store_true", \
                            help="Flag to generate overlapping activity")
    argparse.add_argument("--json_mode", action="store_true", \
                            help="Flag to generate roidb using training data")
    argparse.add_argument("--window_length", type=int,
                        help="Number of frames used for buffering")
    argparse.add_argument("--window_divisor", type=int)
    argparse.add_argument("--min_length", type=int, 
                            help="Minimum length of window")
    argparse.add_argument("--sampling_frequency", type=int, 
                            help="Strides for the window")
    argparse.add_argument("--use_flipped", action="store_true", 
                            help="Allow flipping images")
    argparse.add_argument("--frame_roots", nargs="+",
            help="Relative path from root directory where VIRAT frames are stored")
    argparse.add_argument("--annotation_dirs", nargs="+", 
            help="Relative path from root directiory to the kpf directories")
    argparse.add_argument("--json_index_path", default=None, 
            help="Relative path from data root where file index for test set is stored (Used only for testing)")
    argparse.add_argument("--class_index", 
                            help="path to class index relative to root")
    argparse.add_argument("--pkl_file", 
                            help="path to the pkl file wrt to project root")
    argparse.set_defaults(test_mode=False)
    argparse.set_defaults(kpf_mode=False)
    argparse.set_defaults(use_flipped=False)
    argparse.set_defaults(use_json_mode=False)
    args = argparse.parse_args()
    generate_roidb_pkl(test_mode=args.test_mode,
                        kpf_mode=args.kpf_mode,
                        json_mode=args.json_mode,
                        window_length=args.window_length,
                        window_divisor=args.window_divisor,
                        min_length=args.min_length,
                        sampling_frequency=args.sampling_frequency,
                        use_flipped=args.use_flipped,
                        frame_roots = args.frame_roots,
                        annotation_dirs = args.annotation_dirs,
                        class_index = args.class_index,
                        pickle_file = args.pkl_file,
                        json_index_path = args.json_index_path)
