# global dependencies
import gc
import random

import psutil
import torch as th
import pandas as pd
import numpy as np
import cv2
import os
import traceback

from time import sleep
from collections import OrderedDict

# local dependencies
from naeural_core import DecentrAIObject
from decentra_vision.draw_utils import DrawUtils
from naeural_core.serving.serving_manager import ServingManager
from naeural_core.serving.model_testing.utils import Dataset
from naeural_core.serving.model_testing.utils import ServingUtils


class Base(DecentrAIObject):
  def __init__(self,
               model_name,
               # Previously named test_files
               test_datasets,
               sleep_time=10,
               save_plots=True,
               show_plots=False,
               label_extension='json',
               nr_warmup=2,
               nr_predicts=20,
               inprocess=False,
               serving_manager=None,
               print_errors=True,
               subdir_keys=None,
               **kwargs):
    """

    :param model_name: tuple(CLASS, NAME)
      Passed to ServingManager.start_server server_name

    :param test_datasets: dict(dataset_name: [list of paths or list of raw inference items])
      Datasets to run inference on

      !! Previously named test_files - before non-image inference support was added.

    :param sleep_time: int
      Sleep time at startup

    :param save_plots: bool
      Whether to save ploted images to disk

    :param show_plots: bool
      Whether to show plotted images

    :param nr_warmup: int
      How many warmup steps to run

    :param nr_predicts: int
      How many predicts to run in order to get mean resource consumption

    :param inprocess: bool
      Whether to run inprocess serving processes

    :param serving_manager: ServingManager object

    :param print_errors: bool
      Whether to display inference errors

    :param subdir_keys: list
      keys used to better organize the saved results (mainly plots) in subdirectories

    :param kwargs:
    """
    super().__init__(**kwargs)
    self._model_name = model_name
    self._initial_test_datasets = test_datasets
    self._sleep_time = sleep_time
    self._save_plots = save_plots
    self._show_plots = show_plots
    # self._gpu_device = gpu_device
    self._serving_manager = serving_manager
    # if self._gpu_device is None:
    #   self._gpu_device = 1 if th.cuda.device_count() > 1 else 0
    self._nr_predicts = nr_predicts
    self._nr_warmup = nr_warmup
    self._inprocess = inprocess
    self._print_errors = print_errors
    self.subdir_keys = [] if subdir_keys is None else subdir_keys
    self.current_dct_params = {}

    if self._inprocess is True:
      self.P("WARNING !!! 'inprocess' is set to True which may yield invalid values because of the inability to do a complete cleanup", color='r')
    #
    # self._lst_imgs = None
    # self._lst_img_names = None

    self._dct_res = {
      'MODEL': [],
      'TIME': [],
      'TOTAL_MEM': [],
      'INITIALIZATION_MEM': [],
      'TIME_PER_INPUT': [],
      'NUMBER_OF_INPUTS': [],
      'SUCCESS': []
    }
    self._test_datasets = {}
    self._test_datasets_paths = {}
    self._dataset_types = {}
    self.label_extension = label_extension
    self._preds = OrderedDict()
    self._init()
    return

  def set_loop_stage(*args, **kwargs):
    return

  def _init(self):
    """Starting serving manager, starting utils. Loading data."""

    if self._serving_manager is None:
      self._serving_manager = ServingManager(
        shmem={},
        log=self.log,
        prefix_log='[TFMGR]',
        owner=self,
      )
    self._painter = DrawUtils(log=self.log)
    self._dataset = Dataset(log=self.log)
    self._serving_utils = ServingUtils(log=self.log)
    self._load_data()
    return

  def _load_data(self):
    """Loading testing dataset from indicated list of files or inference items"""

    self.log.P("Preparing data...", color='r')
    self.load_data()
    self.log.P("Done preparing data.")
    return

  def load_data(self):
    for dataset_name, dataset_paths in self._initial_test_datasets.items():
      if not isinstance(dataset_paths, list):
        dataset_paths = [dataset_paths]
      loaded_data, data_type = self._dataset.load(dataset_paths)
      self._test_datasets[dataset_name] = loaded_data
      self._dataset_types[dataset_name] = data_type
      if all(isinstance(dataset_path, str) for dataset_path in dataset_paths):
        self._test_datasets_paths[dataset_name] = [
          f'{os.path.splitext(dataset_path)[0]}.{self.label_extension}'
          for dataset_path in dataset_paths
        ]
      else:
        self._test_datasets_paths[dataset_name] = []
    # self._lst_imgs = list(dct.values())
    # self._lst_img_names = list(dct.keys())
    return

  def get_device_free_memory(
      self, device, device_id: int = None,
      show_logs: bool = True
  ):
    if device == 'cpu':
      res = psutil.virtual_memory().available
    else:
      gpus_info = self.log.gpu_info(show=show_logs, mb=True)
      if device_id is None:
        device_id = [
          i for i, gpu_info in enumerate(gpus_info)
          if gpus_info['NAME'] == th.cuda.get_device_name(device)
        ][0]
      # endif device_id not provided
      res = gpus_info[device_id]['FREE_MEM']
    return res

  def run(self, dct_test={}, dct_params={}, _lst_images=None, kill_serving_manager=True):
    """
    Execute inference testing
    Video file testing params:
      _lst_images = force the use of the given _lst_images
      kill_serving_manager = set to false in video file testing in order to not create a new manager for every movie batch

    Parameters
    ----------
    dct_test : dict
      Test specific parameters
    dct_params : dict
      Common parameters for all tests
    _lst_images : list
      List of images to use for testing
    kill_serving_manager : bool
      Whether to kill the serving manager after test

    Returns
    -------
    preds : list
      List of predictions
    """
    dct_aggregated = {
      **dct_params,
      **dct_test,
    }
    device_id = None
    preds = None
    default_device = dct_aggregated.get('DEFAULT_DEVICE', None)
    dataset_name = dct_aggregated.get('dataset_name')
    dataset_input_type = self.get_dataset_input_type(dataset_name) if dataset_name is not None else None
    input_type = dct_aggregated.get("INPUT_TYPE") or dataset_input_type or "IMG"

    if "MAX_BATCH_FIRST_STAGE" not in dct_aggregated:
      self.P("WARNING! Parameter 'MAX_BATCH_FIRST_STAGE' not set. Using default value 5", color='r')

    batch_first_stage = dct_aggregated.get("MAX_BATCH_FIRST_STAGE", 5)
    batch_second_stage = dct_aggregated.get("MAX_BATCH_SECOND_STAGE", 1)

    batch_size = max(batch_first_stage, batch_second_stage)

    run_dct = {
      'DEFAULT_DEVICE': None,
      'MODEL': None,
      'TIME': None,
      'TOTAL_MEM': None,
      'INITIALIZATION_MEM': None,
      'TIME_PER_INPUT': None,
      'NUMBER_OF_INPUTS': None,
      'SUCCESS': None,
      'PREDICTS': self._nr_predicts,
      'WARMUPS': self._nr_warmup,
      'INPUT_TYPE': input_type,
      **dct_aggregated
    }
    try:
      gpus_info = self.log.gpu_info(show=True, mb=True)
      if len(gpus_info) == 0:
        self.P(f"WARNING! No GPUs found. Running on CPU only.", color='r')
        default_device = "cpu"
      else:
        if default_device is not None:
          if default_device in ['gpu', 'cuda']:
            default_device = 'cuda:0'
          device_id = [
            i for i, gpu_info in enumerate(gpus_info)
            if gpus_info['NAME'] == th.cuda.get_device_name(default_device)
          ][0]
        else:
          default_device = "cuda:0"
          device_id = 0
        # endif default_device
      # endif no gpus
      run_dct['DEFAULT_DEVICE'] = default_device

      self.log.P("Running test for {}".format(dct_test), color='g')
      mem1 = self.get_device_free_memory(
        device=default_device,
        device_id=device_id
      )
      self.log.P(f"Free memory before start: {mem1}")

      str_model_name = self._serving_manager.start_server(
        server_name=self._model_name,
        inprocess=self._inprocess,
        upstream_config=dct_aggregated,
      )
      run_dct['MODEL'] = str_model_name

      mem15 = self.get_device_free_memory(
        device=default_device,
        device_id=device_id
      )
      initialization_mem = mem1 - mem15
      run_dct['INITIALIZATION_MEM'] = initialization_mem

      self.log.P("Model is ready. Check nvidia-smi in next {} seconds...".format(self._sleep_time))
      sleep(self._sleep_time)

      self.log.P("Running warm-up predict...")
      if _lst_images is None:
        _lst_images = self.get_inputs(dataset_name)

      no_imgs = len(_lst_images)
      _lst_images = (_lst_images * np.ceil(batch_size / no_imgs).astype(int))[:max(batch_size, no_imgs)]
      run_dct['NUMBER_OF_INPUTS'] = len(_lst_images)

      yolo_model_inputs = self._serving_utils.get_model_input(
        _lst_images,
        input_type=input_type
      )
      for _ in range(self._nr_warmup):
        preds = self._serving_manager.predict(self._model_name, yolo_model_inputs)
      # endfor

      if preds is None and self._nr_warmup > 0:
        raise Exception("Error raised in predict")

      self.log.reset_timers()

      self.log.P("Running batch predict...")
      yolo_model_inputs = self._serving_utils.get_model_input(
        _lst_images,
        input_type=input_type
      )
      preds = self._serving_manager.predict(self._model_name, yolo_model_inputs)

      if preds is None:
        raise Exception("Error raised in predict")

      for _ in range(self._nr_predicts - 1):
        random.shuffle(_lst_images)
        yolo_model_inputs = self._serving_utils.get_model_input(
          _lst_images,
          input_type=input_type
        )
        self._serving_manager.predict(self._model_name, yolo_model_inputs)
      # endfor

      self._preds[len(self._preds)] = preds

      mem2 = self.get_device_free_memory(
        device=default_device,
        device_id=device_id
      )
      self.log.P(f"Free memory after run: {mem2}")
      total_model_mem = mem1 - mem2
      run_dct['TOTAL_MEM'] = total_model_mem
      self.log.P("Model memory {:.0f} MB".format(total_model_mem), color='m')

      self.log.P("Showing results & timers...", color='g')
      self.log.show_timers(title="on {} for '{}' with {}".format(
          self.log.get_machine_name(),
          self._model_name, dct_test
        ),
        div=max(len(_lst_images), 1)
      )

      if dataset_name is not None:
        plot_kwargs = {**dct_test, **dct_params, "INPUT_TYPE": input_type}
        if input_type == "IMG":
          self.plot(**plot_kwargs)
        elif self._show_plots or self._save_plots:
          self.log.P(f"Skipping plot for dataset '{dataset_name}' with non-image input type '{input_type}'")

        res = self.score(**plot_kwargs)
        if res is not None:
          run_dct['SCORE'] = res
        # endif res None
      # endif dataset available

      if not self._show_plots and not self._save_plots:
        self.log.P("Skipped showing images & inferences.")
      # endif

      mem3 = self.get_device_free_memory(
        device=default_device,
        device_id=device_id
      )
      self.log.P(f"Free memory after shutdown: {mem3}")

      if kill_serving_manager:
        self._serving_manager.stop_server(server_name=self._model_name)

      timer_prefix = 'local_pred_' if self._inprocess else 'remote_pred_'
      run_dct['TIME'] = self.log.get_timer(timer_prefix + str_model_name.upper())['MEAN']
      run_dct['TIME_PER_INPUT'] = self.log.get_timer(timer_prefix + str_model_name.upper())['MEAN'] / len(_lst_images)
      run_dct['SUCCESS'] = True
    except:
      run_dct['SUCCESS'] = False
      if self._print_errors:

        info = traceback.format_exc()
        self.log.P("Inference error: {}".format(info), color='r')
    # endtry

    for k, v in run_dct.items():
      if k not in self._dct_res:
        self._dct_res[k] = []
      self._dct_res[k].append(v)

    gc.collect()
    th.cuda.empty_cache()
    return preds

  def get_results(self):
    df = pd.DataFrame(self._dct_res)
    return df

  def show_results(self):
    df = self.get_results()
    self.log.set_nice_prints()
    self.log.P("Results on {}:\n{}".format(self.log.get_machine_name(), df))
    return

  @property
  def _testing_subfolder_path(self):
    return self._testing_subfolder_path_helper()

  def _testing_subfolder_path_helper(self, skip_subdirs=False):
    res = os.path.join('testing', self.log.file_prefix, self.__name__, self._model_name)
    if not skip_subdirs and len(self.subdir_keys) > 0:
      values = [self.current_dct_params.get(key) for key in self.subdir_keys if key in self.current_dct_params]
      str_values = [str(value) for value in values]
      res = os.path.join(res, *str_values)
    # endif subdir_keys
    return res

  def save_results(self):
    df = self.get_results()
    df_name = 'results.csv'
    self.log.save_dataframe(df, df_name, folder='output',
                            subfolder_path=self._testing_subfolder_path_helper(skip_subdirs=True))

  def run_tests(self, lst_tests, dct_params, save_results=True):
    """
    Run a series of tests over all datasets
    Parameters
    ----------
    lst_tests : list of dicts
      Each dict contains a test configuration
    dct_params : dict
      Common parameters for all tests
    save_results : bool
      Whether to save results to disk

    Returns
    -------
    pd.DataFrame
      Dataframe with results
    """
    # TODO: big ambiguity between lists of test dicts and dct_params
    for dataset_name in self.get_dataset_names():
      for dct_test in lst_tests:
        test_params = {'dataset_name': dataset_name, **dct_test}
        self.current_dct_params = test_params
        self.run(test_params, dct_params)
    # endfor
    self.show_results()
    if save_results:
      self.save_results()
    return self.get_results()

  def plot(self, dataset_name, **kwargs):
    """Draw plots specific to each kind of model"""

    raise NotImplementedError()
    return

  def score(self, dataset_name, **kwargs):
    """Score computing for specific model"""

    return None

  def plot_video(self, cv2_video_writer, lst_imgs, preds, **kwargs):
    """Based on a lst of imgs and preds save a video"""
    raise NotImplementedError()
    return

  def get_inputs(self, dataset_name):
    return self._test_datasets[dataset_name]

  def get_images(self, dataset_name):
    return self._test_datasets[dataset_name]

  def get_dataset_input_type(self, dataset_name):
    return self._dataset_types.get(dataset_name, "IMG")

  def get_dataset_names(self):
    return list(self._test_datasets.keys())

  def get_dataset_output_folder(self, dataset_name):
    return os.path.join(self._testing_subfolder_path, dataset_name)

  def get_dataset_labels_path(self, dataset_name):
    return self._test_datasets_paths[dataset_name]

  def get_preds(self, idx=None):
    preds = self._preds
    if idx is not None:
      preds = self._preds.get(idx, None)
    return preds

  def get_last_preds(self):
    dct_preds = self.get_preds()
    keys = list(dct_preds.keys())
    preds = dct_preds[keys[-1]]
    return preds

  def get_model_name(self):
    return self._model_name

  def run_video_test(self, video_fn, test_params={}, flip_video=False, movie_batch_size=50, plot_video_kwargs={}):
    """

    :param video_fn:
    :param flip_video:
      False - nothing happens
      Rotate code - passed to cv2.rotate
    :param movie_batch_size:
    :return:
    """
    cap = cv2.VideoCapture(video_fn)
    while not cap.isOpened():
      cap = cv2.VideoCapture(video_fn)

    frames = []
    fps = cap.get(cv2.CAP_PROP_FPS)
    pos_frame = cap.get(cv2.CAP_PROP_POS_FRAMES)
    while True:
      flag, frame = cap.read()
      if flip_video is not False:
        frame = cv2.rotate(frame, flip_video)
      frame_rgb = frame[:, :, ::-1]
      frame_rgb = np.ascontiguousarray(frame_rgb)
      frames.append(frame_rgb)

      if cap.get(cv2.CAP_PROP_POS_FRAMES) == cap.get(cv2.CAP_PROP_FRAME_COUNT):
          # If the number of captured frames is equal to the total number of frames,
          # we stop
        break

    output_video_file = '{}_{}'.format(self.log.file_prefix, os.path.split(video_fn)[-1])

    output_video_folder = os.path.join(
      self.log.get_output_folder(),
      'video_testing',
    )
    if not os.path.exists(output_video_folder):
      os.makedirs(output_video_folder)

    output_video_fn = os.path.join(output_video_folder, output_video_file)

    cv2_video_writer = cv2.VideoWriter(
      output_video_fn,
      cv2.VideoWriter_fourcc(*"XVID"),
      fps,
      (frames[0].shape[1], frames[0].shape[0])
    )

    for batch_id in range(len(frames) // movie_batch_size):
      self.P("Processing movie batch {}/{}".format(batch_id + 1, len(frames) // movie_batch_size))
      batch = frames[batch_id * movie_batch_size: (batch_id + 1) * movie_batch_size]
      batch_results = self.run(dct_test={**test_params, "MAX_BATCH_FIRST_STAGE": 5},
                               _lst_images=batch, kill_serving_manager=False)
      self.plot_video(
        cv2_video_writer=cv2_video_writer,
        lst_imgs=batch,
        preds=batch_results,
        **plot_video_kwargs
      )
    cv2_video_writer.release()

    a = 0


if __name__ == '__main__':
  from naeural_core import Logger
  from plugins.serving.testing.th_yolov5l6s.data import files as TEST_FILES

  log = Logger('MPTF', base_folder='.', app_folder='_solisbox', TF_KERAS=False)

  # defining script constants
  GPU_DEVICE = 1 if th.cuda.device_count() > 1 else 0
  TESTS = [
    # {
    #   'USE_AMP'         : False,
    #   'USE_FP16'        : False,
    #   'GPU_PREPROCESS'  : False,
    # },

    {
      'USE_AMP': True,
      'USE_FP16': False,
      'GPU_PREPROCESS': True,
      'DEFAULT_DEVICE': 'cuda:' + str(GPU_DEVICE)
    },

    # {
    #   'USE_AMP'         : True,
    #   'USE_FP16'        : False,
    #   'GPU_PREPROCESS'  : True,
    # },

    # {
    #   'USE_AMP'         : False,
    #   'USE_FP16'        : True,
    #   'GPU_PREPROCESS'  : True,
    # },
  ]
  COVERED = 'eff_det0'
  MODEL_NAME = ('th_y5l6s', 'default')

  test = Base(
    log=log,
    model_name=MODEL_NAME,
    test_datasets=TEST_FILES,
    gpu_device=GPU_DEVICE,
  )

  dct_params = {
    "NEW_PARAM": 0,
    "COVERED_SERVERS": ['th_y5l6', COVERED],
    "NMS_CONF_THR": 0.05,
    "WARMUP_BATCH": len(test.get_images())
  }
  test.run_tests(
    lst_tests=TESTS,
    dct_params=dct_params
  )
