# SPDX-License-Identifier: Apache-2.0

from onnx import onnx_pb as onnx_proto
from ...common._apply_operation import (
    apply_add, apply_mul, apply_sum, apply_div, apply_sub,
    apply_concat, apply_cast)
from ...common._registration import register_converter, register_shape_calculator
from ...common.data_types import FloatTensorType, DoubleTensorType
from ...common.utils import check_input_and_output_numbers, check_input_and_output_types
from .tree_ensemble_common import save_read_sparkml_model_data

MinHashLSH_HASH_PRIME = 2038074743
g_rand_coefficients = None


def get_rand_coefficients(operator):
    global g_rand_coefficients
    if not g_rand_coefficients:
        g_rand_coefficients = save_read_sparkml_model_data(
            operator.raw_params['SparkSession'], operator.raw_operator
        ).first()[0]
    return g_rand_coefficients


def convert_min_hash_lsh(scope, operator, container):
    int_type = onnx_proto.TensorProto.INT64
    rand_coefficients = get_rand_coefficients(operator)
    coeffs = []
    for i in range(0, len(rand_coefficients), 2):
        coeffs.append((rand_coefficients[i], rand_coefficients[i + 1]))
    one = scope.get_unique_variable_name('one_tensor')
    container.add_initializer(one, int_type, [1], [1])
    prime = scope.get_unique_variable_name('prime_tensor')
    container.add_initializer(prime, int_type, [1], [MinHashLSH_HASH_PRIME])

    non_zeros_int = scope.get_unique_variable_name('non_zero_int_tensor')
    container.add_node('NonZero', operator.input_full_names, non_zeros_int, op_version=9)
    non_zeros = scope.get_unique_variable_name('non_zeros_tensor')
    apply_cast(scope, non_zeros_int, non_zeros, container, to=int_type)
    remainders = []
    for coeff in coeffs:
        one_added = scope.get_unique_variable_name('one_added_tensor')
        apply_add(scope, [one, non_zeros], one_added, container)
        a = scope.get_unique_variable_name('a_coeff_tensor')
        container.add_initializer(a, int_type, [1], [coeff[0]])
        b = scope.get_unique_variable_name('b_coeff_tensor')
        container.add_initializer(b, int_type, [1], [coeff[1]])
        coeff_0_times = scope.get_unique_variable_name('coeff_0_times_tensor')
        apply_mul(scope, [a, one_added], coeff_0_times, container)
        coeff_1_added = scope.get_unique_variable_name('coeff_1_added_tensor')
        apply_add(scope, [b, coeff_0_times], coeff_1_added, container)
        # this is for Mod
        div_by_prime = scope.get_unique_variable_name('div_by_prime_tensor')
        apply_div(scope, [coeff_1_added, prime], div_by_prime, container)
        prime_x_floor = scope.get_unique_variable_name('prime_x_floor_tensor')
        apply_mul(scope, [div_by_prime, prime], prime_x_floor, container)
        remainder = scope.get_unique_variable_name('remainder_tensor')
        apply_sub(scope, [coeff_1_added, prime_x_floor], remainder, container)
        float_remainder = scope.get_unique_variable_name('float_remainder_tensor')
        apply_cast(scope, remainder, float_remainder, container, to=onnx_proto.TensorProto.FLOAT)
        min_remainder = scope.get_unique_variable_name('min_remainder_tensor')
        container.add_node('ReduceMin', float_remainder, min_remainder,
                           op_version=1,
                           axes=[1],
                           keepdims=1)
        remainders.append(min_remainder)
    apply_concat(scope, remainders, operator.output_full_names, container, axis=1)


register_converter('pyspark.ml.feature.MinHashLSHModel', convert_min_hash_lsh)


def calculate_min_hash_lsh_output_shapes(operator):
    check_input_and_output_numbers(operator, output_count_range=1)
    check_input_and_output_types(
        operator, good_input_types=[FloatTensorType, DoubleTensorType])

    N = operator.inputs[0].type.shape[0]
    C = len(get_rand_coefficients(operator)) // 2
    operator.outputs[0].type = FloatTensorType([N, C])


register_shape_calculator('pyspark.ml.feature.MinHashLSHModel', calculate_min_hash_lsh_output_shapes)
