#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.
#https://github.com/PRBonn/semantic-kitti-ap
import argparse
import os
import yaml
import sys
import numpy as np
import time
import json

from eval_np import PanopticEval

# possible splits
splits = ["train", "valid", "test"]

if __name__ == '__main__':
  parser = argparse.ArgumentParser("./evaluate_panoptic.py")
  parser.add_argument(
      '--dataset',
      '-d',
      type=str,
      required=True,
      help='Dataset dir. No Default',
  )
  parser.add_argument(
      '--predictions',
      '-p',
      type=str,
      required=None,
      help='Prediction dir. Same organization as dataset, but predictions in'
      'each sequences "prediction" directory. No Default. If no option is set'
      ' we look for the labels in the same directory as dataset')
  parser.add_argument(
      '--split',
      '-s',
      type=str,
      required=False,
      choices=["train", "valid", "test"],
      default="valid",
      help='Split to evaluate on. One of ' + str(splits) + '. Defaults to %(default)s',
  )
  parser.add_argument(
      '--data_cfg',
      '-dc',
      type=str,
      required=False,
      default="config/semantic-kitti.yaml",
      help='Dataset config file. Defaults to %(default)s',
  )
  parser.add_argument(
      '--limit',
      '-l',
      type=int,
      required=False,
      default=None,
      help='Limit to the first "--limit" points of each scan. Useful for'
      ' evaluating single scan from aggregated pointcloud.'
      ' Defaults to %(default)s',
  )
  parser.add_argument(
      '--min_inst_points',
      type=int,
      required=False,
      default=50,
      help='Lower bound for the number of points to be considered instance',
  )
  parser.add_argument(
      '--output',
      type=str,
      required=False,
      default=None,
      help='Output directory for scores.txt and detailed_results.html.',
  )

  start_time = time.time()

  FLAGS, unparsed = parser.parse_known_args()

  # fill in real predictions dir
  if FLAGS.predictions is None:
    FLAGS.predictions = FLAGS.dataset

  # print summary of what we will do
  print("*" * 80)
  print("INTERFACE:")
  print("Data: ", FLAGS.dataset)
  print("Predictions: ", FLAGS.predictions)
  print("Split: ", FLAGS.split)
  print("Config: ", FLAGS.data_cfg)
  print("Limit: ", FLAGS.limit)
  print("Min instance points: ", FLAGS.min_inst_points)
  print("Output directory", FLAGS.output)
  print("*" * 80)

  # assert split
  assert (FLAGS.split in splits)

  # open data config file
  DATA = yaml.safe_load(open(FLAGS.data_cfg, 'r'))

  # get number of interest classes, and the label mappings
  # class
  class_remap = DATA["learning_map"]
  class_inv_remap = DATA["learning_map_inv"]
  class_ignore = DATA["learning_ignore"]
  nr_classes = len(class_inv_remap)
  class_strings = DATA["labels"]

  # make lookup table for mapping
  # class
  maxkey = max(class_remap.keys())

  # +100 hack making lut bigger just in case there are unknown labels
  class_lut = np.zeros((maxkey + 100), dtype=np.int32)
  class_lut[list(class_remap.keys())] = list(class_remap.values())

  # class
  ignore_class = [cl for cl, ignored in class_ignore.items() if ignored]

  print("Ignoring classes: ", ignore_class)

  # create evaluator
  class_evaluator = PanopticEval(nr_classes, None, ignore_class, min_points=FLAGS.min_inst_points)

  # get test set
  test_sequences = DATA["split"][FLAGS.split]

  # get label paths
  label_names = []
  for sequence in test_sequences:
    sequence = '{0:02d}'.format(int(sequence))
    label_paths = os.path.join(FLAGS.dataset, "sequences", sequence, "labels")
    # populate the label names
    seq_label_names = sorted([os.path.join(label_paths, fn) for fn in os.listdir(label_paths) if fn.endswith(".label")])
    label_names.extend(seq_label_names)
  # print(label_names)

  # get predictions paths
  pred_names = []
  for sequence in test_sequences:
    sequence = '{0:02d}'.format(int(sequence))
    pred_paths = os.path.join(FLAGS.predictions, "sequences", sequence, "predictions")
    # populate the label names
    seq_pred_names = sorted([os.path.join(pred_paths, fn) for fn in os.listdir(pred_paths) if fn.endswith(".label")])
    pred_names.extend(seq_pred_names)
  # print(pred_names)


  # check that I have the same number of files
  assert (len(label_names) == len(pred_names))

  print("Evaluating sequences: ", end="", flush=True)
  # open each file, get the tensor, and make the iou comparison

  complete = len(label_names)
  count = 0
  percent = 10
  for label_file, pred_file in zip(label_names, pred_names):
    count = count + 1
    if 100 * count / complete > percent:
      print("{}% ".format(percent), end="", flush=True)
      percent = percent + 10
    # print("evaluating label ", label_file, "with", pred_file)
    # open label

    label = np.fromfile(label_file, dtype=np.uint32)

    u_label_sem_class = class_lut[label & 0xFFFF]  # remap to xentropy format
    u_label_inst = label >> 16
    if FLAGS.limit is not None:
      u_label_sem_class = u_label_sem_class[:FLAGS.limit]
      u_label_sem_cat = u_label_sem_cat[:FLAGS.limit]
      u_label_inst = u_label_inst[:FLAGS.limit]

    label = np.fromfile(pred_file, dtype=np.uint32)

    u_pred_sem_class = class_lut[label & 0xFFFF]  # remap to xentropy format
    u_pred_inst = label >> 16
    if FLAGS.limit is not None:
      u_pred_sem_class = u_pred_sem_class[:FLAGS.limit]
      u_pred_sem_cat = u_pred_sem_cat[:FLAGS.limit]
      u_pred_inst = u_pred_inst[:FLAGS.limit]

    class_evaluator.addBatch(u_pred_sem_class, u_pred_inst, u_label_sem_class, u_label_inst)

  print("100%")

  complete_time = time.time() - start_time

  # when I am done, print the evaluation
  class_PQ, class_SQ, class_RQ, class_all_PQ, class_all_SQ, class_all_RQ = class_evaluator.getPQ()
  class_IoU, class_all_IoU = class_evaluator.getSemIoU()

  # now make a nice dictionary
  output_dict = {}

  # make python variables
  class_PQ = class_PQ.item()
  class_SQ = class_SQ.item()
  class_RQ = class_RQ.item()
  class_all_PQ = class_all_PQ.flatten().tolist()
  class_all_SQ = class_all_SQ.flatten().tolist()
  class_all_RQ = class_all_RQ.flatten().tolist()
  class_IoU = class_IoU.item()
  class_all_IoU = class_all_IoU.flatten().tolist()

  # fill in with the raw values
  # output_dict["raw"] = {}
  # output_dict["raw"]["class_PQ"] = class_PQ
  # output_dict["raw"]["class_SQ"] = class_SQ
  # output_dict["raw"]["class_RQ"] = class_RQ
  # output_dict["raw"]["class_all_PQ"] = class_all_PQ
  # output_dict["raw"]["class_all_SQ"] = class_all_SQ
  # output_dict["raw"]["class_all_RQ"] = class_all_RQ
  # output_dict["raw"]["class_IoU"] = class_IoU
  # output_dict["raw"]["class_all_IoU"] = class_all_IoU

  things = ['car', 'truck', 'bicycle', 'motorcycle', 'other-vehicle', 'person', 'bicyclist', 'motorcyclist']
  stuff = [
      'road', 'sidewalk', 'parking', 'other-ground', 'building', 'vegetation', 'trunk', 'terrain', 'fence', 'pole',
      'traffic-sign'
  ]
  all_classes = things + stuff

  # class

  output_dict["all"] = {}
  output_dict["all"]["PQ"] = class_PQ
  output_dict["all"]["SQ"] = class_SQ
  output_dict["all"]["RQ"] = class_RQ
  output_dict["all"]["IoU"] = class_IoU

  classwise_tables = {}

  for idx, (pq, rq, sq, iou) in enumerate(zip(class_all_PQ, class_all_RQ, class_all_SQ, class_all_IoU)):
    class_str = class_strings[class_inv_remap[idx]]
    output_dict[class_str] = {}
    output_dict[class_str]["PQ"] = pq
    output_dict[class_str]["SQ"] = sq
    output_dict[class_str]["RQ"] = rq
    output_dict[class_str]["IoU"] = iou

  PQ_all = np.mean([float(output_dict[c]["PQ"]) for c in all_classes])
  PQ_dagger = np.mean([float(output_dict[c]["PQ"]) for c in things] + [float(output_dict[c]["IoU"]) for c in stuff])
  RQ_all = np.mean([float(output_dict[c]["RQ"]) for c in all_classes])
  SQ_all = np.mean([float(output_dict[c]["SQ"]) for c in all_classes])

  PQ_things = np.mean([float(output_dict[c]["PQ"]) for c in things])
  RQ_things = np.mean([float(output_dict[c]["RQ"]) for c in things])
  SQ_things = np.mean([float(output_dict[c]["SQ"]) for c in things])

  PQ_stuff = np.mean([float(output_dict[c]["PQ"]) for c in stuff])
  RQ_stuff = np.mean([float(output_dict[c]["RQ"]) for c in stuff])
  SQ_stuff = np.mean([float(output_dict[c]["SQ"]) for c in stuff])
  mIoU = output_dict["all"]["IoU"]

  codalab_output = {}
  codalab_output["pq_mean"] = float(PQ_all)
  codalab_output["pq_dagger"] = float(PQ_dagger)
  codalab_output["sq_mean"] = float(SQ_all)
  codalab_output["rq_mean"] = float(RQ_all)
  codalab_output["iou_mean"] = float(mIoU)
  codalab_output["pq_stuff"] = float(PQ_stuff)
  codalab_output["rq_stuff"] = float(RQ_stuff)
  codalab_output["sq_stuff"] = float(SQ_stuff)
  codalab_output["pq_things"] = float(PQ_things)
  codalab_output["rq_things"] = float(RQ_things)
  codalab_output["sq_things"] = float(SQ_things)

  print("Completed in {} s".format(complete_time))

  if FLAGS.output is not None:
    table = []
    for cl in all_classes:
      entry = output_dict[cl]
      table.append({
          "class": cl,
          "pq": "{:.3}".format(entry["PQ"]),
          "sq": "{:.3}".format(entry["SQ"]),
          "rq": "{:.3}".format(entry["RQ"]),
          "iou": "{:.3}".format(entry["IoU"])
      })

    print("Generating output files.")
    # save to yaml
    output_filename = os.path.join(FLAGS.output, 'scores.txt')
    with open(output_filename, 'w') as outfile:
      yaml.dump(codalab_output, outfile, default_flow_style=False)

    ## producing a detailed result page.
    output_filename = os.path.join(FLAGS.output, "detailed_results.html")
    with open(output_filename, "w") as html_file:
      html_file.write("""
<!doctype html>
<html lang="en" style="scroll-behavior: smooth;">
<head>
  <script src='https://cdnjs.cloudflare.com/ajax/libs/tabulator/4.4.3/js/tabulator.min.js'></script>
  <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/tabulator/4.4.3/css/bulma/tabulator_bulma.min.css">
</head>
<body>
  <div id="classwise_results"></div>

<script>
  let table_data = """ + json.dumps(table) + """


  table = new Tabulator("#classwise_results", {
    layout: "fitData",
    data: table_data,
    columns: [{title: "Class", field:"class", width:200}, 
              {title: "PQ", field:"pq", width:100, align: "center"},
              {title: "SQ", field:"sq", width:100, align: "center"},
              {title: "RQ", field:"rq", width:100, align: "center"},
              {title: "IoU", field:"iou", width:100, align: "center"}]
  });
</script>
</body>
</html>""")
