# -*- coding: utf-8 -*-
"""
Created on Fri May 31 14:50:21 2024

@author: ym
"""
import cv2
import numpy as np
import torch
from scipy.spatial.distance import cdist
from tracking.trackers.reid.config import config as ReIDConfig
from tracking.trackers.reid.reid_interface import ReIDInterface 
ReIDEncoder = ReIDInterface(ReIDConfig)

def read_data_file(datapath):

    with open(datapath, 'r') as file:  
        lines = file.readlines()
    Videos = []
    FrameBoxes, FrameFeats = [], []
    boxes, feats = [], []
    
    bboxes, ffeats = [], []
    timestamp = []
    t1 = None
    for line in lines: 
        if line.find('CameraId') >= 0: 
            t = int(line.split(',')[1].split(':')[1])
            timestamp.append(t)
     
            if len(boxes) and len(feats):
                FrameBoxes.append(np.array(boxes, dtype = np.float32))
                FrameFeats.append(np.array(feats, dtype = np.float32))
                
                boxes, feats = [], []
            
            if t1 and t - t1 > 1e4:
                Videos.append((FrameBoxes, FrameFeats))    
                FrameBoxes, FrameFeats = [], []
            t1 = int(line.split(',')[1].split(':')[1])
            
        if line.find('box') >= 0:
            box = line.split(':', )[1].split(',')[:-1]
            boxes.append(box)
            bboxes.append(boxes)
            
        
        if line.find('feat') >= 0:
            feat = line.split(':', )[1].split(',')[:-1]
            feats.append(feat)
            ffeats.append(feat)
    
    
    
    
    FrameBoxes.append(np.array(boxes, dtype = np.float32))
    FrameFeats.append(np.array(feats, dtype = np.float32))  
    Videos.append((FrameBoxes, FrameFeats)) 
    
    TimeStamp = np.array(timestamp, dtype = np.float32) 
    DimesDiff = np.diff((timestamp))      

    return Videos

def inference_image(image, detections):
    H, W, _ = np.shape(image)
    imgs = []
    batch_patches = []
    patches = []
    for d in range(np.size(detections, 0)):
        tlbr = detections[d, :4].astype(np.int_)
        tlbr[0] = max(0, tlbr[0])
        tlbr[1] = max(0, tlbr[1])
        tlbr[2] = min(W - 1, tlbr[2])
        tlbr[3] = min(H - 1, tlbr[3])
        img1 = image[tlbr[1]:tlbr[3], tlbr[0]:tlbr[2], :]

        img = img1[:, :, ::-1].copy()    # the model expects RGB inputs
        patch = ReIDEncoder.transform(img)
        
        imgs.append(img1)
        # 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 imgs, features

def test_dog():
    
    
    datapath = r"D:\datasets\ym\Img_ResnetData\dog_224x224\dog_224x224.txt"
    with open(datapath, 'r') as file:  
        lines = file.readlines()
    dlist = lines[0].split(',')
    dfloat = [float(d) for d in dlist]
    
    afeat = np.array(dfloat).reshape(1, -1)
    
    
    
    imgpath = r"D:\datasets\ym\Img_ResnetData\dog_224x224\dog_224x224.jpg"
    image = cv2.imread(imgpath)
    patches = []
    
    img = image[:, :, ::-1].copy()    # the model expects RGB inputs
    patch = ReIDEncoder.transform(img)
    
    patch = patch.to(device=ReIDEncoder.device)
    
    patches.append(patch)
    patches = torch.stack(patches, dim=0)
    pred = ReIDEncoder.model(patches)
    pred[torch.isinf(pred)] = 1.0
    bfeat = pred.cpu().data.numpy()
    
    
    
    
    aafeat = afeat / np.linalg.norm(afeat, ord=2, axis=1, keepdims=True)
    bbfeat = bfeat / np.linalg.norm(bfeat, ord=2, axis=1, keepdims=True)
    
    
    cost_matrix = 1 - np.maximum(0.0, cdist(aafeat, bbfeat, 'cosine'))
    
    
    
    
    print("Done!!!")
    
    
    
    
    


def main():
    imgpath = r"D:\datasets\ym\Img_ResnetData\20240531-103547_0354b1cb-53fa-48de-86cd-ac3c5b127ada_6921168593576\3568800050000_0.jpeg"
    datapath = r"D:\datasets\ym\Img_ResnetData\0_tracker_inout.data"
    savepath = r"D:\datasets\ym\Img_ResnetData\result"
    
    image = cv2.imread(imgpath)
    
    
    Videos = read_data_file(datapath)
    
    bboxes, afeats = Videos[0][0][0],  Videos[0][1][0]
    imgs, bfeats = inference_image(image, bboxes)
    
    
    aafeats = afeats / np.linalg.norm(afeats, ord=2, axis=1, keepdims=True)
    bbfeats = bfeats / np.linalg.norm(bfeats, ord=2, axis=1, keepdims=True)
    
    
    cost_matrix = 1 - np.maximum(0.0, cdist(aafeats, bbfeats, 'cosine'))
    
    
    for i, img in enumerate(imgs):
        cv2.imwrite(savepath + f"\{i}.png", img)
    
    
    
    print("Done!!!!")
    
    
    
    
if __name__ == '__main__':    
    main()
    
    # test_dog()