# SPDX-License-Identifier: Apache-2.0

from ...common._registration import register_converter

_LINK_FUNCTION_TO_POST_TRANSFORM = {
    "identity": "NONE",
    "logit": "LOGISTIC",
    "ologit": "LOGISTIC",
}


def _get_post_transform(params):
    link_function = params["link_function"]
    family = params["family"]
    if family == "multinomial":
        return "SOFTMAX"
    elif link_function not in _LINK_FUNCTION_TO_POST_TRANSFORM.keys():
        raise ValueError("Link function %s not supported." % link_function)
    else:
        return _LINK_FUNCTION_TO_POST_TRANSFORM[link_function]


def _get_default_tree_attribute_pairs(is_classifier, params):
    attrs = {"post_transform": _get_post_transform(params)}
    nclasses = params["nclasses"]
    if is_classifier:
        predicted_classes = nclasses if nclasses > 2 else 1
        attrs["base_values"] = [
            params["base_score"] for _ in range(0, predicted_classes)
        ]
    else:
        attrs["n_targets"] = 1
        attrs["base_values"] = [params["base_score"]]
    for k in {
        "nodes_treeids",
        "nodes_nodeids",
        "nodes_featureids",
        "nodes_modes",
        "nodes_values",
        "nodes_truenodeids",
        "nodes_falsenodeids",
        "nodes_missing_value_tracks_true",
    }:
        attrs[k] = []
    node_attr_prefix = _node_attr_prefix(is_classifier)
    for k in {"_treeids", "_nodeids", "_ids", "_weights"}:
        attrs[node_attr_prefix + k] = []
    return attrs


def _add_node(
    attr_pairs,
    is_classifier,
    tree_id,
    node_id,
    feature_id,
    mode,
    value,
    true_child_id,
    false_child_id,
    weights,
    missing,
):
    attr_pairs["nodes_treeids"].append(tree_id)
    attr_pairs["nodes_nodeids"].append(node_id)
    attr_pairs["nodes_featureids"].append(feature_id)
    attr_pairs["nodes_modes"].append(mode)
    attr_pairs["nodes_values"].append(float(value))
    attr_pairs["nodes_truenodeids"].append(true_child_id)
    attr_pairs["nodes_falsenodeids"].append(false_child_id)
    attr_pairs["nodes_missing_value_tracks_true"].append(missing)
    if mode == "LEAF":
        node_attr_prefix = _node_attr_prefix(is_classifier)
        for i, w in enumerate(weights):
            attr_pairs[node_attr_prefix + "_treeids"].append(tree_id)
            attr_pairs[node_attr_prefix + "_nodeids"].append(node_id)
            attr_pairs[node_attr_prefix + "_ids"].append(i)
            attr_pairs[node_attr_prefix + "_weights"].append(float(w))


def _node_attr_prefix(is_classifier):
    return "class" if is_classifier else "target"


def _fill_node_attributes(tree_id, node, attr_pairs, is_classifier):
    if "leftChild" in node:
        if node["isCategorical"]:
            raise ValueError("categorical splits not supported, use one_hot_explicit")
        else:
            operator = "BRANCH_GTE"
            value = node["splitValue"]
        _add_node(
            attr_pairs=attr_pairs,
            is_classifier=is_classifier,
            tree_id=tree_id,
            mode=operator,
            value=value,
            node_id=node["id"],
            feature_id=node["colId"],
            true_child_id=node["rightChild"]["id"],
            false_child_id=node["leftChild"]["id"],
            weights=None,
            missing=(0 if node["leftward"] else 1),
        )
        _fill_node_attributes(tree_id, node["leftChild"], attr_pairs, is_classifier)
        _fill_node_attributes(tree_id, node["rightChild"], attr_pairs, is_classifier)
    else:  # leaf
        weights = [node["predValue"]]
        _add_node(
            attr_pairs=attr_pairs,
            is_classifier=is_classifier,
            tree_id=tree_id,
            value=0.0,
            node_id=node["id"],
            feature_id=0,
            mode="LEAF",
            true_child_id=0,
            false_child_id=0,
            weights=weights,
            missing=False,
        )


def assign_node_ids(node, next_id):
    if node is None:
        return next_id
    node["id"] = next_id
    next_id += 1
    next_id = assign_node_ids(node.get("leftChild", None), next_id)
    return assign_node_ids(node.get("rightChild", None), next_id)


def fill_tree_attributes(model, attr_pairs, node_attr_prefix):
    for tree in model["trees"]:
        assign_node_ids(tree["root"], 0)
        _fill_node_attributes(tree["index"], tree["root"], attr_pairs, node_attr_prefix)


def convert_regression(scope, operator, container, params):
    model = operator.raw_operator

    attr_pairs = _get_default_tree_attribute_pairs(False, params)
    fill_tree_attributes(model, attr_pairs, False)

    # add nodes
    container.add_node(
        "TreeEnsembleRegressor",
        operator.input_full_names,
        operator.output_full_names,
        op_domain="ai.onnx.ml",
        name=scope.get_unique_operator_name("TreeEnsembleRegressor"),
        **attr_pairs
    )


def convert_classifier(scope, operator, container, params):
    if params["family"] == "multinomial" and params["nclasses"] == 2:
        raise ValueError(
            "Multinomial distribution with two classes "
            "not supported, use binomial distribution."
        )
    model = operator.raw_operator

    attr_pairs = _get_default_tree_attribute_pairs(True, params)
    fill_tree_attributes(model, attr_pairs, True)

    n_trees_in_group = params["n_trees_in_group"]
    attr_pairs["class_ids"] = [
        v % n_trees_in_group for v in attr_pairs["class_treeids"]
    ]
    attr_pairs["classlabels_strings"] = params["class_labels"]

    container.add_node(
        "TreeEnsembleClassifier",
        operator.input_full_names,
        operator.output_full_names,
        op_domain="ai.onnx.ml",
        name=scope.get_unique_operator_name("TreeEnsembleClassifier"),
        **attr_pairs
    )


def convert_h2o(scope, operator, container):
    params = operator.raw_operator["params"]
    is_classifier = params["classifier"]
    if is_classifier:
        convert_classifier(scope, operator, container, params)
    else:
        convert_regression(scope, operator, container, params)


register_converter("H2OTreeMojo", convert_h2o)
