# -*- coding: utf-8 -*-
"""
Created on Wed Jul  5 10:01:11 2023

@author: ym
"""

import numpy as np
import os
import cv2
import sys

from scipy.spatial.distance import cdist
import matplotlib
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
# from ultralytics.utils.plotting import Annotator, colors


from .annotator import TrackAnnotator
# from .processboxes import array2list

# boxes Format: [x1, y1, x2, y2, track_id, score, cls, frame_index]
pth = r"D:/DeepLearning/yolov5/tracking/"
colors = np.array([[255, 255, 255], [255, 255, 255], [255, 255, 255], [255, 255, 255], [255, 255, 255],
                              [0, 0, 255], [0, 255, 0], [255, 51, 255], [102, 178, 255], [51, 153, 255],
                              [255, 153, 153], [255, 102, 102], [255, 51, 51], [153, 255, 153], [102, 255, 102],
                              [51, 255, 51], [255, 102, 255], [153, 204, 255], [255, 0, 0], [255, 255, 255]],
                             dtype=np.uint8)

def array2list(bboxes):
    '''
        将 bboxes 变换为 track 列表
        bboxes: [x1, y1, x2, y2, track_id, score, cls, frame_index, box_index]
    Return：
        lboxes：列表，列表中元素具有同一 track_id，x1y1x2y2 格式
                [x1, y1, x2, y2, track_id, score, cls, frame_index, box_index]        
    '''
    tids = set(bboxes[:, 4])
    track_ids = bboxes[:, 4].astype(int)
    lboxes = []
    for t_id in tids: 
        # print(f"The ID is: {t_id}")
        idx = np.where(track_ids == t_id)[0]
        box = bboxes[idx, :]     

        lboxes.append(box)
    
    return lboxes 



def draw5points(bboxes, file):
    """
    显示中心点、4角点的轨迹，以及轨迹 features
    """
    image = cv2.imread(pth + r"/shopcart/cart_tempt/edgeline.png")
    
    imgx = image.copy()
    annotator = TrackAnnotator(imgx, line_width=2)

    lboxes = array2list(bboxes)
    
    
    for k in range(len(lboxes)):
        boxes = lboxes[k]
        
        cls = int(boxes[0, 6])
        tid = int(boxes[0, 4])
        
        # print(tid)
        
        frnum = boxes.shape[0]
        cornpoints = np.zeros((frnum, 10))

        cornpoints[:,0] = (boxes[:, 0] + boxes[:, 2]) / 2
        cornpoints[:,1] = (boxes[:, 1] + boxes[:, 3]) / 2
        cornpoints[:,2], cornpoints[:,3] = boxes[:, 0], boxes[:, 1]
        cornpoints[:,4], cornpoints[:,5] = boxes[:, 2], boxes[:, 1]
        cornpoints[:,6], cornpoints[:,7] = boxes[:, 0], boxes[:, 3]
        cornpoints[:,8], cornpoints[:,9] = boxes[:, 2], boxes[:, 3]
    
        
        x1, y1, x2, y2 = cornpoints[:,2],cornpoints[:,3],cornpoints[:,8],cornpoints[:,9]
        BoundPixel = 10
        BoundThresh = 0.4
        cont1 = sum(abs(x1)<BoundPixel) / frnum > BoundThresh
        cont2 = sum(abs(y1)<BoundPixel) / frnum > BoundThresh
        cont3 = sum(abs(x2-1024)<BoundPixel) / frnum > BoundThresh
        cont4 = sum(abs(y2-1280)<BoundPixel) / frnum > BoundThresh
        isImgBorder = False
        if cont1 or cont2 or cont3 or cont4:
            isImgBorder = True

# =============================================================================
#         '''情况1: 在商品运动过程中，商品检测框始终左下角点和图像左下角点重合, 用中心点代替''' 
#         lfcn_dist = np.linalg.norm(cornpoints[:, 6:8] - [0, 1280], axis=1)
#         idx1 = lfcn_dist<10
#         if sum(idx1)/frnum > 0.5:
#             cornpoints[:, 6:8] = cornpoints[:, 0:2]
#         
#         '''情况2: 在商品运动过程中，商品检测框始终右下角点和图像右下角点重合, 用中心点代替''' 
#         rtcn_dist = np.linalg.norm(cornpoints[:, 8:10] - [1024, 1280], axis=1)
#         idx2 = rtcn_dist<10
#         if sum(idx2)/frnum > 0.5:
#             cornpoints[:, 8:10] = cornpoints[:, 0:2]
# =============================================================================
        
    
        mwh = (np.mean(boxes[:, 2]) + np.mean(boxes[:, 3]))/2
        
        trajectory = []
        trajlens = []
        trajdist = []
        for k in range(5):
            traj = np.linalg.norm(np.diff(cornpoints[:, 2*k:2*(k+1)], axis = 0), axis=1)
            trajlen = np.sum(traj)
            ptdist = np.max(cdist(cornpoints[:, 2*k:2*(k+1)], cornpoints[:, 2*k:2*(k+1)]))
            
            trajectory.append(traj)
            trajlens.append(trajlen)
            trajdist.append(ptdist)
        
        if not isImgBorder:
            idx = trajlens.index(min(trajlens))
            
            trajmin = trajectory[idx]
            trajlen_min = min(trajlens)
            trajdist_min = min(trajdist)
        else:
            trajmin = trajectory[0]
            trajlen_min = trajlens[0]
            trajdist_min = trajdist[0]
  
        
        '''最小轨迹长度/最大轨迹长度，越小，代表运动幅度越小'''
        trajlen_rate = trajlen_min/(max(trajlens)+0.0001)    
        
        '''最小轨迹欧氏距离/目标框尺度均值'''
        trajdist_rate = trajdist_min/(mwh+0.0001)
        
        
        
        
        # idx = trajlens.index(min(trajlens))
        # trajmin = trajectory[idx]
        
        # '''最小轨迹长度/最大轨迹长度，越小，代表运动幅度越小'''
        # trajlen_rate = min(trajlens)/(max(trajlens)+0.0001)   
        
        # '''最小轨迹欧氏距离，越小，代表运动幅度越小'''
        # trajdist_min = min(trajdist)
        
        # '''最小轨迹欧氏距离 / 目标框尺度均值'''
        # mindist_rate = min(trajdist)/(mwh+0.0001)
        

        
        
        img = image.copy()
        for i in range(boxes.shape[0]):
            cv2.circle(img, (int(cornpoints[i, 0]), int(cornpoints[i, 1])), 6, (255, 255, 255), 2)         
            cv2.circle(img, (int(cornpoints[i, 2]), int(cornpoints[i, 3])), 6, (255, 0, 255), 2)
            cv2.circle(img, (int(cornpoints[i, 4]), int(cornpoints[i, 5])), 6, (0, 255, 0), 2)
            cv2.circle(img, (int(cornpoints[i, 6]), int(cornpoints[i, 7])), 6, (64, 128, 255), 2)
            cv2.circle(img, (int(cornpoints[i, 8]), int(cornpoints[i, 9])), 6, (255, 128, 64), 2)
        
        # if frnum>=3:
        #     cntpoints = cornpoints[:, 0:2].astype(np.int64)
        #     rect = cv2.minAreaRect(cntpoints)
        #     box = cv2.boxPoints(rect)
        #     box = np.int0(box)
        #     cv2.drawContours(img, [box], 0, (255, 0, 0), 2)
        
        # img1 = image.copy()
        # for i in range(boxes.shape[0]-1):
        #     pt1 = cornpoints[i, :].astype(np.int64)
        #     pt2 = cornpoints[i+1, :].astype(np.int64)
        #     cv2.line(img1, pt1, pt2, color=(255, 255, 255), thickness=2)
        # gray = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
        # _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
        # contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            
            
         
        color = ((0, 0, 255), (255, 128, 0))
        label_6 = "PCA(singular_values_)    : "
        label_7 = "Rect    : "
        
        if frnum>=3:
            if isImgBorder:
                X = cornpoints[:, 0:2]
            else:
                X = cornpoints[:, 2*idx:2*(idx+1)]
            
            pca = PCA()
            pca.fit(X)
            label_6 = "PCA(variance_ratio)  : {:.2f}".format(pca.explained_variance_ratio_[0])
            for i, (comp, var) in enumerate(zip(pca.components_, pca.explained_variance_ratio_)):
                pt1 = (pca.mean_ - comp*var*200).astype(np.int64)
                pt2 = (pca.mean_ + comp*var*200).astype(np.int64)
                cv2.line(img, pt1, pt2, color=color[i], thickness=2)
                
            rect = cv2.minAreaRect(X.astype(np.int64))
            box = cv2.boxPoints(rect)
            box = np.int0(box)
            cv2.drawContours(img, [box], 0, (0, 255, 0), 2)
            label_7 = "Rect W&H&Ratio: {}, {}, {:.2f}".format(int(rect[1][0]), int(rect[1][1]), min(rect[1])/(max(rect[1])+0.001))

        
        
        '''撰写专利需要，生成黑白图像'''
        # imgbt = cv2.bitwise_not(img)
        # for i in range(box.shape[0]):
        #     cv2.circle(imgbt, (int(cornpoints[i, 0]), int(cornpoints[i, 1])), 14, (0, 0, 0), 2)         
        #     cv2.drawMarker(imgbt, (int(cornpoints[i, 2]), int(cornpoints[i, 3])), color= (0, 0, 0), markerType=3, markerSize = 30, thickness=2)
        #     cv2.drawMarker(imgbt, (int(cornpoints[i, 4]), int(cornpoints[i, 5])), color= (0, 0, 0), markerType=4, markerSize = 30, thickness=2)
        #     cv2.drawMarker(imgbt, (int(cornpoints[i, 6]), int(cornpoints[i, 7])), color= (0, 0, 0), markerType=5, markerSize = 30, thickness=2)
        #     cv2.drawMarker(imgbt, (int(cornpoints[i, 8]), int(cornpoints[i, 9])), color= (0, 0, 0), markerType=6, markerSize = 30, thickness=2)
        # cv2.imwrite(pth + f"/zhuanli/{file}_{tid}.png", imgbt)

        if len(trajmin):
            trajstd = np.std(trajmin)
        else:
            trajstd = 0
            
                
        
        trajlens = [int(t) for t in trajlens]
        trajdist = [int(t) for t in trajdist]
        label_1 = f"trajlens: {trajlens}, trajlen_min: {int(trajlen_min)}"
        label_2 = f"trajdist: {trajdist}: trajdist_min: {int(trajdist_min)}"
        label_3 = "trajlen_min/max(trajlens): {:.2f}/{} = {:.2f}".format(trajlen_min, max(trajlens), trajlen_rate)
        label_4 = "trajdist_min/mwh         : {:.2f}/{} = {:.2f}".format(trajdist_min, int(mwh), trajdist_rate)
        label_5 = "std(trajmin)             : {:.2f}".format(trajstd)
        
        label = [label_1, label_2, label_3, label_4, label_5, label_6, label_7]
        word = 'abc'
        w, h = cv2.getTextSize('abc', 0, fontScale=2, thickness=1)[0] 
        for i in range(len(label)):
            # color = [int(x) for x in colors[i]]
            cv2.putText(img,
                        label[i], 
                        (20, int(50+(i+1)*1.2*h)),
                        0,
                        1,
                        [int(x) for x in colors[i]],
                        2,
                        lineType=cv2.LINE_AA)
        cv2.imwrite(pth + f"/result/cls11_80212/{file}_{tid}.png", img)


       
        
        
def drawtracks(bboxes, imgshow=None):
    
    """
    Inputs
        bboxes: 原始检测跟踪后的结果，变换为 tboxes
        image：只用于获取图像的（Width, Height）
    Outputs:
        imgshow
    """
    if imgshow is None:
        edgeline = cv2.imread(pth + r"/shopcart/cart_tempt/edgeline.png")
        # edgeline = cv2.bitwise_not(edgeline)
        H, W = edgeline.shape[0:2]
        imgshow= np.zeros((H, W, 3), np.uint8)
    
    if 'edgeline' in locals().keys():
        imgshow = cv2.add(imgshow, edgeline)
    
    ## ==== list，其中元素格式: [x, y, w, h, track_id, score, cls, frame_index]

    tboxes = array2list(bboxes) 
    
    # imgshow = cv2.bitwise_not(imgshow)
    
    annotator = TrackAnnotator(imgshow, line_width=2)
    for boxes in tboxes:
        annotator.plotting_track(boxes)
        
    imgshow = annotator.result()
        
    return imgshow



def writefilename():
    npydir = r"D:\DeepLearning\yolov5\runs\boxes"
    files = [f.split('.')[0] for f in os.listdir(npydir)]
    with open('data.txt', 'w') as f:
        [f.write(f"{file}:\n") for file in files]
    
    print("len(files)")
    # for filename in os.listdir(npydir):
    #     file, ext = os.path.splitext(filename)
    

def main():
    npydir = r"D:\DeepLearning\yolov5\runs\boxes"
    k = 0
    fields = []
    
    for filename in os.listdir(npydir):
        # filename = "加购_快速置入_12.npy"
        print(filename)
        
        file, ext = os.path.splitext(filename)
        filepath = os.path.join(npydir, filename)
        
        try:
            bboxes = np.load(filepath)
            
            imgshow = drawtracks(bboxes, file)
            draw5points(bboxes, file)
            
            cv2.imwrite(pth + f"/result/cls11_80212/{file}_show.png", imgshow)
        except Exception as e:
            # print(str(e))
            pass
        
        # k += 1
        # if k == 1:
        #     break  
            

if __name__ == "__main__": 
    
    main()
    
    # writefilename()
    
        
        

    
    
    





















