"""Export the skating, rolling and learned landing actions as one ONNX policy."""

import learn
import numpy as np
import onnx
from onnx import TensorProto, helper, numpy_helper


def export(weights, destination):
    nodes, constants = [], []

    def constant(name, value, dtype=np.float32):
        constants.append(numpy_helper.from_array(np.asarray(value, dtype=dtype), name))
        return name

    def node(kind, name, *inputs, **attributes):
        nodes.append(helper.make_node(kind, inputs, [name], name=name, **attributes))
        return name

    def column(name, ix):
        return node(
            "Gather", name, "obs", constant(name + "_idx", [ix], np.int64), axis=1
        )

    def scaled_clip(value, factor, name):
        scaled = node(
            "Mul", name + "_scaled", value, constant(name + "_factor", factor)
        )
        return node("Clip", name, scaled, neg_one, one)

    zero = constant("zero", 0)
    one = constant("one", 1)
    neg_one = constant("neg_one", -1)
    cue = column("cue", 60)
    phase = [
        node("Greater", "after_" + str(i), cue, constant("boundary_" + str(i), i + 0.5))
        for i in range(5)
    ]
    models = {}
    for name in learn.BASE_MODELS:
        file = learn.model_file(name)
        source = onnx.compose.add_prefix(onnx.load(str(file)), name + "/")
        scale = 1.12 if name == "roll" else 1.0
        divisor = np.ones(61, np.float32)
        divisor[34:48] = scale
        mask = np.ones(61, np.float32)
        mask[48 if name == "roll" else 49 :] = 0
        prepared = node(
            "Div", name + "_obs", "obs", constant(name + "_divisor", divisor)
        )
        node(
            "Mul", source.graph.input[0].name, prepared, constant(name + "_mask", mask)
        )
        nodes.extend(source.graph.node)
        constants.extend(source.graph.initializer)
        models[name] = node(
            "Mul",
            name + "_action",
            source.graph.output[0].name,
            constant(name + "_gain", scale),
        )
    base = node("Where", "recovery_base", phase[3], models["rise"], models["roll"])
    arr = np.asarray(weights, np.float32)
    biases = arr[:24].reshape(4, 6)
    bias = constant("phase2", biases[0])
    for n in range(1, 4):
        bias = node(
            "Where",
            "bias" + str(n),
            phase[n + 1],
            constant("phase" + str(n + 2), biases[n]),
            bias,
        )
    pitch_gravity = column("pitch_gravity", 3)
    roll_gravity = column("roll_gravity", 4)
    pitch_rate = scaled_clip(column("pitch_rate_body", 1), 0.14, "pitch_rate")
    roll_rate = scaled_clip(column("roll_rate_body", 0), 0.15, "roll_rate")

    def vector_term(obs, name, arr):
        return node("Mul", name, obs, constant(name + "_weights", arr))

    # Sum the feedback before adding the offsets, as in the trained policy.
    pitch_terms = node(
        "Add",
        "pitch_feedback",
        vector_term(pitch_gravity, "pitch_gravity_term", arr[24:27]),
        vector_term(pitch_rate, "pitch_rate_term", arr[27:30]),
    )
    roll_terms = node(
        "Add",
        "roll_feedback",
        vector_term(roll_gravity, "roll_gravity_term", arr[30]),
        vector_term(roll_rate, "roll_rate_term", arr[31]),
    )
    zeros = node("Mul", "zeros", pitch_rate, zero)
    feedback = node("Concat", "feedback", pitch_terms, zeros, zeros, roll_terms, axis=1)
    activated = node("Tanh", "bounded", node("Add", "raw", bias, feedback))
    adjustment = node(
        "MatMul", "joint_adjustment", activated, constant("joint_map", learn.MAP)
    )
    corrected = node("Add", "corrected", base, adjustment)
    roll_or_recovery = node("Where", "trick", phase[1], corrected, base)
    node("Where", "actions", phase[0], roll_or_recovery, models["skate"])
    graph = helper.make_graph(
        nodes,
        "skate_land_skate",
        [helper.make_tensor_value_info("obs", TensorProto.FLOAT, [1, 61])],
        [helper.make_tensor_value_info("actions", TensorProto.FLOAT, [1, 14])],
        constants,
    )
    model = helper.make_model(
        graph, ir_version=8, opset_imports=[helper.make_opsetid("", 18)]
    )
    onnx.checker.check_model(model, full_check=True)
    onnx.save(model, str(destination))