"""HNSW build mechanics: maintenance_work_mem, parallel workers, index-before-load vs load-then-index."""
import numpy as np, time, psycopg
from common import *

N = 100_000
base = corpus(N)
out = {"n": N, "dim": 1536, "index": "hnsw m=16 ef_construction=64, vector_cosine_ops"}
notices = []

with connect() as c:
    c.add_notice_handler(lambda d: notices.append(d.message_primary))
    load_table(c, "items_b2", base)

    def build(label, settings):
        for s in settings:
            c.execute(s)
        notices.clear()
        c.execute("DROP INDEX IF EXISTS items_b2_h")
        t = time.perf_counter()
        c.execute("CREATE INDEX items_b2_h ON items_b2 USING hnsw (embedding vector_cosine_ops)")
        r = {"label": label, "settings": settings, "build_s": round(time.perf_counter() - t, 1), "notices": list(notices)}
        c.execute("RESET maintenance_work_mem"); c.execute("RESET max_parallel_maintenance_workers")
        print(r, flush=True)
        return r

    out["builds"] = [
        build("4 GB, 3 parallel workers", ["SET maintenance_work_mem = '4GB'", "SET max_parallel_maintenance_workers = 3"]),
        build("4 GB, 1 parallel worker", ["SET maintenance_work_mem = '4GB'", "SET max_parallel_maintenance_workers = 1"]),
        build("4 GB, no parallel workers", ["SET maintenance_work_mem = '4GB'", "SET max_parallel_maintenance_workers = 0"]),
        build("64 MB (PostgreSQL default), 3 workers", ["SET maintenance_work_mem = '64MB'", "SET max_parallel_maintenance_workers = 3"]),
        build("256 MB, 3 workers", ["SET maintenance_work_mem = '256MB'", "SET max_parallel_maintenance_workers = 3"]),
    ]
    out["index_bytes"] = rel_size(c, "items_b2_h")

    # Index first, then load: every row is inserted into the graph one at a time (no parallelism).
    c.execute("DROP TABLE items_b2")
    c.execute("CREATE TABLE items_b2 (id int PRIMARY KEY, embedding vector(1536))")
    c.execute("CREATE INDEX items_b2_h ON items_b2 USING hnsw (embedding vector_cosine_ops)")
    t = time.perf_counter()
    with c.cursor().copy("COPY items_b2 (id, embedding) FROM STDIN WITH (FORMAT BINARY)") as cp:
        cp.set_types(["int4", "vector"])
        for i, v in enumerate(base):
            cp.write_row([i, v])
    out["copy_into_indexed_table_s"] = round(time.perf_counter() - t, 1)
    out["index_bytes_incremental"] = rel_size(c, "items_b2_h")

    # Insert rate of single-row INSERTs into an already-indexed 100K table.
    extra = corpus(101_000)[100_000:101_000]
    cur = c.cursor()
    t = time.perf_counter()
    for i, v in enumerate(extra):
        cur.execute("INSERT INTO items_b2 VALUES (%s, %s)", (200_000 + i, v), prepare=True)
    el = time.perf_counter() - t
    out["single_insert_ms_indexed"] = round(el / len(extra) * 1000, 3)

    c.execute("DROP INDEX items_b2_h")
    t = time.perf_counter()
    for i, v in enumerate(extra):
        cur.execute("INSERT INTO items_b2 VALUES (%s, %s)", (300_000 + i, v), prepare=True)
    out["single_insert_ms_unindexed"] = round((time.perf_counter() - t) / len(extra) * 1000, 3)
    print(out, flush=True)
    c.execute("DROP TABLE items_b2")
save("build", out)