"""Generate a mock model for LLVM tests for Register Allocation.
The generated model is not a neural net - it is just a tf.function with the
correct input and output parameters. By construction, the mock model will always
output the first liverange that can be evicted.
"""
import os
import sys
import tensorflow as tf

POLICY_DECISION_LABEL = "index_to_evict"
POLICY_OUTPUT_SPEC = """
[
    {
        "logging_name": "index_to_evict",
        "tensor_spec": {
            "name": "StatefulPartitionedCall",
            "port": 0,
            "type": "int64_t",
            "shape": [
                1
            ]
        }
    }
]
"""
PER_REGISTER_FEATURE_LIST = ["mask"]
NUM_REGISTERS = 33


def get_input_signature():
    """Returns (time_step_spec, action_spec) for LLVM register allocation."""
    inputs = dict(
        (key, tf.TensorSpec(dtype=tf.int64, shape=(NUM_REGISTERS), name=key))
        for key in PER_REGISTER_FEATURE_LIST
    )
    return inputs


def get_output_spec_path(path):
    return os.path.join(path, "output_spec.json")


def build_mock_model(path):
    """Build and save the mock model with the given signature."""
    module = tf.Module()
    # We have to set this useless variable in order for the TF C API to correctly
    # intake it
    module.var = tf.Variable(0, dtype=tf.int64)

    def action(*inputs):
        result = (
            tf.math.argmax(tf.cast(inputs[0]["mask"], tf.int32), axis=-1) + module.var
        )
        return {POLICY_DECISION_LABEL: result}

    module.action = tf.function()(action)
    action = {"action": module.action.get_concrete_function(get_input_signature())}
    tf.saved_model.save(module, path, signatures=action)
    output_spec_path = get_output_spec_path(path)
    with open(output_spec_path, "w") as f:
        print(f"Writing output spec to {output_spec_path}.")
        f.write(POLICY_OUTPUT_SPEC)


def main(argv):
    assert len(argv) == 2
    model_path = argv[1]
    build_mock_model(model_path)


if __name__ == "__main__":
    main(sys.argv)