# -*- 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:\yolov5track")
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 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):
    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 = vts.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'
    alltracks = []
    k = 0
    gt = Profile()
    for filename in os.listdir(trackdict):
        filename = 'test_20240402-173935_6920152400975_front_174037379.pkl'
        file, ext = os.path.splitext(filename)
        filepath = os.path.join(trackdict, filename)
        TracksDict = np.load(filepath, allow_pickle=True)

        bboxes = TracksDict['TrackBoxes']
        with gt:
            if filename.find("front") >= 0:
                vts = doFrontTracks(bboxes, TracksDict)
                vts.classify()
                
                save_subimgs(vts, file)

                plt = plot_frameID_y2(vts)
                
                savedir = save_dir.joinpath(f'{file}_y2.png') 
                
                plt.savefig(savedir)  
                plt.close()
            else:
                vts = doBackTracks(bboxes, TracksDict)
                vts.classify()
                alltracks.append(vts)
                
                save_subimgs(vts, file)
                edgeline = cv2.imread("./shopcart/cart_tempt/edgeline.png")
                draw_all_trajectories(vts, edgeline, save_dir, filename)
        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()