# -*- coding: utf-8 -*-
"""
Created on Fri Feb 23 11:04:48 2024

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

# TracksDict
def readDict(boxes, TracksDict):
    feats = []
    for i in range(boxes.shape[0]):   
        tid, fid, bid = int(boxes[i, 4]), int(boxes[i, 7]), int(boxes[i, 8])
        
        trackdict = TracksDict[f"frame_{fid}"]
        if "feats" in trackdict:
            feat = trackdict["feats"][bid]
            feats.append(feat)
        
        if "boxes" in trackdict:
            box = trackdict["boxes"][bid]
            assert (box[:4].astype(int) == boxes[i, :4].astype(int)).all(), f"Please check: frame_{fid}"
        
        if "imgs" in trackdict:
            img = trackdict["imgs"][bid]
            cv2.imwrite(f'./data/imgs/{tid}_{fid}_{bid}.png', img)
    
    return np.asarray(feats, dtype=np.float32)



def track_equal_track(atrack, btrack):
    # boxes: [x, y, w, h, track_id, score, cls, frame_index, box_index]
    #         0  1  2  3       4      5     6        7          8
    aboxes = atrack.boxes
    bboxes = btrack.boxes
    
    afeat = atrack.features
    bfeat = btrack.features
    
    # afeat = readDict(aboxes, TracksDict)
    # bfeat = readDict(bboxes, TracksDict)

    
    ''' 1. 判断轨迹在时序上是否有交集 '''
    afids = aboxes[:, 7].astype(np.int_)
    bfids = bboxes[:, 7].astype(np.int_)
    # 帧索引交集
    interfid = set(afids).intersection(set(bfids))
    
    # 或者直接判断帧索引是否有交集，返回 Ture or False
    # interfid = set(afids).isdisjoint(set(bfids))
    
    if len(interfid):
        return False
    
    ''' 2. 轨迹特征相似度判断'''
    feat = np.concatenate((afeat, bfeat), axis=0)
    
    emb_simil = 1 - np.maximum(0.0, cdist(feat, feat, 'cosine'))
    emb_ = 1 - np.maximum(0.0, cdist(np.mean(afeat, axis=0)[None, :], np.mean(bfeat, axis=0)[None, :], 'cosine'))/2
    
    if emb_[0, 0]<0.66:
        return False
    
    ''' 3. 轨迹空间iou'''
    alabel = np.array([0] * afids.size, dtype=np.int_)
    blabel = np.array([1] * bfids.size, dtype=np.int_)
    
    label = np.concatenate((alabel, blabel), axis=0)
    fids = np.concatenate((afids, bfids), axis=0)
    indices = np.argsort(fids)
    idx_pair = []
    for i in range(len(indices)-1):
        idx1, idx2 = indices[i], indices[i+1]
        if label[idx1] != label[idx2] and fids[idx2] - fids[idx1] <= 3:
            if label[idx1] == 0:
                a_idx = idx1
                b_idx = idx2-alabel.size
            else:
                a_idx = idx2
                b_idx = idx1-alabel.size 
            
            idx_pair.append((a_idx, b_idx))
    
    ious = []
    embs = []
    for a, b in idx_pair:
        abox, bbox = aboxes[a, :], bboxes[b, :]
        
        af, bf = afeat[a, :], bfeat[b, :]
        
        emb_ab = 1 - np.maximum(0.0, cdist(af[None, :], bf[None, :], 'cosine'))
        
        
        xa1, ya1 = abox[0] - abox[2]/2, abox[1] - abox[3]/2 
        xa2, ya2 = abox[0] + abox[2]/2, abox[1] + abox[3]/2
        
        xb1, yb1 = bbox[0] - bbox[2]/2, bbox[1] - bbox[3]/2
        xb2, yb2 = bbox[0] + bbox[2]/2, bbox[1] + bbox[3]/2 
        
        
        inter = (np.minimum(xb2, xa2) - np.maximum(xb1, xa1)).clip(0) * \
                     (np.minimum(yb2, ya2) - np.maximum(yb1, ya1)).clip(0)
                     
        # Union Area
        box1_area = abox[2] * abox[3]
        box2_area = bbox[2] * bbox[3]
        union = box1_area + box2_area - inter + 1e-6

        ious.append(inter/union)
        embs.append(emb_ab[0, 0])
    
    ''' 4. 和同一手部关联，如何将该代码和 iou 部分相融合,需进一步完善'''   
    # ahands = np.array(atrack.Hands)
    # bhands = np.array(btrack.Hands)
    # ahids = ahands[:, 0]
    # bhids = bhands[:, 0]
    # interhid = set(ahids).intersection(set(bhids))
    # for hid in interhid:
    #     aidx = ahands[:, 0] == hid
    #     bidx = bhands[:, 0] == hid
        
    #     ahfids = ahids[aidx, 1]
    #     bhfids = bhids[bidx, 1]
    
    
    

    cont = False if len(interfid) else True    # fid 无交集
    cont1 = all(emb > 0.5 for emb in embs)
    cont2 = all(iou > 0.5 for iou in ious)
    # cont = cont and cont2 and cont3
    
    cont = cont and cont1 and cont2

    
    return cont



def track_equal_str(atrack, btrack):
    if atrack == btrack:
        return True
    else:
        return False


def merge_track(Residual):
    out_list = []
    alist = [t for t in Residual]
    while alist:
        atrack = alist[0]
        cur_list = []
        cur_list.append(atrack)
        alist.pop(0)
        
        blist = [b for b in alist]
        alist = []
        for btrack in blist:
            if track_equal_str(atrack, btrack):
                cur_list.append(btrack)
            else:
                alist.append(btrack)
                
        out_list.append(cur_list)
    return out_list
    
def main():
    Residual = ['a', 'b', 'c', 'd', 'a', 'b', 'c', 'b', 'c', 'd']
    out_list = merge_track(Residual)
    
    print(Residual)
    print(out_list)

if __name__ == "__main__": 
    
    main()




# =============================================================================
# for i, atrack in enumerate(input_list):
#     cur_list = []
#     cur_list.append(atrack)
#     del input_list[i]
#     
#     for j, btrack in enumerate(input_list):
#         if track_equal(atrack, btrack):
#             cur_list.append(btrack)
#             del input_list[j]
#     
#     out_list.append(cur_list)
# =============================================================================