#!/usr/bin/env python
# -*- coding: utf-8 -*-
import json
import os

import tensorflow as tf
import numpy as np

from collections import Counter
from multiprocessing import Queue, Process, Value, get_logger
from time import sleep
from ctypes import c_bool

import backend

from hailo_platform.common.jlf import readjlf
from hailo_platform.drivers.hw_object import HailoUdpControllerObject
from hailo_platform.drivers.hailo_controller.hailo_controller import HailoBroadcastController

from hailo_sdk_client import ClientRunner

# from hailo_ml_utils.ssd_utils.ssd_postproc import add_ssd_postproc

from demos_backbone.demo_wrapper import DemoWrapper
from demos_ui.control import UIControl

IDLE_POWER_FACTOR = 0.04

def _get_jlfs_paths(model):
    my_dir = os.path.dirname(__file__)

    if not isinstance(model, str):
        raise TypeError('model must be a string not {}'.format(model.__class__.__name__))

    return os.path.join('models_files', model, 'results', 'jlfs', '*.jlf')

    if model.lower() in ['resnet50', 'resnet_v1_50']:
        jlfs_paths = os.path.join(my_dir, 'models_files', 'resnet_v1_50', 'jlfs', '*')
    elif model.lower() in ['mobilenet', 'mobilenet_v1']:
        jlfs_paths = os.path.join(my_dir, 'models_files', 'mobilenet_v1', 'jlfs', '*')
        jlfs_paths = os.path.join(my_dir, 'trial', '*')
    elif model.lower() in ['mobilenet_ssd_nms', 'mobilenet_ssd', 'ssd']:
        jlfs_paths = os.path.join(my_dir, 'models_files', 'mobilenet_ssd_nms', 'jlfs', '*')
        #  jlfs_paths = os.path.join(my_dir, 'models_files', 'mobilenet_ssd_nms.new', 'jlfs', '*')
        #  jlfs_paths = os.path.join(my_dir, 'models_files', 'mobilenet_ssd_nms.old', 'jlfs', '*')
        #  jlfs_paths = os.path.join(my_dir, 'models_files', 'mobilenet_ssd_nms_ibc_eq_8bit_bias_16img_90_classes', 'jlfs', '*')
        #  jlfs_paths = os.path.join(my_dir, 'auto', '*')
    else:
        raise ValueError(f'Unsupported model: {model}')

    return jlfs_paths

def get_jlfs(model):
    from glob import glob
    jlfs_paths = _get_jlfs_paths(model)
    jlfs = []

    for jlfp in glob(jlfs_paths):
        with open(jlfp, 'rb') as jlf_file:
            jlfs.append(jlf_file.read())
    print(jlfs_paths)
    return jlfs


def batch_nms_post_processing(net_output,
                              score_threshold=0.3,
                              #max_proposals=100,
                             ):
    '''
    process batched nms output (post-infer).
    
    Args:
    net_output(np.array) - the output of the post infer. should be of size(batch, classes, 5, max_proposals)
    score_threshold(float) - the score threshold by which to select boxes.
    
    Returns:
    boxes(list) - a list of length batches. each element in the list is a np array of shape (detections,4)
    scores(list) - a list of length batches. each element in the list is a np array of shape (detections,)
    classes(list) - a list of length batches. each element in the list is a np array of shape (detections,)
    num_detection(np.array) - each element gives the #detections in each batch
    '''
    indices = np.argwhere(net_output[:, :, 4, :] > score_threshold)
    boxes = net_output[indices[:,0], indices[:,1] ,:4, indices[:,2]]
    scores = net_output[indices[:,0], indices[:,1] ,4, indices[:,2]]
    classes = indices[:,1] + 1
    batch = indices[:,0]
    d = Counter(batch)
    batch_size = net_output.shape[0]
    idx = np.zeros(shape=(batch_size-1),dtype=np.uint16)
    num_detections = np.zeros(shape=batch_size,dtype=np.uint16)
    last_ind = 0
    for i in range(batch_size-1):
        if i in d.keys():
            last_ind += d[i]
            num_detections[i] = d[i]
        else:
            pass
        idx[i] = last_ind
    num_detections[batch_size-1] = 0 if not batch_size-1 in d else d[batch_size-1]
    classes = np.split(classes, idx)
    boxes = np.split(boxes, idx)
    scores = np.split(scores, idx)
    return num_detections, boxes, scores, classes