# -*- coding:utf-8 -*-
import cv2,os,torch,argparse
import queue
import time
import threading
from System_I import Check_start_fordataselect,Density,Detection_Y5,Fine_grained,Fine_grained_T,shishi,save_sigle_image,RealDeal
from common import global_para as gp
r_db = None
import torch.multiprocessing as multiprocessing
import sort.sort as SORT
from imatovideo import images_to_video
from multiprocessing import Array


class DetFineVid():
    def __init__(self, VideoPath,SavePath):
        super(DetFineVid,self).__init__()
        # name
        self.VideoPath = VideoPath
        self.VideoResultSavePathroot = "/data1/Hatcher/model/project/data_MPKFNet/result_vedio/"
        self.VideoName = os.path.basename(self.VideoPath).split(".")[0]
        self.VideoResultSavePath = os.path.join(self.VideoResultSavePathroot,"{}.avi".format(self.VideoName))
        self.SavePath = os.path.join(SavePath,self.VideoName)
        if not os.path.exists(self.SavePath):
            os.makedirs(self.SavePath)

        self.root = "/home/wanghc/output_jiankong"
        self.DensityClass = None  # Density(gp.devices, root)
        self.DetectionClass = Detection_Y5(gp.devices, self.root)
        self.Fine_grianedClass = Fine_grained_T(gp.devices, self.root)
        self.FineRestartFlag = True
        self.numrestart = 0

    def conVideo(self):
        images_to_video( self.SavePath,self.VideoResultSavePath)
    def StartRec(self):


        t1 = time.time()
        cap = cv2.VideoCapture(self.VideoPath)
        t2 = time.time()
        print("capture is ",(t2-t1)*1000)
        num_detection = 0
        num_cuntupian=0
        SortRestart = SORT.Sort(max_age=gp.max_age,
                       min_hits=gp.min_hits,
                       iou_threshold=gp.sortiou_threshold)
        realtimedeal = RealDeal("check_start",self.DetectionClass,self.Fine_grianedClass)
        if cap.isOpened():
            res, frame = cap.read()
            if res:
                shapex = frame.shape[1]
                shapey = frame.shape[0]
                gp.radioShowX = shapex / gp.radioShowX
                gp.radioShowY = shapey / gp.radioShowY
                gp.radioShowtext = gp.radioShowX
                if frame.shape[1] == 3840:
                    gp.fontScale = 3
                elif frame.shape[1] == 2560:
                    gp.fontScale = 2
                else:
                    gp.fontScale = 1


        num_image = 0
        while True:
            if cap.isOpened():
                res, frame = cap.read()
                if res:
                    closed = False
                    failCnt = 0
                    num_image += 1
                    data_today = 8
                    print(num_detection)
                    # if num_detection <12974:
                    #     a=0
                    #     num_detection+=1
                    #     continue
                        #cv2.imwrite("./test.jpg",frame)
                    if (int(data_today) > 1) & (int(data_today) < 20):
                        if (0xFF == ord('q')):
                            print("restart")
                            self.FineRestartFlag = True
                            self.numrestart = 0
                            SortRestart = SORT.Sort(max_age=gp.max_age,
                                                    min_hits=gp.min_hits,
                                                    iou_threshold=gp.sortiou_threshold)
                        if self.FineRestartFlag :
                            realtimedeal.realtimerestart(frame,num_image=num_image,num_detection=num_detection,SortRestart = SortRestart)
                            self.numrestart += 1
                            if self.numrestart >= 5:
                                self.numrestart = 0
                                self.FineRestartFlag = False
                                FirstDetFineResult = realtimedeal.getrestartresult(gp.IndexMax)
                                print("µÚÒ»´Îʶ±ðÄ¿±êÊýÁ¿ ",len(FirstDetFineResult))
                                gp.SortLabel_Index = FirstDetFineResult
                                # SortRealTime = SORT.Sort(max_age=gp.max_age,
                                #                          min_hits=gp.min_hits,
                                #                          iou_threshold=gp.sortiou_threshold)

                        else:
                            tm1 = time.time()
                            #frame = cv2.imread("./test.jpg")
                            img_show,num_image = realtimedeal.realtime(frame,num_image=num_image,num_detection=num_detection,SortRestart = SortRestart)
                            tm2 = time.time()
                            if gp.DebugPrint_flag:
                                print("main1 time is  ", (tm2 - tm1) * 1000)
                            # cv2.imshow("AI-ECO", img_show)
                            # #cv2.imwrite("../data_MPKFNet/RTSP3/" + "{}.jpg".format(time.strftime('%Y%m%d%H%M%S', time.localtime(time.time()))),img_show)
                            # if cv2.waitKey(1) & 0xFF == ord('r'):
                            #     break
                            save_path =os.path.join(self.SavePath,"{}.jpg".format(str(num_detection).zfill(6)))
                            cv2.imwrite(save_path,img_show)
                            tm3 = time.time()
                            if gp.DebugPrint_flag:
                                print("main2 time is  ", (tm3 - tm2) * 1000)
                    else:
                        print("time is out")
                    if (num_detection >1000000) :
                        # save_image = save_sigle_image(frame)
                        # save_image.start()
                        # num_detection=0
                        num_detection = 0
                    num_detection+=1
                    num_cuntupian+=1
                else:
                    break


def commanddeal():
    SavePath = "/data1/Hatcher/model/project/data_MPKFNet/result/"
    DFS = DetFineVid(gp.VideoPath,SavePath)
    DFS.StartRec()
    DFS.conVideo()
def filedeal(rootPath):  #deal all file in a folder
    SavePath = "/data1/Hatcher/model/project/data_MPKFNet/result/"
    i = 0
    for video_path in os.listdir(rootPath):
        if i == 1:
            print("111")
        # print("starting deal {}".format(str(video_path)))
        # VideoPath = os.path.join(rootPath,video_path)
        # DFS = DetFineVid(VideoPath,SavePath)
        # DFS.StartRec()
        # DFS.conVideo()
        P = multiprocessing.Process(target=func,args=(video_path,SavePath,rootPath))
        P.start()
        P.join()
        #globalInit()
        i += 1
def func(video_path,SavePath,rootPath):
    print("starting deal {}".format(str(video_path)))
    VideoPath = os.path.join(rootPath, video_path)
    DFS = DetFineVid(VideoPath, SavePath)
    DFS.StartRec()
    DFS.conVideo()
def globalInit():  #deal all file in a folder

    parser1 = argparse.ArgumentParser(description="CUI")
    parser1.add_argument("--video_num", default='/data1/Hatcher/model/project/data_MPKFNet/rawdata/PYH-DEMO.mp4',
                         type=str,
                         help='Video name.')
    parser1.add_argument("--GPU", default='6', type=str,
                         help='GPU.')
    parser1.add_argument("--thresh", default=0.4, type=float,
                         help='GPU.')
    parser1.add_argument("--red", default=0, type=int,
                         help='red.')
    parser1.add_argument("--add", default=0, type=str,
                         help='add.')

    args = parser1.parse_args()
    # print(args)
    gp.root = "/home/gyd/output1"
    # os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    os.environ['CUDA_VISIBLE_DEVICES'] = args.GPU
    gp.devices = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    gp.camera_address = args.add
    gp.VideoPath = args.video_num
    gp.red_flag = args.red
    gp.detect_thresh = args.thresh

    # fine_thresh2 = torch.tensor(fine_thresh2).cuda()
    gp.birdIndex = 99999
    # birdIndex = torch.tensor(birdIndex).cuda()
    gp.is_ex_bird = False  # Ä¿±ê¼ì²â±ê־λ£¬Ã¿¸öÒ»Ãë¼ì²âÒ»´ÎÄ¿±ê¼ì²â£¬Èç¹ûÓÐÄñ£¬ÔòÖÃΪTrue£¬ÎªTrueµÄʱºò˵Ã÷£¬ ´¦Àí30sµÄÏ̻߳¹ÔÚ¼ÌÐø£¬Ö±µ½±ê־λ±äΪÁËFalseÔÙ¿ªÆôеÄÏß³Ì
    gp.Density_thread_flag = False  # ÃܶȹÀ¼ÆÏ̱߳ê־λ£¬ÎªTrue˵Ã÷¿ªÆôÁËÃܶȹÀ¼ÆỊ̈߳¬´¦ÀíÍê750֡Ϊֹ£¬±ê־λ±êΪFalse£¬
    gp.Detection_thread_flag = False  # Ä¿±ê¼ì²âºÍϸÁ£¶Èʶ±ðÏ̱߳ê־λ£¬ÎªTrue˵Ã÷¿ªÆôÁËÃܶȹÀ¼ÆỊ̈߳¬´¦ÀíÍê750֡Ϊֹ£¬±ê־λ±êΪFalse£¬
    gp.Vedio_thread_flag = False  # ´æ´¢ÊÓƵÏ̱߳ê־룬´æ´¢30s
    gp.SaveImage_thread_flag = False  # ±£´æԭͼÏ̱߳ê־룬´æ´¢30ÕÅԭͼ
    gp.ALL_thread_flag = False  # ËùÓÐÏ̵߳ıê־룬µ¥¶À¿ªÉèÒ»¸öỊ̈߳¬¼à¿ØËùÓÐÏß³ÌÊÇ·ñ¶¼½áÊø£¬Èç¹û¶¼½áÊø£¬ÔòÖÃΪFalse£¬ºóÐøÌí¼ÓÏ̹߳ÜÀí£¬¿ÉÒÔ·ÅÔÚÕâÀïÃæ

    gp.ShowAllNumFlag = False

    gp.DebugPrint_flag = False  # ´òÓ¡±ê־λ

    gp.num_count_old = 0
    gp.boxes_fine = None
    gp.labels_fine = None

    gp.showThresh = 25
    gp.destory_image = False
    gp.Camera_device = args.video_num  # ÉãÏñÍ·±àºÅ£¬0001
    gp.image_time = None  # µ±Ç°´¦ÀíµÄ30ÊÓƵµÄ±àºÅ
    gp.ImgDealQueue = []  # 750Ö¡ÊÓƵ£¬»º´æ¿Õ¼ä
    gp.img_show = None

    # event = threading.Event()
    gp.eventDensity = threading.Event()
    gp.eventDetection = threading.Event()
    gp.eventSaveimage = threading.Event()
    gp.max_age = 2
    gp.min_hits = 3
    gp.sortiou_threshold = 0.3
    gp.Finerestartboxes = []
    gp.IndexMax = 0
    gp.radioShowX = 2560.0
    gp.radioShowY = 1440.0
    gp.radioShowtext = gp.radioShowX
    gp.SortLabel_Index = None  # ±êÇ©ºÍË÷Òý¶ÔÓ¦¹Øϵ[[4,1],[3,2]]
    gp.SortLabel_Indexpre = []  # ±êÇ©ºÍË÷ÒýµÄ¶ÔÓ¦¹Øϵ£¬ÓÃÔÚ×·×Ù¹ý³ÌÖÐгöÏÖµÄÄ¿±ê£¬[[label,Index,Count]]
    # Detection2 = Detection(devices,root)
    gp.FineCountThresh = 5

    gp.img_show = None

    gp.Fine_thresh = 0.4

    # DetWeight = "/data1/Hatcher/model/Detection/yolov5/runs/train/exp21/weights/best.pt"
    # FineWeight = "/data/Hatcher/data/weigth/areaproject/resnet50-PYH-3-20.7/_42522_PYH-3/best.pth"
    gp.DetWeight = "/data1/Hatcher/model/Detection/yolov5/runs/train/KunMingCuiHu-Fish2/weights/best.pt"
    gp.FineWeight = "/data/Hatcher/data/weigth/areaproject/resnet50-DianChiFish-3-3-0.7/_4289_DianChiFish-4/best.pth"
    # DetWeight = "/data1/Hatcher/model/Detection/yolov5/runs/train/BeijingCuiHu-Bird26/weights/best.pt"
    # FineWeight = "/data/Hatcher/data/weigth/areaproject/resnet50-cui90_1model0.7/_42111_CUI-94-1/best.pth"
    # BHG_chineseName = ["kong", '白鹤', '白头鹤','白枕鹤']#["kong","鲫鱼","麦穗鱼","鳑鲏"]
    gp.BHG_chineseName = ["kong", "鲫鱼", "麦穗鱼", "鳑鲏"]

    # BHG_chineseName = ['Kong', '白额雁', '白骨顶', '苍鹭', '苍鹭', '白秋沙鸭', '白腰草鹬', '斑头雁', '斑嘴鸭', '苍鹭', '赤膀鸭', '赤颈鸭', '苍鹭', '大麻鳽',
    #                    '豆雁', '反嘴鹬', '凤头鸊鷉', '凤头麦鸡', '凤头潜鸭', '黑鹳', '黑水鸡', '红嘴鸥', '鸿雁', '灰鹤', '灰雁', '绿翅鸭', '绿头鸭', '琵嘴鸭',
    #                    '普通鸬鹚', '普通秋沙鸭', '翘鼻麻鸭', '鹊鸭', '扇尾沙锥', '小白额雁', '小天鹅', '针尾鸭', '苍鹭', '暗绿绣眼鸟', '八哥', '白腹鹞', '白鹡鸰',
    #                    '白眼潜鸭', '草鹭', '池鹭', '戴胜', '黑翅长脚鹬', '黑鸢', '红隼', '花脸鸭', '环颈雉', '灰头绿啄木鸟', '灰头麦鸡', '矶鹬', '栗苇鳽',
    #                    '罗纹鸭', '牛背鹭', '普通翠鸟', '雀鹰', '喜鹊', '苍鹭', '鸳鸯', '棕背伯劳', '红胸秋沙鸭', '普通燕鸥', '小白额燕鸥', '丘鹬', '凤头蜂鹰',
    #                    '大天鹅', '大鸨', '小田鸡', '山斑鸠', '帆背潜鸭', '日本松雀鹰', '普通鵟', '松雀鹰', '毛脚鵟', '灰斑鸠', '白头鹞', '白尾鹞', '白眉鸭',
    #                    '白鹈鹕', '红头潜鸭', '苍鹰', '蒙古银鸥', '西方秧鸡', '赤嘴潜鸭', '阔嘴鹬', '雕鸮', '鹌鹑', '鹗', '黄苇鳽', '黑颈鹤', '小䴙䴘', '赤麻鸭']
    # BHG_chineseName =["kong","绿头鸭","斑嘴鸭","灰雁","鸿雁","kong","小鷿鷈"]

    gp.FineClassNum = 4
    gp.color_show = (255, 255, 255)
    gp.color_box = (255, 0, 0)
    gp.NoneLabel = "鱼"  # "鹤类"
    gp.fine_thresh1 = 6000
    gp.fine_thresh2 = 10000
    gp.fontScale = 1



    torch.zero_(gp.Species)



if __name__ == '__main__':
    multiprocessing.set_start_method("spawn")
    rootPath = "/data/DataManagement/Project/005_kunmingdianchi/0002_20220412_800_200/研究院-鱼视频/"
    filedeal(rootPath)
    #commanddeal()