"""phx_score.py — score real corpus posts with the released mini-Phoenix ranker.

Adapted from phoenix/run_pipeline.py (Apache 2.0, xAI). Skips retrieval; scores an
arbitrary subset of corpus posts across ALL action heads for a given user sequence,
optionally probing the learned post-age response by varying candidate_impr_ts.

Outputs a .npz with post_ids, author_ids, all head probabilities, and (optionally)
an age sweep.

Usage:
  phxenv/Scripts/python phx_score.py --artifacts x-algorithm/phoenix/artifacts/unpacked \
      --repo x-algorithm/phoenix --n 50000 --out scores.npz
  ... --age-sweep --n 512 --out age_sweep.npz
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time

import numpy as np

p = argparse.ArgumentParser()
p.add_argument("--artifacts", required=True)
p.add_argument("--repo", required=True, help="path to phoenix/ source dir")
p.add_argument("--n", type=int, default=50000)
p.add_argument("--seed", type=int, default=7)
p.add_argument("--out", required=True)
p.add_argument("--age-sweep", action="store_true",
               help="score same candidates at multiple impression ages")
p.add_argument("--ids-file", default=None,
               help="score exactly these post ids (overrides --n random selection)")
p.add_argument("--sequence", default=None, help="user sequence json (default example)")
args = p.parse_args()

sys.path.insert(0, args.repo)

import haiku as hk  # noqa: E402
import jax  # noqa: E402
import jax.numpy as jnp  # noqa: E402

from grok import TransformerConfig  # noqa: E402
from recsys_model import (  # noqa: E402
    HashConfig, PhoenixModelConfig, RecsysBatch, RecsysEmbeddings,
)
from runners import load_embedding_table, load_model_params  # noqa: E402

A = args.artifacts
with open(os.path.join(A, "ranker", "config.json")) as f:
    cfg = json.load(f)
print("[cfg] ranker config:", json.dumps({k: v for k, v in cfg.items() if k != "hash_params"})[:400])

num_actions = cfg["num_actions"]
hist_len = cfg["history_seq_len"]
cand_len = cfg["candidate_seq_len"]
emb_size = cfg["emb_size"]

# --- hashing (verbatim logic from run_pipeline.py) ---

def _hash_ids(ids, scales, biases, modulus, num_buckets):
    ids = np.asarray(ids, dtype=np.int64).ravel()
    scales = np.array(scales, dtype=np.int64)
    biases = np.array(biases, dtype=np.int64)
    n, m = len(ids), len(scales)
    out = np.empty((n, m), dtype=np.int32)
    for i in range(n):
        for j in range(m):
            raw = (ids[i] * scales[j] + biases[j]) % np.int64(modulus)
            out[i, j] = 0 if ids[i] == 0 else int((int(raw) % (num_buckets - 1)) + 1)
    return out


hp = cfg["hash_params"]
pad = 65
uv, iv, av = cfg["user_vocab_size"], cfg["item_vocab_size"], cfg["author_vocab_size"]


def hash_user(x):
    h = _hash_ids(x, hp["user_hash_scales"], hp["user_biases"], hp["user_modulus"], uv)
    return np.where(h == 0, 0, h + pad)


def hash_item(x):
    h = _hash_ids(x, hp["item_hash_scales"], hp["item_biases"], hp["item_modulus"], iv)
    return np.where(h == 0, 0, h + pad + uv)


def hash_author(x):
    h = _hash_ids(x, hp["author_hash_scales"], hp["author_biases"], hp["author_modulus"], av)
    return np.where(h == 0, 0, h + pad + uv + iv)


print("[load] ranker params + embeddings...")
params = load_model_params(os.path.join(A, "ranker", "model_params.npz"))
emb_dict = load_embedding_table(os.path.join(A, "ranker", "embedding_tables.npz"))
table = np.zeros((pad + uv + iv + av, emb_size), dtype=np.float32)
table[pad:pad + uv] = emb_dict["user_embeddings"]
table[pad + uv:pad + uv + iv] = emb_dict["item_embeddings"]
table[pad + uv + iv:] = emb_dict["author_embeddings"]

corpus = np.load(os.path.join(A, "sports_corpus.npz"), allow_pickle=True)
post_ids = corpus["post_ids"]
author_ids = corpus["author_ids"]
print(f"[corpus] {len(post_ids)} posts; fields: {list(corpus.keys())}")

if args.ids_file:
    want = set(int(l.strip()) for l in open(args.ids_file) if l.strip())
    mask = np.array([int(p) in want for p in post_ids])
    sel = np.nonzero(mask)[0]
    print(f"[select] {len(sel)} of {len(want)} requested ids found in corpus")
else:
    rng = np.random.default_rng(args.seed)
    sel = rng.choice(len(post_ids), size=min(args.n, len(post_ids)), replace=False)
sel_posts = np.asarray(post_ids[sel], dtype=np.uint64)
sel_authors = np.asarray(author_ids[sel], dtype=np.uint64)

seq_file = args.sequence or os.path.join(A, "example_sequence.json")
with open(seq_file) as f:
    seq = json.load(f)
user_id = seq["user_id"]
history = seq["history"]

history_post_ids = np.zeros(hist_len, dtype=np.uint64)
history_author_ids = np.zeros(hist_len, dtype=np.uint64)
history_actions = np.zeros((hist_len, num_actions), dtype=np.float32)
for i, item in enumerate(history[:hist_len]):
    history_post_ids[i] = item["post_id"]
    history_author_ids[i] = item["author_id"]
    for k, v in item.get("actions", {}).items():
        if int(k) < num_actions:
            history_actions[i, int(k)] = float(v)

user_h = hash_user(np.array([user_id], dtype=np.uint64))
hist_ph = hash_item(history_post_ids).reshape(1, hist_len, -1)
hist_ah = hash_author(history_author_ids).reshape(1, hist_len, -1)

mc_kwargs = dict(
    emb_size=emb_size,
    history_seq_len=hist_len,
    candidate_seq_len=cand_len,
    hash_config=HashConfig(
        num_user_hashes=cfg["num_user_hashes"],
        num_item_hashes=cfg["num_item_hashes"],
        num_author_hashes=cfg["num_author_hashes"],
    ),
    product_surface_vocab_size=cfg.get("product_surface_vocab_size", 16),
    num_actions=num_actions,
    post_age_granularity_mins=cfg.get("post_age_granularity_mins", 60),
    model=TransformerConfig(
        emb_size=emb_size, key_size=cfg["key_size"],
        num_q_heads=cfg["num_heads"], num_kv_heads=cfg["num_heads"],
        num_layers=cfg["num_layers"], widening_factor=2.0,
        attn_output_multiplier=0.125,
    ),
)
mc = PhoenixModelConfig(**mc_kwargs)
mc.initialize()


def fwd(b, e):
    return mc.make()(b, e)


fn = hk.without_apply_rng(hk.transform(fwd))

SNOWFLAKE_EPOCH_MS = 1288834974657


def creation_ts_sec(pid):
    return ((np.asarray(pid, dtype=np.uint64) >> np.uint64(22)).astype(np.float64)
            + SNOWFLAKE_EPOCH_MS) / 1000.0


def score_batch(cand_posts, cand_authors, impr_ts=None, creation_ts=None):
    cs = len(cand_posts)
    cph = hash_item(cand_posts).reshape(1, cs, -1)
    cah = hash_author(cand_authors).reshape(1, cs, -1)
    if cs < cand_len:
        cph = np.pad(cph, ((0, 0), (0, cand_len - cs), (0, 0)))
        cah = np.pad(cah, ((0, 0), (0, cand_len - cs), (0, 0)))
    kw = {}
    if impr_ts is not None:
        it = np.full((1, cand_len), 0.0)
        ct = np.zeros((1, cand_len))
        it[0, :cs] = impr_ts
        ct[0, :cs] = creation_ts
        kw = dict(candidate_impr_ts=jnp.asarray(it),
                  candidate_post_creation_ts=jnp.asarray(ct))
    rb = RecsysBatch(
        user_hashes=jnp.asarray(user_h),
        history_post_hashes=jnp.asarray(hist_ph),
        history_author_hashes=jnp.asarray(hist_ah),
        history_actions=jnp.asarray(history_actions.reshape(1, hist_len, num_actions)),
        history_product_surface=jnp.zeros((1, hist_len), dtype=jnp.int32),
        candidate_post_hashes=jnp.asarray(cph),
        candidate_author_hashes=jnp.asarray(cah),
        candidate_product_surface=jnp.zeros((1, cand_len), dtype=jnp.int32),
        **kw,
    )
    re = RecsysEmbeddings(
        user_embeddings=jnp.asarray(table[user_h]),
        history_post_embeddings=jnp.asarray(table[hist_ph]),
        candidate_post_embeddings=jnp.asarray(table[cph]),
        history_author_embeddings=jnp.asarray(table[hist_ah]),
        candidate_author_embeddings=jnp.asarray(table[cah]),
    )
    out = fn.apply(params, rb, re)
    probs = np.asarray(jax.nn.sigmoid(out.logits))[0, :cs, :]
    cont = None
    if out.continuous_preds is not None:
        cont = np.asarray(out.continuous_preds)[0, :cs, :]
    return probs, cont


t0 = time.time()
if args.age_sweep:
    ages_h = [0.25, 0.5, 1, 2, 4, 8, 16, 24, 36, 48, 64, 72, 80, 96, 120]
    ct = creation_ts_sec(sel_posts)
    sweep = np.zeros((len(ages_h), len(sel_posts), num_actions), dtype=np.float32)
    for ai, ah in enumerate(ages_h):
        it = ct + ah * 3600.0
        for i in range(0, len(sel_posts), cand_len):
            j = min(i + cand_len, len(sel_posts))
            pr, _ = score_batch(sel_posts[i:j], sel_authors[i:j],
                                impr_ts=it[i:j], creation_ts=ct[i:j])
            sweep[ai, i:j, :] = pr
        print(f"[sweep] age {ah}h done ({time.time()-t0:.0f}s)")
    np.savez_compressed(args.out, ages_h=np.array(ages_h), probs=sweep,
                        post_ids=sel_posts, author_ids=sel_authors)
else:
    all_probs = np.zeros((len(sel_posts), num_actions), dtype=np.float32)
    all_cont = None
    for i in range(0, len(sel_posts), cand_len):
        j = min(i + cand_len, len(sel_posts))
        pr, co = score_batch(sel_posts[i:j], sel_authors[i:j])
        all_probs[i:j, :] = pr
        if co is not None:
            if all_cont is None:
                all_cont = np.zeros((len(sel_posts), co.shape[-1]), dtype=np.float32)
            all_cont[i:j, :] = co
        if (i // cand_len) % 100 == 0:
            print(f"[score] {j}/{len(sel_posts)} ({time.time()-t0:.0f}s)")
    save = dict(post_ids=sel_posts, author_ids=sel_authors, probs=all_probs)
    if all_cont is not None:
        save["continuous"] = all_cont
    np.savez_compressed(args.out, **save)
print(f"[done] {args.out} in {time.time()-t0:.0f}s")