from typing import Tuple
from histomics_detect.boxes import filter_edge_boxes
import tensorflow as tf


@tf.function
def create_anchors(anchor_px, field, width, height, filter_boxes=True):
    """Generates anchors given anchor sizes, receptive field size, and input image size.
    
    This function generates the complete set of anchors for an image with dimensions 
    height x width and backbone with given receptive field size, and removes anchors
    that cross the image boundary. Anchors are squares whose centers coincide with the 
    centers of the receptive fields of the backbone network. Each receptive field 
    corresponds to a feature map generated by the backbone, and each field has a set of
    anchors of various sizes. Each size corresponds to an output from the objectness and 
    regression branches of the region proposal network.
        
    Parameters
    ----------
    anchor_px: tensor (int32)
        K-length 1-d tensor containing the anchor width hyperparameter values in pixel 
        units.
    field: float32
        Edge length of the receptive field in pixels. This defines the area of the 
        image that corresponds to 1 feature map and the anchor gridding.
    width: int32
        Input image width in pixels.
    height: int32
        Input image height in pixels.
    filter_boxes: bool
        if true boxes are filtered, if they overlap with the border
        
    Returns
    -------
    anchors: tensor (float32)
        M x 4 tensor anchor positions organized in a dense grid over the image
        space. Each row contains the x,y upper left corner of the anchor in pixel 
        units relative in the image coordinate frame, and the anchor width and 
        height.
    """

    # generate unfiltered anchor locations using smallest anchors size
    px = tf.expand_dims(anchor_px, 1)

    x_pair, y_pair = _generate_x_y_pairs(px, field, width, height)

    # transform from 2D meshgrid format to array format (1 row per anchor)
    single_size_anchors = tf.concat((tf.reshape(x_pair, (tf.size(x_pair), 1)),
                                    tf.reshape(y_pair, (tf.size(x_pair), 1)),
                                    tf.ones((tf.size(x_pair), 1)),
                                    tf.ones((tf.size(x_pair), 1))),
                                    axis=1)

    # replicate array for each anchor size
    anchors = tf.tile(tf.expand_dims(single_size_anchors, axis=0), [tf.size(px), 1, 1])

    # add the width and height to the anchors
    px_expanded = tf.cast(tf.expand_dims(px, axis=1), tf.float32)
    ones = tf.ones((tf.size(px), 1, 1), tf.float32)
    multiplier = tf.concat((ones, ones, px_expanded, px_expanded), axis=2)
    anchors = tf.reshape(anchors * multiplier, (-1, 4))

    # remove anchors that cross the boundary
    if filter_boxes:
        anchors, _ = filter_edge_boxes(anchors, width, height, 0, tf.constant(False, tf.bool))


    return anchors


@tf.function
def _generate_x_y_pairs(px, field, width, height):
    """Generates 2D tensors containing the (x,y) corner positions of anchors for a
    square px x px anchor, height x width image, and given receptive field size.
    
    Each receptive field has a corresponding set of anchors with varying sizes.
    Since anchors must be contained entirely within the image border, some receptive
    fields at the border may lack a complete set of anchors. This function generates
    the complete set of anchor positions for a given anchor size in a 2D format similar
    to 'meshgrid'. The output includes anchors that cross the image border (these 
    anchors are later filtered using box transformations).
        
    Parameters
    ----------
    size: int32
        Anchor size in pixels.
    field: float32
        Edge length of the receptive field in pixels. This defines the area of the 
        image that corresponds to 1 feature map from the backbone network grid spacing
        of anchors.
    width: int32
        Image width in pixels.
    height: int32
        Image height in pixels.
        
    Returns
    -------
    x_pair: tensor (float32)
        A 2D grid of anchor corner horizontal positions. Each position in the grid
        represents a receptive field. Includes anchors that cross image boundaries.
    y_pair: tensor (float32)
        A 2D grid of anchor corner vertical positions. Each position in the grid
        represents a receptive field. Includes anchors that cross image boundaries.
    """    
    
    # the smallest anchor size will determine spacing from the border - larger 
    # anchor sizes will be filtered later
    size = tf.reduce_min(px)

    # get anchor corners
    x_corners = _anchor_corners(size, field, width)
    y_corners = _anchor_corners(size, field, height)

    # replicate for coordinate pairing
    x_pair = tf.tile(tf.expand_dims(x_corners, axis=0), [tf.size(y_corners), 1])
    y_pair = tf.tile(tf.expand_dims(y_corners, axis=1), [1, tf.size(x_corners)])

    return x_pair, y_pair


def first_last_anchor_indexes(size, field, length) -> Tuple[int, int]:
    """
    calculates the first and last indexes of the anchors where a box of the size 'size' does not intersect
    with the boundary of the image

    Parameters
    ----------
    size: int32
        Anchor size.
    field: float32
        Edge length of the receptive field in pixels. This defines the area of the
        image that corresponds to 1 feature map from the backbone network and the
        anchor gridding.
    length: int32
        Image edge length in pixels.

    Returns
    -------
    first: int32
        index of the first anchor that does not intersect with the boundary
    last: int32
        index of the last anchor
    """
    first = tf.math.ceil(tf.cast(size, tf.float32) /
                         (2 * tf.cast(field, tf.float32)) - 1 / 2)
    last = tf.math.floor((tf.cast(length, tf.float32) -
                          tf.cast(size, tf.float32) / 2) /
                         tf.cast(field, tf.float32) - 1 / 2)

    return first, last


@tf.function
def _anchor_corners(size, field, length):
    """Generates a sequence of anchor corner positions along one dimension of an image.
    
    Anchors are required to be entirely with the image domain during training. Called
    twice for each image to generate a horizontal sequence and a vertical sequence.
    These sequences are combined using a cartesian product to generate a dense grid
    of anchors over the entire image.
        
    Parameters
    ----------
    size: int32
        Anchor size.
    field: float32
        Edge length of the receptive field in pixels. This defines the area of the 
        image that corresponds to 1 feature map from the backbone network and the 
        anchor gridding.
    length: int32
        Image edge length in pixels.
        
        
    Returns
    -------
    corners: tensor (float32)
        1D tensor of upper/left anchor corner positions in pixels.
    """

    # first and last anchor index
    first, last = first_last_anchor_indexes(size, field, length)

    # anchor corners
    corners = tf.cast(field, tf.float32) * (tf.range(first, last + 1) + 1 / 2) - tf.cast(size, tf.float32) / 2

    return corners
