# -*- coding: utf-8 -*-
"""
Created on Fri Aug  9 10:36:45 2024

@author: ym
"""
import os
import cv2
import numpy as np
import torch
import sys
from scipy.spatial.distance import cdist
sys.path.append(r"D:\DetectTracking")
from tracking.trackers.reid.reid_interface import ReIDInterface 
from tracking.trackers.reid.config import config as ReIDConfig
ReIDEncoder = ReIDInterface(ReIDConfig)


def inference_image(images):
    batch_patches = []
    patches = []
    for d, img1 in enumerate(images):


        img = img1[:, :, ::-1].copy()    # the model expects RGB inputs
        patch = ReIDEncoder.transform(img)
        
        # patch = patch.to(device=self.device).half()
        if str(ReIDEncoder.device) != "cpu":
            patch = patch.to(device=ReIDEncoder.device).half()
        else:
            patch = patch.to(device=ReIDEncoder.device)
        
        patches.append(patch)
        if (d + 1) % ReIDEncoder.batch_size == 0:
            patches = torch.stack(patches, dim=0)
            batch_patches.append(patches)
            patches = []

    if len(patches):
        patches = torch.stack(patches, dim=0)
        batch_patches.append(patches)

    features = np.zeros((0, ReIDEncoder.embedding_size))
    for patches in batch_patches:
        pred = ReIDEncoder.model(patches)
        pred[torch.isinf(pred)] = 1.0
        feat = pred.cpu().data.numpy()
        features = np.vstack((features, feat))

    return features


def similarity_compare(root_dir):
    '''
    root_dir：包含 "subimgs"字段的文件夹中图像为 subimg子图
    功能：相邻帧子图间相似度比较
    '''
    
    all_files = []
    extensions = ['.png', '.jpg']
    for dirpath, dirnames, filenames in os.walk(root_dir):
        filepaths = []
        for filename in filenames:
            if os.path.basename(dirpath).find('subimgs') < 0: 
                continue
            file, ext = os.path.splitext(filename)
            if ext in extensions:
                imgpath = os.path.join(dirpath, filename)
                filepaths.append(imgpath)
        nf = len(filepaths)
        if nf==0:
            continue

        fnma = os.path.basename(filepaths[0]).split('.')[0]
        imga = cv2.imread(filepaths[0])
        ha, wa = imga.shape[:2]
        
        for i in range(1, nf):
            fnmb = os.path.basename(filepaths[i]).split('.')[0]
            
            imgb = cv2.imread(filepaths[i])
            hb, wb = imgb.shape[:2]
            
            
            feats = inference_image(((imga, imgb)))
            
            similar = 1 - np.maximum(0.0, cdist(feats, feats, metric='cosine')) 
            
            
            h, w  = max((ha, hb)), max((wa, wb))
            img = np.zeros(((h, 2*w, 3)), np.uint8)
            img[0:ha, 0:wa], img[0:hb, w:(w+wb)] = imga, imgb

            linewidth = max(round(((h+2*w))/2 * 0.001), 2)
            cv2.putText(img,
                        text=f'{similar[0,1]:.2f}',     # Text string to be drawn
                        org=(max(w-20, 10), h-10),       # Bottom-left corner of the text string
                        fontFace=0,                     # Font type
                        fontScale=linewidth/3,          # Font scale factor
                        color=(0, 0, 255),              # Text color
                        thickness=linewidth,            # Thickness of the lines used to draw a text
                        lineType=cv2.LINE_AA,           # Line type
                        )                           
            spath = os.path.join(dirpath, 's'+fnma+'-vs-'+fnmb+'.png')
            cv2.imwrite(spath, img)
            
            
            fnma = os.path.basename(filepaths[i]).split('.')[0]
            imga = imgb.copy()
            ha, wa = imga.shape[:2]
                

                
    return


def main():
    root_dir = r"D:\contrast\dataset\result\20240723-112242_6923790709882"
    
    try:
        similarity_compare(root_dir)   
    except Exception as e:
        print(f'Error: {e}') 
    



if __name__ == '__main__':
    main()