"""Build engine 3.0 (hybrid-engine-3.0-bge384): the v2.0 architecture on 384-d BGE content vectors and ONLY features the UniHive `main` app can supply.

Inputs kept (cold, Wide&Deep):  log likes, log comments, log age, title length, body length, 16 content components (BGE-384 -> PCA), 8 tag-match features
                                (our 114 tags, binarised to confidence 1.0 because the app's tagger returns no confidences)
Inputs kept (warm, DeepFM)   :  the same post inputs + history: h_n, h_pos_rate, h_mean_focal (log active time), h_click, auth_seen, auth_rate,
                                aff_recent, aff_neg, 16 user-taste components (uemb)
Dropped everywhere           :  community (incl. the sub_* rates and the `domain` token, which carries subreddit names), flair, language, post_type,
                                has_caption, is_ad, spoiler, awards, vote ratio, membership, hover, user/post ids, feed position, session signals
Label                        :  engaged (click/vote/save/join/reply/detail) OR active time >= 10 s  (prep_recbole LABEL_MODE=active10000, FOCAL_COL=active_ms)
Architecture, hyper-parameters, seeds, calibration and export are identical to v2.0 (build_engine_v2.py).
Prerequisites (run once, in this order, from reco_data):
  python engine_v3/embed_posts_bge384.py
  EMB_NPZ=post_emb_bge384.npz EMB_KEY=fused LABEL_MODE=active10000 FOCAL_COL=active_ms META_DIR=package_meta_bge SUFFIX=_bge python prep_recbole.py
  (same + SUFFIX=_bge_noid_strict MASK_IDS=1 STRICT_PAST=1)
Writes hybrid_engine/artifacts_v3 (v2.0 artifacts are only read)."""
import os, sys, json, glob, shutil
HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(HERE, "engine_v2"))
import numpy as np, pandas as pd
from engine_lib import CHILD, CHILD_NAME, PARENT_OF, BROAD_NAME, C, TF, tag_features, implied_parents, max_sim, load_post_tags
src = open(os.path.join(HERE, "run_models.py"), encoding="utf8").read().split("if __name__")[0].replace("os.path.dirname(os.path.abspath(__file__))", repr(HERE)); exec(src)
from recbole.quick_start import load_data_and_model
import torch.nn as nn, onnx, onnxruntime as ort
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score

OUT = os.path.join(HERE, "out"); DATA = os.path.join(HERE, "recbole_data"); ART = os.path.join(HERE, "hybrid_engine", "artifacts_v3"); CK = os.path.join(OUT, "engine_v3")
os.makedirs(ART, exist_ok=True); os.makedirs(CK, exist_ok=True); SEEDS = range(5); PICKS = 6
rd = lambda name, sp: pd.read_csv(os.path.join(DATA, name, f"{name}.{sp}.inter"), sep="\t")
lg = lambda q: np.log(np.clip(q, 1e-6, 1 - 1e-6) / (1 - np.clip(q, 1e-6, 1 - 1e-6)))[:, None]
POST5 = ["log_score", "log_comments", "log_age_h", "title_len", "body_len"]; PCA16 = [f"pca{i}" for i in range(16)]
HIST = ["h_n", "h_pos_rate", "h_mean_focal", "h_click", "auth_seen", "auth_rate", "aff_recent", "aff_neg"]; UEMB = [f"uemb{i}" for i in range(16)]
COLD_FLT, WARM_FLT = POST5 + PCA16 + TF, POST5 + HIST + UEMB + PCA16

# ---------------- tag similarity in BGE space (one embedding model end to end) ----------------
p_tag = os.path.join(OUT, "engine", "tag_emb_bge.npy")
if not os.path.exists(p_tag):
    from fastembed import TextEmbedding
    m = TextEmbedding(model_name="BAAI/bge-small-en-v1.5", cache_dir=os.path.join(HERE, "bge_cache"), threads=4)
    E = np.asarray(list(m.embed([f"{CHILD_NAME[i]} ({BROAD_NAME[PARENT_OF[i]]})" for i in range(C)])), dtype=np.float64); E = (E / np.linalg.norm(E, axis=1, keepdims=True)).astype("float32")
    np.save(p_tag, E)
E = np.load(p_tag); S = np.clip(E @ E.T, 0, 1); np.fill_diagonal(S, 1.0)

# ---------------- cold dataset: tag-match features on BINARISED tags, inferred picks from the NEW label ----------------
sl = pd.read_csv(os.path.join(OUT, "split_labels_bge.csv"))
ctr = {sp: rd("reco_ctr_bge", sp) for sp in ("train", "valid", "test")}
rows = pd.concat([ctr[sp].assign(split=sp) for sp in ("train", "valid", "test")], ignore_index=True)
sl = pd.concat([sl[sl.split == sp] for sp in ("train", "valid", "test")], ignore_index=True)
assert (sl.post_id.values == rows["item_id:token"].values).all() and np.allclose(sl.label.values, rows["label:float"].values)
tags = {k: (v > 0).astype(np.float32) for k, v in load_post_tags(os.path.join(OUT, "post_tags.jsonl")).items()}      # binary: the app returns no confidences
P = np.stack([tags.get(p, np.zeros(C, np.float32)) for p in sl.post_id]); print("rows without a tag:", round(float((P.sum(1) == 0).mean()), 3))
idf = (np.log((len(tags) + 1) / ((np.stack(list(tags.values())) > 0).sum(0) + 1)) + 1).astype(np.float32)
label, split = sl.label.values, sl.split.values
Q = np.zeros((len(sl), C), np.float32)
for u in sl.user_id.unique():
    mk = (sl.user_id == u).values; trpos = mk & (split == "train") & (label == 1); T = P[trpos].sum(0)
    for i in np.where(mk)[0]:
        t = T - P[i] if (split[i] == "train" and label[i] == 1) else T
        top = np.argsort(-t)[:PICKS]; top = top[t[top] > 0]; Q[i, top] = 1.0
R = implied_parents(Q); MS = np.zeros((len(sl), C), np.float32); uq = {}
for i in range(len(sl)):
    k = Q[i].tobytes()
    if k not in uq: uq[k] = max_sim(Q[i:i + 1], S)[0]
    MS[i] = uq[k]
TFm = tag_features(Q, R, MS, P, idf); tr = split == "train"; mu, sd = TFm[tr].mean(0), TFm[tr].std(0) + 1e-6
lr = LogisticRegression(C=1.0, max_iter=1000).fit((TFm[tr] - mu) / sd, label[tr]); s1 = ((TFm - mu) / sd) @ lr.coef_[0] + lr.intercept_[0]
stage1_auc = {sp: round(float(roc_auc_score(label[split == sp], s1[split == sp])), 3) for sp in ("train", "valid", "test")}
print("stage-1 tag closeness AUC:", stage1_auc)
z = (TFm - mu) / sd; d_cold = os.path.join(DATA, "reco_engine_bge"); os.makedirs(d_cold, exist_ok=True)
for sp in ("train", "valid", "test"):
    mk = split == sp; out = pd.DataFrame({"user_id:token": "unk", "item_id:token": "unk", "label:float": label[mk], "timestamp:float": rows["timestamp:float"].values[mk]})
    for f in POST5 + PCA16: out[f"{f}:float"] = rows[f"{f}:float"].values[mk]
    for j, f in enumerate(TF): out[f"{f}:float"] = z[mk, j].round(5)
    out.to_csv(os.path.join(d_cold, f"reco_engine_bge.{sp}.inter"), sep="\t", index=False)
np.savez_compressed(os.path.join(OUT, "engine_v3", "real_rows_v3.npz"), Q=Q, P=P, s1=s1, split=split, label=label, user=sl.user_id.values, post=sl.post_id.values, ts=sl.ts.values)

# ---------------- export helpers (identical to v2.0) ----------------
class Ens(nn.Module):
    def __init__(s, models, tok, flt): super().__init__(); s.ms = nn.ModuleList(models); s.tok, s.flt = tok, flt
    def forward(s, token_ids, float_features):
        d = {nm: token_ids[:, i] for i, nm in enumerate(s.tok)}; d.update({nm: torch.stack([float_features[:, i], torch.ones_like(float_features[:, i])], -1) for i, nm in enumerate(s.flt)})
        return torch.stack([torch.sigmoid(m.forward(d)).reshape(-1) for m in s.ms], 0).mean(0)
def export(models, ds, path):
    tok, flt = list(models[0].token_field_names), list(models[0].float_field_names)
    assert all(list(m.token_field_names) == tok and list(m.float_field_names) == flt for m in models)
    e = Ens(models, tok, flt).eval(); torch.onnx.export(e, (torch.zeros(2, len(tok), dtype=torch.long), torch.zeros(2, len(flt))), path, input_names=["token_ids", "float_features"], output_names=["probability"],
                                                         dynamic_axes={k: {0: "batch"} for k in ("token_ids", "float_features", "probability")}, opset_version=17, dynamo=False)
    onnx.checker.check_model(onnx.load(path)); return e, tok, flt, {f: ds.field2token_id[f] for f in tok}
def arrays(df, tok, flt, vocab):
    T = np.stack([df[f"{f}:token"].astype(str).map(lambda s: vocab[f].get(s, 0)).values for f in tok], 1).astype(np.int64); return T, df[[f"{f}:float" for f in flt]].values.astype(np.float32)
def check(e, path, df, tok, flt, vocab):
    T, X = arrays(df, tok, flt, vocab); a = ort.InferenceSession(path, providers=["CPUExecutionProvider"]).run(None, {"token_ids": T, "float_features": X})[0]
    with torch.no_grad(): b = e(torch.tensor(T), torch.tensor(X)).numpy()
    return float(np.abs(a - b).max()), a
auc = lambda name, sp, p: round(float(roc_auc_score(rd(name, sp)["label:float"].values, p)), 3)
report = dict(label="engaged OR active_ms >= 10 s", positives={sp: int(label[split == sp].sum()) for sp in ("train", "valid", "test")}, stage1_auc=stage1_auc,
              features=dict(cold=COLD_FLT, warm=WARM_FLT), embedding="BAAI/bge-small-en-v1.5 (384-d) via fastembed; training text includes captions, serving text does not")

# ---------------- cold model: Wide&Deep, 5 seeds ----------------
S2P = json.load(open(os.path.join(HERE, "engine_v1_frozen", "out_engine", "stage2_summary.json")))["WideDeep"]["params"]; cold = []; cold_auc = []
for s in SEEDS:
    c = config_for("WideDeep", s); c.update(dataset="reco_engine_bge", numerical_features=COLD_FLT, load_col={"inter": ["user_id", "item_id", "label", "timestamp"] + COLD_FLT},
                                            mlp_hidden_size=[64, 32], dropout_prob=0.3); c.update(S2P)
    r = run_recbole(model="WideDeep", config_dict=c); ck = max(glob.glob(os.path.join(HERE, "recbole_saved", "WideDeep-*.pth")), key=os.path.getmtime)
    shutil.copy(ck, os.path.join(CK, f"ship3_cold_s{s}.pth")); _, m, ds_c, *_ = load_data_and_model(ck); cold.append(m.cpu().eval()); cold_auc.append(round(float(r["test_result"]["auc"]), 3))
ec, tok_c, flt_c, voc_c = export(cold, ds_c, os.path.join(ART, "cold_stage2.onnx"))
dv, pv = check(ec, os.path.join(ART, "cold_stage2.onnx"), rd("reco_engine_bge", "valid"), tok_c, flt_c, voc_c); dt, pt = check(ec, os.path.join(ART, "cold_stage2.onnx"), rd("reco_engine_bge", "test"), tok_c, flt_c, voc_c)
platt = LogisticRegression(C=1e6).fit(lg(pv), label[split == "valid"])
report["cold"] = dict(onnx_vs_pytorch_max_abs_diff=max(dv, dt), valid_auc=auc("reco_engine_bge", "valid", pv), test_auc=auc("reco_engine_bge", "test", pt), per_seed_test_auc=cold_auc, token_inputs=tok_c, float_inputs=flt_c, params="embedding 4, dropout 0.3, weight decay 1e-2 (as v2.0)")
print("cold:", {k: report["cold"][k] for k in ("valid_auc", "test_auc", "per_seed_test_auc")}, flush=True)

# ---------------- warm model: DeepFM, serving-rule history, 5 seeds ----------------
BESTW = json.load(open(os.path.join(HERE, "out", "best_params_v2.json")))["DeepFM"]; warm = []; warm_auc = []
for s in SEEDS:
    c = config_for("DeepFM", s); c.update(BESTW); c.update(dataset="reco_ctr_bge_noid_strict", numerical_features=WARM_FLT, load_col={"inter": ["user_id", "item_id", "label", "timestamp"] + WARM_FLT})
    r = run_recbole(model="DeepFM", config_dict=c); ck = max(glob.glob(os.path.join(HERE, "recbole_saved", "DeepFM-*.pth")), key=os.path.getmtime)
    shutil.copy(ck, os.path.join(CK, f"ship3_warm_s{s}.pth")); _, m, ds_w, *_ = load_data_and_model(ck); warm.append(m.cpu().eval()); warm_auc.append(round(float(r["test_result"]["auc"]), 3))
ew, tok_w, flt_w, voc_w = export(warm, ds_w, os.path.join(ART, "warm_deepfm.onnx"))
d1, wv = check(ew, os.path.join(ART, "warm_deepfm.onnx"), rd("reco_ctr_bge_noid_strict", "valid"), tok_w, flt_w, voc_w); d2, wt = check(ew, os.path.join(ART, "warm_deepfm.onnx"), rd("reco_ctr_bge_noid_strict", "test"), tok_w, flt_w, voc_w)
report["warm"] = dict(onnx_vs_pytorch_max_abs_diff=max(d1, d2), valid_auc=auc("reco_ctr_bge_noid_strict", "valid", wv), test_auc=auc("reco_ctr_bge_noid_strict", "test", wt), per_seed_test_auc=warm_auc, token_inputs=tok_w, float_inputs=flt_w, params="embedding 16, dropout 0.3, weight decay 1e-3 (as v2.0)")
print("warm:", {k: report["warm"][k] for k in ("valid_auc", "test_auc", "per_seed_test_auc")}, flush=True)
np.savez_compressed(os.path.join(CK, "ship3_offline_scores.npz"), cold_valid=pv, cold_test=pt, warm_valid=wv, warm_test=wt)

# ---------------- meta.json ----------------
st = json.load(open(os.path.join(HERE, "package_meta_bge", "feature_stats.json"))); pca = json.load(open(os.path.join(HERE, "package_meta_bge", "pca.json")))
OLD = json.load(open(os.path.join(HERE, "hybrid_engine", "artifacts", "meta.json"), encoding="utf8")); tax = dict(OLD["taxonomy"]); tax["tag_similarity"] = np.round(S, 4).tolist()
it = pd.read_csv(os.path.join(OUT, "interactions_v2.csv"))
meta = dict(OLD)
meta.update(version="hybrid-engine-3.0-bge384",
            embedding=dict(model="BAAI/bge-small-en-v1.5", runtime="fastembed TextEmbedding.embed (passage mode), L2-normalised", dim=384,
                           text="'\\n'.join(title, body, summary, description)[:12000], as the UniHive main personalization worker builds it; training text additionally puts the image caption in the summary slot"),
            cold=dict(onnx="cold_stage2.onnx", token_inputs=tok_c, float_inputs=flt_c, vocab=voc_c, platt=dict(coef=float(platt.coef_[0][0]), intercept=float(platt.intercept_[0]))),
            warm=dict(onnx="warm_deepfm.onnx", token_inputs=tok_w, float_inputs=flt_w, vocab=voc_w, history_rule="strictly before the request, viewing ended; focal time := active time"),
            feature_stats=dict(zip(st["features"], zip(st["mean"], st["std"]))),
            stage1=dict(features=TF, mean=mu.tolist(), std=sd.tolist(), coef=lr.coef_[0].tolist(), intercept=float(lr.intercept_[0]), idf=idf.tolist()),
            taxonomy=tax,
            pca=dict(mean=pca["pca_mean"], components=pca["pca_components"], component_std=pca["pca_component_std"], embedding_model="BAAI/bge-small-en-v1.5",
                     embedding_text="title + body (+ caption in training only); no community, no flair"),
            constants=dict(global_positive_rate=float(label[tr].mean()), age_hours_median=float(it.age_hours.median()), history_window=20, top_domains=[], click_prior=0.08),
            label=dict(definition="engaged OR active_ms >= 10000", active_ms_threshold=10000))
json.dump(meta, open(os.path.join(ART, "meta.json"), "w", encoding="utf8"), ensure_ascii=False)
report["meta_kb"] = round(os.path.getsize(os.path.join(ART, "meta.json")) / 1024, 1); report["onnx_kb"] = {f: round(os.path.getsize(os.path.join(ART, f)) / 1024, 1) for f in ("cold_stage2.onnx", "warm_deepfm.onnx")}
json.dump(report, open(os.path.join(ART, "build_report.json"), "w"), indent=1); print(json.dumps({k: v for k, v in report.items() if k != "features"}, indent=1)); print("BUILD V3 DONE")