"""
Determine spatial relationships between layers to relate their coordinates.
Coordinates are mapped from input-to-output (forward), but can
be mapped output-to-input (backward) by the inverse mapping too.
This helps crop and align feature maps among other uses.
"""

from __future__ import division
import numpy as np
from caffe import layers as L

PASS_THROUGH_LAYERS = ['AbsVal', 'BatchNorm', 'Bias', 'BNLL', 'Dropout',
                       'Eltwise', 'ELU', 'Log', 'LRN', 'Exp', 'MVN', 'Power',
                       'ReLU', 'PReLU', 'Scale', 'Sigmoid', 'Split', 'TanH',
                       'Threshold']


def conv_params(fn):
    """
    Extract the spatial parameters that determine the coordinate mapping:
    kernel size, stride, padding, and dilation.

    Implementation detail: Convolution, Deconvolution, and Im2col layers
    define these in the convolution_param message, while Pooling has its
    own fields in pooling_param. This method deals with these details to
    extract canonical parameters.
    """
    params = fn.params.get('convolution_param', fn.params)
    axis = params.get('axis', 1)
    ks = np.array(params['kernel_size'], ndmin=1)
    dilation = np.array(params.get('dilation', 1), ndmin=1)
    assert len({'pad_h', 'pad_w', 'kernel_h', 'kernel_w', 'stride_h',
                'stride_w'} & set(fn.params)) == 0, \
        'cropping does not support legacy _h/_w params'
    return (axis, np.array(params.get('stride', 1), ndmin=1),
            (ks - 1) * dilation + 1,
            np.array(params.get('pad', 0), ndmin=1))


def crop_params(fn):
    """
    Extract the crop layer parameters with defaults.
    """
    params = fn.params.get('crop_param', fn.params)
    axis = params.get('axis', 2)  # default to spatial crop for N, C, H, W
    offset = np.array(params.get('offset', 0), ndmin=1)
    return (axis, offset)


class UndefinedMapException(Exception):
    """
    Exception raised for layers that do not have a defined coordinate mapping.
    """
    pass


def coord_map(fn):
    """
    Define the coordinate mapping by its
    - axis
    - scale: output coord[i * scale]