from numpy import linalg as LA
import numpy as np

MAX_SIG_VALUE = 10000000


def get_avg_gradient(gradient_list, num_of_workers):
    summed_gradient = gradient_list[0]
    i = 0
    for gradient in gradient_list:
        i += 1
        if i == 1:
            continue
        j = 0
        for gradient_part in gradient:
            summed_gradient[j] = np.add(summed_gradient[j], gradient_part)
            j += 1

    avg_gradient = []
    for gradient_part in summed_gradient:
        avg_gradient.append(gradient_part / num_of_workers)
    return avg_gradient


def get_gradient_diff(previous_gradient, new_gradient):
    gradient_diff = []
    for i in range(len(previous_gradient)):
        gradient_diff.append(np.subtract(new_gradient[i], previous_gradient[i]))

    return gradient_diff


def get_aggregate_gradient_update(previous_gradient, new_gradient):
    gradient_diff = []
    for i in range(len(previous_gradient)):
        gradient_diff.append(np.subtract(new_gradient[i], previous_gradient[i]))

    aggregate_sum = 0
    for gradient in gradient_diff:
        aggregate_sum += gradient.sum()
    return aggregate_sum


def get_l2_norm_of_gradients(gradient_list):
    l2_norm_list = []
    for gradient in gradient_list:
        raveled_gradient = np.ravel(gradient)
        l2_norm_list.append(LA.norm(raveled_gradient, 2))

    return LA.norm(l2_norm_list, 2)


def get_significance_value(previous_gradient, new_gradient):
    aggregate_update = get_aggregate_gradient_update(previous_gradient, new_gradient)
    print("Aggregate update: " + str(aggregate_update))
    previous_gradient_value = 0
    for gradient in previous_gradient:
        previous_gradient_value += gradient.sum()
    print("Previous gradient value: " + str(previous_gradient_value))
    return abs(aggregate_update / previous_gradient_value)


def get_significance_value_2(previous_gradient, new_gradient):
    gradient_diff = get_gradient_diff(previous_gradient, new_gradient)
    gradient_diff_value = get_l2_norm_of_gradients(gradient_diff)
    previous_gradient_value = get_l2_norm_of_gradients(previous_gradient)
    if previous_gradient_value <= 0:
        return MAX_SIG_VALUE
    else:
        return abs(gradient_diff_value / previous_gradient_value)


def get_significance_value_3(current_gradient, updated_gradient, summed_gradient, grad_counter):
    current_gradient = get_sum_of_two_gradients(current_gradient, summed_gradient)
    updated_gradient = get_sum_of_two_gradients(updated_gradient, summed_gradient)
    current_avg_gradient = divide_gradient_by_constant(current_gradient, grad_counter + 1)
    updated_gradient_val = get_l2_norm_of_gradients(updated_gradient)
    current_avg_gradient_value = get_l2_norm_of_gradients(current_avg_gradient)
    print(current_avg_gradient_value)
    # current_avg_gradient = get_sum_of_two_gradients(current_avg_gradient, summed_gradient)
    if current_avg_gradient_value <= 0:
        return MAX_SIG_VALUE, updated_gradient, current_gradient
    else:
        return abs(updated_gradient_val / current_avg_gradient_value), updated_gradient, current_gradient


def get_significance_value_4(avg_gradient, updated_gradient):
    # current_gradient = get_sum_of_two_gradients(current_gradient, summed_gradient)
    updated_gradient = get_sum_of_two_gradients(updated_gradient, avg_gradient)
    # current_avg_gradient = divide_gradient_by_constant(current_gradient, grad_counter+1)
    updated_gradient_val = get_l2_norm_of_gradients(updated_gradient)
    current_avg_gradient_value = get_l2_norm_of_gradients(avg_gradient)
    # current_avg_gradient = get_sum_of_two_gradients(current_avg_gradient, summed_gradient)
    if current_avg_gradient_value <= 0:
        return MAX_SIG_VALUE, updated_gradient
    else:
        return abs(updated_gradient_val / current_avg_gradient_value), updated_gradient


def get_initial_threshold(initial_gradient):
    l2_norm_list = []
    linf_norm_list = []
    for gradient in initial_gradient:
        raveled_gradient = np.ravel(gradient)
        l2_norm_list.append(LA.norm(raveled_gradient, 2))
        linf_norm_list.append(LA.norm(raveled_gradient, np.inf))

    initial_threshold = LA.norm(l2_norm_list, 2) / LA.norm(linf_norm_list, np.inf)
    return initial_threshold


def get_new_threshold(loss_list, worker_count, lr_change, previous_threshold):
    avg_loss = 0
    for loss in loss_list:
        avg_loss = avg_loss + loss
    avg_loss = avg_loss / worker_count
    # print("Momentum correction")
    # print(avg_loss)
    # return lr_change * avg_loss * previous_threshold
    return previous_threshold * 0.9


def get_new_threshold_2(momentum_correction, lr_change, previous_threshold):
    return lr_change * momentum_correction * previous_threshold


def get_gradient_difference(new_gradient, previous_gradient):
    gradient_diff = []
    for i in range(len(previous_gradient)):
        gradient_diff.append(np.subtract(new_gradient[i], previous_gradient[i]))

    l2_norm_diff = get_l2_norm_of_gradients(gradient_diff)
    print(l2_norm_diff)
    if l2_norm_diff > 0:
        return True, gradient_diff
    else:
        return False, None


def get_momentum_correction(gradient_list, num_of_workers, batch_size):
    summed_gradient = get_gradient_summation(gradient_list)
    for i in range(len(summed_gradient)):
        summed_gradient[i] = summed_gradient[i] / (num_of_workers * batch_size)
    l2_norm = get_l2_norm_of_gradients(summed_gradient) / batch_size
    return l2_norm


def get_momentum_correction_2(summed_gradient, batch_size):
    thresh = []
    print(len(summed_gradient))
    for i in range(len(summed_gradient)):
        raveled_gradient = np.ravel(summed_gradient[i])
        ind = np.argpartition(raveled_gradient, -4, axis=None)[-4:]
        thresh.append(np.amin(raveled_gradient[ind]))
        # summed_gradient[i] = summed_gradient[i] / batch_size
    l2_norm = LA.norm(thresh, 2)
    return l2_norm / batch_size


def get_gradient_summation(gradient_list):
    summed_gradient = gradient_list[0]
    i = 0
    for gradient in gradient_list:
        i += 1
        if i == 1:
            continue
        j = 0
        for gradient_part in gradient:
            summed_gradient[j] = np.add(summed_gradient[j], gradient_part)
            j += 1
    return summed_gradient


def get_max_gradients(gradient_list):
    max_gradient = gradient_list[0]
    max_grad_value = get_l2_norm_of_gradients(max_gradient)
    for i in range(len(gradient_list)):
        if i == 0:
            continue
        new_grad_val = get_l2_norm_of_gradients(gradient_list[i])
        if new_grad_val > max_grad_value:
            max_gradient = gradient_list[i]
            max_grad_value = new_grad_val
    return max_gradient, max_grad_value


def get_sum_of_two_gradients(gradient_1, gradient_2):
    for i in range(len(gradient_1)):
        gradient_1[i] = np.add(gradient_1[i], gradient_2[i])
    return gradient_1


def divide_gradient_by_constant(gradient, constant):
    for i in range(len(gradient)):
        gradient[i] = gradient[i] / constant
    return gradient

# b = np.arange(9).reshape((3, 3))
# added_list = [a, b]
# g_list = [added_list, added_list]
# # print(get_avg_gradient(g_list, 2))
# new_a = np.arange(2, 20).reshape((3, 3, 2))
# new_b = np.arange(2, 11).reshape((3, 3))
# new_list = [new_a, new_b]
# # print(get_aggregate_gradient_update(added_list, new_list))
# # print(get_significance_value(added_list, new_list))
# print(get_initial_threshold(added_list))
