import sys
import torch

from tools.config import cfg as conf
sys.path.append('contrast')
# from config import config as conf
# from model import resnet18, MobileNetV3_Large
from test_logic import similarity_interface
from img_data import queueImgs_add

# import pymilvus


class datacollection:
    barcode_flag = None
    add_flag = None
    queImgsDict = None
    mainMilvus = None
    tempLibList = None
    model = None
    barcode_list = None
    actionModel = True  # 是否是运行模式， False是测试模式 True是运行模式


class similarityResult:
    top10 = None
    top1 = None
    tempLibList = None
    topn = None


class similarity:
    def __init__(self):
        pass

    def getSimilarity(self, model, dataCollection, similarityRes):
        dataCollection.mainMilvus = model.milvusModel
        dataCollection.model = model.similarityModel
        # try:
        if dataCollection.add_flag:
            if dataCollection.barcode_flag:  # 加购 有barcode -> 输出top10和top1
                similarityRes.top10, similarityRes.top1, similarityRes.tempLibList = similarity_interface(
                    dataCollection)
                print(f"top10: {similarityRes.top10}\ntop1: {similarityRes.top1}")
            else:  # 加购  无barcode -> 输出top10
                similarityRes.top10, similarityRes.tempLibList = similarity_interface(dataCollection)
        else:  # 退购 -> 输出top10和topn
            if dataCollection.barcode_flag:
                similarityRes.top10, similarityRes.top1, similarityRes.topn = similarity_interface(dataCollection)
            else:
                similarityRes.top10, similarityRes.topn = similarity_interface(dataCollection)
        return similarityRes
        # except pymilvus.exceptions.SchemaNotReadyException as SchemaNotReadyException:  ###当前特征库不存在
        #     print('pymilvus.exceptions.SchemaNotReadyException', SchemaNotReadyException)

def main():
    data_collection = datacollection()
    similarityRes = similarityResult()
    data_collection.barcode_flag = queueImgs_add['barcode_flag']
    data_collection.add_flag = queueImgs_add['add_flag']
    data_collection.queImgsDict = queueImgs_add
    similarity().getSimilarity(data_collection, similarityRes)


if __name__ == '__main__':
    main()
