# -*- coding: utf-8 -*-
"""
Created on Wed Sep 20 17:33:00 2023

@author: ym
"""

import cv2
import os
import numpy as np
import time
import pickle
import matplotlib.pyplot as plt
import pandas as pd 
from scipy.spatial.distance import cdist

from pathlib import Path
# ================= using for import ultralytics
import sys
sys.path.append(r"D:\DetectTracking")
from utils.gen import Profile

from dotrack.dotracks_back import doBackTracks
from dotrack.dotracks_front import doFrontTracks 
from utils.drawtracks import draw5points, drawTrack, drawtracefeat, plot_frameID_y2, drawFeatures, draw_all_trajectories

from utils.read_data import extract_data_realtime, read_tracking_output_realtime

# from datetime import datetime
# from utils.proBoxes import boxes_add_fid
# from utils.plotting import boxing_img  #, Annotator, colors,
# from utils import Boxes, IterableSimpleNamespace, yaml_load
# from trackers import BOTSORT, BYTETracker
# from utils.mergetrack import track_equal_track
# from utils.basetrack import MoveState, ShoppingCart, doTracks


def detect_start_end(bboxes, features_dict, filename):
    
    boxes = np.empty(shape=(0, 9), dtype = np.float)
    if filename.find("back") >= 0:
        vts = doBackTracks(bboxes, features_dict)
        vtx = [t for t in vts if t.cls != 0]     
        for track in vtx:
            if track.moving_index.size:
                boxes = np.concatenate((boxes, track.moving_index), axis=0)
                
    elif filename.find("front") >= 0:
        vts = doFrontTracks(bboxes, features_dict)
        vtx = [t for t in vts if t.cls != 0]
        
        for track in vtx:
            for start, end in track.dynamic_y2:
                boxes = np.concatenate((boxes, track.boxes[start:end+1, :]), axis=0)    
            for start, end in track.dynamic_y1:
                boxes = np.concatenate((boxes, track.boxes[start:end+1, :]), axis=0)
    
    start = np.min(boxes[:, 0])
    end = np.max(boxes[:, 1])
    
    if start > 5:
        start = start - 5
    else:
        start = 0
    
    return start, end
    
def save_subimgs(vts, file, TracksDict):
    imgdir = Path(f'./result/imgs/{file}')
    if not imgdir.exists():
        imgdir.mkdir(parents=True, exist_ok=True)
    
    for i, track in enumerate(vts.Residual):  
        boxes = track.boxes
        for ii in range(len(boxes)):
            tid, fid, bid = int(boxes[ii, 4]), int(boxes[ii, 7]), int(boxes[ii, 8])
            
            img = TracksDict[f"frame_{fid}"]["imgs"][bid]
            # feat = TracksDict[f"frame_{fid}"]["feats"][bid]
            # box = TracksDict[f"frame_{fid}"]["boxes"][bid]
            
            cv2.imwrite(str(imgdir) + f"/{tid}_{fid}_{bid}.png", img)

def have_tracked():
    # trackdict = r'./data/trackdicts'
    trackdict = r'D:\全实时\202502\20250228-152846-438_6901668934727_6901668934727'
    
    bboxes, alltracks = [], []
    k = 0
    gt = Profile()
    for filename in os.listdir(trackdict):
        # filename = '153112511_0_seek_105.pkl'
        
        file, ext = os.path.splitext(filename)
        filepath = os.path.join(trackdict, filename)
        
        if file.split('_')[0]=='0' or file.find("back") >= 0:
            CamerType = "back"
        if file.split('_')[0]=='1' or file.find("front") >= 0:
            CamerType = "front"
        
        ## 1. 加载 tracker 输出的 pickle 文件
        if ext in ['.pkl', '.pickle']:
            filepath = os.path.join(trackdict, filename)
            TracksDict = np.load(filepath, allow_pickle=True)
            bboxes = TracksDict['TrackBoxes']
        
        ## 2. 加载 data 文件    
        if filename.find('tracker.data')>0:
            bboxes, TracksDict = extract_data_realtime(filepath)
            
        if len(bboxes)==0:
            continue

        with gt:
            if CamerType == "front":
                vts = doFrontTracks(bboxes, TracksDict)
                
                Intrude = vts.isintrude()
                
                vts.classify()
                
                save_subimgs(vts, file, TracksDict)

                plt = plot_frameID_y2(vts)
                
                savedir = save_dir.joinpath(f'{file}_y2.png') 
                
                plt.savefig(savedir)  
                plt.close()
                
                edgeline = cv2.imread("./shopcart/cart_tempt/board_ftmp_line.png")
                img_tracking = draw_all_trajectories(vts, edgeline, save_dir, file, draw5p=True)
                
            if CamerType == "back":
                vts = doBackTracks(bboxes, TracksDict)
                
                Intrude = vts.isintrude()
                
                vts.classify()
                alltracks.append(vts)
                
                save_subimgs(vts, file, TracksDict)
                edgeline = cv2.imread("./shopcart/cart_tempt/edgeline.png")
                img_tracking = draw_all_trajectories(vts, edgeline, save_dir, file)
                
            trackpath = save_dir.joinpath(f'{file}.png')
            cv2.imwrite(str(trackpath), img_tracking)
        print(file+f" need time: {gt.dt:.2f}s")

        k += 1
        if k==1:
            break
    
    if len(alltracks):
        drawFeatures(alltracks, save_dir)
         
    
    
if __name__ == "__main__": 
    # now = datetime.now()
    # time_string = now.strftime("%Y%m%d%H%M%S")[:8]
    
    save_dir = Path('./result/tracks')
    if not save_dir.exists():
        save_dir.mkdir(parents=True, exist_ok=True)

    have_tracked()    
   

    
    
    
    
    
    