from tools.Interface import AiInterface, AiClass
# from Interface import AiInterface, AiClass
from config import cfg
import numpy as np
from PIL import Image
import pdb
from track_ import *
import time

'''
跟踪与对比输入
'''
ai_obj = AiClass()


def deal_similarity_data(message):
    # print('message --- > {}'.format(message))
    car_mac = message['videoIds'].split('_')[-3]
    similar = {'add_flag': True if message['action'] in {'1', '3'} else False,
               'barcode_flag': True if not message['barcode'] == 'null' else False}
    if similar['add_flag'] and similar['barcode_flag']:  # 加购有barcode
        for Id, image_list in zip(message['trackIdList'],
                                  message['images']):
            similar[car_mac + '_' + message['barcode'] + '_' + Id] = image_list
        similar['barcode_list'] = message['barcodeList']
    elif similar['add_flag'] and (not similar['barcode_flag']):  # 加购无barcode
        for Id, image_list in zip(message['trackIdList'],
                                  message['images']):
            similar[car_mac + '_' + Id] = image_list
        similar['barcode_list'] = message['barcodeList']
    elif (not similar['add_flag']) and similar['barcode_flag']:  # 退购有barcode
        for Id, image_list in zip(message['trackIdList'],
                                  message['images']):
            similar[car_mac + '_' + message['barcode'] + '_' + Id] = image_list
        similar['barcode_list'] = message['barcodeList']
    else:  # 退购无barcode
        for Id, image_list in zip(message['trackIdList'],
                                  message['images']):
            similar[car_mac + '_' + Id] = image_list
        similar['barcode_list'] = message['barcodeList']
    similar['sequenceId'] = message['sequenceId']
    return similar


def get_similarity_result(obj: AiInterface, videopth, model, camera_id, message):
    """
    获取相似度结果。

    :param videopth:
    :param obj: AiInterface 对象，用于获取跟踪框和相似度数据。
    :param message: dict, 可选参数，包含跟踪数据和相似度处理结果。
    :return: 相似度结果。
    """
    try:
        track_boxes, features_dict, frame_id_img = run(model, source=videopth)
        allimages, trackIdList = obj.getTrackingBox(track_boxes, features_dict, camera_id, frame_id_img)
        message['trackIdList'] = trackIdList
        message['images'] = allimages
        message = deal_similarity_data(message)
        similarityRes = obj.getSimilarity(model, message)
    except ValueError as ve:
        print('ve >>>> ', ve)
        similarityRes = {'top10': {},
                         'top1': {},
                         'topn': {},
                         'tempLibList': [],
                         'sequenceId': message['sequenceId']
                         }
    return similarityRes


if __name__ == '__main__':
    message = {
        'action': '1',
        'barcode': '084501446314',
        'sequenceId': 'test'
    }
    get_similarity_result(ai_obj, message)