cli.py

from __future__ import annotations
import argparse
import os
import random

from .kg import build_toy_family_kg, KnowledgeGraph
from .rule_mining import mine_binary_path_rules, mine_unary_rules
from .contexts import discover_contexts, rules_that_entail
from .pc import estimate_with_negatives
from .inference import infer_probability, InferenceResult
from .neo4j_utils import load_triples, fetch_triples
from .llm_interface import parse_natural_query, llm_wrap_explanation
from .anyburl import write_triples_tsv, run_anyburl
from .juice_trainer import train_with_juice, score_query_with_juice, JuiceModel
from .datasets import download_dataset, load_triples_from_dir, all_triples
from .model_io import save_model, load_model


def build_pipeline():
    kg = build_toy_family_kg()
    bin_rules = mine_binary_path_rules(kg, min_support=1, min_conf=0.2)
    un_rules = mine_unary_rules(kg, min_support=1, min_conf=0.3)
    binary = list(bin_rules.keys())
    unary = list(un_rules.keys())
    train = list(kg.all_triples())
    contexts, _ = discover_contexts(kg, train, binary, unary)
    params = estimate_with_negatives(contexts)
    return kg, binary, unary, contexts, params


def main():
    ap = argparse.ArgumentParser(description="KG Probabilistic Contexts Demo")
    sub = ap.add_subparsers(dest="cmd", required=True)

    sub.add_parser("demo", help="Run end-to-end toy demo")
    q = sub.add_parser("query", help="Query a triple")
    q.add_argument("triple", help='Triple as "H R T"')

    sub.add_parser("load-neo4j", help="Load toy KG to Neo4j")

    ask = sub.add_parser("ask", help="LLM-flavored QA over KG")
    ask.add_argument("prompt", help="Question like 'Is John grandparent of Alice?' or 'John Grandparent Alice'")

    exp = sub.add_parser("neo4j-export", help="Export triples from Neo4j to TSV")
    exp.add_argument("--out", default="data/triples.tsv")

    ab = sub.add_parser("anyburl-run", help="Run AnyBURL on TSV triples")
    ab.add_argument("--jar", required=True, help="Path to AnyBURL jar")
    ab.add_argument("--input", default="data/triples.tsv")
    ab.add_argument("--out", default="data/anyburl_rules.txt")
    ab.add_argument("--time", type=int, default=60)

    ctxp = sub.add_parser("contexts", help="Discover contexts (fallback without PyClause)")
    ctxp.add_argument("--triples", default="data/triples.tsv")
    ctxp.add_argument("--rules", default="data/anyburl_rules.txt")
    ctxp.add_argument("--out", default="data/contexts.tsv", help="Outputs triple->context mapping")

    trn = sub.add_parser("train", help="Train probabilistic model (uses pyjuice if available)")
    trn.add_argument("--contexts", default="data/contexts.tsv")
    trn.add_argument("--out", default="data/model.tsv")

    inf = sub.add_parser("infer", help="Infer using trained model")
    inf.add_argument("--model", default="data/model.tsv")
    inf.add_argument("--query", required=True, help='Triple as "H R T"')

    dld = sub.add_parser("dataset-download", help="Download a benchmark dataset")
    dld.add_argument("--name", choices=["FB15k-237", "WN18RR"], required=True)
    dld.add_argument("--out", default="data")

    lnd = sub.add_parser("dataset-to-neo4j", help="Load dataset triples into Neo4j")
    lnd.add_argument("--name", choices=["FB15k-237", "WN18RR"], required=True)
    lnd.add_argument("--split", choices=["train", "valid", "test", "all"], default="train")

    # Pure-Python pipeline
    pypl = sub.add_parser("py-pipeline", help="Run pure-Python pipeline on a dataset")
    pypl.add_argument("--name", choices=["FB15k-237", "WN18RR"], required=True)
    pypl.add_argument("--split", choices=["train", "valid"], default="train")
    pypl.add_argument("--samples", type=int, default=5, help="Number of demo queries to print")
    pypl.add_argument("--limit", type=int, default=0, help="If >0, sample this many triples for mining/training")
    pypl.add_argument("--min-support", type=int, default=5)
    pypl.add_argument("--min-conf", type=float, default=0.5)

    pytrain = sub.add_parser("py-train", help="Train and save a pure-Python model on a dataset split")
    pytrain.add_argument("--name", choices=["FB15k-237", "WN18RR"], required=True)
    pytrain.add_argument("--split", choices=["train", "valid"], default="train")
    pytrain.add_argument("--out", default="data/py_model.json")
    pytrain.add_argument("--limit", type=int, default=0, help="If >0, sample this many triples for mining/training")
    pytrain.add_argument("--min-support", type=int, default=5)
    pytrain.add_argument("--min-conf", type=float, default=0.5)

    pyinfer = sub.add_parser("py-query", help="Query a triple using a saved pure-Python model")
    pyinfer.add_argument("--name", choices=["FB15k-237", "WN18RR"], required=True)
    pyinfer.add_argument("--split", choices=["train", "valid"], default="train")
    pyinfer.add_argument("--model", default="data/py_model.json")
    pyinfer.add_argument("--triple", required=True, help='Triple as "H R T" (exact strings from the dataset)')

    pyask = sub.add_parser("py-ask", help="Natural language query using a saved model")
    pyask.add_argument("--name", choices=["FB15k-237", "WN18RR"], required=True)
    pyask.add_argument("--split", choices=["train", "valid"], default="train")
    pyask.add_argument("--model", default="data/py_model.json")
    pyask.add_argument("prompt", help="Question like 'h r t' or 'Is h r of t?' (must use dataset strings)")

    pyeval = sub.add_parser("py-eval", help="Evaluate a saved model with negative sampling")
    pyeval.add_argument("--name", choices=["FB15k-237", "WN18RR"], required=True)
    pyeval.add_argument("--split", choices=["train", "valid", "test"], default="valid")
    pyeval.add_argument("--model", default="data/py_model.json")
    pyeval.add_argument("--neg_k", type=int, default=10, help="Negative samples per positive")

    pyctx = sub.add_parser("py-llm-context", help="Build an LLM-ready context for a query")
    pyctx.add_argument("--name", choices=["FB15k-237", "WN18RR"], required=True)
    pyctx.add_argument("--split", choices=["train", "valid"], default="train")
    pyctx.add_argument("--model", default="data/py_model.json")
    pyctx.add_argument("--triple", help='Triple as "H R T"')
    pyctx.add_argument("--prompt", help="Natural query like 'Is h r of t?' (must use dataset strings)")
    pyctx.add_argument("--limit", type=int, default=4000, help="Sample size for witness search")
    pyctx.add_argument("--call-llm", action="store_true", help="If set, call LLM with the built context")

    # Labels + NL pipeline
    lbl = sub.add_parser("labels-build", help="Build human-friendly labels for a dataset")
    lbl.add_argument("--name", choices=["WN18RR", "FB15k-237"], required=True)
    lbl.add_argument("--out", default="data/wn18rr_labels.json")

    asknl = sub.add_parser("py-ask-nl", help="Ask a natural-language question using labels + saved model")
    asknl.add_argument("--name", choices=["WN18RR", "FB15k-237"], required=True)
    asknl.add_argument("--split", choices=["train", "valid"], default="train")
    asknl.add_argument("--model", default="data/py_model.json")
    asknl.add_argument("--labels", required=True)
    asknl.add_argument("prompt", help="e.g., 'Is dog a hypernym of canine?' or 'dog hypernym canine'")

    args = ap.parse_args()

    if args.cmd == "demo":
        from .toy_demo import run_demo
        run_demo()
        return

    if args.cmd == "load-neo4j":
        kg = build_toy_family_kg()
        load_triples(kg.all_triples())
        print("Loaded toy KG into Neo4j.")
        return

    # Build toy pipeline for ask/query on toy KG
    kg, binary, unary, contexts, params = build_pipeline()

    if args.cmd == "query":
        parts = args.triple.split()
        if len(parts) != 3:
            raise SystemExit("Provide triple as: 'H R T'")
        h, r, t = parts
        res = infer_probability(kg, (h, r, t), contexts, params, binary, unary)
        print(f"({h}, {r}, {t}) -> p={res.probability:.2f}")
        if res.supporting_rules:
            print("Supporting rules:")
            for s in res.supporting_rules:
                print("  -", s)
        return

    if args.cmd == "ask":
        triple = parse_natural_query(args.prompt)
        if not triple:
            raise SystemExit("Could not parse your prompt into a triple.")
        res = infer_probability(kg, triple, contexts, params, binary, unary)
        print(llm_wrap_explanation(args.prompt, res))
        return

    if args.cmd == "neo4j-export":
        triples = fetch_triples()
        write_triples_tsv(triples, args.out)
        print(f"Exported {len(triples)} triples to {args.out}")
        return

    if args.cmd == "anyburl-run":
        run_anyburl(args.jar, args.input, args.out, runtime_seconds=args.time)
        print(f"AnyBURL rules written to {args.out}")
        return

    if args.cmd == "contexts":
        triples = []
        with open(args.triples, "r", encoding="utf-8") as f:
            for line in f:
                h, r, t = line.rstrip("\n").split("\t")
                triples.append((h, r, t))
        kg2 = KnowledgeGraph()
        kg2.add_many(triples)
        bin_rules = list(mine_binary_path_rules(kg2, min_support=args.min_support, min_conf=args.min_conf).keys())
        un_rules = list(mine_unary_rules(kg2, min_support=args.min_support, min_conf=args.min_conf).keys())
        contexts2, triple_to_ctx = discover_contexts(kg2, triples, bin_rules, un_rules)
        os.makedirs(os.path.dirname(args.out), exist_ok=True)
        with open(args.out, "w", encoding="utf-8") as f:
            for t in triples:
                f.write(f"{t[0]}\t{t[1]}\t{t[2]}\t{triple_to_ctx.get(t, -1)}\n")
        print(f"Wrote contexts mapping to {args.out} (fallback discovery)")
        return

    if args.cmd == "train":
        triple_to_ctx = {}
        with open(args.contexts, "r", encoding="utf-8") as f:
            for line in f:
                h, r, t, cid = line.rstrip("\n").split("\t")
                triple_to_ctx[(h, r, t)] = int(cid)
        # Legacy trainer path
        model = train_with_juice({t: [c] for t, c in triple_to_ctx.items()}, set(triple_to_ctx.keys()))
        os.makedirs(os.path.dirname(args.out), exist_ok=True)
        with open(args.out, "w", encoding="utf-8") as f:
            for cid, th in model.theta.items():
                f.write(f"{cid}\t{th}\n")
        print(f"Trained model written to {args.out}")
        return

    if args.cmd == "infer":
        theta = {}
        with open(args.model, "r", encoding="utf-8") as f:
            for line in f:
                cid, th = line.rstrip("\n").split("\t")
                theta[int(cid)] = float(th)
        print("Model contexts and weights:")
        for cid, th in sorted(theta.items()):
            print(f"  ctx {cid}: {th:.3f}")
        qparts = args.query.split()
        if len(qparts) != 3:
            raise SystemExit("--query expects 'H R T'")
        print("Note: For full inference, use the 'py-query' command with a saved model.")
        return

    if args.cmd == "dataset-download":
        ds_dir = download_dataset(args.name, out_dir=args.out)
        print(f"Downloaded {args.name} to {ds_dir}")
        return

    if args.cmd == "dataset-to-neo4j":
        ds_dir = os.path.join("data", args.name)
        if args.split == "all":
            triples = all_triples(ds_dir)
        else:
            triples = load_triples_from_dir(ds_dir, args.split)
        from .neo4j_utils import bulk_load_triples
        bulk_load_triples(triples)
        print(f"Loaded {len(triples)} {args.split} triples from {args.name} into Neo4j.")
        return

    if args.cmd == "py-pipeline":
        ds_dir = os.path.join("data", args.name)
        triples = load_triples_from_dir(ds_dir, args.split)
        if args.limit and args.limit > 0 and len(triples) > args.limit:
            random.seed(0)
            triples = random.sample(triples, args.limit)
        # Optional RDFlib store
        try:
            from .rdflib_store import RDFStore  # type: ignore
            store = RDFStore()
            store.add_triples(triples)
        except Exception:
            pass
        kg2 = KnowledgeGraph()
        kg2.add_many(triples)
        bin_rules = list(mine_binary_path_rules(kg2, min_support=args.min_support, min_conf=args.min_conf).keys())
        un_rules = list(mine_unary_rules(kg2, min_support=args.min_support, min_conf=args.min_conf).keys())
        contexts2, triple_to_ctx = discover_contexts(kg2, triples, bin_rules, un_rules)
        triple_to_candidates = {t: [cid] for t, cid in triple_to_ctx.items()}
        positives = set(triple_to_candidates.keys())
        model = train_with_juice(triple_to_candidates, positives)
        print(f"Pure-Python pipeline: {args.name} [{args.split}] -> {len(triples)} triples, {len(contexts2)} contexts")
        count = 0
        for t in triples:
            cids = triple_to_candidates.get(t, [])
            p = score_query_with_juice(model, cids)
            print(f"  {t} -> p={p:.2f}")
            count += 1
            if count >= args.samples:
                break
        return

    if args.cmd == "py-train":
        ds_dir = os.path.join("data", args.name)
        triples = load_triples_from_dir(ds_dir, args.split)
        if args.limit and args.limit > 0 and len(triples) > args.limit:
            random.seed(0)
            triples = random.sample(triples, args.limit)
        kg2 = KnowledgeGraph()
        kg2.add_many(triples)
        bin_rules = list(mine_binary_path_rules(kg2, min_support=5, min_conf=0.5).keys())
        un_rules = list(mine_unary_rules(kg2, min_support=5, min_conf=0.5).keys())
        contexts2, triple_to_ctx = discover_contexts(kg2, triples, bin_rules, un_rules)
        triple_to_candidates = {t: [cid] for t, cid in triple_to_ctx.items()}
        positives = set(triple_to_candidates.keys())
        model = train_with_juice(triple_to_candidates, positives)
        theta = model.theta
        os.makedirs(os.path.dirname(args.out), exist_ok=True)
        save_model(args.out, bin_rules, un_rules, contexts2, theta)
        print(f"Saved model to {args.out} with {len(contexts2)} contexts.")
        return

    if args.cmd == "py-query":
        ds_dir = os.path.join("data", args.name)
        triples = load_triples_from_dir(ds_dir, args.split)
        kg2 = KnowledgeGraph()
        kg2.add_many(triples)
        bin_rules, un_rules, contexts2, theta = load_model(args.model)
        parts = args.triple.split()
        if len(parts) != 3:
            raise SystemExit("--triple expects 'H R T'")
        h, r, t = parts
        brs, urs = rules_that_entail(kg2, h, r, t, bin_rules, un_rules)
        cids = [cid for cid, ctx in contexts2.items() if ctx.binary_rules == brs and ctx.unary_rules == urs]
        model = JuiceModel(theta=theta, triple_to_cids={})
        p = score_query_with_juice(model, cids)
        rules_str = [str(x) for x in list(brs) + list(urs)]
        print(f"({h}, {r}, {t}) -> p={p:.2f}")
        if rules_str:
            print("Supporting rules:")
            for s in rules_str[:10]:
                print("  -", s)
        else:
            print("No supporting rules fired; probability from priors (likely low).")
        return

    if args.cmd == "py-ask":
        ds_dir = os.path.join("data", args.name)
        triples = load_triples_from_dir(ds_dir, args.split)
        kg2 = KnowledgeGraph()
        kg2.add_many(triples)
        bin_rules, un_rules, contexts2, theta = load_model(args.model)
        triple = parse_natural_query(args.prompt)
        if not triple:
            raise SystemExit("Could not parse your prompt into a triple. Use exact dataset strings.")
        h, r, t = triple
        brs, urs = rules_that_entail(kg2, h, r, t, bin_rules, un_rules)
        cids = [cid for cid, ctx in contexts2.items() if ctx.binary_rules == brs and ctx.unary_rules == urs]
        model = JuiceModel(theta=theta, triple_to_cids={})
        p = score_query_with_juice(model, cids)
        res = InferenceResult(triple=(h, r, t), probability=p, supporting_rules=[str(x) for x in list(brs) + list(urs)])
        print(llm_wrap_explanation(args.prompt, res))
        return

    if args.cmd == "py-eval":
        ds_dir = os.path.join("data", args.name)
        triples = load_triples_from_dir(ds_dir, args.split)
        entities = sorted({h for h, _, _ in triples} | {t for _, _, t in triples})
        kg2 = KnowledgeGraph()
        kg2.add_many(triples)
        bin_rules, un_rules, contexts2, theta = load_model(args.model)
        model = JuiceModel(theta=theta, triple_to_cids={})
        def score_triple(tr):
            h, r, t = tr
            brs, urs = rules_that_entail(kg2, h, r, t, bin_rules, un_rules)
            cids = [cid for cid, ctx in contexts2.items() if ctx.binary_rules == brs and ctx.unary_rules == urs]
            return score_query_with_juice(model, cids)
        pos_scores = []
        neg_scores = []
        for (h, r, t) in triples:
            pos_scores.append(score_triple((h, r, t)))
            for _ in range(args.neg_k):
                nt = t
                while nt == t:
                    nt = random.choice(entities)
                neg_scores.append(score_triple((h, r, nt)))
        mean_pos = sum(pos_scores) / max(1, len(pos_scores))
        mean_neg = sum(neg_scores) / max(1, len(neg_scores))
        combined = [(s, 1) for s in pos_scores] + [(s, 0) for s in neg_scores]
        combined.sort(key=lambda x: x[0])
        rank = 1
        sum_ranks_pos = 0.0
        n_pos = len(pos_scores)
        n_neg = len(neg_scores)
        for s, y in combined:
            if y == 1:
                sum_ranks_pos += rank
            rank += 1
        auc = (sum_ranks_pos - n_pos * (n_pos + 1) / 2) / max(1, (n_pos * n_neg)) if n_neg > 0 else 0.0
        print(f"Pos mean p={mean_pos:.3f}  Neg mean p={mean_neg:.3f}  AUROC={auc:.3f}  (neg_k={args.neg_k})")
        return

    if args.cmd == "py-llm-context":
        from .llm_interface import build_llm_context
        ds_dir = os.path.join("data", args.name)
        triples = load_triples_from_dir(ds_dir, args.split)
        if args.limit and args.limit > 0 and len(triples) > args.limit:
            random.seed(0)
            triples = random.sample(triples, args.limit)
        kg2 = KnowledgeGraph()
        kg2.add_many(triples)
        bin_rules, un_rules, contexts2, theta = load_model(args.model)
        # Parse input
        if args.triple:
            parts = args.triple.split()
            if len(parts) != 3:
                raise SystemExit("--triple expects 'H R T'")
            h, r, t = parts
            qtext = f"({h}, {r}, {t})"
        elif args.prompt:
            triple = parse_natural_query(args.prompt)
            if not triple:
                raise SystemExit("Could not parse your prompt into a triple.")
            h, r, t = triple
            qtext = args.prompt
        else:
            raise SystemExit("Provide either --triple or --prompt")

        # Compute support
        brs, urs = rules_that_entail(kg2, h, r, t, bin_rules, un_rules)
        cids = [cid for cid, ctx in contexts2.items() if ctx.binary_rules == brs and ctx.unary_rules == urs]
        model = JuiceModel(theta=theta, triple_to_cids={})
        p = score_query_with_juice(model, cids)
        res = InferenceResult(triple=(h, r, t), probability=p, supporting_rules=[str(x) for x in list(brs) + list(urs)])

        # Context weights
        ctx_weights = [(cid, theta.get(cid, 0.2)) for cid in cids]

        # Witness extraction: for each binary rule fired, find an example y such that r1(h,y) & r2(y,t)
        witnesses = []
        for rule in brs:
            for (hh, y) in kg2.facts.get(rule.r1, set()):
                if hh != h:
                    continue
                if (y, t) in kg2.facts.get(rule.r2, set()):
                    witnesses.append(f"{rule}: with y={y}")
                    break

        context_text = build_llm_context(qtext, res, ctx_weights, witnesses)
        print(context_text)
        if getattr(args, "call_llm", False):
            # Call LLM with the built context as the user message
            try:
                from openai import OpenAI
                import os as _os
                api_key = _os.environ.get("OPENAI_API_KEY")
                if not api_key:
                    print("OPENAI_API_KEY not set; skipping LLM call.")
                    return
                client = OpenAI(api_key=api_key)
                resp = client.chat.completions.create(
                    model="gpt-4o-mini",
                    messages=[
                        {"role": "system", "content": "You explain KG reasoning with probabilities clearly and concisely."},
                        {"role": "user", "content": context_text},
                    ],
                    temperature=0.2,
                )
                print("\nLLM Answer:\n" + (resp.choices[0].message.content or ""))
            except Exception as e:
                print(f"LLM call failed: {e}")
        return

    if args.cmd == "labels-build":
        ds_dir = os.path.join("data", args.name)
        if args.name == "WN18RR":
            from .labels import build_wn18rr_labels as builder
        else:
            from .labels import build_fb15k_labels as builder
        from .labels import save_labels
        labels = builder(ds_dir)
        save_labels(args.out, labels)
        print(f"Labels saved to {args.out}")
        return

    if args.cmd == "py-ask-nl":
        from .labels import load_labels, parse_nl_query_to_triple, parse_fb_nl_query
        labels = load_labels(args.labels)
        # Resolve NL to triple ids
        if args.name == "WN18RR":
            triple = parse_nl_query_to_triple(args.prompt, labels)
        else:
            triple = parse_fb_nl_query(args.prompt, labels)
        if not triple:
            raise SystemExit("Could not resolve your question to a dataset triple. Try different lemmas/words.")
        h, r, t = triple
        # Load dataset + model
        ds_dir = os.path.join("data", args.name)
        triples = load_triples_from_dir(ds_dir, args.split)
        kg2 = KnowledgeGraph()
        kg2.add_many(triples)
        bin_rules, un_rules, contexts2, theta = load_model(args.model)
        brs, urs = rules_that_entail(kg2, h, r, t, bin_rules, un_rules)
        cids = [cid for cid, ctx in contexts2.items() if ctx.binary_rules == brs and ctx.unary_rules == urs]
        model = JuiceModel(theta=theta, triple_to_cids={})
        p = score_query_with_juice(model, cids)
        res = InferenceResult(triple=(h, r, t), probability=p, supporting_rules=[str(x) for x in list(brs) + list(urs)])
        # Construct a nicer prompt using labels
        from .llm_interface import build_llm_context
        subj_lbl = labels.get("entity_labels", {}).get(h, {}).get("label", h)
        obj_lbl = labels.get("entity_labels", {}).get(t, {}).get("label", t)
        rel_lbl = labels.get("relation_labels", {}).get(r, r)
        question = f"Is {subj_lbl} {rel_lbl} of {obj_lbl}?"
        ctx_text = build_llm_context(question, res, [(cid, theta.get(cid, 0.2)) for cid in cids], [])
        print(ctx_text)
        # Optionally call LLM if key is present
        from .llm_interface import OpenAI
        api_key = os.environ.get("OPENAI_API_KEY")
        if api_key and OpenAI is not None:
            try:
                client = OpenAI(api_key=api_key)
                resp = client.chat.completions.create(
                    model="gpt-4o-mini",
                    messages=[
                        {"role": "system", "content": "You explain KG reasoning with probabilities clearly and concisely."},
                        {"role": "user", "content": ctx_text},
                    ],
                    temperature=0.2,
                )
                print("\nLLM Answer:\n" + (resp.choices[0].message.content or ""))
            except Exception as e:
                print(f"LLM call failed: {e}")
        return


if __name__ == "__main__":
    main()