# SPDX-License-Identifier: Apache-2.0

import itertools
import numpy as np
from ...common._registration import register_converter


def extract_support_vectors_as_dense_tensor(svm_model):
    support_type = svm_model.WhichOneof("supportVectors")
    if support_type == "denseSupportVectors":
        vectors = svm_model.denseSupportVectors.vectors
        support_vectors = np.array([v.values for v in vectors]).flatten()
    elif support_type == "sparseSupportVectors":
        # Since current ONNX doesn't support sparse representations, we have to save them as dense tensors. It may
        # dramatically prolong prediction time and increase memory usage.
        vectors = svm_model.sparseSupportVectors.vectors
        # Search for the maximum dimension over all sparse support vectors
        max_idx = 0
        for v in vectors:
            for n in v.nodes:
                if n.index > max_idx:
                    max_idx = n.index
        # Save sparse vectors to dense vectors
        support_vectors = np.zeros(shape=(len(vectors), max_idx))
        for i, v in enumerate(vectors):
            for n in v.nodes:
                support_vectors[i][n.index - 1] = n.value
        support_vectors = support_vectors.flatten()
    else:
        raise ValueError("Unsupported support vector type: %s" % support_type)
    return len(vectors), support_vectors


def convert_svm_classifier(scope, operator, container):
    params = operator.raw_operator.supportVectorClassifier
    kernel_enum = {
        "linearKernel": "LINEAR",
        "polyKernel": "POLY",
        "rbfKernel": "RBF",
        "sigmoidKernel": "SIGMOID",
        "precomputedKernel": "PRECOMPUTED",
    }
    kernel = params.kernel
    kernel_val = kernel.WhichOneof("kernel")
    svc_kernel = kernel_enum[kernel_val]

    if kernel_val == "rbfKernel":
        svc_kernel_params = [kernel.rbfKernel.gamma, 0.0, 0.0]
    elif kernel_val == "polyKernel":
        svc_kernel_params = [
            kernel.polyKernel.gamma,
            kernel.polyKernel.coef0,
            kernel.polyKernel.degree,
        ]
    elif kernel_val == "sigmoidKernel":
        svc_kernel_params = [
            kernel.sigmoidKernel.gamma,
            kernel.sigmoidKernel.coef0,
            0.0,
        ]
    elif kernel_val == "linearKernel":
        svc_kernel_params = [0.0, 0.0, 0.0]

    prob_a = params.probA
    prob_b = params.probB
    support_vectors_per_class = params.numberOfSupportVectorsPerClass
    n_supports, svc_support_vectors = extract_support_vectors_as_dense_tensor(
        operator.raw_operator.supportVectorClassifier
    )
    chain_coef = list(
        itertools.chain.from_iterable([coef.alpha for coef in params.coefficients])
    )
    svc_coefficients = chain_coef
    svc_rho = [-x for x in params.rho]

    op_type = "SVMClassifier"
    op_name = scope.get_unique_operator_name(op_type)
    attrs = {"name": op_name}
    attrs["kernel_type"] = svc_kernel
    attrs["kernel_params"] = svc_kernel_params
    if prob_a:
        attrs["prob_a"] = prob_a
    if prob_b:
        attrs["prob_b"] = prob_b
    attrs["vectors_per_class"] = support_vectors_per_class
    attrs["support_vectors"] = svc_support_vectors
    attrs["coefficients"] = svc_coefficients
    attrs["rho"] = svc_rho
    zipmap_attrs = {"name": scope.get_unique_operator_name("ZipMap")}
    svc_classes = params.WhichOneof("ClassLabels")
    if svc_classes == "int64ClassLabels":
        class_labels = list(int(i) for i in params.int64ClassLabels.vector)
        attrs["classlabels_ints"] = class_labels
        zipmap_attrs["classlabels_int64s"] = class_labels
    elif svc_classes == "stringClassLabels":
        class_labels = list(
            str(s).encode("utf-8") for s in params.stringClassLabels.vector
        )
        attrs["classlabels_strings"] = class_labels
        zipmap_attrs["classlabels_strings"] = class_labels
    else:
        raise ValueError("Unknown class label type")

    # For classifiers, due to the different representation
    # of classes' probabilities, we need to add some
    # operators for type conversion. It turns out that we
    # have the following topology.
    # input features ---> SupportVectorClassifier ---> label (must present)
    #                               |
    #                               '--> probability tensor

    raw_model = operator.raw_operator
    # Find label name and probability name
    proba_output_name = None
    for variable in operator.outputs:
        if raw_model.description.predictedFeatureName == variable.raw_name:
            label_output_name = variable.full_name
        if (
            raw_model.description.predictedProbabilitiesName != ""
            and raw_model.description.predictedProbabilitiesName == variable.raw_name
        ):
            proba_output_name = variable.full_name

    inputs = [variable.full_name for variable in operator.inputs]
    proba_tensor_name = scope.get_unique_variable_name("ProbabilityTensor")

    if proba_output_name is not None:
        # Add support vector classifier in terms of ONNX node with probability output
        container.add_node(
            op_type,
            inputs,
            [label_output_name, proba_tensor_name],
            op_domain="ai.onnx.ml",
            **attrs
        )

        # Add ZipMap to convert probability tensor into probability map
        container.add_node(
            "ZipMap",
            [proba_tensor_name],
            [proba_output_name],
            op_domain="ai.onnx.ml",
            **zipmap_attrs
        )
    else:
        # Add support vector classifier in terms of ONNX node
        container.add_node(
            op_type,
            inputs,
            [label_output_name, proba_tensor_name],
            op_domain="ai.onnx.ml",
            **attrs
        )


register_converter("supportVectorClassifier", convert_svm_classifier)
