"""Minimal example: score a post with a CTR model and get next-post suggestions from SASRec.
pip install onnxruntime numpy
"""
import json, os, numpy as np, onnxruntime as ort
HERE = os.path.dirname(os.path.abspath(__file__))
meta = lambda f: json.load(open(os.path.join(HERE, "metadata", f), encoding="utf8"))
report, vocab, stats = meta("export_report.json"), meta("ctr_vocab.json"), meta("feature_stats.json")

# ---------- 1) engagement probability: LR / DeepFM / WideDeep / DCNV2 ----------
def score(model_name, tokens: dict, raw_features: dict):
    """tokens: {'user_id': 'P09', 'subreddit': 'r/MachineLearning', ...}  (unknown value -> id 0)
       raw_features: {feature_name: raw value}  (missing -> training mean, i.e. z-score 0)"""
    r = report[model_name]
    tok = np.array([[vocab[f].get(str(tokens.get(f, "none")), 0) for f in r["token_inputs"]]], dtype=np.int64)
    z = [(raw_features.get(n, m) - m) / s for n, m, s in zip(stats["features"], stats["mean"], stats["std"])]
    flt = np.array([z], dtype=np.float32)
    sess = ort.InferenceSession(os.path.join(HERE, "onnx", f"{model_name}.onnx"), providers=["CPUExecutionProvider"])
    return float(sess.run(None, {"token_ids": tok, "float_features": flt})[0][0])

example_tokens = {"user_id": "P09", "subreddit": "r/MachineLearning", "post_type": "text", "page_kind": "subreddit_feed"}
example_raw = {"rank": np.log1p(3), "log_score": np.log1p(250), "log_comments": np.log1p(40), "has_caption": 0.0}
for m in ("LR", "DeepFM", "WideDeep", "DCNV2"):
    print(f"{m:9s} P(engage) = {score(m, example_tokens, example_raw):.3f}")

# ---------- 2) next-post suggestion: SASRec ----------
sv = meta("sasrec_vocab.json")["item_id"]; inv = {i: t for t, i in sv.items()}
def suggest(history_post_ids, k=10, L=20):
    ids = [sv[p] for p in history_post_ids if p in sv][-L:]
    seq = np.zeros((1, L), dtype=np.int64); seq[0, :len(ids)] = ids          # right-padded with 0
    sess = ort.InferenceSession(os.path.join(HERE, "onnx", "SASRec.onnx"), providers=["CPUExecutionProvider"])
    scores = sess.run(None, {"item_seq": seq, "item_seq_len": np.array([len(ids)], dtype=np.int64)})[0][0]
    scores[0] = -1e9                                                         # never suggest [PAD]
    for i in set(ids): scores[i] = -1e9                                      # skip posts already seen
    return [inv[i] for i in np.argsort(-scores)[:k]]
some = [t for t in list(sv)[1:6]]
print("history:", some, "\nSASRec suggests:", suggest(some, k=5))