import tensorflow as tf


def fast_rcnn(features, units=[4096, 4096], activations=['relu', 'relu'], pool=2, tiles=3):
    """
    Produces the terminal network used to refine regressions from the region-proposal network
    and to optionally classify the objects.
    
    Parameters
    ----------
    features : int
        The number of features in the 3D feature map generated by the backbone. Required to
        set input shape properties of network.
    units : array_like (int)
        The number of units in each dense layer. Default value [4096, 4096].
    activations : array_like (string)
        A list of strings defining the activations to use in the dense layers.
        Must be same length as dense. Default value ['relu', 'relu'].    
    pool : int
        pool^2 is the number of locations to interpolate features at within 
        each tile. Default value 2.
    tiles : int
        tile^2 is the number of tiles that each regressed bounding box is divided
        into. Default value 3.

    Returns
    -------
    fastrcnn : tf.keras.Model
        .
    """
    
    #generates fast RCNN model used in second regression that follows region proposal network.
    
    #create input layer
    fastrcnn_input = tf.keras.Input(shape=(pool*tiles, pool*tiles, features))

    #pooling over each roialign tile
    pooled = tf.keras.layers.MaxPool2D(pool_size=(pool, pool), padding='valid')(fastrcnn_input)

    #stack the pooled feature vectors from each tile into a single feature vector
    pooled = tf.reshape(pooled, (tf.shape(pooled)[0], tiles*tiles*features))

    #fully connected layers
    for i, (unit, activation) in enumerate(zip(units, activations)):
        if i == 0:
            layer_input = pooled
        else:
            layer_input = x
        x = tf.keras.layers.Dense(unit, activation=activation)(layer_input)

    #final layer to generate parameterized regression prediction
    regression_align = tf.keras.layers.Dense(4, activation='linear')(x)
    
    #create model
    fastrcnn = tf.keras.Model(inputs=fastrcnn_input,
                              outputs=[regression_align])
    
    return fastrcnn
