# -*- coding: utf-8 -*-
"""
Created on Sun Sep 29 08:59:21 2024

@author: ym
"""
import os
import sys
import cv2
import pickle
import argparse
import numpy as np
from pathlib import Path
from track_reid import parse_opt
from track_reid import yolo_resnet_tracker
# FILE = Path(__file__).resolve()
# ROOT = FILE.parents[0]  # YOLOv5 root directory
# if str(ROOT) not in sys.path:
#     sys.path.append(str(ROOT))  # add ROOT to PATH
# ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

from tracking.dotrack.dotracks_back import doBackTracks
from tracking.dotrack.dotracks_front import doFrontTracks 
from tracking.utils.drawtracks import plot_frameID_y2, draw_all_trajectories
from utils.getsource import get_image_pairs, get_video_pairs

def get_interbcd_inputenents():
    bcdpath = r"\\192.168.1.28\share\测试_202406\contrast\std_barcodes_2192"
    eventpath = r"\\192.168.1.28\share\测试_202406\0918"

    barcodes = []
    eventpaths = []
    for featname in os.listdir(bcdpath):
        barcode, ext = os.path.splitext(featname)
        barcodes.append(barcode)
    
    input_enents = []    
    for root, dirs, files in os.walk(eventpath):
        input_enent = [os.path.join(root, d) for d in dirs if d.split('_')[-1] in barcodes]
        
        input_enents.extend(input_enent)
        
        
    return input_enents

def pipeline(        
        eventpath,
        savepath,
        SourceType = "image",   # video
        stdfeat_path = None 
        ):
    if SourceType == "video":
        vpaths = get_video_pairs(eventpath)
    elif SourceType == "image":
        vpaths = get_image_pairs(eventpath)

    '''======== 函数 yolo_resnet_tracker() 的参数字典 ========'''
    opt = parse_opt()
    optdict = vars(opt)
    optdict["is_save_img"] = True
    optdict["is_save_video"] = True
    
    event_tracks = []
    for vpath in vpaths:
        '''事件结果文件夹'''
        save_dir_event = Path(savepath) / Path(os.path.basename(eventpath))
        if isinstance(vpath, list):
            save_dir_video = save_dir_event / Path("images")
        else:
            save_dir_video = save_dir_event / Path(str(Path(vpath).stem))
        
        if not save_dir_video.exists():
            save_dir_video.mkdir(parents=True, exist_ok=True)
        
        '''Yolo + Resnet + Tracker'''
        optdict["source"] = vpath
        optdict["save_dir"] = save_dir_video        
        
        
        tracksdict = yolo_resnet_tracker(**optdict)
        
        bboxes = tracksdict['TrackBoxes']
        
        bname = os.path.basename(vpath[0]) if isinstance(vpath, list) else os.path.basename(vpath)
        if bname.split('_')[0] == "0" or bname.find('back')>=0:
            vts = doBackTracks(bboxes, tracksdict)
            vts.classify()
            
            event_tracks.append(("back", vts))
            
        if  bname.split('_')[0] == "1" or bname.find('front')>=0:
            vts = doFrontTracks(bboxes, tracksdict)
            vts.classify()
            event_tracks.append(("front", vts))
    
    
    '''轨迹显示模块'''
    illus = [None, None]
    for CamerType, vts in event_tracks:
        if CamerType == 'front':
            edgeline = cv2.imread("./tracking/shopcart/cart_tempt/board_ftmp_line.png")
            
            h, w = edgeline.shape[:2]
            nh, nw = h//2, w//2
            edgeline = cv2.resize(edgeline, (nw, nh), interpolation=cv2.INTER_AREA)
            
            img_tracking = draw_all_trajectories(vts, edgeline, save_dir_event, CamerType, draw5p=True)
            illus[0] = img_tracking
            
            
            plt = plot_frameID_y2(vts)
            plt.savefig(os.path.join(save_dir_event, "front_y2.png"))
        
        if CamerType == 'back':
            edgeline = cv2.imread("./tracking/shopcart/cart_tempt/edgeline.png")
            
            h, w = edgeline.shape[:2]
            nh, nw = h//2, w//2
            edgeline = cv2.resize(edgeline, (nw, nh), interpolation=cv2.INTER_AREA)
            
            img_tracking = draw_all_trajectories(vts, edgeline, save_dir_event, CamerType, draw5p=True)
            illus[1] = img_tracking
    
    illus = [im for im in illus if im is not None]
    if len(illus):
        img_cat = np.concatenate(illus, axis = 1)
        if len(illus)==2:
            H, W = img_cat.shape[:2]
            cv2.line(img_cat, (int(W/2), 0), (int(W/2), int(H)), (128, 128, 255), 3)
            
        trajpath = os.path.join(save_dir_event, "traj.png")
        cv2.imwrite(trajpath, img_cat)
    
    
        
    '''前后摄轨迹选择'''
    
    if stdfeat_path is not None:
        with open(stdfeat_path, 'rb') as f:  
            featDict = pickle.load(f) 
        

def main_loop():
    bcdpath = r"\\192.168.1.28\share\测试_202406\contrast\std_barcodes_2192"
    eventpath = r"\\192.168.1.28\share\测试_202406\0918\images1"
    SourceType = "image"   # video, image

    barcodes = []
    input_enents = []  
    output_events = []


    '''1. 获得barcode标准特征集列表'''    
    for featname in os.listdir(bcdpath):
        barcode, ext = os.path.splitext(featname)
        if not barcode.isdigit() or len(barcode)<=8 or ext != ".pickle" :
            continue
        barcodes.append(barcode)
    
    '''2. 构造（放入事件，标准特征）对'''
    for filename in os.listdir(eventpath):
        '''barcode为时间文件夹的最后一个字段'''
        bcd = filename.split('_')[-1]
        
        event_path = os.path.join(eventpath, filename)
        stdfeat_path = None
        if bcd in barcodes:
            stdfeat_path = os.path.join(bcdpath, f"{bcd}.pickle")            
            input_enents.append((event_path, stdfeat_path))

    parmDict = {}
    parmDict["SourceType"] = "image"
    parmDict["savepath"] = r"D:\contrast\detect"
    for eventpath, stdfeat_path in input_enents:
        parmDict["eventpath"] = eventpath
        parmDict["stdfeat_path"] = stdfeat_path
        
        pipeline(**parmDict)
    

def main():
    '''
    函数：pipeline()，遍历事件文件夹，选择类型 image 或 video,
    '''
    parmDict = {}
    parmDict["eventpath"] = r"\\192.168.1.28\share\测试视频数据以及日志\各模块测试记录\展厅测试\1120_展厅模型v801测试\扫A放A\20241121-144855-dce94b09-1100-43f1-92e8-33a1b538b159_6924743915848_6924743915848"
    
    parmDict["savepath"] = r"D:\contrast\detect"
    parmDict["SourceType"] = "image"   # video, image
    parmDict["stdfeat_path"] = None

    pipeline(**parmDict)
    

if __name__ == "__main__": 
    main()