###################################################################
# File Name: box.py
# Author: Zhongdao Wang
# mail: wcd17@mails.tsinghua.edu.cn
# Created Time: Fri Jan 29 15:16:53 2021
###################################################################

import torch
from torchvision import ops

from .basetrack import STrack
from .multitracker import AssociationTracker
from utils.box import scale_box, scale_box_input_size, xywh2xyxy, tlbr_to_tlwh


class BoxAssociationTracker(AssociationTracker):
    def __init__(self, opt):
        super(BoxAssociationTracker, self).__init__(opt)
        self.opt = opt

    def extract_emb(self, img, obs, img_seq): # test
        input_img = img_seq.unsqueeze(0) if self.opt.use_img_seq else img.unsqueeze(0)
        
        if self.opt.use_diff_feat:
            feat_t, feat_d = self.app_model(input_img.to(self.opt.device).float())
            assert(feat_t.shape[-2:] == feat_d.shape[-2:])
            scale = [feat_t.shape[-1]/self.opt.img_size[0],
                     feat_t.shape[-2]/self.opt.img_size[1]]
            obs_feat = scale_box(scale, obs).to(self.opt.device)
            obs_feat = [obs_feat[:, :4], ]
            ret_t = ops.roi_align(feat_t, obs_feat, self.opt.feat_size).detach().cpu()
            ret_d = ops.roi_align(feat_d, obs_feat, self.opt.feat_size).detach().cpu()

            return ret_t, ret_d

        else:
            feat = self.app_model(input_img.to(self.opt.device).float())
            scale = [feat.shape[-1]/self.opt.img_size[0],
                     feat.shape[-2]/self.opt.img_size[1]]
            obs_feat = scale_box(scale, obs).to(self.opt.device)
            obs_feat = [obs_feat[:, :4], ]
            ret = ops.roi_align(feat, obs_feat, self.opt.feat_size).detach().cpu()
            return ret

    def prepare_obs(self, img, img0, obs, img_seq): # test
        if len(obs) > 0:
            obs = torch.from_numpy(obs[obs[:, 4] > self.opt.conf_thres]).float()
            obs = xywh2xyxy(obs)
            obs = scale_box(self.opt.img_size, obs)
            if self.opt.use_diff_feat:
                embs_t, embs_d = self.extract_emb(img, obs, img_seq) # test
            else:
                embs = self.extract_emb(img, obs, img_seq) # test
            obs = scale_box_input_size(self.opt.img_size, obs, img0.shape)

            if obs.shape[1] == 5:
                if self.opt.use_diff_feat:
                    detections_t = [STrack(tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f,
                                  self.buffer_size, use_kalman=self.opt.use_kalman)
                                  for (tlbrs, f) in zip(obs, embs_t)]
                    detections_d = [STrack(tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f,
                                  self.buffer_size, use_kalman=self.opt.use_kalman)
                                  for (tlbrs, f) in zip(obs, embs_d)]

                    return detections_t, detections_d

                else:
                    detections = [STrack(tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f,
                                  self.buffer_size, use_kalman=self.opt.use_kalman)
                                  for (tlbrs, f) in zip(obs, embs)]
            elif obs.shape[1] == 6:
                detections = [STrack(tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f,
                              self.buffer_size, category=tlbrs[5],
                              use_kalman=self.opt.use_kalman)
                              for (tlbrs, f) in zip(obs, embs)]
            else:
                raise ValueError(
                        'Shape of observations should be [n, 5] or [n, 6].')
        else:
            detections = []
        return detections
