#
# Copyright (c) 2018 cTuning foundation.
# See CK COPYRIGHT.txt for copyright details.
#
# SPDX-License-Identifier: BSD-3-Clause.
# See CK LICENSE.txt for licensing details.
#

import os
import json

TOP1 = 0
TOP5 = 0
IMAGES_COUNT = 0 # to be assigned
SEPARATOR = '---------------------------------------'

def ck_postprocess(i):
  def my_env(var): return i['env'].get(var)
  def dep_env(dep, var): return i['deps'][dep]['dict']['env'].get(var)

  # Init variables from environment
  BATCH_COUNT = int(my_env('CK_BATCH_COUNT'))
  BATCH_SIZE = int(my_env('CK_BATCH_SIZE'))
  global IMAGES_COUNT
  IMAGES_COUNT = BATCH_COUNT * BATCH_SIZE
  SKIP_IMAGES = int(my_env('CK_SKIP_IMAGES'))
  COUNT_SKIPPED_IMAGES = (my_env('CK_CONTINUE') or 'no').lower() in ('true', 'yes', 'on', '1')
  RESULTS_DIR = 'predictions'
  NUM_CLASSES = 1000
  PREPROCESSED_EXT = dep_env('images', 'CK_ENV_DATASET_IMAGENET_PREPROCESSED_NEW_EXTENSION') or 'JPEG'
  AUX_DIR = dep_env('imagenet-aux', 'CK_ENV_DATASET_IMAGENET_AUX')
  CLASSES_FILE = os.path.join(AUX_DIR, 'synset_words.txt')
  VALUES_FILE = os.path.join(AUX_DIR, 'val.txt')
  CLASSES_LIST = []
  VALUES_MAP = {}
  IMAGE_FILE = my_env('CK_IMAGE_FILE')
  FULL_REPORT = my_env('CK_SILENT_MODE') not in ('TRUE', 'True', 'true', 'YES', 'Yes', 'yes', 'ON', 'On', 'on', '1', 1)
  NO_LOADGEN = my_env('CK_LOADGEN_SUPPORT') not in ('TRUE', 'True', 'true', 'YES', 'Yes', 'yes', 'ON', 'On', 'on', '1', 1)

  # Loads ImageNet classes and correct predictions
  def load_ImageNet_classes():
    global IMAGES_COUNT
    classes_list = []
    with open(CLASSES_FILE, 'r') as classes_file:
      classes_list = classes_file.read().splitlines()

    values_map = {}
    with open(VALUES_FILE, 'r') as values_file:
      if IMAGE_FILE:
        # Single file mode: try to find this file in values
        for line in values_file:
          file_name, file_class = line.split()
          if file_name == IMAGE_FILE:
            values_map[file_name] = int(file_class)
            break
      else:
        # Directory mode: load only required number of values
        for _ in range(SKIP_IMAGES):
          val = values_file.readline().split()
          if COUNT_SKIPPED_IMAGES:
              values_map[val[0]] = int(val[1])
        for _ in range(IMAGES_COUNT):
          val = values_file.readline().split()
          values_map[val[0]] = int(val[1])
        if COUNT_SKIPPED_IMAGES:
            IMAGES_COUNT += SKIP_IMAGES

    return classes_list, values_map


  # Returns printable string for ImageNet specific class
  def get_class_str(class_index):
    obj_class = CLASSES_LIST[class_index]
    if len(obj_class) > 50:
        obj_class = obj_class[:50] + '...'
    return '(%d) %s' % (class_index, obj_class)


  # Shows prediction results for image file
  # top5 - list of pairs (prob, class_index)
  def print_predictions(top5, img_file):
    print(SEPARATOR)
    if img_file in VALUES_MAP:
      class_correct = VALUES_MAP[img_file]
      print('%s - %s' % (img_file, get_class_str(class_correct)))
    else:
      print(img_file)
    for prob, class_index in top5:
      print('%.2f - %s' % (prob, get_class_str(class_index)))


  # Returns list of pairs (prob, class_index)
  def get_topn(all_probs, n=5):
    probs_with_classes = []
    for class_index in range(len(all_probs)):
      prob = all_probs[class_index]
      probs_with_classes.append((prob, class_index))
    sorted_probs = sorted(probs_with_classes, key = lambda pair: pair[0], reverse=True)
    return sorted_probs[0:n]


  # Calculates if prediction was correct for specified image file
  # topn - list of pairs (prob, class_index)
  def check_predictions(img_file, topn):
    if img_file not in VALUES_MAP:
      print('Correctness information is not available')
      return {}

    class_correct = VALUES_MAP[img_file]
    classes = [c[1] for c in topn]
    is_top1 = class_correct == classes[0]
    is_top5 = class_correct in classes[0:5]
    is_topn = class_correct in classes
    if is_top1:
      global TOP1
      TOP1 += 1
    if is_top5:
      global TOP5
      TOP5 += 1
    if is_topn:
      global TOPN
      TOPN += 1
    return {
      'accuracy_top1': 'yes' if is_top1 else 'no',
      'accuracy_top5': 'yes' if is_top5 else 'no',
      'accuracy_topn': 'yes' if is_topn else 'no',
      'class_correct': class_correct,
      'class_topmost': classes[0],
      'file_name': img_file
    }

  frame_predictions = {}

  def calculate_precision():
    print('\nProcessed results from {} :'.format(RESULTS_DIR))

    def load_probes(filename):
      probes = []
      with open(os.path.join(RESULTS_DIR, filename), 'r') as f:
        for line in f:
          s = line.strip()
          if s: probes.append(float(s))
      return probes

    checked_files = 0

    for res_file in sorted(os.listdir(RESULTS_DIR)):
      # remove trailing suffix '.txt'
      img_file = res_file[:-4]
      # replace preprocesed extension '.rgb8' or '.rgbf32' with original '.JPEG'
      img_file = img_file.replace(PREPROCESSED_EXT, 'JPEG')
      checked_files += 1

      all_probes = load_probes(res_file)
      if len(all_probes) != NUM_CLASSES:
        print('WARNING: {} is invalid probes count in file {}, results ignored'.format(len(all_probes), res_file))
        global IMAGES_COUNT
        IMAGES_COUNT -= 1
        continue

      topn = get_topn(all_probes, n=5)
      if FULL_REPORT:
        print_predictions(topn, img_file)
      elif checked_files % 100 == 0:
        print('Predictions checked: {}'.format(checked_files))

      res = check_predictions(img_file, topn)

      # Augment the result with actual TopN class probabilities and names.
      res['topn'] = []
      for prob_index in topn:
        res['topn'].append({
          'index': prob_index[1],
          'prob': prob_index[0],
          'name': get_class_str(prob_index[1])
        })

      frame_predictions[img_file] = res


  global TOP1
  global TOP5
  global TOPN
  TOP1 = 0
  TOP5 = 0
  TOPN = 0
  CLASSES_LIST, VALUES_MAP = load_ImageNet_classes()
  calculate_precision()

  accuracy_top1 = TOP1 / float(IMAGES_COUNT) if IMAGES_COUNT > 0 else 0
  accuracy_top5 = TOP5 / float(IMAGES_COUNT) if IMAGES_COUNT > 0 else 0
  accuracy_topn = TOPN / float(IMAGES_COUNT) if IMAGES_COUNT > 0 else 0

  # Store benchmark results
  openme = {}

  # Print metrics
  print(SEPARATOR)
  print('\nSummary:')
  print(SEPARATOR)

  if NO_LOADGEN:
    # Preserve values stored by program
    with open('tmp-ck-timer.json', 'r') as o:
      old_values = json.load(o)
    for key in old_values:
      # xopenmp c++ writes this section, copy it into root object
      if key == 'run_time_state':
        for key1 in old_values[key]:
          openme[key1] = old_values[key][key1]
      else:
        openme[key] = old_values[key]

    setup_time = openme.get('setup_time_s', 0.0)
    test_time = openme.get('test_time_s', 0.0)
    total_load_images_time = openme.get('images_load_time_total_s', 0.0)
    total_prediction_time = openme.get('prediction_time_total_s', 0.0)
    avg_prediction_time = openme.get('prediction_time_avg_s', 0.0)

    openme['execution_time'] = total_prediction_time
    openme['execution_time_sum'] = setup_time + test_time

    print('Graph loaded in {:.6f}s'.format(setup_time))
    print('All images loaded in {:.6f}s'.format(total_load_images_time))
    print('All images classified in {:.6f}s'.format(total_prediction_time))
    print('Average classification time: {:.6f}s'.format(avg_prediction_time))


  print('Accuracy top 1: {} ({} of {})'.format(accuracy_top1, TOP1, IMAGES_COUNT))
  openme['accuracy_top1'] = accuracy_top1
  print('Accuracy top 5: {} ({} of {})'.format(accuracy_top5, TOP5, IMAGES_COUNT))
  openme['accuracy_top5'] = accuracy_top5
  # TODO: Read from env
  N=5
  if N!=5:
    print('Accuracy top {}: {} ({} of {})'.format(N, accuracy_topn, TOPN, IMAGES_COUNT))
    openme['accuracy_topn'] = accuracy_topn
  openme['frame_predictions'] = frame_predictions

  with open('tmp-ck-timer.json', 'w') as o:
    json.dump(openme, o, indent=2, sort_keys=True)

  print(SEPARATOR+"\n")
  return {'return': 0}