########################################################################################
# 
# Hierarchical Attentive Recurrent Tracking
# Copyright (C) 2017  Adam R. Kosiorek, Oxford Robotics Institute, University of Oxford
# email:   adamk@robots.ox.ac.uk
# webpage: http://ori.ox.ac.uk
# 
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
# 
########################################################################################

import numpy as np
import tensorflow as tf
from tensorflow.contrib.rnn.python.ops.core_rnn_cell_impl import LSTMCell
from tensorflow.python.ops.rnn_cell_impl import _RNNCell as RNNCell
from tensorflow.python.util import nest

from hart.model import tensor_ops
from hart.model.nn import DynamicFilterModel
from hart.model.rnn import ZoneoutWrapper, IdentityLSTMCell
from neurocity.component.layer import AffineLayer, ConvLayer
from neurocity.tensor_ops import convert_shape


def gaussian_mask(params, R, C):
    """Define a mask of size RxC given by one 1-D Gaussian per row.

    u, s and d must be 1-dimensional vectors"""
    u, s, d = (params[..., i] for i in xrange(3))

    for i in (u, s, d):
        assert len(u.get_shape()) == 1, i

    batch_size = tf.to_int32(tf.shape(u)[0])

    R = tf.range(tf.to_int32(R))
    C = tf.range(tf.to_int32(C))
    R = tf.to_float(R)[tf.newaxis, tf.newaxis, :]
    C = tf.to_float(C)[tf.newaxis, :, tf.newaxis]
    C = tf.tile(C, (batch_size, 1, 1))

    u, d = u[:, tf.newaxis, tf.newaxis], d[:, tf.newaxis, tf.newaxis]
    s = s[:, tf.newaxis, tf.newaxis]

    ur = u + (R - 0.) * d
    sr = tf.ones_like(ur) * s

    mask = C - ur
    mask = tf.exp(-.5 * (mask / sr) ** 2)

    mask /= tf.reduce_sum(mask, 1, keep_dims=True) + 1e-8
    return mask


def extract_glimpse(inpt, attention_params, glimpse_size):
    """Extracts an attention glimpse

    :param inpt: tensor of shape == (batch_size, img_height, img_width)
    :param attention_params: tensor of shape = (batch_size, 6) as
        [uy, sy, dy, ux, sx, dx] with u - mean, s - std, d - stride"
    :param glimpse_size: 2-tuple of ints as (height, width),
        size of the extracted glimpse
    :return: tensor
    """

    ap = attention_params
    shape = inpt.get_shape()
    rank = len(shape)

    assert rank in (3, 4), "Input must be 3 or 4 dimensional tensor"

    inpt_H, inpt_W = shape[1:3]
    if rank == 3:
        inpt = inpt[..., tf.newaxis]
        rank += 1

    Fy = gaussian_mask(ap[..., 0::2], glimpse_size[0], inpt_H)
    Fx = gaussian_mask(ap[..., 1::2], glimpse_size[1], inpt_W)

    gs = []
    for channel in tf.unstack(inpt, axis=rank - 1):
        g = tf.matmul(tf.matmul(Fy, channel, adjoint_a=True), Fx)
        gs.append(g)
    g = tf.stack(gs, axis=rank - 1)

    g.set_shape([shape[0]] + list(glimpse_size))
    return g


class Attention(object):
    n_params = None

    def __init__(self, inpt_size, glimpse_size):

        self.inpt_size = np.asarray(inpt_size)
        self.glimpse_size = np.asarray(glimpse_size)

    def extract_glimpse(self, inpt, raw_att, return_all=False):
        raw_att_flat = tf.reshape(raw_att, (-1, self.n_params), 'flat_raw_att')
        att_flat = self._to_attention(raw_att_flat)

        shape = raw_att.get_shape().as_list()
        n_glimpses = int(raw_att.get_shape()[-1]) // self.n_params

        att = tf.reshape(att_flat, shape[:-1] + [n_glimpses, int(att_flat.get_shape()[-1])])
        glimpse = []
        for a in tf.unstack(att, axis=1):
            glimpse.append(self._extract_glimpse(inpt, a))

        glimpse = tf.stack(glimpse, 1)
        glimpse = tf.reshape(glimpse, (-1,) + tuple(self.glimpse_size))

        if return_all:
            return raw_att_flat, att_flat, glimpse
        else:
            return glimpse

    def attention_to_bbox(self, att):
        with tf.variable_scope('attention_to_bbox'):
            yx = att[..., :2] * self.inpt_size[np.newaxis, :2]
            hw = att[..., 2:4] * (self.inpt_size[np.newaxis, :2] - 1)
            bbox = tf.concat(axis=tf.rank(att) - 1, values=(yx, hw))
            bbox.set_shape(att.get_shape()[:-1].concatenate((4,)))
        return bbox

    def attention_region(self, att):
        return self.attention_to_bbox(att)

    def _extract_glimpse(self, inpt, att_flat):
        return extract_glimpse(inpt, att_flat, self.glimpse_size)


class RATMAttention(Attention):
    """Implemented after https://arxiv.org/abs/1510.08660"""
    n_params = 6

    def bbox_to_attention(self, bbox):
        with tf.variable_scope('ratm_bbox_to_attention'):
            us = bbox[..., :2] / self.inpt_size[np.newaxis, :2]
            ss = 0.5 * bbox[..., 2:] / self.inpt_size[np.newaxis, :2]
            ds = bbox[..., 2:] / (self.inpt_size[np.newaxis, :2] - 1.)

            att = tf.concat(axis=tf.rank(bbox) - 1, values=(us, ss, ds))
        return att

    @staticmethod
    def _to_axis_attention(params, glimpse_dim, inpt_dim):
        u, s, d = (params[..., i] for i in xrange(RATMAttention.n_params // 2))
        u = u * inpt_dim
        s = (s + 1e-5) * float(inpt_dim) / glimpse_dim
        d = d * float(inpt_dim - 1) / (glimpse_dim - 1)
        return u, s, d

    def _to_attention(self, params):
        (y, x), (u, v) = self.inpt_size[:2], self.glimpse_size[:2]
        uy, sy, dy = self._to_axis_attention(params[..., ::2], u, y)
        ux, sx, dx = self._to_axis_attention(params[..., 1::2], v, x)

        ap = (uy, ux, sy, sx, dy, dx)
        ap = tf.transpose(tf.stack(ap), name='attention')
        assert ap.get_shape()[-1] == self.n_params, 'Invalid attention shape={}!'.format(ap.get_shape())
        return ap


class FixedStdAttention(Attention):
    """Like RATM but std for the gaussian mask depends directly and exclusively on the stride
    between gaussians. I used a small neural net to compute std to approximate bicubic
    interpolation and then fitted a 4th order surface to predictions of the neural net.

    There's also an additive (learnt) bias in pixels to the upper left corner of the attention
     window."""
    n_params = 4
    offset_bias = np.asarray([0.00809737, 0.50086582], dtype=np.float32).reshape(1, 2)
    weights = np.asarray([[6.12598441e-01, 9.25613308e-01],
                          [-1.05801568e-02, -2.18224973e-03],
                          [1.32131897e-04, -6.09307166e-06],
                          [-2.87635530e-07, 9.08051012e-08],
                          [1.94529164e-10, -9.47235313e-11],
                          [1.44468477e-04, -1.19733592e-02],
                          [-4.30590720e-06, 7.71485474e-05],
                          [1.05376852e-08, -1.05474865e-07],
                          [-6.49625282e-12, 3.43567810e-11],
                          [9.85685680e-06, 6.57580098e-05],
                          [-1.41991381e-08, -9.46024867e-08],
                          [-9.81123812e-09, -1.56167932e-07],
                          [-3.61024557e-12, 7.68954027e-11],
                          [2.65848501e-11, 4.39530612e-11],
                          [-1.01850187e-11, 9.85289183e-11]], dtype=np.float32)

    def bbox_to_attention(self, bbox):
        with tf.variable_scope('fixed_std_bbox_to_attention'):
            us = bbox[..., :2] / self.inpt_size[np.newaxis, :2]
            ds = bbox[..., 2:] / (self.inpt_size[np.newaxis, :2] - 1.)

            att = tf.concat(axis=tf.rank(bbox) - 1, values=(us, ds))
            att.set_shape(bbox.get_shape()[:-1].concatenate([4]))
        return att

    def _stride_to_std(self, stride):
        shape = convert_shape(stride.get_shape())
        stride_flat = tf.reshape(stride, (-1, shape[-1]))
        y, x = stride_flat[..., 0], stride_flat[..., 1]
        features = [
            tf.ones_like(y),
            y, y ** 2, y ** 3, y ** 4,
            x, x ** 2, x ** 3, x ** 4,
               y * x, y * x ** 2, y ** 2 * x,
               y * x ** 3, y ** 2 * x ** 2, y ** 3 * x
        ]

        features = tf.concat(axis=1, values=[f[..., tf.newaxis] for f in features])
        sigma_flat = tf.matmul(features, self.weights)
        return tf.reshape(sigma_flat, shape)

    def _to_attention(self, raw_att, with_bias=True):
        bbox = FixedStdAttention.attention_to_bbox(self, raw_att)
        us = bbox[..., :2]
        if with_bias:
            us += self.offset_bias

        ds = bbox[..., 2:4] / (self.glimpse_size[np.newaxis, :2] - 1)
        ss = self._stride_to_std(ds)

        ap = tf.concat(axis=tf.rank(raw_att) - 1, values=(us, ss, ds), name='attention')
        ap.set_shape(raw_att.get_shape()[:-1].concatenate((6,)))
        return ap


class AttentionCell(RNNCell):
    def __init__(self, feature_extractor, n_units, att_gain, glimpse_size,
                 input_size=None, batch_size=None,
                 zoneout_prob=0., attention_module=RATMAttention,
                 normalize_glimpse=False, identity_init=True, debug=False,
                 predict_appearance=False, feature_shape=None, is_training=True):

        assert len(glimpse_size) in (2, 3), 'Invalid size'
        assert input_size is None or len(input_size) == len(glimpse_size), 'Invalid size'

        self.feature_extractor = feature_extractor
        self.n_units = n_units
        self.att_gain = att_gain
        self.glimpse_size = glimpse_size
        self.input_size = input_size
        self.batch_size = batch_size
        self.normalize_glimpse = normalize_glimpse
        self.identity_init = identity_init
        self.debug = debug
        self.predict_appearance = predict_appearance

        self.attention = attention_module(self.input_size, self.glimpse_size)

        if not isinstance(zoneout_prob, (tuple, list)):
            zoneout_prob = (zoneout_prob, 0.)

        self.zoneout_prob = zoneout_prob

        self.cell = self._make_cell(is_training)
        self._rec_init = tf.random_uniform_initializer(-1e-3, 1e-3)

        self._att_size = self.att_size
        self._state_size = (self._att_size, 1, self.cell.state_size)
        self._output_size = (self.cell.output_size, self._att_size, 1)

        if self.debug:
            self._output_size += (np.prod(self.glimpse_size),)

        if self.predict_appearance:
            self._state_size += (self.n_units,)
            self._output_size += (np.prod(feature_shape), 10, 1)

    @property
    def att_size(self):
        return self.attention.n_params

    @property
    def state_size(self):
        return self._state_size

    @property
    def output_size(self):
        return self._output_size

    def zero_state(self, batch_size, dtype, bbox0=None, presence=None, img0=None,
                   transform_featuers=False, transform_state=False):
        if bbox0 is None:
            return super(AttentionCell, self).zero_state(batch_size, dtype)

        self.att0 = self.attention.bbox_to_attention(bbox0)
        att0 = tf.reshape(self.att0, (-1, self.att_size), 'shape_first_att')
        presence = tf.to_float(tf.reshape(presence, (-1, 1)), 'shape_first_presence') * 1e3

        state = att0, presence
        zero_state = self.cell.zero_state(batch_size, dtype)
        if img0 is not None:
            zero_state, rnn_outputs = self._zero_state(img0, att0, presence, zero_state,
                                                       transform_featuers, transform_state)

            att_bias = tf.get_variable('att_bias', (1, 1, 1, self.att_size), initializer=self._rec_init)
            self.att_bias = .25 * tf.nn.tanh(att_bias)
            att0 += tf.reshape(tf.tile(self.att_bias, (1, 1, 1, 1)), (1, -1))
            self.att0 += self.att_bias[0]
            state += (zero_state,)
            if self.predict_appearance:
                rnn_outputs = tf.reshape(rnn_outputs, (-1, self.n_units))
                state += (rnn_outputs,)
        else:
            state += (zero_state,)

        return state

    def _zero_state(self, img, att, presence, state, transform_features, transform_state=False):

        with tf.variable_scope(self.__class__.__name__) as vs:
            features = self.extract_features(img, att)[1]

            if transform_features:
                features_flat = tf.reshape(features, (-1, self.n_units))
                features_flat = AffineLayer(features_flat, self.n_units, name='init_feature_transform').output
                features = tf.reshape(features_flat, tf.shape(features))

            rnn_outputs, hidden_state = self._propagate(features, state)

            hidden_state = nest.flatten(hidden_state)

            if transform_state:
                for i, hs in enumerate(hidden_state):
                    name = 'init_state_transform_{}'.format(i)
                    hidden_state[i] = AffineLayer(hs, self.n_units, name=name).output

            state = nest.pack_sequence_as(structure=state, flat_sequence=hidden_state)
        self.rnn_vs = vs
        return state, rnn_outputs

    def _make_cell(self, is_training):

        raw_cell = IdentityLSTMCell if self.identity_init else LSTMCell

        if self.zoneout_prob[0] > 0.:
            cell = lambda: ZoneoutWrapper(raw_cell(self.n_units), self.zoneout_prob, is_training)
        else:
            cell = lambda: raw_cell(self.n_units)

        return cell()

    def _propagate(self, inpt, state):
        features = tf.reshape(inpt, (self.batch_size, self.n_units))
        outputs, hidden_state = self.cell(features, state)
        return tf.reshape(outputs, (self.batch_size, 1, self.n_units)), hidden_state

    def extract_features(self, inpt, raw_att, apperance_vec=None, reuse=False):

        raw_att_flat, att_flat, glimpse_flat = self.attention.extract_glimpse(inpt, raw_att,
                                                                              return_all=True)
        if self.normalize_glimpse:
            # do not normalize depth
            colour = tensor_ops.normalize_contrast(glimpse_flat[..., :3])

            if glimpse_flat.get_shape()[-1] == 4:
                ax = len(glimpse_flat.get_shape()) - 1
                glimpse_flat = tf.concat(axis=ax, values=(colour, glimpse_flat[..., 3:]))
            else:
                glimpse_flat = colour

        features = self.feature_extractor(glimpse_flat, reuse=reuse)

        def flatten_features(f, name='', more_feats=None):
            with tf.variable_scope(tf.get_variable_scope(), reuse=reuse):
                f_flat = tf.reshape(f, (self.batch_size * 1, -1), 'reshape' + name)
                to_concat = (f_flat, att_flat)
                if more_feats is not None:
                    to_concat += more_feats

                f_flat = tf.concat(axis=1, values=to_concat, name='concat_att' + name)
                return AffineLayer(f_flat, self.n_units, transfer=tf.nn.elu, name='fc_after_conv' + name)

        if apperance_vec is not None:
            before_features = self.feature_extractor.orig_output
            self.dfn_inpt = DynamicFilterModel(before_features, apperance_vec, ksize=(1, 1), n_channels=5,
                                               name='pre_DFN')
            self.dfn = DynamicFilterModel(self.dfn_inpt.features, apperance_vec, ksize=(3, 3), n_channels=5)
            dfn = self.dfn.features

            obj_mask_logit = ConvLayer(dfn, (1, 1), 1, transfer=None, name='obj_mask').output
            obj_mask = tf.nn.sigmoid(obj_mask_logit)
            features *= obj_mask

            flat_mask = tf.reshape(obj_mask, (self.batch_size * 1, -1))

        features_flat = flatten_features(features)

        features = tf.reshape(features_flat, (self.batch_size, 1, self.n_units), 'to_features')

        output = raw_att_flat, features, glimpse_flat
        if apperance_vec is not None:
            output += (obj_mask_logit, flat_mask)
        return output

    def __call__(self, inpt, state, scope=None):
        raw_att, presence, hidden_state = state[:3]
        if self.predict_appearance:
            apperance_vec = tf.reshape(state[3], (self.batch_size * 1, self.n_units))

        if self.batch_size is None:
            self.batch_size = int(inpt.get_shape()[0])

        if self.input_size is None:
            self.input_size = [tf.to_float(i) for i in inpt.get_shape()[-2:]]

        with tf.variable_scope(self.__class__.__name__):

            all_features = self.extract_features(inpt, raw_att, apperance_vec=apperance_vec, reuse=True)
            raw_att_flat, features, glimpses = all_features[:3]

            with tf.variable_scope(self.rnn_vs, initializer=self._rec_init, reuse=True):
                rnn_outputs, hidden_state = self._propagate(features, hidden_state)

            # delta-update of the raw attention params
            zero_init = tf.constant_initializer()
            outputs_flat = tf.reshape(rnn_outputs, (-1, self.n_units), 'outputs_flat')

            att_inpt = outputs_flat
            if self.predict_appearance:
                flat_mask = all_features[-1]
                mask_features = AffineLayer(flat_mask, 10, transfer=tf.nn.elu, name='mask_features')
                att_inpt = tf.concat(axis=1, values=(att_inpt, mask_features))

            att_readout = AffineLayer(att_inpt, self.n_units, transfer=tf.nn.elu, name='att_readout_1')
            att_diff_flat = AffineLayer(att_readout, self.att_size,
                                        transfer=tf.nn.tanh,
                                        weight_init=self._rec_init,
                                        bias_init=zero_init,
                                        name='att_readout')

            att_delta_scale = tf.Variable(self.att_gain, name='att_delta_scale')
            new_att_flat = raw_att_flat + tf.nn.sigmoid(att_delta_scale) * att_diff_flat.output

            new_att = tf.reshape(new_att_flat, (-1, self.att_size), 'new_att_shape')
            rnn_outputs = tf.reshape(rnn_outputs, (-1, self.n_units), 'outputs_shape')

        outputs, state = (rnn_outputs, new_att, presence), (new_att, presence, hidden_state)

        if self.debug:
            glimpse_flat = tf.reshape(glimpses, (self.batch_size, -1))
            outputs += (glimpse_flat,)

        if self.predict_appearance:
            rnn_outputs = tf.reshape(rnn_outputs, (self.batch_size, self.n_units))
            state += (rnn_outputs,)

            # concat flat obj_mask to outputs
            flat_obj_mask = tf.reshape(all_features[-2], (self.batch_size, -1))
            flat_mask_features = tf.reshape(mask_features, (self.batch_size, -1))

            # weight decay
            def weight_decay(w):
                l = tf.reshape(w, (self.batch_size, 1, -1))
                return tf.reduce_sum(l ** 2, (1, 2))[..., tf.newaxis] / (2 * 1)

            dynamic_weights = (
            self.dfn.dynamic_weights, self.dfn.dynamic_bias, self.dfn_inpt.dynamic_weights, self.dfn_inpt.dynamic_bias)
            dfn_weight_decay = sum((weight_decay(i) for i in dynamic_weights))

            outputs += (flat_obj_mask, flat_mask_features, dfn_weight_decay)

        return outputs, state
