# -*- coding: utf-8 -*-
"""
Created on Mon Mar  4 18:38:20 2024

@author: ym
"""
import cv2
import copy
import numpy as np
from pathlib import Path

curpath = Path(__file__).resolve().parents[0]
curpath = Path(curpath)
parpath = curpath.parent
# from tracking.utils.mergetrack import track_equal_track
from .dotracks import doTracks
from .track_front import frontTrack

class doFrontTracks(doTracks):
    def __init__(self, bboxes, frameDictList):
        super().__init__(bboxes, frameDictList)
        
        # self.tracks = [frontTrack(b) for b in self.lboxes]
        self.tracks = [frontTrack(b, f) for b, f in zip(self.lboxes, self.lfeats)]
        
        self.incart = self.getincart()
        
    def getincart(self):
        img = cv2.imread(str(parpath/'shopcart/cart_tempt/incart_ftmp.png'), cv2.IMREAD_GRAYSCALE)
        ret, binary = cv2.threshold(img, 250, 255, cv2.THRESH_BINARY)
        
        return binary
            

    def classify(self):
        '''功能：对 tracks 中元素分类 '''
        
        tracks = self.tracks

        '''提取手的 tracks'''
        hand_tracks = [t for t in tracks if t.cls==0]
        
                
        self.Hands.extend(hand_tracks)
        tracks = self.sub_tracks(tracks, hand_tracks)
        
        
        
        '''提取小孩的 tracks'''
        kid_tracks = [t for t in tracks if t.cls==9]        
        tracks = self.sub_tracks(tracks, kid_tracks)
        
        out_trcak = [t for t in tracks if t.isWholeOutCart]
        tracks = self.sub_tracks(tracks, out_trcak)

        '''静态 tracks'''        
        static_tracks = [t for t in tracks if t.frnum>1 and t.is_static()] 

        
        '''剔除静止目标后的 tracks'''
        tracks = self.sub_tracks(tracks, static_tracks)
        
        tracks_free = [t for t in tracks if t.frnum>1 and t.is_freemove()]
        self.FreeMove.extend(tracks_free)
        tracks = self.sub_tracks(tracks, tracks_free)
        
        # [self.associate_with_hand(htrack, gtrack) for htrack in hand_tracks for gtrack in tracks]
        '''轨迹循环归并'''
        merged_tracks = self.merge_tracks_loop(tracks)
        
        [self.associate_with_hand(htrack, gtrack) for htrack in hand_tracks for gtrack in merged_tracks]
        
        tracks = [t for t in merged_tracks if t.frnum > 1]

        # for gtrack in tracks:
        #     # print(f"Goods ID:{gtrack.tid}")
        #     for htrack in hand_tracks:
        #         hand_ious = self.associate_with_hand(htrack, gtrack)
        #         if len(hand_ious):
        #             gtrack.Hands.append(htrack)
        #             gtrack.HandsIou.append(hand_ious)
                    
        '''静止 tracks 判断与剔除静止 tracks'''
        static_tracks = [t for t in tracks if t.frnum>1 and t.is_static()] 
        tracks = self.sub_tracks(tracks, static_tracks)            

        freemoved_tracks = [t for t in tracks if t.is_free_move()]
        tracks = self.sub_tracks(tracks, freemoved_tracks)
        
        self.Residual = tracks
        self.Confirmed = self.confirm_track()
        
    def confirm_track(self):
        Confirmed = None
        mindist = 0
        for track in self.Residual:
            md = min(track.trajrects_wh)
            if md > mindist:
                mindist = copy.deepcopy(md)
                Confirmed = copy.deepcopy(track)
                
        if Confirmed is not None:
            return [Confirmed]

        return []
  
    def associate_with_hand(self, htrack, gtrack):
        ''' 
        迁移至基类：
            手部 Track、商品 Track 建立关联的依据：
            a. 运动帧的帧索引有交集
            b. 帧索引交集部分iou均大于0
        '''
        assert htrack.cls==0 and gtrack.cls!=0 and gtrack.cls!=9, 'Track cls is Error!'
        
        hboxes = np.empty(shape=(0, 9), dtype = np.float)
        gboxes = np.empty(shape=(0, 9), dtype = np.float)
        
        # start, end 为索引值，需要 start:(end+1)
        for start, end in htrack.dynamic_y2:
            hboxes = np.concatenate((hboxes, htrack.boxes[start:end+1, :]), axis=0)    
        for start, end in gtrack.dynamic_y1:
            gboxes = np.concatenate((gboxes, gtrack.boxes[start:end+1, :]), axis=0)
        
        hfids, gfids = hboxes[:, 7], gboxes[:, 7]
        fids = sorted(set(hfids).intersection(set(gfids)))
        
        if len(fids)==0:
            return None

        # print(f"Goods ID: {gtrack.tid}, Hand ID: {htrack.tid}")
        for f in fids:
            h = np.where(hfids==f)[0][0]
            g = np.where(gfids==f)[0][0]
            
            x11, y11, x12, y12 = hboxes[h, 0:4]
            x21, y21, x22, y22 = gboxes[g, 0:4]
            
            x1, y1 = max((x11, x21)), max((y11, y21))
            x2, y2 = min((x12, x22)), min((y12, y22))
            
            union = (x2 - x1).clip(0) * (y2 - y1).clip(0)
            area1 = (x12 - x11) * (y12 - y11)
            area2 = (x22 - x21) * (y22 - y21)
            
            iou = union / (area1 + area2 - union + 1e-6)
            
            if iou >= 0.01: 
                gtrack.Hands.append((htrack.tid, f, iou))

        return gtrack.Hands
    
    
    
    def merge_tracks(self, Residual):
        """
        对不同id，但可能是同一商品的目标进行归并
        和 dotrack_back.py中函数相同，可以合并至基类
        """ 
        mergedTracks = self.base_merge_tracks(Residual)
        
        oldtracks, newtracks = [], []
        for tracklist in mergedTracks:
            if len(tracklist) > 1:
                boxes = np.empty((0, 9), dtype=np.float32)
                feats = np.empty((0, 256), dtype=np.float32)
                for i, track in enumerate(tracklist):
                    if i==0: ntid, ncls=track.boxes[0, 4], track.boxes[0, 6]
                    iboxes = track.boxes.copy()
                    ifeats = track.features.copy()
                    
                    # iboxes[:, 4], iboxes[:, 6] = ntid, ncls
                    
                    boxes = np.concatenate((boxes, iboxes), axis=0)
                    feats = np.concatenate((feats, ifeats), axis=0)
                    
                    oldtracks.append(track)
                
                fid_indices = np.argsort(boxes[:, 7])
                boxes_fid = boxes[fid_indices]
                feats_fid = feats[fid_indices]
                
                newtracks.append(frontTrack(boxes_fid, feats_fid))
            elif len(tracklist) == 1:
                oldtracks.append(tracklist[0])
                newtracks.append(tracklist[0])
                

        redu = self.sub_tracks(Residual, oldtracks)
        merged = self.join_tracks(redu, newtracks)

        return merged