/**
 * Benchmark filtered semantic search to find the optimal crossover point
 * between filter-first and KNN-first strategies.
 *
 * Measures:
 * 1. filterArtIds() latency for representative filters
 * 2. searchFiltered() (chunked vec_distance_cosine) at various candidate set sizes up to full collection
 * 3. Pure KNN + post-filter at various pool sizes
 * 4. Head-to-head: filter-first vs KNN-first across diverse queries and filters
 *
 * Run:  node scripts/tests/bench-filtered-search.mjs
 */
import { fileURLToPath } from "node:url";
import path from "node:path";

const PROJECT_DIR = path.resolve(path.dirname(fileURLToPath(import.meta.url)), "../..");

// Dynamic import of compiled modules
const { EmbeddingsDb } = await import(path.join(PROJECT_DIR, "dist/api/EmbeddingsDb.js"));
const { VocabularyDb } = await import(path.join(PROJECT_DIR, "dist/api/VocabularyDb.js"));
const { EmbeddingModel } = await import(path.join(PROJECT_DIR, "dist/api/EmbeddingModel.js"));

// ── Setup ────────────────────────────────────────────────────────

const embeddingsDb = new EmbeddingsDb();
const vocabDb = new VocabularyDb();
const embeddingModel = new EmbeddingModel();

if (!embeddingsDb.available) {
  console.error("Embeddings DB not available — cannot benchmark");
  process.exit(1);
}
if (!vocabDb.available) {
  console.error("Vocabulary DB not available — cannot benchmark");
  process.exit(1);
}

console.log("Loading embedding model...");
await embeddingModel.init(
  process.env.EMBEDDING_MODEL_ID || "Xenova/multilingual-e5-small",
  embeddingsDb.vectorDimensions,
);

// ── Queries ──────────────────────────────────────────────────────

const QUERIES = [
  "portrait of a woman",            // high overlap ambiguity in Section 5
  "still life with flowers",        // moderate ambiguity (12/15 drawing)
  "battle scene",                   // moderate ambiguity (12/15 paper)
  "map of the world",               // moderate ambiguity (14/15 drawing)
  "landscape with windmill",        // baseline: no ambiguity
];

console.log(`Embedding ${QUERIES.length} queries...`);
const queryVecs = new Map();
for (const q of QUERIES) {
  queryVecs.set(q, await embeddingModel.embed(q));
}
console.log();

// ── Helpers ──────────────────────────────────────────────────────

function timeMs(fn) {
  const start = performance.now();
  const result = fn();
  return { result, ms: performance.now() - start };
}

function fmt(ms) {
  if (ms < 1) return `${(ms * 1000).toFixed(0)}µs`;
  if (ms >= 1000) return `${(ms / 1000).toFixed(2)}s`;
  return `${ms.toFixed(1)}ms`;
}

function median(arr) {
  const s = [...arr].sort((a, b) => a - b);
  const mid = Math.floor(s.length / 2);
  return s.length % 2 ? s[mid] : (s[mid - 1] + s[mid]) / 2;
}

// ── Section 1: filterArtIds latency ─────────────────────────────

console.log("═══════════════════════════════════════════════════════════════════════════");
console.log("  Section 1: filterArtIds() latency + true counts");
console.log("═══════════════════════════════════════════════════════════════════════════\n");

const FILTER_TESTS = [
  { label: "type: 'painting'",                          params: { type: "painting" } },
  { label: "type: 'print'",                             params: { type: "print" } },
  { label: "type: 'photograph'",                        params: { type: "photograph" } },
  { label: "type: 'drawing'",                           params: { type: "drawing" } },
  { label: "type: 'furniture'",                         params: { type: "furniture" } },
  { label: "material: 'oil paint'",                     params: { material: "oil paint" } },
  { label: "material: 'paper'",                         params: { material: "paper" } },
  { label: "creator: 'Rembrandt'",                      params: { creator: "Rembrandt" } },
  { label: "creator: 'anonymous'",                      params: { creator: "anonymous" } },
  { label: "subject: 'landscape'",                      params: { subject: "landscape" } },
  { label: "subject: 'portrait'",                       params: { subject: "portrait" } },
  { label: "depictedPerson: 'Jesus'",                   params: { depictedPerson: "Jesus" } },
  { label: "depictedPerson: 'Maria'",                   params: { depictedPerson: "Maria" } },
  { label: "productionPlace: 'Amsterdam'",              params: { productionPlace: "Amsterdam" } },
  { label: "productionPlace: 'Japan'",                  params: { productionPlace: "Japan" } },
  { label: "type: 'painting' + creator: 'Rembrandt'",   params: { type: "painting", creator: "Rembrandt" } },
  { label: "type: 'print' + subject: 'landscape'",      params: { type: "print", subject: "landscape" } },
  { label: "type: 'photograph' + place: 'Amsterdam'",   params: { type: "photograph", productionPlace: "Amsterdam" } },
  { label: "material: 'paper' + subject: 'portrait'",   params: { material: "paper", subject: "portrait" } },
];

// Get true counts via collection_stats for the broad filters
const trueCountCache = new Map();
for (const test of FILTER_TESTS) {
  // Use computeCollectionStats to get the true count (no LIMIT)
  // We pass dimension: 'type' just to trigger the count — we only need `total`
  const stats = vocabDb.computeCollectionStats({ dimension: "type", ...test.params, topN: 1 });
  trueCountCache.set(test.label, stats.total);
}

const filterData = [];
for (const test of FILTER_TESTS) {
  vocabDb.filterArtIds(test.params); // warm up
  const { result, ms } = timeMs(() => vocabDb.filterArtIds(test.params));
  const count = result?.length ?? 0;
  const trueCount = trueCountCache.get(test.label);
  const capped = count >= 200000;
  filterData.push({ ...test, count, trueCount, ms, capped });
  const flag = capped ? ` ⚠ CAPPED (true: ${trueCount.toLocaleString()})` : "";
  console.log(`  ${test.label.padEnd(52)} → ${String(count).padStart(7)} artworks  ${fmt(ms).padStart(10)}${flag}`);
}

// ── Section 2: searchFiltered() scaling ─────────────────────────

console.log("\n═══════════════════════════════════════════════════════════════════════════");
console.log("  Section 2: searchFiltered() — scaling with candidate set size");
console.log("═══════════════════════════════════════════════════════════════════════════\n");

// Build an uncapped pool: get ALL art_ids from the DB
const allArtIdsSql = vocabDb.db.prepare("SELECT art_id FROM artworks").all();
const allArtIds = allArtIdsSql.map(r => r.art_id);
console.log(`  Full art_id pool: ${allArtIds.length.toLocaleString()}\n`);

// Shuffle deterministically (Fisher-Yates with seed)
function seededShuffle(arr, seed = 42) {
  const copy = [...arr];
  let s = seed;
  for (let i = copy.length - 1; i > 0; i--) {
    s = (s * 1664525 + 1013904223) & 0x7fffffff;
    const j = s % (i + 1);
    [copy[i], copy[j]] = [copy[j], copy[i]];
  }
  return copy;
}
const shuffledIds = seededShuffle(allArtIds);

const CANDIDATE_SIZES = [
  500, 1_000, 2_500, 5_000, 10_000, 25_000, 50_000,
  100_000, 200_000, 400_000, 600_000, 832_000,
];

const K = 15;
const RUNS = 3;
const benchQuery = queryVecs.get("landscape with windmill");

console.log(`  Query: "landscape with windmill", k=${K}, ${RUNS} runs each\n`);
console.log(`  ${"Candidates".padStart(12)} │ ${"Median".padStart(9)} │ ${"Min".padStart(9)} │ ${"Max".padStart(9)} │ ${"ms/1K".padStart(7)} │ Results`);
console.log(`  ${"─".repeat(12)} ┼ ${"─".repeat(9)} ┼ ${"─".repeat(9)} ┼ ${"─".repeat(9)} ┼ ${"─".repeat(7)} ┼ ${"─".repeat(7)}`);

for (const size of CANDIDATE_SIZES) {
  if (size > shuffledIds.length) break;
  const candidates = shuffledIds.slice(0, size);
  const times = [];

  for (let r = 0; r < RUNS; r++) {
    const { ms } = timeMs(() => embeddingsDb.searchFiltered(benchQuery, candidates, K));
    times.push(ms);
  }

  const med = median(times);
  const min = Math.min(...times);
  const max = Math.max(...times);
  const perK = (med / (size / 1000)).toFixed(2);
  const lastResult = embeddingsDb.searchFiltered(benchQuery, candidates, K);
  const fallback = lastResult.warning ? " (KNN fallback)" : "";

  console.log(
    `  ${size.toLocaleString().padStart(12)} │ ${fmt(med).padStart(9)} │ ${fmt(min).padStart(9)} │ ${fmt(max).padStart(9)} │ ${perK.padStart(7)} │ ${lastResult.results.length}${fallback}`
  );
}

// ── Section 3: Pure KNN + post-filter ───────────────────────────

console.log("\n═══════════════════════════════════════════════════════════════════════════");
console.log("  Section 3: search() (pure KNN) + post-filter — baseline");
console.log("═══════════════════════════════════════════════════════════════════════════\n");

const KNN_POOL_SIZES = [512, 1024, 2048, 4096];
const KNN_FILTERS = [
  { label: "painting (~0.6%)",   params: { type: "painting" } },
  { label: "print (~51%)",       params: { type: "print" } },
  { label: "photograph (~15%)",  params: { type: "photograph" } },
  { label: "Rembrandt (~0.5%)",  params: { creator: "Rembrandt" } },
  { label: "Amsterdam (~12%)",   params: { productionPlace: "Amsterdam" } },
];

// Pre-resolve filter ID sets
const filterIdSets = new Map();
for (const f of KNN_FILTERS) {
  const ids = vocabDb.filterArtIds(f.params);
  if (ids) filterIdSets.set(f.label, new Set(ids));
}

console.log(`  ${"Pool".padStart(6)} │ ${"Filter".padEnd(22)} │ ${"Median".padStart(9)} │ ${"Matches".padStart(7)} │ ${"Fill".padStart(5)}`);
console.log(`  ${"─".repeat(6)} ┼ ${"─".repeat(22)} ┼ ${"─".repeat(9)} ┼ ${"─".repeat(7)} ┼ ${"─".repeat(5)}`);

for (const poolSize of KNN_POOL_SIZES) {
  for (const filter of KNN_FILTERS) {
    const idSet = filterIdSets.get(filter.label);
    if (!idSet) continue;
    const times = [];
    let lastMatches = 0;

    for (let r = 0; r < RUNS; r++) {
      const { result, ms } = timeMs(() => {
        const knnResults = embeddingsDb.search(benchQuery, poolSize);
        return knnResults.filter(r => idSet.has(r.artId)).slice(0, K);
      });
      times.push(ms);
      lastMatches = result.length;
    }

    const med = median(times);
    const fillRate = ((lastMatches / K) * 100).toFixed(0) + "%";
    console.log(
      `  ${String(poolSize).padStart(6)} │ ${filter.label.padEnd(22)} │ ${fmt(med).padStart(9)} │ ${String(lastMatches).padStart(7)} │ ${fillRate.padStart(5)}`
    );
  }
}

// ── Section 4: Multi-query head-to-head ─────────────────────────

console.log("\n═══════════════════════════════════════════════════════════════════════════");
console.log("  Section 4: Head-to-head across queries × filters");
console.log("═══════════════════════════════════════════════════════════════════════════\n");

const H2H_FILTERS = [
  { label: "print",       params: { type: "print" } },
  { label: "photograph",  params: { type: "photograph" } },
  { label: "drawing",     params: { type: "drawing" } },
  { label: "Amsterdam",   params: { productionPlace: "Amsterdam" } },
  { label: "paper",       params: { material: "paper" } },
];

// Pre-resolve filter sets for both strategies
const h2hFilterData = new Map();
for (const f of H2H_FILTERS) {
  const ids = vocabDb.filterArtIds(f.params);
  const trueCount = trueCountCache.get(
    FILTER_TESTS.find(t => JSON.stringify(t.params) === JSON.stringify(f.params))?.label
  ) ?? ids?.length ?? 0;
  h2hFilterData.set(f.label, { ids: ids || [], idSet: new Set(ids || []), trueCount });
}

console.log(`  ${"Query".padEnd(30)} │ ${"Filter".padEnd(12)} │ ${"True#".padStart(7)} │ ${"A:filt-1st".padStart(10)} │ ${"B:KNN-1st".padStart(10)} │ ${"A dist".padStart(7)} │ ${"B dist".padStart(7)} │ ${"Overlap".padStart(7)}`);
console.log(`  ${"─".repeat(30)} ┼ ${"─".repeat(12)} ┼ ${"─".repeat(7)} ┼ ${"─".repeat(10)} ┼ ${"─".repeat(10)} ┼ ${"─".repeat(7)} ┼ ${"─".repeat(7)} ┼ ${"─".repeat(7)}`);

for (const query of QUERIES) {
  const qVec = queryVecs.get(query);

  for (const filter of H2H_FILTERS) {
    const { ids, idSet, trueCount } = h2hFilterData.get(filter.label);
    if (ids.length === 0) continue;

    // Strategy A: filter-first (capped at 200K by filterArtIds)
    const timesA = [];
    for (let r = 0; r < RUNS; r++) {
      const { ms } = timeMs(() => embeddingsDb.searchFiltered(qVec, ids, K));
      timesA.push(ms);
    }
    const medA = median(timesA);
    const resA = embeddingsDb.searchFiltered(qVec, ids, K);

    // Strategy B: KNN-first + post-filter (pool=4096)
    const timesB = [];
    for (let r = 0; r < RUNS; r++) {
      const { ms } = timeMs(() => {
        const knn = embeddingsDb.search(qVec, 4096);
        return knn.filter(r => idSet.has(r.artId)).slice(0, K);
      });
      timesB.push(ms);
    }
    const medB = median(timesB);
    const knn = embeddingsDb.search(qVec, 4096);
    const resB = knn.filter(r => idSet.has(r.artId)).slice(0, K);

    // Overlap
    const setA = new Set(resA.results.map(r => r.objectNumber));
    const setB = new Set(resB.map(r => r.objectNumber));
    const overlap = [...setA].filter(x => setB.has(x)).length;

    const distA = resA.results[0]?.distance.toFixed(4) ?? "N/A";
    const distB = resB[0]?.distance.toFixed(4) ?? "N/A";
    const qLabel = query.length > 28 ? query.slice(0, 26) + "…" : query;

    console.log(
      `  ${qLabel.padEnd(30)} │ ${filter.label.padEnd(12)} │ ${String(trueCount).padStart(7)} │ ${fmt(medA).padStart(10)} │ ${fmt(medB).padStart(10)} │ ${distA.padStart(7)} │ ${distB.padStart(7)} │ ${`${overlap}/${K}`.padStart(7)}`
    );
  }
}

// ── Section 5: Quality deep-dive — does the cap affect results? ─

console.log("\n═══════════════════════════════════════════════════════════════════════════");
console.log("  Section 5: Quality — slice stability for capped filters");
console.log("═══════════════════════════════════════════════════════════════════════════\n");
console.log("  Compares results from first half vs last half of capped candidate sets.\n");

// For filters that hit the 200K cap, compare first-half results with last-half results.
// This measures whether the arbitrary row order affects which results appear.

const SLICE_SIZE = 100_000; // half of FILTER_ART_IDS_LIMIT
const QUALITY_FILTERS = filterData.filter(f => f.capped);

if (QUALITY_FILTERS.length === 0) {
  console.log("  No capped filters — nothing to compare.\n");
} else {
  console.log(`  ${"Query".padEnd(30)} │ ${"Filter".padEnd(16)} │ ${"True#".padStart(7)} │ ${"1st half dist".padStart(13)} │ ${"2nd half dist".padStart(13)} │ ${"Overlap".padStart(7)}`);
  console.log(`  ${"─".repeat(30)} ┼ ${"─".repeat(16)} ┼ ${"─".repeat(7)} ┼ ${"─".repeat(13)} ┼ ${"─".repeat(13)} ┼ ${"─".repeat(7)}`);

  for (const query of QUERIES.slice(0, 6)) {
    const qVec = queryVecs.get(query);
    for (const filter of QUALITY_FILTERS) {
      const filterIds = vocabDb.filterArtIds(filter.params);
      if (!filterIds || filterIds.length < SLICE_SIZE * 2) continue;

      const firstHalf = filterIds.slice(0, SLICE_SIZE);
      const secondHalf = filterIds.slice(SLICE_SIZE, SLICE_SIZE * 2);

      const resFirst = embeddingsDb.searchFiltered(qVec, firstHalf, K);
      const resSecond = embeddingsDb.searchFiltered(qVec, secondHalf, K);

      const setFirst = new Set(resFirst.results.map(r => r.objectNumber));
      const setSecond = new Set(resSecond.results.map(r => r.objectNumber));
      const overlap = [...setFirst].filter(x => setSecond.has(x)).length;

      const distFirst = resFirst.results[0]?.distance.toFixed(4) ?? "N/A";
      const distSecond = resSecond.results[0]?.distance.toFixed(4) ?? "N/A";
      const qLabel = query.length > 28 ? query.slice(0, 26) + "…" : query;

      console.log(
        `  ${qLabel.padEnd(30)} │ ${filter.label.padEnd(16)} │ ${String(filter.trueCount).padStart(7)} │ ${distFirst.padStart(13)} │ ${distSecond.padStart(13)} │ ${`${overlap}/${K}`.padStart(7)}`
      );
    }
  }
}

// ── Cleanup ─────────────────────────────────────────────────────

embeddingModel.dispose?.();
console.log("\nDone.");
