"""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 = "priority"
POLICY_OUTPUT_SPEC = """
[
    {
        "logging_name": "priority", 
        "tensor_spec": {
            "name": "StatefulPartitionedCall", 
            "port": 0, 
            "type": "float", 
            "shape": [
                1
            ]
        }
    }
]
"""
PER_LIVEINTERVAL_INT64_FEATURE_LIST = ["li_size", "stage"]
PER_LIVEINTERVAL_FLOAT32_FEATURE_LIST = ["weight"]
PER_LIVEINTERVAL_FEATURE_LIST = (
    PER_LIVEINTERVAL_FLOAT32_FEATURE_LIST + PER_LIVEINTERVAL_INT64_FEATURE_LIST
)
CONTEXT_FEATURE_LIST = ("discount", "reward", "step_type")


def get_input_signature():
    """Returns (time_step_spec, action_spec) for LLVM register allocation."""
    inputs = dict(
        (key, tf.TensorSpec(dtype=tf.int64, shape=(), name=key))
        for key in PER_LIVEINTERVAL_INT64_FEATURE_LIST
    )
    inputs.update(
        dict(
            (key, tf.TensorSpec(dtype=tf.float32, shape=(), name=key))
            for key in PER_LIVEINTERVAL_FLOAT32_FEATURE_LIST
        )
    )
    inputs.update(
        dict(
            (key, tf.TensorSpec(dtype=tf.float32, shape=(), name=key))
            for key in ["discount", "reward"]
        )
    )
    inputs.update(
        dict(
            (key, tf.TensorSpec(dtype=tf.int32, shape=(), name=key))
            for key in ["step_type"]
        )
    )
    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.float32)

    def action(*inputs):
        s1 = tf.reduce_sum(
            [
                tf.cast(inputs[0][key], tf.float32)
                for key in PER_LIVEINTERVAL_FEATURE_LIST
            ],
            axis=0,
        )
        s2 = tf.reduce_sum(
            [tf.cast(inputs[0][key], tf.float32) for key in CONTEXT_FEATURE_LIST]
        )
        # Add a large number so s won't be 0.
        s = s1 + s2
        result = s + 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)