from imandra.u.agents.code_logician.graph import GraphState
from imandra.u.agents.code_logician.command import Command
from imandra.u.agents import create_thread_sync, get_remote_graph

## 
graph = get_remote_graph("code_logician", api_key="**YOUR_IMANDRA_UNIVERSE_API_KEY**") # add your own Imandra Universe key here
create_thread_sync(graph)

## Let's now create a sample Python program which we'll use to convert into IML.
src_code = """
def g(x: int) -> int:
    if x > 22:
        return 9
    else:
        return 100 + x
    
def f(x: int) -> int:
    if x > 99:
        return 100
    elif 70 > x > 23:
        return 89 + x
    elif x > 20:
        return g(x) + 20
    elif x > -2:
        return 103
    else:
        return 99
"""

# Create an empty agent state (`GraphState`), set the Python source program to the `src_code` element of the state and run the agent.
gs = GraphState()

# Create a list of commands for CL to execute
gs = gs.add_commands([
    Command(type="init_state", src_code=src_code, src_lang="python"), # Will initialize the state with specified Python source program
    Command(type="gen_formalization_data"), # Will gather relevant formalization data (required to create the model)
    Command(type="gen_model"), # Will attempt to generate the model
])

res = await gs.run(graph) # Run the agent

# from imandra.u.agents.code_logician.command import FormalizationState # this is the type of the state variable we'll be accessing.
print(res[0].last_fstate.iml_code) # this contains the resulting IML code