"""ZeroGPU Space: ZeroEntropy zembed-1 (embed) + zerank-2 (rerank) on a free H200 — for the
ZeroEntropy-vs-current A/B benchmark. Both are 4B (Qwen3-based), so they need the GPU.

Endpoints (call via gradio_client):
- `/embed`  (texts, mode)  -> L2-normalised vectors. mode='query' -> encode_query (search-query
                              task prompt); anything else -> encode_document. Optional Matryoshka
                              truncation via EMBED_DIM.
- `/rerank` (groups)        -> per-group relevance logits, groups = [[query, [passages]], ...].

Models load to CPU once at startup, `.to('cuda')` per call (no re-download). `@spaces.GPU(duration=)`
is generous because 4B inference is much slower than a 300M embedder.
"""
import os

import gradio as gr
import numpy as np
import spaces
from sentence_transformers import CrossEncoder, SentenceTransformer

EMB_ID = os.environ.get("EMBED_MODEL_ID", "zeroentropy/zembed-1-embedding")
RR_ID = os.environ.get("RERANK_MODEL_ID", "zeroentropy/zerank-2-reranker")
EMB_DIM = int(os.environ.get("EMBED_DIM", "0"))  # 0 = full (2560); else Matryoshka-truncate + renorm
TOK = os.environ.get("HF_TOKEN")

EMB = SentenceTransformer(EMB_ID, trust_remote_code=True,
                          model_kwargs={"torch_dtype": "bfloat16"}, device="cpu", token=TOK)
RR = CrossEncoder(RR_ID, trust_remote_code=True, device="cpu", token=TOK)
_LOG = {"emb": 0, "rr": 0}
print(f"[startup] loaded embed={EMB_ID} rerank={RR_ID} dim={EMB_DIM or 'full'}", flush=True)


def _emb_dur(texts, mode="document"):
    return min(300, 30 + len(texts or []) // 4)  # 4B embed is slow


def _rr_dur(groups):
    n = sum(len(g[1]) for g in (groups or []))
    return min(300, 30 + n // 20)


@spaces.GPU(duration=_emb_dur)
def embed(texts, mode="document"):
    if not texts:
        return []
    _LOG["emb"] += len(texts)
    EMB.to("cuda")
    enc = EMB.encode_query if mode == "query" else EMB.encode_document
    print(f"[embed:{mode}] +{len(texts)} on cuda -> {_LOG['emb']} total", flush=True)
    v = np.asarray(enc([(t or "")[:6000] for t in texts], batch_size=16,
                       convert_to_numpy=True, device="cuda"), dtype=np.float32)
    if EMB_DIM and EMB_DIM < v.shape[1]:
        v = v[:, :EMB_DIM]
    v = v / (np.linalg.norm(v, axis=1, keepdims=True) + 1e-9)  # (re)normalise (and after truncation)
    return v.tolist()


@spaces.GPU(duration=_rr_dur)
def rerank(groups):
    """groups: list of [query, [passages]]. Returns one list of relevance logits per group."""
    if not groups:
        return []
    _LOG["rr"] += len(groups)
    RR.model.to("cuda")
    pairs, spans = [], [0]
    for q, ps in groups:
        pairs += [(q, (p or "")[:4000]) for p in ps]
        spans.append(len(pairs))
    print(f"[rerank] +{len(groups)} q / {len(pairs)} pairs -> {_LOG['rr']} total", flush=True)
    scores = RR.predict(pairs, batch_size=16, convert_to_numpy=True) if pairs else np.array([])
    return [list(map(float, scores[spans[i]:spans[i + 1]])) for i in range(len(groups))]


with gr.Blocks() as demo:
    gr.Markdown("ZeroGPU zembed-1 + zerank-2 — call /embed (texts, mode) and /rerank (groups) via gradio_client.")
    ein, emode, eout = gr.JSON(label="texts"), gr.Textbox(label="mode", value="document"), gr.JSON(label="vectors")
    gr.Button("embed").click(embed, [ein, emode], eout, api_name="embed")
    rin, rout = gr.JSON(label="groups [[query,[passages]],...]"), gr.JSON(label="scores")
    gr.Button("rerank").click(rerank, rin, rout, api_name="rerank")

if __name__ == "__main__":
    demo.queue(max_size=256, default_concurrency_limit=1).launch()