
import timeit
import argparse
import zen
from zen import ZenEngine, ZenDecision
graph_content = """
{
  "contentType": "application/vnd.gorules.decision",
  "nodes": [
    {
      "type": "inputNode",
      "content": {
        "schema": ""
      },
      "id": "d4c5fdef-96c0-4c5e-8e7b-e1e4f9797251",
      "name": "request",
      "position": {
        "x": 245,
        "y": 355
      }
    },
    {
      "type": "outputNode",
      "content": {
        "schema": ""
      },
      "id": "a2cb6433-140b-4d2c-89f2-0e4f3140dd1a",
      "name": "response",
      "position": {
        "x": 1120,
        "y": 365
      }
    },
    {
      "type": "expressionNode",
      "content": {
        "expressions": [
          {
            "id": "8bf2c094-3115-4608-9416-1c854c68cc89",
            "key": "result",
            "value": "num * 2"
          }
        ],
        "passThrough": false,
        "inputField": null,
        "outputPath": null,
        "executionMode": "single"
      },
      "id": "45bbd693-493b-415d-8ad6-0d57c9a40d3e",
      "name": "expression1",
      "position": {
        "x": 715,
        "y": 400
      }
    }
  ],
  "edges": [
    {
      "id": "84fa9a9a-e90a-4322-be6f-94d3efd84edf",
      "sourceId": "d4c5fdef-96c0-4c5e-8e7b-e1e4f9797251",
      "type": "edge",
      "targetId": "45bbd693-493b-415d-8ad6-0d57c9a40d3e"
    },
    {
      "id": "0cd7ec7d-05d9-4497-abc7-d51b2202ddb4",
      "sourceId": "45bbd693-493b-415d-8ad6-0d57c9a40d3e",
      "type": "edge",
      "targetId": "a2cb6433-140b-4d2c-89f2-0e4f3140dd1a"
    }
  ]
}
"""


def init_zen():
    engine = zen.ZenEngine()
    decision = engine.create_decision(graph_content)
    return graph_content, engine, decision


con, engine, decision = init_zen()


def demo1(con, engine, decision):
    decision = engine.create_decision(con)
    input = {
      "num": 10,
    }
    result = decision.evaluate(input, {"trace": True})
    return result


def demo2(con, engine, decision):
    input = {
      "num": 10,
    }
    result = decision.evaluate(input, {"trace": True})
    return result


def demo3(con, engine, decision):
    input = {
      "num": 10,
    }
    result = decision.evaluate(input, {"trace": False})
    return result


# 1. Create the ArgumentParser object
parser = argparse.ArgumentParser(description="test different zen engine usage demo performace")

# 2. Add arguments
parser.add_argument("name", help="name of test function.", choices=['demo1', 'demo2', 'demo3', 'bench'])
# parser.add_argument("bench", choices=['d', 's', 'b'], help="Include a greeting.")
# parser.print_help()

# 3. Parse the arguments
args = parser.parse_args()

# 4. Access and use the arguments
from pprint import pprint
print(args)
if args.name in globals():
  print(globals()[args.name](con, engine, decision))
else:
  if args.name == 'bench':
    number = 10000
    r1 = timeit.timeit('demo1(con, engine, decision)', number=number, globals=globals())
    print(f"demo1: {r1} s/{number}ops")
    r2 = timeit.timeit('demo2(con, engine, decision)', number=number, globals=globals())
    print(f"demo2: {r2} s/{number}ops")
    r3 = timeit.timeit('demo3(con, engine, decision)', number=number, globals=globals())
    print(f"demo3: {r3} s/{number}ops")
  else:
    print(f"{args.name} function not found")


