import jittor as jt 
import numpy as np 
from jittor import nn

def generate_anchor_base(base_size=16, ratios=[0.5, 1, 2],scales=[8, 16, 32],scales_per_octave=None,octave_base_scale=None):
    """Generate anchor base windows by enumerating aspect ratio and scales.
    Generate anchors that are scaled and modified to the given aspect ratios.
    Area of a scaled anchor is preserved when modifying to the given aspect
    ratio.
    :obj:`R = len(ratios) * len(scales)` anchors are generated by this
    function.
    The :obj:`i * len(scales) + j` th anchor corresponds to an anchor
    generated by :obj:`ratios[i]` and :obj:`scales[j]`.
    For example, if the scale is :math:`8` and the ratio is :math:`0.25`,
    the width and the height of the base window will be stretched by :math:`8`.
    For modifying the anchor to the given aspect ratio,
    the height is halved and the width is doubled.
    Args:
        base_size (number): The width and the height of the reference window.
        ratios (list of floats): This is ratios of width to height of
            the anchors.
        scales (list of numbers): This is areas of anchors.
            Those areas will be the product of the square of an element in
            :obj:`scales` and the original area of the reference
            window.
        octave_base_scale (int): The base scale of octave.
        scales_per_octave (int): Number of scales for each octave.
            `octave_base_scale` and `scales_per_octave` are usually used in
            retinanet and the `scales` should be None when they are set.
    Returns:
        ~numpy.ndarray:
        An array of shape :math:`(R, 4)`.
        Each element is a set of coordinates of a bounding box.
        The second axis corresponds to
        :math:`(y_{min}, x_{min}, y_{max}, x_{max})` of a bounding box.
    """
    assert ((octave_base_scale is not None and scales_per_octave is not None) ^ (scales is not None)),'scales and octave_base_scale with scales_per_octave cannot be set at the same time'
    
    if scales is None:
        octave_scales = np.array(
                [2**(i / scales_per_octave) for i in range(scales_per_octave)])
        scales = octave_scales * octave_base_scale

    py = base_size / 2.
    px = base_size / 2.

    anchor_base = np.zeros((len(ratios) * len(scales), 4),
                           dtype=np.float32)
    for i in range(len(ratios)):
        for j in range(len(scales)):
            h = base_size * scales[j] * np.sqrt(ratios[i])
            w = base_size * scales[j] * np.sqrt(1. / ratios[i])

            index = i * len(scales) + j
            anchor_base[index, 1] = py - h / 2.
            anchor_base[index, 0] = px - w / 2.
            anchor_base[index, 3] = py + h / 2.
            anchor_base[index, 2] = px + w / 2.
    return anchor_base

def grid_anchors(anchor_base, feat_stride, feat_size):
    # Enumerate all shifted anchors:
    #
    # add A anchors (1, A, 4) to
    # cell K shifts (K, 1, 4) to get
    # shift anchors (K, A, 4)
    # reshape to (K*A, 4) shifted anchors
    # return (K*A, 4)
    
    height, width = feat_size
    shift_y = np.arange(0, height * feat_stride, feat_stride)
    shift_x = np.arange(0, width * feat_stride, feat_stride)
    shift_x, shift_y = np.meshgrid(shift_x, shift_y)
    shift = np.stack((shift_x.ravel(), shift_y.ravel(),
                      shift_x.ravel(), shift_y.ravel()), axis=1)

    A = anchor_base.shape[0]
    K = shift.shape[0]
    anchor = anchor_base.reshape((1, A, 4)) + \
             shift.reshape((1, K, 4)).transpose((1, 0, 2))
    anchor = anchor.reshape((K * A, 4)).astype(np.float32)
    return anchor

def loc2bbox(src_bbox,loc,mean=[0.,0.,0.,0.],std=[1.,1.,1.,1.]):
    if src_bbox.shape[0] == 0:
        return jt.zeros((0, 4), dtype=loc.dtype)
    

    mean = jt.array(mean)
    std = jt.array(std)
    loc = loc*std+mean 

    src_width = src_bbox[:, 2:3] - src_bbox[:, 0:1]
    src_height = src_bbox[:, 3:4] - src_bbox[:, 1:2]
    src_center_x = src_bbox[:, 0:1] + 0.5 * src_width
    src_center_y = src_bbox[:, 1:2] + 0.5 * src_height

    dx = loc[:, 0:1]
    dy = loc[:, 1:2]
    dw = loc[:, 2:3]
    dh = loc[:, 3:4]

    center_x = dx*src_width+src_center_x
    center_y = dy*src_height+src_center_y
        
    w = jt.exp(dw) * src_width
    h = jt.exp(dh) * src_height
        
    x1,y1,x2,y2 = center_x-0.5*w, center_y-0.5*h, center_x+0.5*w, center_y+0.5*h
        
    dst_bbox = jt.contrib.concat([x1,y1,x2,y2],dim=1)

    return dst_bbox

def loc2bbox_r(src_bbox,loc,mean=[0.,0.,0.,0.,0.],std=[1.,1.,1.,1.,1.]):
    if src_bbox.shape[0] == 0:
        return jt.zeros((0, 4), dtype=loc.dtype)
    

    mean = jt.array(mean)
    std = jt.array(std)
    loc = loc*std+mean 

    src_center_x = src_bbox[:, 0:1]
    src_center_y = src_bbox[:, 1:2]
    src_width = src_bbox[:, 2:3]
    src_height = src_bbox[:, 3:4]

    dx = loc[:, 0:1]
    dy = loc[:, 1:2]
    dw = loc[:, 2:3]
    dh = loc[:, 3:4]

    center_x = dx*src_width+src_center_x
    center_y = dy*src_height+src_center_y
        
    w = jt.exp(dw) * src_width
    h = jt.exp(dh) * src_height
        
    x1,y1,x2,y2 = center_x-0.5*w, center_y-0.5*h, center_x+0.5*w, center_y+0.5*h
    theta = loc[:, 4:5] + src_bbox[:, 4:5]
    dst_bbox = jt.contrib.concat([(x1+x2)/2,(y1+y2)/2,x2-x1,y2-y1,theta],dim=1)
    return dst_bbox
    
def bbox2loc_r(src_bbox,dst_bbox,mean=[0.,0.,0.,0.,0.],std=[1.,1.,1.,1.,1.]):        
    center_x, center_y, width, height = src_bbox[:, 0:1], src_bbox[:, 1:2], src_bbox[:, 2:3], src_bbox[:, 3:4]
    base_center_x, base_center_y, base_width, base_height = dst_bbox[:, 0:1], dst_bbox[:, 1:2], dst_bbox[:, 2:3], dst_bbox[:, 3:4]

    eps = 1e-5

    dx = (base_center_x - center_x) / (width + 1)
    dy = (base_center_y - center_y) / (height + 1)

    dw = jt.log(base_width / (width + 1) + eps)
    dh = jt.log(base_height / (height + 1) + eps)

    da = dst_bbox[:, 4:5] - src_bbox[:, 4:5]
        
    loc = jt.contrib.concat([dx,dy,dw,dh,da],dim=1)

    mean = jt.array(mean)
    std = jt.array(std)
    loc = (loc-mean)/std 

    return loc
    
def bbox2loc(src_bbox,dst_bbox,mean=[0.,0.,0.,0.],std=[1.,1.,1.,1.]):        
    width = src_bbox[:, 2:3] - src_bbox[:, 0:1]
    height = src_bbox[:, 3:4] - src_bbox[:, 1:2]
    center_x = src_bbox[:, 0:1] + 0.5 * width
    center_y = src_bbox[:, 1:2] + 0.5 * height

    base_width = dst_bbox[:, 2:3] - dst_bbox[:, 0:1]
    base_height = dst_bbox[:, 3:4] - dst_bbox[:, 1:2]
    base_center_x = dst_bbox[:, 0:1] + 0.5 * base_width
    base_center_y = dst_bbox[:, 1:2] + 0.5 * base_height

    eps = 1e-5
    height = jt.maximum(height, eps)
    width = jt.maximum(width, eps)

    dy = (base_center_y - center_y) / height
    dx = (base_center_x - center_x) / width

    dw = jt.safe_log(base_width / width)
    dh = jt.safe_log(base_height / height)
        
    loc = jt.contrib.concat([dx,dy,dw,dh],dim=1)

    mean = jt.array(mean)
    std = jt.array(std)
    loc = (loc-mean)/std 

    return loc
    
def bbox_iou(bbox_a, bbox_b):
    assert bbox_a.shape[1]==4 and bbox_b.shape[1]==4
    if bbox_a.numel()==0 or bbox_b.numel()==0:
        return jt.zeros((bbox_a.shape[0],bbox_b.shape[0]))
    # top left
    tl = jt.maximum(bbox_a[:, :2].unsqueeze(1), bbox_b[:, :2])
    # bottom right
    br = jt.minimum(bbox_a[:,2:].unsqueeze(1), bbox_b[:, 2:])

    area_i = jt.prod(br - tl, dim=2) * (tl < br).all(dim=2)
    area_a = jt.prod(bbox_a[:, 2:] - bbox_a[:, :2], dim=1)
    area_b = jt.prod(bbox_b[:, 2:] - bbox_b[:, :2], dim=1)
    return area_i / (area_a.unsqueeze(1) + area_b - area_i)

def _unmap(data, count, index, fill=0):
    # Unmap a subset of item (data) back to the original set of items (of
    # size count)
    ret_shape = list(data.shape)
    ret_shape[0]=count
    ret = jt.ones(ret_shape,dtype=data.dtype)*fill
    ret[index]=data
    return ret


class AnchorTargetCreator(nn.Module):
    """Assign the ground truth bounding boxes to anchors.
    
    Args:
        n_sample (int): The number of regions to produce.
        pos_iou_thresh (float): Anchors with IoU above this
            threshold will be assigned as positive.
        neg_iou_thresh (float): Anchors with IoU below this
            threshold will be assigned as negative.
        pos_ratio (float): Ratio of positive regions in the
            sampled regions.
    """

    def __init__(self,
                 n_sample=256,
                 pos_iou_thresh=0.7, 
                 neg_iou_thresh=0.3,
                 pos_ratio=0.5):
        super(AnchorTargetCreator,self).__init__()
        self.n_sample = n_sample
        self.pos_iou_thresh = pos_iou_thresh
        self.neg_iou_thresh = neg_iou_thresh
        self.pos_ratio = pos_ratio

    def execute(self, bbox, anchor, img_size):
        """Assign ground truth supervision to sampled subset of anchors.
        Types of input arrays and output arrays are same.
        Here are notations.
        * :math:`S` is the number of anchors.
        * :math:`R` is the number of bounding boxes.
        Args:
            bbox (array): Coordinates of bounding boxes. Its shape is
                :math:`(R, 4)`.
            anchor (array): Coordinates of anchors. Its shape is
                :math:`(S, 4)`.
            img_size (tuple of ints): A tuple :obj:`H, W`, which
                is a tuple of height and width of an image.
        Returns:
            (array, array):
            #NOTE: it's scale not only  offset
            * **loc**: Offsets and scales to match the anchors to \
                the ground truth bounding boxes. Its shape is :math:`(S, 4)`.
            * **label**: Labels of anchors with values \
                :obj:`(1=positive, 0=negative, -1=ignore)`. Its shape \
                is :math:`(S,)`.
        """

        img_W, img_H = img_size

        n_anchor = len(anchor)
        inside_index = jt.where(
                (anchor[:, 0] >= 0) &
                (anchor[:, 1] >= 0) &
                (anchor[:, 2] <= img_W) &
                (anchor[:, 3] <= img_H)
                )[0]
        anchor = anchor[inside_index]
        argmax_ious, label = self._create_label(anchor, bbox)

        # compute bounding box regression targets
        loc = bbox2loc(anchor, bbox[argmax_ious])

        # map up to original set of anchors
        label = _unmap(label, n_anchor, inside_index, fill=-1)
        loc = _unmap(loc, n_anchor, inside_index, fill=0)

        return loc, label

    def _create_label(self, anchor, bbox):
        # label: 1 is positive, 0 is negative, -1 is dont care
        label = -jt.ones((anchor.shape[0],)).int32()

        argmax_ious, max_ious, gt_argmax_ious = self._calc_ious(anchor, bbox)

        # assign negative labels first so that positive labels can clobber them
        label[max_ious < self.neg_iou_thresh] = 0

        # positive label: for each gt, anchor with highest iou
        label[gt_argmax_ious] = 1

        # positive label: above threshold IOU
        label[max_ious >= self.pos_iou_thresh] = 1
        
        # subsample positive labels if we have too many
        n_pos = int(self.pos_ratio * self.n_sample)
        pos_index = jt.where(label == 1)[0]
        if len(pos_index) > n_pos:
            tmp_index = np.arange(0,pos_index.shape[0])
            np.random.shuffle(tmp_index)
            disable_index = tmp_index[:pos_index.shape[0] - n_pos]
            disable_index = pos_index[disable_index]
            label[disable_index] = -1

        # subsample negative labels if we have too many
        n_neg = self.n_sample - jt.sum(label == 1).item() if label.numel()>0 else 0
        neg_index = jt.where(label == 0)[0]
        if len(neg_index) > n_neg:
            tmp_index = np.arange(0,neg_index.shape[0])
            np.random.shuffle(tmp_index)
            disable_index = tmp_index[:neg_index.shape[0] - n_neg]
            disable_index = neg_index[disable_index]
            label[disable_index] = -1
        return argmax_ious, label

    def _calc_ious(self, anchor, bbox):
        # ious between the anchors and the gt boxes
        ious = bbox_iou(anchor, bbox)
        argmax_ious,max_ious = ious.argmax(dim=1)
        gt_argmax_ious,gt_max_ious = ious.argmax(dim=0)
        gt_argmax_ious = jt.where(ious == gt_max_ious)[0]

        return argmax_ious, max_ious, gt_argmax_ious


    
class ProposalCreator(nn.Module):
    """Proposal regions are generated by calling this object.
    Args:
        nms_thresh (float): Threshold value used when calling NMS.
        n_train_pre_nms (int): Number of top scored bounding boxes
            to keep before passing to NMS in train mode.
        n_train_post_nms (int): Number of top scored bounding boxes
            to keep after passing to NMS in train mode.
        n_test_pre_nms (int): Number of top scored bounding boxes
            to keep before passing to NMS in test mode.
        n_test_post_nms (int): Number of top scored bounding boxes
            to keep after passing to NMS in test mode.
        force_cpu_nms (bool): If this is :obj:`True`,
            always use NMS in CPU mode. If :obj:`False`,
            the NMS mode is selected based on the type of inputs.
        min_size (int): A paramter to determine the threshold on
            discarding bounding boxes based on their sizes.
    """

    def __init__(self,
                 nms_thresh=0.7,
                 n_train_pre_nms=12000,
                 n_train_post_nms=2000,
                 n_test_pre_nms=6000,
                 n_test_post_nms=300,
                 min_size=16
                 ):
        super(ProposalCreator,self).__init__()
        self.nms_thresh = nms_thresh
        self.n_train_pre_nms = n_train_pre_nms
        self.n_train_post_nms = n_train_post_nms
        self.n_test_pre_nms = n_test_pre_nms
        self.n_test_post_nms = n_test_post_nms
        self.min_size = min_size

    def execute(self, loc, score,anchor, img_size):
        """input should  be ndarray
        Propose RoIs.
        Inputs :obj:`loc, score, anchor` refer to the same anchor when indexed
        by the same index.
        On notations, :math:`R` is the total number of anchors. This is equal
        to product of the height and the width of an image and the number of
        anchor bases per pixel.
        Type of the output is same as the inputs.
        Args:
            loc (array): Predicted offsets and scaling to anchors.
                Its shape is :math:`(R, 4)`.
            score (array): Predicted foreground probability for anchors.
                Its shape is :math:`(R,)`.
            anchor (array): Coordinates of anchors. Its shape is
                :math:`(R, 4)`.
            img_size (tuple of ints): A tuple :obj:`height, width`,
                which contains image size after scaling.
            scale (float): The scaling factor used to scale an image after
                reading it from a file.
        Returns:
            array:
            An array of coordinates of proposal boxes.
            Its shape is :math:`(S, 4)`. :math:`S` is less than
            :obj:`self.n_test_post_nms` in test time and less than
            :obj:`self.n_train_post_nms` in train time. :math:`S` depends on
            the size of the predicted bounding boxes and the number of
            bounding boxes discarded by NMS.
        """
        # NOTE: when test, remember
        if self.is_training():
            n_pre_nms = self.n_train_pre_nms
            n_post_nms = self.n_train_post_nms
        else:
            n_pre_nms = self.n_test_pre_nms
            n_post_nms = self.n_test_post_nms

        # Convert anchors into proposal via bbox transformations.
        roi = loc2bbox(anchor, loc)

        # Clip predicted boxes to image.
        roi[:,0] = jt.clamp(roi[:,0],min_v=0,max_v=img_size[0])
        roi[:,2] = jt.clamp(roi[:,2],min_v=0,max_v=img_size[0])
        
        roi[:,1] = jt.clamp(roi[:,1],min_v=0,max_v=img_size[1])
        roi[:,3] = jt.clamp(roi[:,3],min_v=0,max_v=img_size[1])

        # Remove predicted boxes with either height or width < threshold.
        min_size = self.min_size
        ws = roi[:, 2] - roi[:, 0]
        hs = roi[:, 3] - roi[:, 1]
        keep = jt.where((hs >= min_size) & (ws >= min_size))[0]
        roi = roi[keep, :]
        score = score[keep]

        # Sort all (proposal, score) pairs by score from highest to lowest.
        # Take top pre_nms_topN (e.g. 6000).
        order,_ = jt.argsort(score, descending=True)
        if n_pre_nms > 0:
            order = order[:n_pre_nms]
        roi = roi[order, :]
        score = score[order]

        # Apply nms (e.g. threshold = 0.7).
        # Take after_nms_topN (e.g. 300).
        
        dets = jt.contrib.concat([roi,score.unsqueeze(1)],dim=1)
        keep = jt.nms(dets,self.nms_thresh)
        if n_post_nms > 0:
            keep = keep[:n_post_nms]
        roi = roi[keep]
        return roi



class ProposalTargetCreator(nn.Module):
    """Assign ground truth bounding boxes to given RoIs.
    Args:
        n_sample (int): The number of sampled regions.
        pos_ratio (float): Fraction of regions that is labeled as a
            foreground.
        pos_iou_thresh (float): IoU threshold for a RoI to be considered as a
            foreground.
        neg_iou_thresh_hi (float): RoI is considered to be the background
            if IoU is in
            [:obj:`neg_iou_thresh_hi`, :obj:`neg_iou_thresh_hi`).
        neg_iou_thresh_lo (float): See above.
    """

    def __init__(self,
                 n_sample=128,
                 pos_ratio=0.25, 
                 pos_iou_thresh=0.5,
                 neg_iou_thresh_hi=0.5, 
                 neg_iou_thresh_lo=0.0
                 ):
        super(ProposalTargetCreator,self).__init__()
        self.n_sample = n_sample
        self.pos_ratio = pos_ratio
        self.pos_iou_thresh = pos_iou_thresh
        self.neg_iou_thresh_hi = neg_iou_thresh_hi
        self.neg_iou_thresh_lo = neg_iou_thresh_lo  # NOTE:default 0.1 in py-faster-rcnn
        

    def execute(self, roi, bbox, label):
        """Assigns ground truth to sampled proposals.
        This function samples total of :obj:`self.n_sample` RoIs
        from the combination of :obj:`roi` and :obj:`bbox`.
        The RoIs are assigned with the ground truth class labels as well as
        bounding box offsets and scales to match the ground truth bounding
        boxes. As many as :obj:`pos_ratio * self.n_sample` RoIs are
        sampled as foregrounds.
        Offsets and scales of bounding boxes are calculated using
        :func:`model.utils.bbox_tools.bbox2loc`.
        Also, types of input arrays and output arrays are same.
        Here are notations.
        * :math:`S` is the total number of sampled RoIs, which equals \
            :obj:`self.n_sample`.
        * :math:`L` is number of object classes possibly including the \
            background.
        Args:
            roi (array): Region of Interests (RoIs) from which we sample.
                Its shape is :math:`(R, 4)`
            bbox (array): The coordinates of ground truth bounding boxes.
                Its shape is :math:`(R', 4)`.
            label (array): Ground truth bounding box labels. Its shape
                is :math:`(R',)`. Its range is :math:`[0, L - 1]`, where
                :math:`L` is the number of foreground classes.
        Returns:
            (array, array, array):
            * **sample_roi**: Regions of interests that are sampled. \
                Its shape is :math:`(S, 4)`.
            * **gt_roi_loc**: Offsets and scales to match \
                the sampled RoIs to the ground truth bounding boxes. \
                Its shape is :math:`(S, 4)`.
            * **gt_roi_label**: Labels assigned to sampled RoIs. Its shape is \
                :math:`(S,)`. Its range is :math:`[0, L]`. The label with \
                value 0 is the background.
        """
        pos_roi_per_image = np.round(self.n_sample * self.pos_ratio)
        iou = bbox_iou(roi, bbox)
        gt_assignment,max_iou = iou.argmax(dim=1)
        # Offset range of classes from [0, n_fg_class - 1] to [1, n_fg_class].
        # The label with value 0 is the background.
        gt_roi_label = label[gt_assignment]

        # Select foreground RoIs as those with >= pos_iou_thresh IoU.
        pos_index = jt.where(max_iou >= self.pos_iou_thresh)[0]
        pos_roi_per_this_image = int(min(pos_roi_per_image, pos_index.shape[0]))
        if pos_index.shape[0] > 0:
            tmp_indexes = np.arange(0,pos_index.shape[0])
            np.random.shuffle(tmp_indexes)
            tmp_indexes = tmp_indexes[:pos_roi_per_this_image]
            pos_index = pos_index[tmp_indexes]

        # Select background RoIs as those within
        # [neg_iou_thresh_lo, neg_iou_thresh_hi).
        neg_index = jt.where((max_iou < self.neg_iou_thresh_hi) &
                             (max_iou >= self.neg_iou_thresh_lo))[0]
        neg_roi_per_this_image = self.n_sample - pos_roi_per_this_image
        neg_roi_per_this_image = int(min(neg_roi_per_this_image,
                                         neg_index.shape[0]))
        if neg_index.shape[0] > 0:
            tmp_indexes = np.arange(0,neg_index.shape[0])
            np.random.shuffle(tmp_indexes)
            tmp_indexes = tmp_indexes[:neg_roi_per_this_image]
            neg_index = neg_index[tmp_indexes]
        

        # The indices that we're selecting (both positive and negative).
        keep_index = jt.contrib.concat((pos_index, neg_index),dim=0)
        gt_roi_label = gt_roi_label[keep_index]
        gt_roi_label[pos_roi_per_this_image:] = 0  # negative labels --> 0
        sample_roi = roi[keep_index]

        # Compute offsets and scales to match sampled RoIs to the GTs.
        gt_roi_loc = bbox2loc(sample_roi, bbox[gt_assignment[keep_index]])

        return sample_roi, gt_roi_loc, gt_roi_label



# TODO: delete this class using above functions.
class AnchorGenerator(object):

    def __init__(self, base_size, scales, ratios, scale_major=True, ctr=None):
        self.base_size = base_size
        self.scales = jt.array(scales).float32()
        self.ratios = jt.array(ratios).float32()
        self.scale_major = scale_major
        self.ctr = ctr
        self.base_anchors = self.gen_base_anchors()

    @property
    def num_base_anchors(self):
        return self.base_anchors.size(0)

    def gen_base_anchors(self):
        w = self.base_size
        h = self.base_size
        if self.ctr is None:
            x_ctr = 0.5 * (w - 1)
            y_ctr = 0.5 * (h - 1)
        else:
            x_ctr, y_ctr = self.ctr

        h_ratios = jt.sqrt(self.ratios)
        w_ratios = 1 / h_ratios
        if self.scale_major:
            ws = (w * w_ratios[:, None] * self.scales[None, :]).view(-1)
            hs = (h * h_ratios[:, None] * self.scales[None, :]).view(-1)
        else:
            ws = (w * self.scales[:, None] * w_ratios[None, :]).view(-1)
            hs = (h * self.scales[:, None] * h_ratios[None, :]).view(-1)

        base_anchors = jt.stack(
            [
                x_ctr - 0.5 * (ws - 1), y_ctr - 0.5 * (hs - 1),
                x_ctr + 0.5 * (ws - 1), y_ctr + 0.5 * (hs - 1)
            ],
            dim=-1).round()

        return base_anchors.float32()

    def _meshgrid(self, x, y, row_major=True):
        xx = x.repeat(len(y))
        yy = y.view(-1, 1).repeat(1, len(x)).view(-1)
        if row_major:
            return xx, yy
        else:
            return yy, xx

    def grid_anchors(self, featmap_size, stride=16):
        base_anchors = self.base_anchors

        feat_h, feat_w = featmap_size
        shift_x = jt.arange(0, feat_w) * stride
        shift_y = jt.arange(0, feat_h) * stride
        shift_xx, shift_yy = self._meshgrid(shift_x, shift_y)
        shifts = jt.stack([shift_xx, shift_yy, shift_xx, shift_yy], dim=-1)
        shifts = shifts.type_as(base_anchors)
        # first feat_w elements correspond to the first row of shifts
        # add A anchors (1, A, 4) to K shifts (K, 1, 4) to get
        # shifted anchors (K, A, 4), reshape to (K*A, 4)

        all_anchors = base_anchors[None, :, :] + shifts[:, None, :]
        all_anchors = all_anchors.view(-1, 4)
        # first A rows correspond to A anchors of (0, 0) in feature map,
        # then (0, 1), (0, 2), ...
        return all_anchors

    def valid_flags(self, featmap_size, valid_size):
        feat_h, feat_w = featmap_size
        valid_h, valid_w = valid_size
        assert valid_h <= feat_h and valid_w <= feat_w
        valid_x = jt.zeros(feat_w, dtype=jt.bool)
        valid_y = jt.zeros(feat_h, dtype=jt.bool)
        valid_x[:valid_w] = 1
        valid_y[:valid_h] = 1
        valid_xx, valid_yy = self._meshgrid(valid_x, valid_y)
        valid = valid_xx & valid_yy
        valid = valid[:, None].expand((valid.shape[0], self.num_base_anchors)).view(-1)
        return valid