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

@author: ym
"""
import os
import cv2
import pickle
import numpy as np
from pathlib import Path
from track_reid import parse_opt, yolo_resnet_tracker

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





std_feature_path = r"\\192.168.1.28\share\测试_202406\contrast\std_features_2192_ft32vsft16"

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, stdfeat_path=None, SourceType = "image"): 
    '''
    inputs:
        eventpath: 事件文件夹
        stdfeat_path: 标准特征文件地址
    outputs:
        
    '''
    SourceType = "image"   # image
    
    # eventpath = r"\\192.168.1.28\share\测试_202406\0918\images1\20240918-110822-1bc3902e-5a8e-4e23-8eca-fb3f02738551_6938314601726"
    
    savepath = r"D:\contrast\detect"
    
    opt = parse_opt()
    optdict = vars(opt)
    optdict["project"] = savepath
    
    eventname = os.path.basename(eventpath)
    # barcode = eventname.split('_')[-1]
    
    
    if SourceType == "video":
        vpaths = get_video_pairs(eventpath)
    elif SourceType == "image":
        vpaths = get_image_pairs(eventpath)



    event_tracks = []
    for vpath in vpaths:  
        
        '''事件结果文件夹'''
        save_dir_event = Path(savepath) / Path(eventname)
        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        
        optdict["is_save_img"] = True
        optdict["is_save_video"] = True
        
        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")
            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")
            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))

    
    for eventpath, stdfeat_path in input_enents:
        pipeline(eventpath, stdfeat_path, SourceType)
        
            
def main():
    eventpath = r"D:\datasets\ym\exhibition\175836"
    SourceType = 'image'
    stdfeat_path = None
   
    pipeline(eventpath, stdfeat_path, SourceType)
    
       
        
        

        
        
        


if __name__ == "__main__": 
    main()