# -*- coding: utf-8 -*-
"""
Created on Mon Oct 14 10:01:24 2024

@author: ym
"""
import numpy as np
import cv2
from scipy.spatial.distance import cdist


class TrackFrag:
    
    def __init__(self, boxes, imgshape=(1280, 1024)):
        self.boxes = boxes
        self.cls = int(boxes[0, 6])
        self.tid = int(boxes[0, 4])
        self.imgshape = imgshape
        
        
        '''轨迹的持续时间，以帧ID表征, （或考虑用时间戳）'''
        self.during = (np.min(boxes[:, 7]), np.max(boxes[:, 7]))

        self.groups = [set(np.unique(boxes[:, 7].astype(int)))]
        
        # '''5个关键点（中心点、左上点、右上点、左下点、右下点 ）坐标'''
        self.isCornpoint = self.is_cornpoint(10)
        self.compute_cornpoints()
        

    def is_cornpoint(self, edge=10):
        
        isleft = min(self.boxes[:, 0]) < edge
        istop = min(self.boxes[:, 1]) < edge
        isright = max(self.boxes[:, 2]) > self.imgshape[0] - edge
        isbottom = max(self.boxes[:, 3]) > self.imgshape[0] - edge

        isCornpoint = isbottom or istop or isleft or isright
        return isCornpoint
    


    def compute_cornpoints(self):
            '''
            cornpoints 共10项，分别是个点的坐标值（x, y）
                (center, top_left, top_right, bottom_left, bottom_right)
            '''
            boxes = self.boxes
            cornpoints = np.zeros((len(boxes), 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]
    
            trajdist = []
            for k in range(5):            
                X = cornpoints[:, 2*k:2*(k+1)]
                trajdist.append(np.max(cdist(X, X)))
            
            idx = trajdist.index(min(trajdist))
            
            self.trajdist_min = trajdist[idx]
            self.cornpoints = cornpoints
            
        

    
        
    def update_groups(self, THRESH=18):
        '''
        对 self.groups 重新赋值
        '''

        boxes = self.boxes
        nbox = len(boxes)
        
        X = np.zeros((len(boxes), 2))
        X[:,0] = (boxes[:, 0] + boxes[:, 2]) / 2
        X[:,1] = (boxes[:, 1] + boxes[:, 3]) / 2
        
        dist2 = cdist(X, X)
        # label = np.zeros(nbox, dtype=np.int)
        
        marked, groups = set(), []        
        for k in range(nbox):
            if k in marked:
                continue 
            group = set()
            
            dt = dist2[k, :]
            idx = np.where(dt < THRESH)[0]

            if len(idx) == 1: 
                groups.append({k})
                marked.add(k)
                continue
            '''初始近邻样本点集合, 并移除当前点'''
            seeds = set(idx) 
            seeds.remove(k)
            
            group.add(k)
            marked.add(k)
            while len(seeds) !=0:
                pt = seeds.pop()
                dt = dist2[pt, :]
                
                seed = set(np.where(dt < THRESH)[0])
                seed.remove(pt)
                
                seed.difference_update(marked)
                seeds.update(seed)
                
                group.add(pt)
                marked.add(pt)
            
            groups.append(group)
        
        self.groups = groups

    
    def jump_boxes(self):
        gpboxes = []    
        for group in self.groups:
            box = self.boxes[list(group), :]    
            gpboxes.append(box)

        return gpboxes
        
    
            
    def is_moving(self):
        if len(self.groups)>=3:
            return True
        
        return False

    
    def is_static(self, THRESH=50):
        box1 = self.boxes[0, :4]
        box2 = self.boxes[-1, :4]
        
        ''' 第1帧、最后一帧 boxes 四个角点间的距离 '''
        ptd = box2 - box1
        ptd1 = np.linalg.norm((ptd[0], ptd[1]))
        ptd2 = np.linalg.norm((ptd[2], ptd[1]))
        ptd3 = np.linalg.norm((ptd[0], ptd[3]))
        ptd4 = np.linalg.norm((ptd[2], ptd[3]))
        condt1 = ptd1<THRESH and ptd2<THRESH and ptd3<THRESH and ptd4<THRESH
        
        if not self.isCornpoint:
            self.trajdist_min < 120
            
        
        # condt2 = self.TrajFeat[3] < 50        
        # condt = condt1 or condt2
        return condt1
    
    
class MoveDetect:
    def __init__(self, bboxes, imgshape=(1280, 1024)):
        self.bboxes = bboxes
        self.shape = imgshape
        
        self.temp = np.zeros(imgshape, np.uint8)
        
        
        self.trackIDs = np.unique(bboxes[:, 4].astype(int))
        # self.frameID = np.unique(bboxes[:, 7].astype(int))
        # self.fnum = len(self.frameID)
        
        self.lboxes = self.array2list() 
        
        self.tracks = [TrackFrag(b) for b in self.lboxes]
        
    def classify(self):

        tracks = self.tracks
        
        '''1. 提取手部轨迹轨迹'''
        hand_tracks = [t for t in tracks if t.cls==0]
        tracks = self.sub_tracks(tracks, hand_tracks)
        
        '''2. 提取静止轨迹'''
        tracks_static = [t for t in tracks if t.is_static()]
        tracks = self.sub_tracks(tracks, tracks_static)

        
        '''3. 更新轨迹点聚类'''
        for track in tracks:
            track.update_groups(18)
            
        self.hand_tracks = hand_tracks    
        self.track_motion = [t for t in tracks if len(t.groups)>=3]
        

    
    def draw(self):
        pass
    
    
    
    def array2list(self):
        '''
            将 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]        
        '''
        track_ids = self.bboxes[:, 4].astype(int)
        lboxes = []
        for t_id in self.trackIDs: 
            # print(f"The ID is: {t_id}")
            idx = np.where(track_ids == t_id)[0]
            box = self.bboxes[idx, :]
            
            lboxes.append(box)
        
        return lboxes

    @staticmethod
    def sub_tracks(tlista, tlistb):
        track_ids_b = {t.tid for t in tlistb}
        return [t for t in tlista if t.tid not in track_ids_b]
    
        
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    