import types, json
from stokenizer import STokenizer
from dataset import get_graph_finalonly_dataset
from graph_metrics import _build_finalonly_items

tok = STokenizer()
def show(ids):
    return " ".join(tok.convert_ids_to_tokens(ids))

d = json.load(open("data/star_d1_valid_fo_coconut.json"))
print("=== RAW GRAPH (first 2 held-out val samples) ===")
for s in d[:2]:
    print(f"root={s['root']}  target={s['target']}  neg_target={s['neg_target']}")
    print(f"  edges (a->b) = {s['edges']}")
    print(f"  neighbor_k(reachable frontier from root)   = {s['neighbor_k']}")
    print(f"  neg_neighbor_k(frontier from neg_root)      = {s['neg_neighbor_k']}")
    print()

cfg = types.SimpleNamespace(debug=False, uniform_prob=0.0)
ds = get_graph_finalonly_dataset("data/star_d1_valid_fo_coconut.json", 0, cfg, tok)  # stage 0 -> depth 1
print("=== TRAINING examples (what the model sees + is supervised on) ===")
for ex in ds[:3]:
    ids, labels = ex["input_ids"], ex["labels"]
    lbl = [t for t in labels if t != -100]
    print("  input :", show(ids))
    print("  label :", show(lbl), "   (all other positions = -100, ignored)")
    print()

edl, meta = _build_finalonly_items("data/star_d1_valid_fo_coconut.json", tok, max_samples=None)
print("=== EVAL examples (prompt fed at test; generate 1 token after [A]) ===")
for i in range(3):
    k, reach, *_ = meta[i]
    print("  prompt:", show(edl[i]["input_ids"]))
    print(f"  -> correct answer = {reach} (the root-reachable candidate); scored on the token emitted after [A]")
    print()