# /// script
# requires-python = ">=3.11"
# dependencies = ["duckdb", "pyarrow", "numpy"]
# ///
import duckdb, numpy as np, pyarrow.parquet as pq
con = duckdb.connect()
P = "embeddings.parquet"
print(con.execute(f"SELECT count(*) n, count(DISTINCT id) uid, count(emb_caption) c, count(emb_scene) s, count(emb_events) e FROM '{P}'").fetchall())
print("types:", con.execute(f"SELECT typeof(id), typeof(emb_caption) FROM '{P}' LIMIT 1").fetchall())
print("sorted by id:", con.execute(f"SELECT bool_and(id >= lag) FROM (SELECT id, lag(id) OVER (ORDER BY rowid) lag FROM (SELECT row_number() OVER () rowid, id FROM '{P}'))").fetchall())
# self-join: identical id must score exactly 1.0 on every column
print("self-join:", con.execute(f"""
  SELECT min(list_cosine_similarity(a.emb_caption, b.emb_caption)) cap,
         min(list_cosine_similarity(a.emb_scene,  b.emb_scene))  scn,
         min(list_cosine_similarity(a.emb_events, b.emb_events)) evt
  FROM '{P}' a JOIN '{P}' b USING (id)""").fetchall())
# cross-similarity must be finite everywhere (a NULL element would throw here)
q = con.execute(f"SELECT emb_caption FROM '{P}' LIMIT 1").fetchone()[0]
r = con.execute(f"SELECT count(*) n, min(s), max(s) FROM (SELECT list_cosine_similarity(emb_caption, ?::FLOAT[]) s FROM '{P}') WHERE s IS NOT NULL", [q]).fetchall()
print("full caption scan:", r)
r = con.execute(f"SELECT count(*) FROM (SELECT list_cosine_similarity(emb_events, ?::FLOAT[]) s FROM '{P}') WHERE s IS NULL", [q]).fetchall()
print("null-scoring events rows (expect 19):", r)
# fp16 fidelity against the fp32 source of the same vector
t = pq.read_table(P, columns=["emb_caption"]).column("emb_caption").combine_chunks()
v = np.asarray(t.flatten()[:1024], dtype=np.float32)
print("norm of a stored fp16 vector:", float(np.linalg.norm(v)))