"""Home-page embedding map: 2,000 DBpedia articles (OpenAI text-embedding-3-large, 1,536-d).
Neighbours: exact top-10 by cosine distance computed BY PGVECTOR (ORDER BY embedding <=> ... LIMIT 11).
Layout: PCA(50) -> t-SNE(2) for display only. Colours: k-means(12) in the full 1,536-d space."""
import json, sys
import numpy as np
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from sklearn.cluster import KMeans
from common import corpus, connect, NP

OUT = sys.argv[1]
rng = np.random.default_rng(2026)
idx = np.sort(rng.choice(500_000, 2000, replace=False))
V = corpus(500_000)[idx]
meta = []
with open(f"{NP}/corpus_meta.jsonl") as fh:
    want = set(idx.tolist())
    for i, line in enumerate(fh):
        if i in want:
            meta.append(json.loads(line)["title"])
titles = meta

with connect() as c:
    c.execute("DROP TABLE IF EXISTS mapdemo")
    c.execute("CREATE TABLE mapdemo (id int PRIMARY KEY, embedding vector(1536))")
    with c.cursor().copy("COPY mapdemo (id, embedding) FROM STDIN WITH (FORMAT BINARY)") as cp:
        cp.set_types(["int4", "vector"])
        for i, v in enumerate(V):
            cp.write_row([i, v])
    nn, dd = [], []
    cur = c.cursor()
    for i, v in enumerate(V):
        rows = cur.execute("SELECT id, embedding <=> %s AS d FROM mapdemo WHERE id <> %s ORDER BY embedding <=> %s LIMIT 10",
                           (v, i, v), prepare=True).fetchall()
        nn.append([r[0] for r in rows]); dd.append([round(float(r[1]), 4) for r in rows])
    c.execute("DROP TABLE mapdemo")

X = PCA(n_components=50, random_state=0).fit_transform(V)
Y = TSNE(n_components=2, perplexity=35, init="pca", random_state=0, max_iter=1500).fit_transform(X)
Y = (Y - Y.min(0)) / (Y.max(0) - Y.min(0))
km = KMeans(n_clusters=12, n_init=10, random_state=0).fit(V)
# demo points: articles whose neighbours are tight (small mean distance) and spread across clusters
tight = np.argsort([np.mean(d) for d in dd])
demo, seen = [], set()
for i in tight:
    if km.labels_[i] not in seen:
        demo.append(int(i)); seen.add(km.labels_[i])
    if len(demo) == 12:
        break
out = {"n": 2000, "source": "DBpedia × OpenAI text-embedding-3-large (1,536-d); neighbours by pgvector exact search",
       "x": [round(float(a), 4) for a in Y[:, 0]], "y": [round(float(a), 4) for a in Y[:, 1]],
       "c": [int(l) for l in km.labels_], "t": titles, "nn": nn, "d": dd, "demo": demo}
json.dump(out, open(OUT, "w"), separators=(",", ":"))
for k in range(12):
    members = [titles[i] for i in np.where(km.labels_ == k)[0][:12]]
    print(k, members)