import json, random, hashlib
from collections import deque
from dataclasses import dataclass
from pathlib import Path

@dataclass(frozen=True)
class Unit:
    kind: str
    ident: str
    hop: int
    text: str
    length: int


def build_graph(n=36, seed=17):
    rng = random.Random(seed)
    names = [f"Entity_{i:02d}" for i in range(n)]
    desc = {x: f"{x} is a synthetic entity with property {((i*7)%19)+3} and category C{i%5}." for i, x in enumerate(names)}
    edges = []
    # A connected backbone plus deterministic extra links.
    for i in range(n - 1):
        edges.append((names[i], "related_to", names[i+1]))
    seen = {(u, w) for u, _, w in edges}
    for _ in range(42):
        u, w = rng.sample(names, 2)
        if (u, w) not in seen and (w, u) not in seen:
            edges.append((u, "linked_with", w)); seen.add((u, w))
    edges.sort(key=lambda x: (x[0], x[1], x[2]))
    return names, desc, edges


def name_match(question, names):
    # deterministic exact entity-name matching, with stable identifier order
    return sorted([x for x in names if x in question])


def distances(names, edges, seeds, k):
    adj = {x: [] for x in names}
    for u, _, w in edges:
        adj[u].append(w); adj[w].append(u)
    d = {s: 0 for s in seeds}
    q = deque(sorted(seeds))
    while q:
        v = q.popleft()
        if d[v] >= k: continue
        for w in sorted(adj[v]):
            if w not in d:
                d[w] = d[v] + 1; q.append(w)
    return d


def units_for(names, desc, edges, d, separator="\n"):
    units = []
    for v in sorted(d):
        text = f"NODE {v}: {desc[v]}"
        units.append(Unit("node", v, d[v], text, len(text) + len(separator)))
    for i, (u, r, w) in enumerate(edges):
        if u in d and w in d:
            ident = f"edge:{u}|{r}|{w}|{i:03d}"
            text = f"EDGE {u} -[{r}]-> {w}"
            units.append(Unit("edge", ident, max(d[u], d[w]), text, len(text) + len(separator)))
    return units


def render(units, budget, mode="anchored", rng=None):
    if mode == "anchored":
        ordered = sorted(units, key=lambda x: (x.hop, x.ident))
    elif mode == "global":
        ordered = sorted(units, key=lambda x: x.ident)
    elif mode == "random":
        ordered = list(units); rng.shuffle(ordered)
    else:
        raise ValueError(mode)
    chosen, used = [], 0
    for u in ordered:
        if used + u.length > budget: break
        chosen.append(u); used += u.length
    return separator_join(chosen), chosen, used


def separator_join(chosen):
    return "\n".join(u.text for u in chosen) + ("\n" if chosen else "")


def evaluate(units, budgets, gold_ids):
    out = {}
    for mode in ("anchored", "global", "random"):
        rows = []
        for b in budgets:
            rng = random.Random(1000 + b)
            context, chosen, used = render(units, b, mode, rng)
            ids = {u.ident for u in chosen}
            rows.append({"budget": b, "recall": len(ids & gold_ids) / len(gold_ids),
                         "units": len(chosen), "used": used, "context_sha256": hashlib.sha256(context.encode()).hexdigest()})
        out[mode] = rows
    return out


def main():
    names, desc, edges = build_graph()
    question = "Compare Entity_03 and Entity_11 and explain their local connections."
    seeds = name_match(question, names)
    k = 2
    d = distances(names, edges, seeds, k)
    units = units_for(names, desc, edges, d)
    D = sum(u.length for u in units)
    # The candidate set itself is the annotated gold region for the formal claim.
    gold = {u.ident for u in units}
    # Include values around the mathematically predicted transition and a wide range.
    budgets = sorted(set([1, D//4, D//2, max(1, D-20), D-1, D, D+1, D+100, 2*D]))
    results = evaluate(units, budgets, gold)
    # Verify exact implication and stable contexts under repeated execution.
    full_below = all(next(r for r in results["anchored"] if r["budget"] == b)["recall"] < 1.0 for b in budgets if b < D)
    full_at = next(r for r in results["anchored"] if r["budget"] == D)["recall"] == 1.0
    full_above = all(next(r for r in results["anchored"] if r["budget"] == b)["recall"] == 1.0 for b in budgets if b >= D)
    c1 = render(units, max(1, D//2), "anchored", random.Random(4))[0]
    c2 = render(units, max(1, D//2), "anchored", random.Random(999))[0]
    stable = c1 == c2
    # Small useful summary at budgets where truncation is meaningful.
    summary = []
    for b in [D//4, D//2, D-1, D, D+100]:
        row = {"budget": b}
        for mode in ("anchored", "global", "random"):
            row[mode] = next(x["recall"] for x in results[mode] if x["budget"] == b)
        summary.append(row)
    report = {"seeds": seeds, "k": k, "candidate_units": len(units), "D_chars": D,
              "math_check": {"all_budgets_below_D_incomplete": full_below,
                             "at_D_full": full_at, "all_tested_budgets_at_or_above_D_full": full_above,
                             "repeated_context_identical": stable},
              "summary": summary, "all_results": results}
    Path("results.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == "__main__":
    main()