import sys
from concurrent.futures import ThreadPoolExecutor

sys.path.extend(['..','../AIlib' ])

from AI import AI_process,AI_process_forest,get_postProcess_para
import cv2,os,time
from segutils.segmodel import SegModel
from models.experimental import attempt_load
from utils.torch_utils import select_device
from utilsK.queRiver import get_labelnames,get_label_arrays
import numpy as np
import torch
import tensorrt as trt
from utilsK.masterUtils import get_needed_objectsIndex

'''
    多线程
'''

def process_v1(frame):
    #try:
  
    time00 = time.time()
    H,W,C = frame[0][0].shape
    p_result,timeOut = AI_process(frame[0],frame[1],frame[2],frame[3],frame[4],frame[5],frame[6],frame[7],frame[8], frame[9],frame[10],font=frame[11],trtFlag_det=frame[13],trtFlag_seg=frame[14],segPar=frame[15])
    time11 = time.time()
    image_array = p_result[1]
    cv2.imwrite(os.path.join('images/results/',frame[12] ) ,image_array)
    time22 = time.time()
    print('%s,%d*%d,AI-process: %.1f,image save:%.1f , %s'%(frame[12],H,W, (time11 - time00) * 1000.0, (time22-time11)*1000.0,timeOut))
    return 'success'
    #except Exception as e:
    #    return 'failed:'+str(e)        

def river_demo_v3():
    ##预先设置的参数
    device_='0'  ##选定模型，可选 cpu,'0','1'
    ###注意TRT模型生成时，就需要对应cuda device，下面的trt文件是cuda:0生成的，device只能是0
    ##以下参数目前不可改
    labelnames = "../AIlib/weights/yolov5/class8/labelnames.json" ##对应类别表
    gpuname='3090';
    max_workers=1;
    trtFlag_det=True;trtFlag_seg=True
    device = select_device(device_)
    names=get_labelnames(labelnames)
    
    half = device.type != 'cpu'  # half precision only supported on CUDA
    if trtFlag_det:
        Detweights = "../AIlib/weights/yolov5/class8/bestcao_%s_fp16.engine"%(gpuname) ##升级后的检测模型
        logger = trt.Logger(trt.Logger.ERROR)
        with open(Detweights, "rb") as f, trt.Runtime(logger) as runtime:
            model=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
        print('############locad det model trt success#######')
    else:     
        Detweights = "../AIlib/weights/yolov5/class8/bestcao.pt"
        model = attempt_load(Detweights, map_location=device)  # load FP32 model
        print('############locad det model pth success#######')
        if half: model.half()
        

    seg_nclass = 2
    segPar={'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'numpy':False, 'RGB_convert_first':True}
    if trtFlag_seg:
        Segweights = '../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640_%s_fp16.engine'%(gpuname)
        logger = trt.Logger(trt.Logger.ERROR)
        with open(Segweights, "rb") as f, trt.Runtime(logger) as runtime:
            segmodel=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
        print('############locad seg model trt success#######')
    else:
        Segweights = '../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640.pth'
        segmodel = SegModel(nclass=seg_nclass,weights=Segweights,device=device)
        print('############locad seg model pth success#######')

    postFile= '../AIlib/conf/para.json'
    digitFont= { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3}
    conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile)

    ####模型选择参数用如下：
    mode_paras=[
        {"id":"0","config":{"k1":"v1","k2":"v2"}},
        {"id":"1","config":{"k1":"v1","k2":"v2"}},
        {"id":"2","config":{"k1":"v1","k2":"v2"}},
        {"id":"3","config":{"k1":"v1","k2":"v2"}},
        {"id":"4","config":{"k1":"v1","k2":"v2"}},
        {"id":"5","config":{"k1":"v1","k2":"v2"}},
        {"id":"6","config":{"k1":"v1","k2":"v2"}},
        {"id":"7","config":{"k1":"v1","k2":"v2"}},
    ]
    allowedList,allowedList_string=get_needed_objectsIndex(mode_paras)
    #allowedList=[0,1,2,3]
    ##加载模型，准备好显示字符
    label_arraylist = get_label_arrays(names,rainbows,outfontsize=40,fontpath="../AIlib/conf/platech.ttf")
    
  
    ##图像测试
    #impth = 'images/slope/'
    impth = '../../../data/无人机起飞测试图像/'
    outpth = 'images/results/'
    folders = os.listdir(impth)
    frames=[]
     
    for i in range(len(folders)):
        imgpath = os.path.join(impth, folders[i])
        im0s=[cv2.imread(imgpath)]
        frame=(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList,digitFont,folders[i],trtFlag_det,trtFlag_seg,segPar)
        frames.append(frame)
    t1=time.time()     
    if max_workers==1:
        for i in range(len(folders)):   
            t5=time.time()        
            process_v1(frames[i])
            t6=time.time()       
            print('#######%s, ms:%.1f , accumetate time:%.1f, avage:%1.f '%(folders[i], (t6-t5)*1000.0,(t6-t1)*1000.0, (t6-t1)*1000.0/(i+1)))
    else:
        with ThreadPoolExecutor(max_workers=max_workers) as t:
            for result in t.map(process_v1, frames):
                #print(result)
                t=result
    
    t2=time.time()
    print('All %d images time:%.1f ms ,each:%.1f ms, with %d threads'%(len(folders),(t2-t1)*1000, (t2-t1)*1000.0/len(folders) , max_workers) )   

        
def river_demo():
    ##预先设置的参数
    device_='1'  ##选定模型，可选 cpu,'0','1'
     
    ##以下参数目前不可改
    #Detweights = "../AIlib/weights/yolov5/class5/best_5classes.pt" 
    #labelnames = "../AIlib/weights/yolov5/class5/labelnames.json"
    
    Detweights = "../AIlib/weights/yolov5/class8/bestcao.pt" ##升级后的检测模型
    labelnames = "../AIlib/weights/yolov5/class8/labelnames.json" ##对应类别表
    
    seg_nclass = 2 
    #Segweights = "../AIlib/weights/BiSeNet/checkpoint.pth"
    Segweights = '../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640.pth' ##升级的分割模型
    
    postFile= '../AIlib/conf/para.json'
    digitFont= { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3} 
    conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile)
   
    ####模型选择参数用如下：
    mode_paras=[
        {"id":"0","config":{"k1":"v1","k2":"v2"}},
        {"id":"1","config":{"k1":"v1","k2":"v2"}},
        {"id":"2","config":{"k1":"v1","k2":"v2"}},
        {"id":"3","config":{"k1":"v1","k2":"v2"}},
        {"id":"4","config":{"k1":"v1","k2":"v2"}},
        {"id":"5","config":{"k1":"v1","k2":"v2"}},
        {"id":"6","config":{"k1":"v1","k2":"v2"}},
        {"id":"7","config":{"k1":"v1","k2":"v2"}},
    ]

    allowedList,allowedList_string=get_needed_objectsIndex(mode_paras)
    #allowedList=[0,1,2,3]
    
    

    ##加载模型，准备好显示字符
    device = select_device(device_)
    names=get_labelnames(labelnames)
    
    label_arraylist = get_label_arrays(names,rainbows,outfontsize=40,fontpath="../AIlib/conf/platech.ttf")
    half = device.type != 'cpu'  # half precision only supported on CUDA
    model = attempt_load(Detweights, map_location=device)  # load FP32 model
    if half: model.half()
    
    segmodel = SegModel(nclass=seg_nclass,weights=Segweights,device=device)
    
    
    
    ##图像测试
    #url='images/examples/20220624_响水河_12300_1621.jpg'
    impth = 'images/slope/'
    outpth = 'images/results/'
    folders = os.listdir(impth)
    for i in range(len(folders)):
        imgpath = os.path.join(impth, folders[i])
        im0s=[cv2.imread(imgpath)]
        H,W,C = im0s[0].shape
        time00 = time.time()
        p_result,timeOut = AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList,font=digitFont)
        time11 = time.time()
        image_array = p_result[1]
        cv2.imwrite(  os.path.join( outpth,folders[i] ) ,image_array )
        print('%s,%d*%d,AI-process: %.1f, %s'%(folders[i],H,W, (time11 - time00) * 1000,timeOut))
        
def road_forest_demo(business ):

    ##使用森林，道路模型,business 控制['forest','road']
    ##预先设置的参数
    gpuname='3090'#如果用trt就需要此参数，只能是"3090" "2080Ti"
    device_='0'  ##选定模型，可选 cpu,'0','1'
    device = select_device(device_)
    half = device.type != 'cpu'  # half precision only supported on CUDA
   
    trtFlag_det=False ###是否采用TRT模型加速
    
    ##以下参数目前不可改
    #business='forest';imageW=4916  ####森林模型
    #business='road';
    imageW=1536     ####道路模型
    digitFont= { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3 }
    
    
    if trtFlag_det:
        Detweights="../AIlib/weights/%s/best_%s_fp16.engine"%(business,gpuname)
        logger = trt.Logger(trt.Logger.ERROR)
        with open(Detweights, "rb") as f, trt.Runtime(logger) as runtime:
            model=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
        print('####load TRT model :%s'%(Detweights))    
    else:
        Detweights="../AIlib/weights/%s/best.pt"%(business)
        model = attempt_load(Detweights, map_location=device)  # load FP32 model
        if half: model.half()
        
    labelnames = "../AIlib/weights/%s/labelnames.json"%(business)
    postFile= '../AIlib/weights/%s/para.json'%(business)
    print( Detweights,labelnames  ) 
    conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile)
    
    ####模型选择参数用如下：
    if business == 'road':
        mode_paras=[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,3,4,5,6]  ]###类别2为“修补”，不输出
    else:
        mode_paras=[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6]  ]
    allowedList,allowedList_string=get_needed_objectsIndex(mode_paras)    
    ##只加载检测模型，准备好显示字符
    
    names=get_labelnames(labelnames)
    #imageW=4915;###默认是1920，在森林巡检的高清图像中是4920
    outfontsize=int(imageW/1920*40);###
    label_arraylist = get_label_arrays(names,rainbows,outfontsize=outfontsize,fontpath="../AIlib/conf/platech.ttf")
    
    segmodel = None 
    ##图像测试
    #url='images/examples/20220624_响水河_12300_1621.jpg'
    impth = 'images/%s/'%(business)
    outpth = 'images/results/'
    folders = os.listdir(impth)
    folders.sort()
    for i in range(len(folders)):
    #for i in range(2):
        imgpath = os.path.join(impth, folders[i])
        im0s=[cv2.imread(imgpath)]
        time00 = time.time()
        #使用不同的函数。每一个领域采用一个函数
        p_result,timeOut = AI_process_forest(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList,font=digitFont,trtFlag_det=trtFlag_det)
        time11 = time.time()
        image_array = p_result[1]
        cv2.imwrite(  os.path.join( outpth,folders[i] ) ,image_array )
        
        print('----image:%s, process:%s ,save:%s, %s'%(folders[i],(time11-time00) * 1000, (time.time() - time11) * 1000,timeOut) )

    
def jkm_demo():
    from utilsK.jkmUtils import pre_process,post_process,get_return_data
    img_type = 'plate' ## code,plate
    par={'code':{'weights':'../AIlib/weights/jkm/health_yolov5s_v3.jit','img_type':'code','nc':10 },
        'plate':{'weights':'../AIlib/weights/jkm/plate_yolov5s_v3.jit','img_type':'plate','nc':1 },
        'conf_thres': 0.4,
        'iou_thres':0.45,
        'device':'cuda:0',
        'plate_dilate':(0.5,0.1)
    }
     ###加载模型
    device = torch.device(par['device'])
    jit_weights = par['code']['weights']
    model = torch.jit.load(jit_weights)

    jit_weights =  par['plate']['weights']
    model_plate = torch.jit.load(jit_weights)

    imgd='images/plate'
    imgpaths = os.listdir(imgd)
    for imgp in imgpaths[0:]:
        #imgp = 'plate_IMG_20221030_100612.jpg'
        imgpath = os.path.join(imgd,imgp)
        im0 = cv2.imread(imgpath) #读取数据
        img ,padInfos = pre_process(im0,device) ##预处理
        if img_type=='code':  pred = model(img) ##模型推理
        else: pred = model_plate(img)
        boxes = post_process(pred,padInfos,device,conf_thres= par['conf_thres'], iou_thres= par['iou_thres'],nc=par[img_type]['nc'])  #后处理
        dataBack=get_return_data(im0,boxes,modelType=img_type,plate_dilate=par['plate_dilate'])
        print(imgp,boxes,dataBack['type'])
        for key in dataBack.keys():
            if isinstance(dataBack[key],list):
                cv2.imwrite(  'images/results/%s_%s.jpg'%( imgp.replace('.jpg','').replace('.png',''),key),dataBack[key][0] )    ###返回值： dataBack

        1
if __name__=="__main__":
    #river_demo_v3()
    road_forest_demo('road'  )
    #jkm_demo()