from __future__ import division
from math import sqrt as sqrt
from itertools import product as product
import torch
import numpy as np
import cv2

from lib.utils.visualize_utils import TBWriter


def vis(func):
    """tensorboard visualization if has writer as input"""

    def wrapper(*args, **kw):
        return func(*args, **kw) if kw['tb_writer'] is not None else None

    return wrapper


class PriorBoxBase(object):
    """Compute priorbox coordinates in center-offset form for each source
    feature map.
    """

    def __init__(self, cfg):
        super(PriorBoxBase, self).__init__()
        self.image_size = cfg.MODEL.IMAGE_SIZE
        self._steps = cfg.MODEL.STEPS
        self._cfg_list = []
        self._prior_cfg = {}
        self._clip = cfg.MODEL.CLIP
        self._variance = cfg.MODEL.VARIANCE
        for v in self._variance:
            if v <= 0:
                raise ValueError('Variances must be greater than 0')

    def _setup(self, cfg):
        num_feat = len(self._steps)
        for item in self._cfg_list:
            if item not in cfg.MODEL:
                raise Exception("wrong anchor config!")
            if len(cfg.MODEL[item]) != num_feat and len(cfg.MODEL[item]) != 0:
                raise Exception("config {} length does not match step length!".format(item))
            self._prior_cfg[item] = cfg.MODEL[item]

    @property
    def num_priors(self):
        """allow prior num calculation before knowing feature map size"""
        assert self._prior_cfg is not {}
        return [int(len(self._create_prior(0, 0, k)) / 4) for k in range(len(self._steps))]

    def _create_prior(self, cx, cy, k):
        raise NotImplementedError

    @vis
    def _image_proc(self, image=None, tb_writer=None):
        # TODO test with image
        if isinstance(image, type(None)):
            image = np.ones((self.image_size[1], self.image_size[0], 3))
        elif isinstance(image, str):
            image = cv2.imread(image, -1)
        image = cv2.resize(image, (self.image_size[1], self.image_size[0]))
        return image

    @vis
    def _prior_vis(self, anchor, image_ori, feat_idx, tb_writer=None):
        # TODO add output path to the signature
        writer = tb_writer.writer
        prior_num = self.num_priors[feat_idx]

        # transform coordinates
        scale = [self.image_size[1], self.image_size[0], self.image_size[1], self.image_size[0]]
        bboxs = np.array(anchor).reshape((-1, 4))
        box_centers = bboxs[:, :2] * scale[:2]  # [x, y]
        # bboxs: [xmin, ymin, xmax, ymax]
        bboxs = np.hstack((bboxs[:, :2] - bboxs[:, 2:4] / 2, bboxs[:, :2] + bboxs[:, 2:4] / 2)) * scale
        box_centers = box_centers.astype(np.int32)
        bboxs = bboxs.astype(np.int32)
        # visualize each anchor box on a feature map
        for prior_idx in range(prior_num):
            image = image_ori.copy()
            bboxs_ = bboxs[prior_idx::prior_num, :]
            box_centers_ = box_centers[4 * prior_idx::prior_num, :]
            for archor, bbox in zip(box_centers_, bboxs_):
                cv2.circle(image, (archor[0], archor[1]), 1, (0, 0, 255), -1)
                if archor[0] == archor[1]:  # only show diagnal anchors
                    cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 1)
            image = image[..., ::-1]
            image = image.transpose((2,0,1))
            writer.add_image('base/feature_map_{}_{}'.format(feat_idx, prior_idx), image, 2)

    def forward(self, layer_dims, tb_writer=None, image=None):
        priors = []
        image = self._image_proc(image=image, tb_writer=tb_writer)

        for k in range(len(layer_dims)):
            prior = []
            for i, j in product(range(layer_dims[k][0]), range(layer_dims[k][1])):
                steps_x = self.image_size[1] / self._steps[k]
                steps_y = self.image_size[0] / self._steps[k]
                cx = (j + 0.5) / steps_x  # unit center x,y
                cy = (i + 0.5) / steps_y
                prior += self._create_prior(cx, cy, k)
            priors += prior
            self._prior_vis(prior, image, k, tb_writer=tb_writer)

        output = torch.Tensor(priors).view(-1, 4)
        # TODO this clip is meanless, should clip on [xmin, ymin, xmax, ymax]
        if self._clip:
            output.clamp_(max=1, min=0)
        return output


class PriorBoxSSD(PriorBoxBase):
    def __init__(self, cfg):
        super(PriorBoxSSD, self).__init__(cfg)
        # self.image_size = cfg['image_size']
        self._cfg_list = ['MIN_SIZES', 'MAX_SIZES', 'ASPECT_RATIOS']
        self._flip = cfg.MODEL.FLIP
        self._setup(cfg)

    def _create_prior(self, cx, cy, k):
        # as the original paper do
        prior = []
        min_sizes = self._prior_cfg['MIN_SIZES'][k]
        min_sizes = [min_sizes] if not isinstance(min_sizes, list) else min_sizes
        for ms in min_sizes:
            # min square
            s_i = ms / self.image_size[0]
            s_j = ms / self.image_size[1]
            prior += [cx, cy, s_j, s_i]
            # min max square
            if len(self._prior_cfg['MAX_SIZES']) != 0:
                assert type(self._prior_cfg['MAX_SIZES'][k]) is not list  # one max size per layer
                s_i_prime = sqrt(s_i * (self._prior_cfg['MAX_SIZES'][k] / self.image_size[0]))
                s_j_prime = sqrt(s_j * (self._prior_cfg['MAX_SIZES'][k] / self.image_size[1]))
                prior += [cx, cy, s_j_prime, s_i_prime]
            # rectangles by min and aspect ratio
            for ar in self._prior_cfg['ASPECT_RATIOS'][k]:
                prior += [cx, cy, s_j * sqrt(ar), s_i / sqrt(ar)]  # a vertical box
                if self._flip:
                    prior += [cx, cy, s_j / sqrt(ar), s_i * sqrt(ar)]
        return prior


# PriorBox = PriorBoxSSD


def test_no_vis(cfg, tb_writer):
    cfg = copy.deepcopy(cfg)
    cfg['feature_maps'] = [38, 19, 10, 5, 3, 1]
    cfg['min_sizes'] = [[30], [60], 111, 162, 213, 264]
    cfg['flip'] = True
    feat_dim = [list(a) for a in zip(cfg['feature_maps'], cfg['feature_maps'])]
    p = PriorBoxSSD(cfg)
    print(p.num_priors)
    p1 = p.forward(feat_dim)
    print(p1)


def test_filp(cfg, tb_writer):
    cfg = copy.deepcopy(cfg)
    cfg['feature_maps'] = [38, 19, 10, 5, 3, 1]
    cfg['flip'] = True
    feat_dim = [list(a) for a in zip(cfg['feature_maps'], cfg['feature_maps'])]
    p = PriorBoxSSD(cfg)
    p1 = p.forward(feat_dim, tb_writer=tb_writer)

    cfg['flip'] = False
    cfg['aspect_ratios'] = [[2, 1 / 2], [2, 1 / 2, 3, 1 / 3], [2, 1 / 2, 3, 1 / 3],
                            [2, 1 / 2, 3, 1 / 3], [2, 1 / 2], [2, 1 / 2]]
    p = PriorBox(cfg)
    p2 = p.forward(feat_dim, tb_writer=tb_writer)
    # print(p2)
    assert (p2 - p1).sum() < 1e-8


def test_rectangle(cfg, tb_writer):
    cfg = copy.deepcopy(cfg)
    cfg['feature_maps'] = [38, 19, 10, 5, 3, 1]
    cfg['min_sizes'] = [30, 60, 111, 162, 213, 264]
    cfg['flip'] = True
    # feat_dim = [list(a) for a in zip(cfg['feature_maps'], cfg['feature_maps'])]
    # cfg['image_size'] = [300, 300]
    # feat_dim = [list(a) for a in zip(cfg['feature_maps'], [item * 2 for item in cfg['feature_maps']])]
    # cfg['image_size'] = [300, 600]
    feat_dim = [list(a) for a in zip([item * 2 for item in cfg['feature_maps']], cfg['feature_maps'])]
    cfg['image_size'] = [600, 300]
    p = PriorBoxSSD(cfg)
    p1 = p.forward(feat_dim, tb_writer=tb_writer)
    print(p1.shape)


if __name__ == '__main__':
    import copy
    # from lib.datasets.config import ssd_voc_vgg as cfg
    # from lib.utils.visualize_utils import TBWriter
    # tb_writer = TBWriter(log_dir, {'epoch': 50})
    #
    # test_no_vis(cfg, tb_writer)
    # test_filp(cfg, tb_writer)
    # test_rectangle(cfg, tb_writer)
    print('haha')
    from lib.utils.config import cfg
    print(cfg)


[((150, 10, 150, 28), 'copy.deepcopy', 'copy.deepcopy', ({(150, 24, 150, 27): 'cfg'}, {}), '(cfg)', False, 'import copy\n'), ((162, 10, 162, 28), 'copy.deepcopy', 'copy.deepcopy', ({(162, 24, 162, 27): 'cfg'}, {}), '(cfg)', False, 'import copy\n'), ((179, 10, 179, 28), 'copy.deepcopy', 'copy.deepcopy', ({(179, 24, 179, 27): 'cfg'}, {}), '(cfg)', False, 'import copy\n'), ((62, 16, 62, 75), 'cv2.resize', 'cv2.resize', ({(62, 27, 62, 32): 'image', (62, 34, 62, 74): '(self.image_size[1], self.image_size[0])'}, {}), '(image, (self.image_size[1], self.image_size[0]))', False, 'import cv2\n'), ((59, 20, 59, 72), 'numpy.ones', 'np.ones', ({(59, 28, 59, 71): '(self.image_size[1], self.image_size[0], 3)'}, {}), '((self.image_size[1], self.image_size[0], 3))', True, 'import numpy as np\n'), ((76, 16, 76, 95), 'numpy.hstack', 'np.hstack', ({(76, 26, 76, 94): '(bboxs[:, :2] - bboxs[:, 2:4] / 2, bboxs[:, :2] + bboxs[:, 2:4] / 2)'}, {}), '((bboxs[:, :2] - bboxs[:, 2:4] / 2, bboxs[:, :2] + bboxs[:, 2:4] / 2))', True, 'import numpy as np\n'), ((61, 20, 61, 41), 'cv2.imread', 'cv2.imread', ({(61, 31, 61, 36): 'image', (61, 38, 61, 40): '-1'}, {}), '(image, -1)', False, 'import cv2\n'), ((73, 16, 73, 32), 'numpy.array', 'np.array', ({(73, 25, 73, 31): 'anchor'}, {}), '(anchor)', True, 'import numpy as np\n'), ((85, 16, 85, 77), 'cv2.circle', 'cv2.circle', ({(85, 27, 85, 32): 'image', (85, 34, 85, 56): '(archor[0], archor[1])', (85, 58, 85, 59): '(1)', (85, 61, 85, 72): '(0, 0, 255)', (85, 74, 85, 76): '(-1)'}, {}), '(image, (archor[0], archor[1]), 1, (0, 0, 255), -1)', False, 'import cv2\n'), ((107, 17, 107, 37), 'torch.Tensor', 'torch.Tensor', ({(107, 30, 107, 36): 'priors'}, {}), '(priors)', False, 'import torch\n'), ((135, 28, 135, 94), 'math.sqrt', 'sqrt', ({(135, 33, 135, 93): "s_i * (self._prior_cfg['MAX_SIZES'][k] / self.image_size[0])"}, {}), "(s_i * (self._prior_cfg['MAX_SIZES'][k] / self.image_size[0]))", True, 'from math import sqrt as sqrt\n'), ((136, 28, 136, 94), 'math.sqrt', 'sqrt', ({(136, 33, 136, 93): "s_j * (self._prior_cfg['MAX_SIZES'][k] / self.image_size[1])"}, {}), "(s_j * (self._prior_cfg['MAX_SIZES'][k] / self.image_size[1]))", True, 'from math import sqrt as sqrt\n'), ((87, 20, 87, 96), 'cv2.rectangle', 'cv2.rectangle', ({(87, 34, 87, 39): 'image', (87, 41, 87, 59): '(bbox[0], bbox[1])', (87, 61, 87, 79): '(bbox[2], bbox[3])', (87, 81, 87, 92): '(0, 255, 0)', (87, 94, 87, 95): '(1)'}, {}), '(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 1)', False, 'import cv2\n'), ((140, 40, 140, 48), 'math.sqrt', 'sqrt', ({(140, 45, 140, 47): 'ar'}, {}), '(ar)', True, 'from math import sqrt as sqrt\n'), ((140, 56, 140, 64), 'math.sqrt', 'sqrt', ({(140, 61, 140, 63): 'ar'}, {}), '(ar)', True, 'from math import sqrt as sqrt\n'), ((142, 44, 142, 52), 'math.sqrt', 'sqrt', ({(142, 49, 142, 51): 'ar'}, {}), '(ar)', True, 'from math import sqrt as sqrt\n'), ((142, 60, 142, 68), 'math.sqrt', 'sqrt', ({(142, 65, 142, 67): 'ar'}, {}), '(ar)', True, 'from math import sqrt as sqrt\n')]