"""Load graph-reasoner-typed and answer one synthetic question."""
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE / "code"))
sys.path.insert(0, str(HERE / "code" / "scripts"))

import torch  # noqa: E402

from graph.graph import GraphStore  # noqa: E402
from graph.synthetic.generator import SyntheticGraphGenerator  # noqa: E402
from graph.typed_readout import TypedReadout  # noqa: E402

device = "cuda" if torch.cuda.is_available() else "cpu"
readout = TypedReadout.from_checkpoint(HERE / "readout.pt", device=device)
example = SyntheticGraphGenerator().generate_example(seed=7, stage=2)
result = readout.reason(GraphStore.from_example(example.graph), example.query)
print("question:", example.query.text)
print("status:", result.status.value, "| confidence: %.3f" % result.confidence)
print("answers:", [a.entity_id for a in result.answers], "| gold:", list(example.answers))