# -*- coding: utf-8 -*-
"""
Created on Sun Mar  2 14:15:57 2025

@author: ym
"""
import numpy as np
import cv2
import os
from pathlib import Path

import sys
sys.path.append(r"D:\DetectTracking")

# from tracking.utils.read_data import extract_data_realtime, read_tracking_output_realtime
from tracking.utils.plotting import Annotator, colors
from tracking.utils import Boxes, IterableSimpleNamespace, yaml_load, boxes_add_fid
from tracking.trackers import BOTSORT, BYTETracker
from tracking.utils.showtrack import drawtracks
from hands.hand_inference import hand_pose
from tracking.utils.read_data import read_weight_sensor, extract_data_realtime, read_tracking_output_realtime
from contrast.feat_extract.config import config as conf
from contrast.feat_extract.inference import FeatsInterface
from tracking.utils.drawtracks import drawTrack


ReIDEncoder = FeatsInterface(conf)


W, H = 1024, 1280
Mode = 'front'  #'back'
ImgFormat = ['.jpg', '.jpeg', '.png', '.bmp']


'''调用tracking()函数，利用本地跟踪算法获取各目标轨迹，可以比较本地跟踪算法与现场跟踪算法的区别。'''
def init_tracker(tracker_yaml = None, bs=1):
    """
    Initialize tracker for object tracking during prediction.
    """  
    TRACKER_MAP = {'bytetrack': BYTETracker, 'botsort': BOTSORT}   
    cfg = IterableSimpleNamespace(**yaml_load(tracker_yaml))

    tracker = TRACKER_MAP[cfg.tracker_type](args=cfg, frame_rate=30)

    return tracker


def init_trackers(tracker_yaml = None, bs=1):
    """
    Initialize trackers for object tracking during prediction.
    """      
    # tracker_yaml = r"./tracking/trackers/cfg/botsort.yaml"
    
    TRACKER_MAP = {'bytetrack': BYTETracker, 'botsort': BOTSORT}
    
    cfg = IterableSimpleNamespace(**yaml_load(tracker_yaml))
    trackers = []
    for _ in range(bs):
        tracker = TRACKER_MAP[cfg.tracker_type](args=cfg, frame_rate=30)
        trackers.append(tracker)
    
    return trackers


def draw_box(img, tracks):
    annotator = Annotator(img.copy(), line_width=2)
    # for *xyxy, conf, cls in reversed(tracks):
    #     name = f'{int(cls)} {conf:.2f}'
    #     color = colors(int(cls), True)  
    #     annotator.box_label(xyxy, name, color=color)
        
    for  *xyxy, id, conf, cls, fid, bid in reversed(tracks):
        name = f'ID:{int(id)} {int(cls)} {conf:.2f}'
        color = colors(int(cls), True)  
        annotator.box_label(xyxy, name, color=color)
    
    
    
    im0 = annotator.result()
    
    return im0




def tracking(bboxes, ffeats):
    tracker_yaml = "./tracking/trackers/cfg/botsort.yaml"
    
    tracker = init_tracker(tracker_yaml)
    
    TrackBoxes = np.empty((0, 9), dtype = np.float32)
    TracksDict = {}
    
    frmIds = []
    '''========================== 执行跟踪处理 ============================='''
    # dets 与 feats 应保持严格对应
    k=0
    for dets, feats in zip(bboxes, ffeats):
        
        frmIds.append(np.unique(dets[:, 6]).astype(np.int64)[0])
        
        boxes = dets[:, :6]
        det_tracking = Boxes(boxes).cpu().numpy()
        tracks, outfeats = tracker.update(det_tracking, features=feats)
        '''tracks: [x1, y1, x2, y2, track_id, score, cls, frame_index, box_index]
                    0   1   2   3      4        5     6       7           8
                  这里，frame_index 也可以用视频的 帧ID 代替, box_index 保持不变
        '''
        k += 1
        imgpath = r"D:\全实时\202502\tracker\Yolos_Tracking\tracker\1_1740891284792\1_1740891284792_{}.png".format(int(k))
        img = cv2.imread(imgpath)
        
        im0 = draw_box(img, tracks)
        savepath = r"D:\全实时\202502\tracker\Yolos_Tracking\tracker\1_1740891284792\b\1_1740891284792_{}_b.png".format(k)
        cv2.imwrite(savepath, im0)
        
        
        if len(tracks):
            TrackBoxes = np.concatenate([TrackBoxes, tracks], axis=0)
            
# =============================================================================
#             FeatDict = {}
#             for track in tracks:
#                 tid = int(track[8])
#                 FeatDict.update({tid: feats[tid, :]})
#                 
#             frameID = tracks[0, 7]
#             
#             # print(f"frameID: {int(frameID)}")
#             assert len(tracks) == len(FeatDict), f"Please check the func: tracker.update() at frameID({int(frameID)})"
# 
#             TracksDict[f"frame_{int(frameID)}"] = {"feats":FeatDict}
# =============================================================================
            
            
    return TrackBoxes, TracksDict

def dotrack():
    
    datapath = r"D:\全实时\202502\tracker\1_tracker_in.data"
    
    bboxes, ffeats = extract_data_realtime(datapath)
    trackerboxes, tracker_feat_dict = tracking(bboxes, ffeats)
    
    print("done!")


# def plotbox():
    
#     fpath = r"\\192.168.1.28\share\测试视频数据以及日志\全实时测试\V12\2025-3-3\20250303-103833-338_6928804010091_6928804010091\1_tracking_output.data"
#     imgpath = r"D:\全实时\202502\result\Yolos_Tracking\20250303-103833-338_6928804010091_6928804010091\1_1740969517953"
#     trackingboxes, trackingfeats, tracking_outboxes, tracking_outfeats = read_tracking_output_realtime(fpath)
   
#     for  *xyxy, id, conf, cls, fid, bid in tracking_outboxes[0]:
#         imgname = f"1_1740969517953_{int(fid)}.png"
        
#         img_path = os.path.join(imgpath, imgname)
#         img = cv2.imread(img_path)
#         annotator = Annotator(img.copy(), line_width=2)
        
#         name = f'ID:{int(id)} {int(cls)} {conf:.2f}'
#         color = colors(int(cls), True)  
#         annotator.box_label(xyxy, name, color=color)
#         im0 = annotator.result()
#         cv2.imwrite(os.path.join(imgpath, f"1_1740969517953_{int(fid)}_.png"), im0)
        
#         print(f"1_1740969676295_{int(fid)}_.png")
    
#     print("done")

def video2imgs(videopath):
    cap = cv2.VideoCapture(str(videopath))
    k = 0
    while True:
        ret, frame = cap.read()
        if frame is None:
            break
        k += 1
        imgpath = videopath.parent / f"{videopath.stem}_{k}.png"
        cv2.imwrite(str(imgpath), frame)        
        
 

def extract_evtimgs(evtpath):
    vidpaths = [v for v in evtpath.iterdir() if v.suffix == '.mp4']
    for vidpath in vidpaths:
        video2imgs(vidpath)
        
    stamps = [name.stem.split('_')[1] for name in vidpaths]
    
    if len(set(stamps)==1):
        return stamps[0]
    return None

def draw_tracking_boxes(evtpath, stamp):
    for datapath in evtpath.iterdir():
        if datapath.name.find('_tracking_output.data')<=0:
            continue
        
        camera = datapath.stem.split('_')[0]
        trackingboxes, trackingfeats, tracking_outboxes, tracking_outfeats = read_tracking_output_realtime(str(datapath))
        
        ## 该模块先读取轨迹数据，再根据帧ID读取相应图像
        for  *xyxy, id, conf, cls, fid, bid in tracking_outboxes[0]:
            imgpath = evtpath /  f"{camera}_{stamp}_{int(fid)}.png"
            
            img = cv2.imread(str(imgpath))
            annotator = Annotator(img.copy(), line_width=2)
            
            name = f'ID:{int(id)} {int(cls)} {conf:.2f}'
            color = colors(int(cls), True)  
            annotator.box_label(xyxy, name, color=color)
            im0 = annotator.result()
            cv2.imwrite(imgpath, im0)
            
        print(datapath.name)

def draw_traj(evtpath):
    for datapath in evtpath.iterdir():
        if datapath.name.find('_tracking_output.data')<=0:
            continue
    
        fname = datapath.name
        trackingboxes, trackingfeats, tracking_outboxes, tracking_outfeats = read_tracking_output_realtime(datapath)
            
        CamerType = fname.split('_')[0]
        if CamerType == '1':
            edgeline = cv2.imread("./CartTemp/board_ftmp_line.png")
        if CamerType == '0':
            edgeline = cv2.imread("./CartTemp/edgeline.png")
        edgeline = drawTrack(tracking_outboxes, edgeline)
        
        
        imgpath = datapath.parent / f"{datapath.stem}.png"
        cv2.imwrite(str(imgpath), edgeline)    
    
    
def main():
    
    path = r"\\192.168.1.28\share\测试视频数据以及日志\全实时测试\V12\2025-3-3\20250303-104225-381_6920459958674"
    evtpaths = [p for p in Path(path).iterdir() if p.is_dir()]
    for evtpath in evtpaths:
        #1. 从事件的前后摄视频提取图像
        stamp = extract_evtimgs(evtpath)
        
        #2. 根据 0/1_tracking_output.data 中提取的轨迹在img中绘制box
        draw_tracking_boxes(evtpath, stamp)
        
        #3. 根据 0/1_tracking_output.data 中提取的轨迹在edgeline中绘制box
        draw_traj(evtpath)
        

    
    
if __name__ == '__main__':
    # dotrack()  
    # plotbox()
    
    vpath = r"D:\datasets\ym\VID_20250307_105606"
    extract_evtimgs(Path(vpath))
    
    
    


















