import numpy as np
import cv2
import os
from utils import linear_mapping, pre_process, random_warp

"""
This module implements the basic correlation filter based tracking algorithm -- MOSSE

Date: 2018-05-28

"""

# flags: 是否手动
# dir_path: datasets/
# sub_dir_path: datasets目录下各个目录
class mosse:
    def __init__(self, args, dir_path, sub_dir_path, flags):
        # get arguments..
        self.args = args
        self.img_path = os.path.join(dir_path, sub_dir_path)
        self.sub_dir_path = sub_dir_path
        self.flags = flags
        # get the img lists...
        self.frame_lists = self._get_img_lists(self.img_path)
        self.frame_lists.sort()

    # start to do the object tracking...
    def start_tracking(self):
        # 添加预测ground-truth.txt
        target_path = './predict/' + self.sub_dir_path + '.txt'
        gfile = open(target_path, "w", encoding='utf-8')

        # get the image of the first frame... (read as gray scale image...)
        init_img = cv2.imread(self.frame_lists[0])
        init_frame = cv2.cvtColor(init_img, cv2.COLOR_BGR2GRAY)
        init_frame = init_frame.astype(np.float32)
        # get the init ground truth.. [x, y, width, height]

        # TODO 使用ROI形式获取第一帧识别框；或者使用ground-truth.txt
        if self.flags:
            init_gt = cv2.selectROI('demo', init_img, False, False)
        else:
            # get the first ground-truth by the manual click
            # get the first ground-truth by the ground-truth.txt
            init_gt = self._get_init_ground_truth(self.img_path)
        print('%s is running...' % self.sub_dir_path)

        init_gt = np.array(init_gt).astype(np.int32)

        # start to draw the gaussian response...
        response_map = self._get_gauss_response(init_frame, init_gt)
        # start to create the training set ...
        # get the goal..
        g = response_map[init_gt[1]:init_gt[1] + init_gt[3], init_gt[0]:init_gt[0] + init_gt[2]]
        fi = init_frame[init_gt[1]:init_gt[1] + init_gt[3], init_gt[0]:init_gt[0] + init_gt[2]]
        # print('g:%s' %g)
        G = np.fft.fft2(g)
        # start to do the pre-training...
        Ai, Bi = self._pre_training(fi, G)
        # start the tracking...
        for idx in range(len(self.frame_lists)):
            current_frame = cv2.imread(self.frame_lists[idx])
            frame_gray = cv2.cvtColor(current_frame, cv2.COLOR_BGR2GRAY)
            frame_gray = frame_gray.astype(np.float32)
            if idx == 0:
                Ai = self.args.lr * Ai
                Bi = self.args.lr * Bi
                pos = init_gt.copy()
                clip_pos = np.array([pos[0], pos[1], pos[0] + pos[2], pos[1] + pos[3]]).astype(np.int64)
            else:
                Hi = Ai / Bi
                fi = frame_gray[clip_pos[1]:clip_pos[3], clip_pos[0]:clip_pos[2]]
                # TODO add condition
                try:
                    fi = pre_process(cv2.resize(fi, (init_gt[2], init_gt[3])))
                except Exception as e:
                    pass
                else:
                    Gi = Hi * np.fft.fft2(fi)
                    gi = linear_mapping(np.fft.ifft2(Gi))
                # find the max pos...
                max_value = np.max(gi)
                max_pos = np.where(gi == max_value)
                dy = int(np.mean(max_pos[0]) - gi.shape[0] / 2)
                dx = int(np.mean(max_pos[1]) - gi.shape[1] / 2)

                # update the position...
                pos[0] = pos[0] + dx
                pos[1] = pos[1] + dy

                # trying to get the clipped position [xmin, ymin, xmax, ymax]
                clip_pos[0] = np.clip(pos[0], 0, current_frame.shape[1])
                clip_pos[1] = np.clip(pos[1], 0, current_frame.shape[0])
                clip_pos[2] = np.clip(pos[0] + pos[2], 0, current_frame.shape[1])
                clip_pos[3] = np.clip(pos[1] + pos[3], 0, current_frame.shape[0])
                clip_pos = clip_pos.astype(np.int64)

                # get the current fi..
                fi = frame_gray[clip_pos[1]:clip_pos[3], clip_pos[0]:clip_pos[2]]

                # TODO add condition
                try:
                    fi = pre_process(cv2.resize(fi, (init_gt[2], init_gt[3])))
                except Exception as e:
                    pass
                # online update...
                else:
                    Ai = self.args.lr * (G * np.conjugate(np.fft.fft2(fi))) + (1 - self.args.lr) * Ai
                    Bi = self.args.lr * (np.fft.fft2(fi) * np.conjugate(np.fft.fft2(fi))) + (1 - self.args.lr) * Bi

            # TODO visualize the tracking process...
            gfile.write("%f,%f,%f,%f," % (pos[0], pos[1], pos[0] + pos[2], pos[1]))
            gfile.write("%f,%f,%f,%f\n" % (pos[0] + pos[2], pos[1] + pos[3], pos[0], pos[1] + pos[3]))

            # show images/videos
            if self.flags:
                cv2.rectangle(current_frame, (pos[0], pos[1]), (pos[0] + pos[2], pos[1] + pos[3]), (255, 0, 0), 2)
                cv2.imshow('demo', current_frame)
                cv2.waitKey(100)

            # if record... save the frames..
            if self.args.record:
                frame_path = 'record_frames/' + self.img_path.split('/')[1] + '/'
                if not os.path.exists(frame_path):
                    os.mkdir(frame_path)
                cv2.imwrite(frame_path + str(idx).zfill(5) + '.png', current_frame)
        gfile.close()

    # pre train the filter on the first frame...
    def _pre_training(self, init_frame, G):
        height, width = G.shape
        fi = cv2.resize(init_frame, (width, height))
        # pre-process img..
        fi = pre_process(fi)
        Ai = G * np.conjugate(np.fft.fft2(fi))
        Bi = np.fft.fft2(init_frame) * np.conjugate(np.fft.fft2(init_frame))
        for _ in range(self.args.num_pretrain):
            if self.args.rotate:
                fi = pre_process(random_warp(init_frame))
            else:
                fi = pre_process(init_frame)
            Ai = Ai + G * np.conjugate(np.fft.fft2(fi))
            Bi = Bi + np.fft.fft2(fi) * np.conjugate(np.fft.fft2(fi))

        return Ai, Bi

    # get the ground-truth gaussian reponse...
    def _get_gauss_response(self, img, gt):
        # get the shape of the image..
        height, width = img.shape
        # get the mesh grid...
        # 每一个像素的组合
        xx, yy = np.meshgrid(np.arange(width), np.arange(height))
        # get the center of the object...
        # gt: [xmin, ymin, width, height]
        center_x = gt[0] + 0.5 * gt[2]
        center_y = gt[1] + 0.5 * gt[3]
        # cal the distance...
        dist = (np.square(xx - center_x) + np.square(yy - center_y)) / (2 * self.args.sigma)
        # get the response map...
        response = np.exp(-dist)
        # normalize...
        response = linear_mapping(response)
        return response

    # it will extract the image list 
    def _get_img_lists(self, img_path):
        frame_list = []
        for frame in os.listdir(img_path):
            if os.path.splitext(frame)[1] == '.jpg':
                frame_list.append(os.path.join(img_path, frame))
        return frame_list

    # it will get the first ground truth of the video..
    def _get_init_ground_truth(self, img_path):
        gt_path = os.path.join(img_path, 'groundtruth.txt')
        with open(gt_path, 'r') as f:
            # just read the first frame...
            line = f.readline()
            gt_pos = line.split(',')

        # TODO ground-truth
        # 左上, 右上, 右下, 左下
        # xmin = float(gt_pos[0])
        xmin = min(float(gt_pos[0]), float(gt_pos[6]))
        # ymin = float(gt_pos[1])
        ymin = min(float(gt_pos[1]), float(gt_pos[3]))

        width = float(gt_pos[4]) - xmin
        height = float(gt_pos[5]) - ymin
        if width < 0:
            width = float(gt_pos[2]) - xmin
        if height < 0:
            height = float(gt_pos[7]) - ymin
        return [xmin, ymin, width, height]

        # return [float(element) for element in gt_pos]
