"""Feature builders for the hybrid engine. They reproduce, from raw post/user data, exactly the inputs the models were trained on
(verified against the training files by verify_hybrid_engine.py)."""
import math, re
import numpy as np

def clean(s):
    return " ".join(str(s).split()) if isinstance(s, str) else ""

def token(v):
    """Category value -> the token string used in training ('none' for missing/empty, whitespace -> '_')."""
    if v is None or (isinstance(v, float) and math.isnan(v)):
        return "none"
    s = re.sub(r"\s+", "_", str(v))
    return s if s != "" else "none"

def unit(v):
    return v / (np.linalg.norm(v) + 1e-9)

EMB_DIM = 1024          # Qwen/Qwen3-Embedding-0.6B

def content_vector(emb):
    """A usable Qwen3 vector as a unit numpy array, or None when the embedding is missing, has the wrong size
    (for example the backend's 512-d CLIP vectors, which live in a different space) or is all zeros / not finite."""
    if emb is None:
        return None
    v = np.asarray(emb, dtype=float).ravel()
    if v.shape[0] != EMB_DIM or not np.isfinite(v).all() or not v.any():
        return None
    return unit(v)

def rank01(x):
    """Rank-normalise scores inside one request: lowest -> 1/n, highest -> 1."""
    x = np.asarray(x, dtype=float)
    return (np.argsort(np.argsort(x)) + 1) / len(x) if len(x) else x


class FeatureBuilder:
    def __init__(self, meta):
        self.m = meta
        self.stats = meta["feature_stats"]
        c = meta["constants"]
        self.gl, self.age_median, self.W, self.top_domains, self.click_prior = c["global_positive_rate"], c["age_hours_median"], c["history_window"], set(c["top_domains"]), c["click_prior"]
        p = meta["pca"]
        self.pca_mean, self.pca_comp, self.pca_std = np.array(p["mean"]), np.array(p["components"]), np.array(p["component_std"])
        t = meta["taxonomy"]
        self.child = t["child"]; self.cidx = {s: i for i, s in enumerate(self.child)}; self.parent_of = np.array(t["parent_of"]); self.n_broad = len(t["broad"])
        self.S = np.array(t["tag_similarity"], dtype=float)
        s1 = meta["stage1"]
        self.s1_mu, self.s1_sd, self.s1_coef, self.s1_b, self.idf = np.array(s1["mean"]), np.array(s1["std"]), np.array(s1["coef"]), float(s1["intercept"]), np.array(s1["idf"])

    # ---------- post content ----------
    def pca16(self, emb):
        """1024-d Qwen3-Embedding-0.6B vector (unit length) -> the 16 content components the models use."""
        return ((np.asarray(emb, dtype=float) - self.pca_mean) @ self.pca_comp.T) / self.pca_std

    def z(self, name, raw):
        mu, sd = self.stats[name]
        return (raw - mu) / sd

    def post_floats(self, post):
        """Raw post statistics + content components (before standardisation)."""
        score = max(float(post.get("score") or 0), 0.0)
        age = post.get("age_hours")
        age = self.age_median if age is None or (isinstance(age, float) and math.isnan(age)) else max(float(age), 0.0)
        ur = post.get("upvote_ratio")
        missing = ur is None or (isinstance(ur, float) and math.isnan(ur))
        raw = {
            "log_score": math.log1p(score), "log_comments": math.log1p(float(post.get("comments") or 0)), "log_age_h": math.log1p(age),
            "title_len": math.log1p(len(clean(post.get("title")))), "body_len": math.log1p(len(clean(post.get("body")))),
            "has_caption": float(bool(clean(post.get("image_caption")))), "is_ad": float(bool(post.get("is_ad"))),
            "upvote_ratio": 1.0 if missing else float(ur), "upvote_missing": float(missing),
            "log_awards": math.log1p(float(post.get("awards") or 0)), "is_spoiler": float(bool(post.get("is_spoiler"))),
        }
        emb = post.get("embedding")      # None = content unknown -> 0 on every component (the average post)
        for i, v in enumerate(self.pca16(emb) if emb is not None else np.zeros(16)):
            raw[f"pca{i}"] = float(v)
        return raw

    def post_tokens(self, post):
        dom = post.get("domain")
        dom = dom if isinstance(dom, str) and dom in self.top_domains else "other"
        mem = post.get("is_member")
        return {"user_id": "unk", "item_id": "unk", "subreddit": token(post.get("community")), "post_type": token(post.get("post_type")), "flair": token(post.get("flair")),
                "domain": token(dom), "post_language": token(post.get("language")), "member": {True: "yes", False: "no"}.get(mem, "unk")}

    # ---------- tags ----------
    def tag_vector(self, tags):
        """{slug: confidence} or [slug, ...] -> confidence vector over the taxonomy's child tags (unknown slugs ignored)."""
        v = np.zeros(len(self.child))
        items = tags.items() if isinstance(tags, dict) else ((t, 1.0) for t in (tags or []))
        for s, c in items:
            if s in self.cidx:
                v[self.cidx[s]] = max(v[self.cidx[s]], float(c))
        return v

    def interest_profile(self, picks):
        q = np.zeros(len(self.child))
        for s in picks or []:
            if s in self.cidx:
                q[self.cidx[s]] = 1.0
        r = np.array([q[self.parent_of == b].max() if (self.parent_of == b).any() else 0.0 for b in range(self.n_broad)])
        ms = (self.S * q[None, :]).max(1) if q.any() else np.zeros(len(self.child))
        return q, r, ms

    def tag_match(self, P, q, r, ms):
        """The 8 tag-match features (exact, parent, rarity-weighted exact, soft, top-tag hit, #matches, #tags, mean confidence)."""
        sumc = P.sum(); safe = sumc if sumc > 0 else 1.0; nt = int((P > 0).sum())
        exact = (P * q).sum() / safe
        parent = (P * (1 - q) * r[self.parent_of]).sum() / safe
        den = (P * self.idf).sum(); idf_e = (P * q * self.idf).sum() / (den if den > 0 else 1.0)
        soft = (P * ms).sum() / safe
        top = float(q[int(P.argmax())]) if nt > 0 else 0.0
        return np.array([exact, parent, idf_e, soft, top, ((P > 0) * q).sum(), nt, sumc / max(nt, 1) if nt > 0 else 0.0])

    def tag_closeness(self, tm):
        """Stage-1 tag closeness (fitted logistic score) from the raw tag-match features."""
        return float(((tm - self.s1_mu) / self.s1_sd) @ self.s1_coef + self.s1_b)

    # ---------- user history ----------
    def history_floats(self, history, post, now_ms, post_X, strict_prefix=False):
        """Causal history features for one candidate: only past posts whose viewing had ended before `now_ms` count.
        strict_prefix=True means `history` is already exactly the earlier interactions (used by the parity test).
        history: list of dicts sorted by ts with keys ts_ms, visible_ms, focal_ms, hover_ms, clicked, engaged, community, author, embedding (unit 1024-d), pca16."""
        J = [h for h in history if (h["ts_ms"] < now_ms or strict_prefix) and h["ts_ms"] + (h.get("visible_ms") or 0) <= now_ms]   # past interactions whose viewing had ended
        R = J[-self.W:]; n = len(R); gl = self.gl
        lab = lambda xs: sum(1 for h in xs if h["engaged"])
        out = {"h_n": math.log1p(len(J)), "h_pos_rate": (lab(R) + 3 * gl) / (n + 3),
               "h_mean_focal": sum(math.log1p((h.get("focal_ms") or 0) / 1000) for h in R) / (n + 3),
               "h_hover": sum(math.log1p(h.get("hover_ms") or 0) for h in R) / (n + 3),
               "h_click": (sum(1 for h in R if h.get("clicked")) + 3 * self.click_prior) / (n + 3)}
        ur = (lab(J) + 3 * gl) / (len(J) + 3)
        com = post.get("community")
        sm = [h for h in J if isinstance(com, str) and h.get("community") == com]
        out["sub_seen"], out["sub_rate"] = math.log1p(len(sm)), (lab(sm) + 3 * ur) / (len(sm) + 3)
        au = post.get("author")
        am = [h for h in J if isinstance(au, str) and h.get("author") == au]
        out["auth_seen"], out["auth_rate"] = math.log1p(len(am)), (lab(am) + 3 * ur) / (len(am) + 3)
        # content features use only past posts that have an embedding; without one (post or history) they stay 0, as for a user with no history
        Rp = [h for h in R if h["engaged"] and h.get("embedding") is not None]
        Rn = [h for h in R if not h["engaged"] and h.get("embedding") is not None]
        out["aff_recent"] = float(post_X @ unit(np.mean([h["embedding"] for h in Rp], 0))) if Rp and post_X is not None else 0.0
        out["aff_neg"] = float(post_X @ unit(np.mean([h["embedding"] for h in Rn], 0))) if Rn and post_X is not None else 0.0
        ue = np.mean([h["pca16"] for h in Rp], 0) if Rp else np.zeros(16)
        for i in range(16):
            out[f"uemb{i}"] = float(ue[i])
        return out, lab(J)