# -*- coding: utf-8 -*-
"""
Created on Thu Mar 27 16:09:07 2025

@author: ym
"""
import os
import sys
import cv2
import pickle
import numpy as np
from pathlib import Path
from scipy.spatial.distance import cdist
import copy

from dotrack.dotracks_back import doBackTracks
from dotrack.dotracks_front import doFrontTracks 
from utils.drawtracks import plot_frameID_y2, draw_all_trajectories
from utils.read_data import read_similar


def get_trail(ShoppingDict, ppath):
    
    evtname = ShoppingDict["eventName"]
    
    back_yrt = ShoppingDict["backCamera"]["yoloResnetTracker"]
    front_yrt = ShoppingDict["frontCamera"]["yoloResnetTracker"]
    
    back_vts = ShoppingDict["frontCamera"]["tracking"]
    front_vts = ShoppingDict["backCamera"]["tracking"]
    
    
    event_tracks = [("back", back_yrt, back_vts), ("front", front_yrt, front_vts)]
    
    
    savepath = ppath / "alltrail" 
    if not savepath.exists():
        savepath.mkdir()
    
    savepath = str(savepath)
    evtime = evtname[:15]
        
    illus = [None, None]
    for camera_type, yrtOut, vts in event_tracks:
        if len(vts.Residual)==1: continue
        
        if camera_type == 'front':
            edgeline = cv2.imread("./shopcart/cart_tempt/board_ftmp_line.png")
            
            img_tracking = draw_all_trajectories(vts, edgeline, savepath, camera_type, draw5p=False)
            illus[0] = img_tracking

            plt = plot_frameID_y2(vts)
            plt.savefig(os.path.join(savepath, f"{evtime}_front.png"))
            
        if camera_type == 'back':
            edgeline = cv2.imread("./shopcart/cart_tempt/edgeline.png")
 
            img_tracking = draw_all_trajectories(vts, edgeline, savepath, camera_type, draw5p=False)
            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(savepath, f"{evtime}.png")
        cv2.imwrite(trajpath, img_cat)
        
        return evtime
    
    return None
    
    
    
    






def track_opt(ShoppingDict, ppath):
    '''
    将一个对象读取，修改其中一个属性
    
    '''
    evtname = ShoppingDict["eventName"]
    
    shopping = copy.deepcopy(ShoppingDict)
    
    ## only need to init item: tracking for each Camera
    shopping["frontCamera"]["tracking"] = []
    shopping["backCamera"]["tracking"] = []

    back_camera = ShoppingDict["backCamera"]["cameraType"]        
    back_yrt = ShoppingDict["backCamera"]["yoloResnetTracker"]
    front_camera = ShoppingDict["frontCamera"]["cameraType"] 
    front_yrt = ShoppingDict["frontCamera"]["yoloResnetTracker"]
    yrts = [(back_camera, back_yrt), (front_camera, front_yrt)]
    
    event_tracks = []
    errtrail = ''
    for camera_type, yrtOut in yrts:        
        '''================= 1. tracking ================='''
        '''(1) 生成用于 tracking 模块的 boxes、feats'''  
        # bboxes = np.empty((0, 6), dtype=np.float64)
        trackerboxes = np.empty((0, 9), dtype=np.float64)
        trackefeats = {}
        for frameDict in yrtOut:            
            tboxes = frameDict["tboxes"]
            ffeats = frameDict["feats"]
            
            # boxes = frameDict["bboxes"]
            # bboxes = np.concatenate((bboxes, np.array(boxes)), axis=0)
            trackerboxes = np.concatenate((trackerboxes, np.array(tboxes)), axis=0)
            for i in range(len(tboxes)):
                fid, bid = int(tboxes[i, 7]), int(tboxes[i, 8])
                trackefeats.update({f"{fid}_{bid}": ffeats[f"{fid}_{bid}"]})

        '''(2) tracking, 后摄'''
        if camera_type == "back":
            vts = doBackTracks(trackerboxes, trackefeats)
            vts.classify()
            shopping["backCamera"]["tracking"] = vts
            
            if len(vts.Residual)!=1:
                errtrail = evtname

        '''(3) tracking, 前摄'''    
        if camera_type == "front":
            vts = doFrontTracks(trackerboxes, trackefeats)
            vts.classify()
            shopping["frontCamera"]["tracking"] = vts
            
            if len(vts.Residual)!=1:
                errtrail = evtname

        event_tracks.append((camera_type, yrtOut, vts))

    
    pckpath = ppath / "track_optim"
    if not pckpath.exists():
        pckpath.mkdir()
        
    fpath = pckpath / "{}_new.pickle".format(evtname)
    with open(str(fpath), 'wb') as f:
        pickle.dump(shopping, f)
     
    
    savepath = ppath / "yolos_tracking" / evtname
        
    illus = [None, None]
    for camera_type, yrtOut, vts in event_tracks:
        if len(vts.tracks)==0: continue
        
        if camera_type == 'front':
            edgeline = cv2.imread("./shopcart/cart_tempt/board_ftmp_line.png")
            
            img_tracking = draw_all_trajectories(vts, edgeline, savepath, camera_type, draw5p=False)
            illus[0] = img_tracking

            plt = plot_frameID_y2(vts)
            plt.savefig(os.path.join(savepath, "front_y2_new.png"))
            
        if camera_type == 'back':
            edgeline = cv2.imread("./shopcart/cart_tempt/edgeline.png")
 
            img_tracking = draw_all_trajectories(vts, edgeline, savepath, camera_type, draw5p=False)
            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(savepath, "trajectory_new.png")
        cv2.imwrite(trajpath, img_cat)
        
    return errtrail
        

def main():
    # evttypes = ["single_event_V10", "single_event_V5", "performence_V10", "performence_V5"]
    evttypes = ["single_event_V10"]
    
    k = 0
    error_trail = []
    for evttype in evttypes:
        ppath = Path("/home/wqg/dataset/pipeline/yrt/{}".format(evttype))
    
        pkpath =  ppath / "shopping_pkl"
        for fp in pkpath.iterdir():
            # fp = pkpath / "{}.pickle".format("20250305-152917-635_6970209860221_6970209860221")
            print(fp)
            
            if fp.suffix != '.pickle': continue
            with open(str(fp), 'rb') as f:
                ShoppingDict = pickle.load(f)
                
            # errtrail = track_opt(ShoppingDict, ppath)
            # error_trail.append(errtrail)
            
            errtrail = get_trail(ShoppingDict, ppath)
            if errtrail is not None:
                error_trail.append(errtrail)
            
            # k+=1
            # if k==100:
            #     break
            
    errfile = ppath / 'error_trail.txt'
    with open(errfile, 'w', encoding='utf-8') as f:
        for line in error_trail:
            f.write(line + '\n')
            
            


if __name__ == "__main__": 
    main()
    
    
    
    
    
    
    