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)