from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import gin
import numpy as np
import tensorflow as tf
from weighpoint.layers import bias
from weighpoint.layers import query as _q
from weighpoint.layers import utils
from weighpoint.layers import core


@gin.configurable
def optimizer(factory=tf.keras.optimizers.Adam, lr=1e-3, decay=0.0):
    return factory(lr=lr, decay=decay)


@gin.configurable
def model_dir(model_id, base_directory='/tmp'):
    return os.path.join(base_directory, model_id)


@gin.configurable(blacklist=['flat_features', 'units'])
def mlp_layer(
        flat_features, units, activation='relu',
        use_batch_normalization=False,
        drop_rate=None, init_kernel_scale=1.0):
    kernel_initializer = utils.scaled_glorot_uniform(scale=init_kernel_scale)
    flat_features = core.Dense(
        units, activation=activation,
        use_bias=not use_batch_normalization,
        kernel_initializer=kernel_initializer)(flat_features)
    if drop_rate is not None:
        flat_features = core.dropout(flat_features, rate=drop_rate)
    if use_batch_normalization:
        flat_features = core.batch_norm(
            flat_features, scale=(activation != 'relu'))
    return flat_features


@gin.configurable
def dropout(features, dropout_rate=0.5):
    if dropout_rate is None:
        return features
    return tf.keras.layers.Dropout(dropout_rate)(features)


@gin.configurable
def generalized_activation(
        features, activation='relu', add_bias=False,
        use_batch_normalization=False,
        dropout_rate=None):
    if add_bias and not use_batch_normalization:
        features = bias.add_bias(features)
    features = tf.keras.layers.Lambda(
        tf.keras.activations.get(activation))(features)
    if use_batch_normalization:
        features = core.batch_norm(features)
    return dropout(features, dropout_rate)


@gin.configurable(whitelist=['max_neighbors'])
def query_pairs(
        coords, radius, name=None, include_sample_rate=True,
        max_neighbors=None):
    neighbors = _q.query_pairs(
        coords, radius, name=name, max_neighbors=max_neighbors)
    if include_sample_rate:
        sample_rate = utils.reciprocal(utils.cast(
                utils.row_lengths(neighbors), tf.float32))
        return neighbors, sample_rate
    else:
        return neighbors


@gin.configurable(whitelist=['k'])
def query_knn(coords, radius, name=None, include_sample_rate=True, k=8):
    neighbors, dists = _q.query_knn(coords, k, name=name)
    if include_sample_rate:
        if dists is None:
            sample_rate = None
        else:
            sample_rate = utils.reduce_sum(dists, axis=-1)
        return neighbors, sample_rate
    else:
        return neighbors


@gin.configurable(blacklist=['num_res'])
def constant_radii(num_res):
    return np.power(4.0, np.arange(num_res)).astype(np.float32)


def _cost(radii2, density):
    return tf.reduce_sum(radii2 * density)


def _rescale_for_budget(radii2, density, budget):
    return radii2 * budget / _cost(radii2, density)


@gin.configurable(blacklist=['num_res'])
def fixed_budget_radii(num_res):
    r20 = constant_radii(num_res)
    radii = utils.variable(
        shape=(num_res,), dtype=tf.float32,
        initializer=tf.constant_initializer(np.sqrt(r20)))
    check = tf.debugging.assert_all_finite(radii, 'radii finite')
    if not tf.executing_eagerly() and check is not None:
        with tf.compat.v1.control_dependencies([check]):
            radii = tf.keras.layers.Lambda(tf.identity)(radii)
    radii2 = tf.keras.layers.Lambda(tf.square)(radii)
    density = np.power(4.0, -np.arange(num_res)).astype(np.float32)
    budget = utils.lambda_call(_cost, r20, density)
    return utils.lambda_call(_rescale_for_budget, radii2, density, budget)


def convolve(
        features, radius2, filters, neighborhood, coords_transform,
        weights_transform, convolver_fn, global_features=None):
    rel_coords = neighborhood.rel_coords.flat_values
    offset_batched_neighbors = neighborhood.offset_batched_neighbors

    coord_features = coords_transform(rel_coords, radius2)
    weights = weights_transform(rel_coords, radius2)

    features = convolver_fn(
        features, coord_features, offset_batched_neighbors, filters,
        weights)
    if global_features is not None:
        features = tf.RaggedTensor.from_row_splits(
            offset_batched_neighbors.row_splits)
        features = add_local_global(features, global_features).flat_values
    return features, offset_batched_neighbors.nested_row_splits


def _add_local_global(args):
    local_features, global_features = args
    assert(local_features.shape.ndims == 3)
    assert(global_features.shape.ndims == 2)
    assert(local_features.shape[-1] == global_features.shape[-1])
    return local_features + tf.expand_dims(global_features, axis=-2)


def add_local_global(local_features, global_features):
    if global_features is None:
        return local_features
    assert(global_features.shape[-1] == local_features.shape[-1])
    return utils.ragged_lambda(_add_local_global)(
        [local_features, global_features])
