import os
from collections import defaultdict
import cv2
import numpy as np
from argparse import ArgumentParser
from tqdm import tqdm

def analyze_results(base_dir, res_file):
    # 查询图像和测试图像的目录
    imgs_dir_query = os.path.join(base_dir, 'image_query')
    imgs_dir_test = os.path.join(base_dir, 'image_test')

    # 读取结果文件
    file_track2 = open(res_file)
    label_images = defaultdict(list)
    for index, line in enumerate(file_track2.readlines()):
        curLine = line.strip().split(" ")
        label_images[index] = curLine

    row_id = 0
    index = 1
    col = []
    # 遍历每个query标签及其匹配的test图像列表
    for label, images in tqdm(label_images.items()):
        # 读取query图像
        img = cv2.imread(os.path.join(imgs_dir_query, str(int(label) + 1).zfill(6) + '.jpg'))
        # print('label:{}'.format(label))
        # print('path:{}'.format(os.path.join(imgs_dir_query, str(int(label) + 1).zfill(6) + '.jpg')))

        img = cv2.resize(img, (100, 100))
        img = cv2.putText(img, str(int(label) + 1), (10, 10), cv2.FONT_HERSHEY_COMPLEX, 0.5, (0, 0, 255), 2)
        row = img
        # 将query图像与匹配的test图像拼接在一起
        for j in range(10):
            if j < len(images):
                # 读取test图像
                img = cv2.imread(os.path.join(imgs_dir_test, images[j].zfill(6) + '.jpg'))
                print("path_test:{}".format(os.path.join(imgs_dir_test, images[j].zfill(6) + '.jpg')))
                img = cv2.resize(img, (100, 100))
            else:
                img = np.zeros((100, 100, 3), dtype=np.uint8)

            row = np.concatenate((row, img), axis=1)

        col.extend(row)

        # 每100个拼接的图像保存为一张大图
        row_id += 1
        if row_id % 100 == 0:
            col = np.array(col)
            cv2.imwrite(os.path.join(base_dir, 'test_{}.jpg'.format(index)), col)
            index += 1
            col = []
            row_id = 0
    col = np.array(col)
    cv2.imwrite(os.path.join(base_dir, 'test_{}.jpg'.format(index)), col)
    print('over')

if __name__ == '__main__':

    parser = ArgumentParser(description='vis txt result Tool')
    parser.add_argument('--base_dir', default="AIC21/AIC21_Track2_ReID", help='dir to the datasets images')
    parser.add_argument('--store_dir', default="runs/test2", help='dir to store the inference result')
    parser.add_argument('--result', default="track2.txt", help='result file (txt format) path')

    args = parser.parse_args()
    analyze_results(args.base_dir, args.result)
    #analyze_results('../AIC20_track2_reid/AIC20_track2/AIC20_ReID/',
    #                os.path.join('../AIC20_track2_reid/AIC20_track2/AIC20_ReID/track2.txt'))