# Copyright (C) 2020-2022 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

from functools import partial

import numpy as np
from ..function_selector import ACTIVATIONS_STATS_FN, PERTENSOR, PERCHANNEL

compute_act_stats_fn_per_tensor = ACTIVATIONS_STATS_FN['compute_statistic'][PERTENSOR]
compute_act_stats_fn_per_channel = ACTIVATIONS_STATS_FN['compute_statistic'][PERCHANNEL]

get_act_stats_fn_per_tensor = ACTIVATIONS_STATS_FN['statistic_in_graph'][PERTENSOR]
get_act_stats_fn_per_channel = ACTIVATIONS_STATS_FN['statistic_in_graph'][PERCHANNEL]


# helper functions to calculate statistics for activations
def calculate_per_channel_stats(acts, fn, axis=1):
    """ Calculates per-channel statistics for activations using a specific function
    :param act: activation
    :param fn: function to calculate per-channel statistics
    :return statistics generated by fn for each activation in the batch
    """
    if len(acts.shape) < 3:
        return acts
    acts = np.moveaxis(acts, axis, 1)
    t = acts.reshape(acts.shape[0], acts.shape[1], -1)
    return fn(t, axis=2)


def calculate_per_tensor_stats(acts, fn):
    """ Calculates statistics by whole tensor for activations using a specific function
    :param act: activation
    :param fn: function to calculate per-tensor statistics
    :return statistics generated by fn for each activation in the batch
    """
    if len(acts.shape) < 2:
        return np.atleast_1d(fn(acts))

    t = acts.reshape(acts.shape[0], -1)
    return fn(t, axis=1)


@compute_act_stats_fn_per_tensor.register('max')
def max_per_tensor(acts, **_):
    return calculate_per_tensor_stats(acts, np.max)


@compute_act_stats_fn_per_tensor.register('min')
def min_per_tensor(acts, **_):
    return calculate_per_tensor_stats(acts, np.min)


@compute_act_stats_fn_per_tensor.register('abs_max')
def abs_max_per_tensor(acts, **_):
    return max_per_tensor(np.abs(acts))


@compute_act_stats_fn_per_tensor.register('quantile')
def quantile_per_tensor(acts, q, **_):
    return calculate_per_tensor_stats(acts, partial(np.quantile, q=q))


@compute_act_stats_fn_per_tensor.register('abs_quantile')
def abs_quantile_per_tensor(acts, q, **_):
    return quantile_per_tensor(np.abs(acts), q)


@compute_act_stats_fn_per_channel.register('mean')
def mean_per_channel(acts, **_):
    return calculate_per_channel_stats(acts, np.mean)


@compute_act_stats_fn_per_channel.register('mean_axis')
def mean_per_channel_axis(acts, layer_key=None, **kwargs):
    axis = kwargs.get('channel', {}).get(layer_key, 1)
    return calculate_per_channel_stats(acts, np.mean, axis=axis)


@compute_act_stats_fn_per_channel.register('quantile')
def quantile_per_channel(acts, q, **_):
    return calculate_per_channel_stats(acts, partial(np.quantile, q=q))


@compute_act_stats_fn_per_channel.register('max')
def max_per_channel(acts, **_):
    return calculate_per_channel_stats(acts, np.max)


@compute_act_stats_fn_per_channel.register('min')
def min_per_channel(acts, **_):
    return calculate_per_channel_stats(acts, np.min)


@compute_act_stats_fn_per_channel.register('abs_max')
def abs_max_per_channel(acts, **_):
    return max_per_channel(np.abs(acts))


@compute_act_stats_fn_per_channel.register('abs_quantile')
def abs_quantile_per_channel(acts, q, **_):
    return quantile_per_channel(np.abs(acts), q)


@get_act_stats_fn_per_tensor.register('max')
def get_max_per_tensor(acts, **_):
    return np.atleast_1d(acts)


@get_act_stats_fn_per_tensor.register('min')
def get_min_per_tensor(acts, **_):
    return np.atleast_1d(acts)


@get_act_stats_fn_per_tensor.register('abs_max')
def get_abs_max_per_tensor(acts, **_):
    return np.atleast_1d(acts)


@get_act_stats_fn_per_channel.register('mean')
def get_mean_per_channel(acts, **_):
    return np.atleast_1d(acts)


@get_act_stats_fn_per_channel.register('mean_axis')
def get_mean_per_channel_axis(acts, _, **__):
    return np.atleast_1d(acts)


@get_act_stats_fn_per_channel.register('max')
def get_max_per_channel(acts, **_):
    return np.atleast_1d(acts)


@get_act_stats_fn_per_channel.register('min')
def get_min_per_channel(acts, **_):
    return np.atleast_1d(acts)


@get_act_stats_fn_per_channel.register('abs_max')
def get_abs_max_per_channel(acts, **_):
    return np.atleast_1d(acts)
