import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from data import v2, v1
from math import sqrt as sqrt
from itertools import product as product
if torch.cuda.is_available():
    torch.set_default_tensor_type('torch.FloatTensor')
import numpy as np
from ..box_utils import match, log_sum_exp, point_form

class PriorBox(object):
    """Compute priorbox coordinates in center-offset form for each source
    feature map.
    Note:
    This 'layer' has changed between versions of the original SSD
    paper, so we include both versions, but note v2 is the most tested and most
    recent version of the paper.

    """
    def __init__(self, cfg):
        super(PriorBox, self).__init__()
        # self.type = cfg.name
        self.image_size = cfg['min_dim']
        # number of priors for feature map location (either 4 or 6)
        self.num_priors = len(cfg['aspect_ratios'])
        self.variance = cfg['variance'] or [0.1]
        self.feature_maps = cfg['feature_maps']
        self.min_sizes = cfg['min_sizes']
        self.max_sizes = cfg['max_sizes']
        self.steps = cfg['steps']
        self.aspect_ratios = cfg['aspect_ratios']
        self.clip = cfg['clip']
        self.version = cfg['name']
        self.two = cfg['two']
        self.cfg = cfg
        for v in self.variance:
            if v <= 0:
                raise ValueError('Variances must be greater than 0')

    def forward(self):
        mean = []
        lmean = 0
        # TODO merge these
        if self.version == 'v':
            for k, f in enumerate(self.feature_maps):
                for i, j in product(range(f), repeat=2):
                    f_k = self.image_size / self.steps[k]
                    # unit center x,y
                    cx = (j + 0.5) / f_k
                    cy = (i + 0.5) / f_k

                    # aspect_ratio: 1
                    # rel size: min_size
                    s_k = self.min_sizes[k]/self.image_size
                    mean += [cx, cy, s_k, s_k]

                    if self.two:
                        # aspect_ratio: 1
                        # rel size: sqrt(s_k * s_(k+1))
                        s_k_prime = sqrt(s_k * (self.max_sizes[k]/self.image_size))
                        mean += [cx, cy, s_k_prime, s_k_prime]

                    # rest of aspect ratios
                    for ar in self.aspect_ratios[k]:
                        mean += [cx, cy, s_k*sqrt(ar), s_k/sqrt(ar)]
                        mean += [cx, cy, s_k/sqrt(ar), s_k*sqrt(ar)]

                #print(k,f,len(mean)//4)
        else:# self.version == 'v2' or self.version == 'v3':
            for k, f in enumerate(self.feature_maps):
                for i, j in product(range(f), repeat=2):
                    f_k = self.image_size / self.steps[k]
                    # unit center x,y
                    cx = (j + 0.5) / (f_k)
                    cy = (i + 0.5) / (f_k)
                    #print("cxy",cx,cy)
                    # aspect_ratio: 1
                    # rel size: min_size
                    s_k = self.min_sizes[k]/self.image_size
                    mean += [cx, cy, s_k, s_k]

                    # aspect_ratio: 1
                    # rel size: sqrt(s_k * s_(k+1))
                    if self.cfg["sqrt"]:
                        s_k_prime = sqrt(s_k * (self.max_sizes[k]/self.image_size))
                        mean += [cx, cy, s_k_prime, s_k_prime]

                    # rest of aspect ratios
                    for ar in self.aspect_ratios[k]:
                        mean += [cx, cy, s_k*sqrt(ar), s_k/sqrt(ar)]
                        mean += [cx, cy, s_k/sqrt(ar), s_k*sqrt(ar)]
                    #print(k,f,len(mean)//4)
                    
        arr = (np.array(mean).reshape(-1,4)*self.image_size).astype('int32')
        print("prior boxes", arr.shape)
        w,h = arr[:,2], arr[:,3]
        wh = set(sorted(list(zip(w.tolist(), h.tolist()))))
        print("prior_box", wh)
        """else:
            # original version generation of prior (default) boxes
            for i, k in enumerate(self.feature_maps):
                step_x = step_y = self.image_size/k
                for h, w in product(range(k), repeat=2):
                    c_x = ((w+0.5) * step_x)
                    c_y = ((h+0.5) * step_y)
                    c_w = c_h = self.min_sizes[i] / 2
                    s_k = self.image_size  # 300
                    # aspect_ratio: 1,
                    # size: min_size
                    mean += [(c_x-c_w)/s_k, (c_y-c_h)/s_k,
                             (c_x+c_w)/s_k, (c_y+c_h)/s_k]
                    if self.max_sizes[i] > 0:
                        # aspect_ratio: 1
                        # size: sqrt(min_size * max_size)/2
                        c_w = c_h = sqrt(self.min_sizes[i] *
                                         self.max_sizes[i])/2
                        mean += [(c_x-c_w)/s_k, (c_y-c_h)/s_k,
                                 (c_x+c_w)/s_k, (c_y+c_h)/s_k]
                    # rest of prior boxes
                    for ar in self.aspect_ratios[i]:
                        if not (abs(ar-1) < 1e-6):
                            c_w = self.min_sizes[i] * sqrt(ar)/2
                            c_h = self.min_sizes[i] / sqrt(ar)/2
                            mean += [(c_x-c_w)/s_k, (c_y-c_h)/s_k,
                                     (c_x+c_w)/s_k, (c_y+c_h)/s_k]
                                     """
        # back to torch land
        output = torch.Tensor(mean).view(-1, 4)
        priors_point = point_form(output)
        
        b = (priors_point * self.image_size).round().numpy()
        print("self.image_size", self.image_size)
        print("wunique", np.unique(b[:,2] - b[:,0]))
        print("hunique", np.unique(b[:,3] - b[:,1]))

        if self.clip:
            output.clamp_(max=1, min=0)
        return output
