"""3,072 dimensions (text-embedding-3-large full size): what fails, what works, and shortening."""
import numpy as np, time, psycopg
from common import *

v = corpus(100_000, dim=3072)
rng = np.random.default_rng(7)
qi = np.sort(rng.choice(len(v), 1000, replace=False))
mask = np.ones(len(v), bool); mask[qi] = False
q, base = v[qi], v[mask]
gt = ground_truth(base, q, k=100)
out = {"n": len(base), "dim": 3072, "k": 10, "queries": len(q)}

# Shortening (Matryoshka): truncate + renormalise, compare against full-3072 exact top-10.
out["shortening"] = []
for d in (256, 512, 768, 1024, 1536, 2048, 3072):
    b = base[:, :d] / np.linalg.norm(base[:, :d], axis=1, keepdims=True)
    qq = q[:, :d] / np.linalg.norm(q[:, :d], axis=1, keepdims=True)
    g = ground_truth(b, qq, k=10)
    r = float(np.mean([len(set(g[i]) & set(gt[i, :10])) / 10 for i in range(len(q))]))
    out["shortening"].append({"dims": d, "recall_vs_3072": round(r, 4), "vector_bytes": 4 * d + 8, "halfvec_bytes": 2 * d + 8})
    print("short", d, r, flush=True)

# Does the API's 1536-d output equal the 3072-d output truncated to 1536 and renormalised?
c1536 = corpus(500_000)
ids1536 = [None]  # rows are aligned by _id order? verify via ids
import json
ids3072 = json.load(open(f"{NP}/corpus3072_ids.json"))
meta_ids = []
with open(f"{NP}/corpus_meta.jsonl") as fh:
    for line in fh:
        meta_ids.append(json.loads(line)["id"])
pos = {k: i for i, k in enumerate(meta_ids)}
sample = [pos[k] for k in ids3072[:2000]]
t = v[:2000, :1536]
t = t / np.linalg.norm(t, axis=1, keepdims=True)
cos = np.sum(t * c1536[sample], axis=1)
out["api1536_vs_truncated3072_cosine"] = {"min": round(float(cos.min()), 5), "mean": round(float(cos.mean()), 5)}
print("api vs truncated", out["api1536_vs_truncated3072_cosine"], flush=True)

with connect() as c:
    out["vector_load_s"] = round(load_table(c, "items3k", base, dim=3072), 1)
    out["vector_table_bytes"] = table_size(c, "items3k")
    try:
        c.execute("CREATE INDEX ON items3k USING hnsw (embedding vector_cosine_ops)")
    except psycopg.Error as e:
        out["vector_hnsw_error"] = str(e).strip()
    try:
        c.execute("CREATE INDEX ON items3k USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100)")
    except psycopg.Error as e:
        out["vector_ivfflat_error"] = str(e).strip()
    print(out.get("vector_hnsw_error"), out.get("vector_ivfflat_error"), flush=True)

    SQLX = "SELECT id FROM items3k ORDER BY embedding <=> %s LIMIT 10"
    gtid = np.where(mask)[0]  # map row -> id used in table is row number within base
    out["exact_parallel"] = run_queries(c, SQLX, q[:100], gt[:100], warmup=5)

    t = time.perf_counter()
    c.execute("CREATE INDEX items3k_h ON items3k USING hnsw ((embedding::halfvec(3072)) halfvec_cosine_ops)")
    out["halfvec_build_s"] = round(time.perf_counter() - t, 1)
    out["halfvec_index_bytes"] = rel_size(c, "items3k_h")
    SQLH = "SELECT id FROM items3k ORDER BY embedding::halfvec(3072) <=> %s::halfvec(3072) LIMIT 10"
    out["halfvec_sweep"] = []
    for ef in (20, 40, 100, 200):
        c.execute(f"SET hnsw.ef_search = {ef}")
        s = run_queries(c, SQLH, q, gt); s["ef_search"] = ef
        out["halfvec_sweep"].append(s); print("3072 halfvec", s, flush=True)
    c.execute("DROP INDEX items3k_h")

    try:
        c.execute("CREATE INDEX items3k_h2 ON items3k USING hnsw ((embedding::halfvec(4001)) halfvec_cosine_ops)")
    except psycopg.Error as e:
        out["halfvec_4001_error"] = str(e).strip()

    t = time.perf_counter()
    c.execute("CREATE INDEX items3k_bq ON items3k USING hnsw ((binary_quantize(embedding)::bit(3072)) bit_hamming_ops)")
    out["bq_build_s"] = round(time.perf_counter() - t, 1)
    out["bq_index_bytes"] = rel_size(c, "items3k_bq")
    out["bq_rerank"] = []
    for cand in (40, 80, 160, 320):
        c.execute(f"SET hnsw.ef_search = {max(cand, 40)}")
        SQLR = f"""SELECT id FROM (SELECT id, embedding FROM items3k
              ORDER BY binary_quantize(embedding)::bit(3072) <~> binary_quantize(%s) LIMIT {cand}) c
            ORDER BY c.embedding <=> %s LIMIT 10"""
        cur = c.cursor(); lat, rec = [], []
        for vv in q[:30]:
            cur.execute(SQLR, (vv, vv), prepare=True).fetchall()
        for i, vv in enumerate(q):
            t = time.perf_counter()
            ids = [r[0] for r in cur.execute(SQLR, (vv, vv), prepare=True).fetchall()]
            lat.append((time.perf_counter() - t) * 1000); rec.append(recall(ids, gt[i]))
        lat.sort()
        s = {"candidates": cand, "recall": round(float(np.mean(rec)), 4), "p50_ms": round(lat[500], 3), "p95_ms": round(lat[950], 3)}
        out["bq_rerank"].append(s); print("3072 bq", s, flush=True)
    c.execute("RESET hnsw.ef_search")
    c.execute("DROP TABLE items3k")
save("dims3072", out)