import multiprocessing as mp
import sys
import time
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, wait, ALL_COMPLETED
from multiprocessing import Queue

import cv2
import tensorrt as trt
sys.path.extend(['/home/th/tuo_heng/dev/tuoheng_alg', '/home/th/tuo_heng/dev/tuoheng_alg/util'])
from util.PlotsUtils import get_label_arrays
from util.TorchUtils import select_device

sys.path.extend(['/home/th/tuo_heng/', '/home/th/tuo_heng/dev', '/home/th/tuo_heng/dev/AIlib2', '/home/th/tuo_heng/dev/AIlib2/segutils'])
from segutils.segmodel import SegModel
from models.experimental import attempt_load
from AI import AI_process
from utilsK.queRiver import riverDetSegMixProcess

COLOR = (
    [0, 0, 255],
    [255, 0, 0],
    [211, 0, 148],
    [0, 127, 0],
    [0, 69, 255],
    [0, 255, 0],
    [255, 0, 255],
    [0, 0, 127],
    [127, 0, 255],
    [255, 129, 0],
    [139, 139, 0],
    [255, 255, 0],
    [127, 255, 0],
    [0, 127, 255],
    [0, 255, 127],
    [255, 127, 255],
    [8, 101, 139],
    [171, 130, 255],
    [139, 112, 74],
    [205, 205, 180])
par = {
    'device': '0',
    'labelnames': ["排口", "水生植被", "其它", "漂浮物", "污染排口", "菜地", "违建", "岸坡垃圾"],
    'seg_nclass': 2,
    'trtFlag_seg': True,
    'trtFlag_det': True,
    'segRegionCnt': 1,
    'segPar': {
        'modelSize': (640, 360),
        'mean': (0.485, 0.456, 0.406),
        'std': (0.229, 0.224, 0.225),
        'numpy': False,
        'RGB_convert_first': True,
        'mixFunction': {
            'function': riverDetSegMixProcess,
            'pars': {
                'slopeIndex': [5, 6, 7],
                'riverIou': 0.1
            }
        }
    },
    'postFile': {
        "name": "post_process",
        "conf_thres": 0.25,
        "iou_thres": 0.45,
        "classes": 5,
        "rainbows": COLOR
    },
    'Detweights': "/home/th/tuo_heng/dev/AIlib2/weights/river/yolov5_2080Ti_fp16.engine",
    'Segweights': '/home/th/tuo_heng/dev/AIlib2/weights/river/stdc_360X640_2080Ti_fp16.engine'
}
mode, postPar, segPar = par.get('mode', 'others'), par.get('postPar'), par.get('segPar')
new_device = select_device(par.get('device'))
names = par['labelnames']
half = new_device.type != 'cpu'
Detweights = par['Detweights']
with open(Detweights, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.ERROR)) as runtime:
    model = runtime.deserialize_cuda_engine(f.read())
Segweights = par['Segweights']
if Segweights:
    with open(Segweights, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.ERROR)) as runtime:
        segmodel = runtime.deserialize_cuda_engine(f.read())
else:
    segmodel = None
postFile = par['postFile']
rainbows = postFile["rainbows"]
objectPar = {
    'half': half,
    'device': new_device,
    'conf_thres': postFile["conf_thres"],
    'ovlap_thres_crossCategory': postFile.get("ovlap_thres_crossCategory"),
    'iou_thres': postFile["iou_thres"],
    'allowedList': [],
    'segRegionCnt': par['segRegionCnt'],
    'trtFlag_det': par['trtFlag_det'],
    'trtFlag_seg': par['trtFlag_seg']
}
Detweights = "/home/th/tuo_heng/dev/AIlib2/weights/river2/yolov5_2080Ti_fp16.engine"
with open(Detweights, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.ERROR)) as runtime:
    model = runtime.deserialize_cuda_engine(f.read())
Segweights = "/home/th/tuo_heng/dev/AIlib2/weights/river2/stdc_360X640_2080Ti_fp16.engine"
with open(Segweights, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.ERROR)) as runtime:
    segmodel = runtime.deserialize_cuda_engine(f.read())


allowedList = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]



def one_label(width, height, model_param):
    names = model_param[4]
    rainbows = model_param[6]
    digitFont, label_arraylist, font_config = get_label_arraylist(width, height, names, rainbows)
    """
    font_config, frame, model, segmodel, names, label_arraylist, rainbows, objectPar, font, segPar, mode, postPar, 
    requestId
    """
    model_param[5] = label_arraylist
    model_param[8] = digitFont
    model_param[0] = font_config


def get_label_arraylist(*args):
    width, height, names, rainbows = args
    # line = int(round(0.002 * (height + width) / 2) + 1)
    line = int(width / 1920 * 3 - 1)
    label = ' 0.95'
    tf = max(line, 1)
    fontScale = line * 0.33
    text_width, text_height = cv2.getTextSize(label, 0, fontScale=fontScale, thickness=tf)[0]
    fontsize = int(width / 1920 * 40)
    numFontSize = float(format(width / 1920 * 1.1, '.1f'))
    digitFont = {'line_thickness': line,
                 'boxLine_thickness': line,
                 'fontSize': numFontSize,
                 'waterLineColor': (0, 255, 255),
                 'segLineShow': False,
                 'waterLineWidth': line}
    label_arraylist = get_label_arrays(names, rainbows, text_height, fontSize=fontsize,
                                       fontPath="/home/th/tuo_heng/dev/AIlib2/conf/platech.ttf")
    return digitFont, label_arraylist, (line, text_width, text_height, fontScale, tf)
image = cv2.imread("/home/th/tuo_heng/dev/ompv2fn94m_1687259193110.jpg")
start_time1 = time.time()
with ThreadPoolExecutor(max_workers=3) as t:
    rs = []
    for i in range(500):
        rr = t.submit(AI_process, [image], model, segmodel, names, None, rainbows, objectPar=objectPar,
                      font={'line_thickness': 1,
                            'boxLine_thickness': 1,
                            'fontSize': 1.1,
                            'waterLineColor': (0, 255, 255),
                            'segLineShow': False,
                            'waterLineWidth': 1}, segPar=segPar, mode=mode, postPar=postPar)
        rs.append(rr)
    for i in rs:
        i.result()
print(time.time() - start_time1)

start_time = time.time()
for i in range(500):
    AI_process([image], model, segmodel, names, None, rainbows, objectPar=objectPar,
               font={'line_thickness': 1,
                     'boxLine_thickness': 1,
                     'fontSize': 1.1,
                     'waterLineColor': (0, 255, 255),
                     'segLineShow': False,
                     'waterLineWidth': 1}, segPar=segPar, mode=mode, postPar=postPar)
print(time.time() - start_time)