// Reads the dataset's Parquet files from the Hub with hyparquet, a few byte ranges at a time:
// a file's footer once, then only the row groups a lookup can need, found by the footer's
// min/max statistics (the files are sorted by the column looked up), adjacent ones in one
// request, straight from the CDN.
//
// The hyparquet functions and compressors are passed in, so the same code runs in the page
// (from a CDN) and in Node (from node_modules).

const TAIL = 1 << 20;  // the bytes fetched from a file's end on opening: its footer, usually
const KEEP_BYTES = 64 << 20;  // the most bytes kept from a file's earlier requests

// A file on the Hub as hyparquet's AsyncBuffer, read straight from the CDN. The Hub's URL
// redirects each request to a signed CDN URL (about 0.15 s a time, an hour's validity), so
// the first request, for the file's last TAIL bytes (its footer), finds the CDN URL and the
// file's length, and later ranges go to the CDN URL, found again when it expires or is
// refused. Spans fetched ahead (`prefetch`) serve the slices inside them, so a run of
// adjacent row groups costs one request.
async function hubBuffer(url) {
  let target = url, expires = 0;
  const spans = [];  // { start, end, data: Promise<ArrayBuffer> }
  const get = async (range) => {
    for (let tries = 0; ; tries++) {
      const res = await fetch(target, { headers: { Range: range } });
      if (res.ok || tries) {
        if (!res.ok) throw new Error(`fetch failed ${res.status}`);
        if (res.redirected || target === url) {
          target = res.url;
          const m = /[?&]Expires=(\d+)/.exec(target);
          expires = m ? +m[1] * 1000 : Date.now() + 30 * 60e3;
        }
        return res;
      }
      target = url;  // the CDN URL expired or was refused: through the Hub again
    }
  };
  const first = await get(`bytes=-${TAIL}`);
  const total = Number((first.headers.get("content-range") ?? "").split("/")[1]);
  const byteLength = total || Number(first.headers.get("content-length"));
  const tailStart = Math.max(0, byteLength - TAIL);
  spans.push({ start: tailStart, end: byteLength, data: first.arrayBuffer() });
  const fetchSpan = (start, end) => {
    if (Date.now() > expires - 60e3) target = url;
    const data = get(`bytes=${start}-${end - 1}`).then(async (res) => {
      const buf = await res.arrayBuffer();
      return res.status === 206 ? buf : buf.slice(start, end);  // a 200 is the whole file
    });
    const span = { start, end, data };
    spans.push(span);
    // The oldest spans let go past KEEP_BYTES (the footer's stays); a slice already waiting on
    // one still gets it
    let kept = 0;
    for (let i = spans.length - 1; i >= 1; i--) {
      kept += spans[i].end - spans[i].start;
      if (kept > KEEP_BYTES && i < spans.length - 1) { spans.splice(1, i); break; }
    }
    return span;
  };
  return {
    byteLength,
    prefetch(start, end) {
      if (!spans.some((x) => x.start <= start && end <= x.end)) fetchSpan(start, end);
    },
    async slice(start, end = byteLength) {
      const span = spans.find((x) => x.start <= start && end <= x.end) ?? fetchSpan(start, end);
      const buf = await span.data;
      return buf.slice(start - span.start, end - span.start);
    },
  };
}

export function makeData({ hyparquet, compressors, base }) {
  const { parquetMetadataAsync, parquetRead, parquetReadObjects, parquetSchema } = hyparquet;
  const files = new Map();
  const postingsOf = new Map();

  // A file's buffer and footer, fetched once
  function open(name) {
    if (!files.has(name)) {
      files.set(name, (async () => {
        const file = await hubBuffer(base + name);
        const metadata = await parquetMetadataAsync(file, { initialFetchSize: TAIL });
        let start = 0;
        const groups = metadata.row_groups.map((rg) => {
          const rows = Number(rg.num_rows);
          const g = { start, end: start + rows, stats: {}, from: Infinity, to: 0, bytes: {} };
          for (const col of rg.columns) {
            const m = col.meta_data;
            if (m?.statistics) g.stats[m.path_in_schema.join(".")] = m.statistics;
            if (!m) continue;
            const at = Number(m.dictionary_page_offset ?? m.data_page_offset);
            const top = m.path_in_schema[0];
            g.bytes[top] = (g.bytes[top] ?? 0) + Number(m.total_compressed_size);
            g.from = Math.min(g.from, at);
            g.to = Math.max(g.to, at + Number(m.total_compressed_size));
          }
          start += rows;
          return g;
        });
        const names = new Set(parquetSchema(metadata).children.map((c) => c.element.name));
        return { file, metadata, groups, names };
      })());
    }
    return files.get(name);
  }

  // The row groups whose [min, max] of `column` overlaps [lo, hi]
  function overlapping(groups, column, lo, hi) {
    return groups.filter((g) => {
      const s = g.stats[column];
      if (!s) return true;
      const min = s.min_value ?? s.min, max = s.max_value ?? s.max;
      return !(max < lo || min > hi);
    });
  }

  // Adjacent row groups' bytes in one request each run, when the columns read are most of
  // their bytes (else hyparquet's own requests, column by column, fetch less)
  function prefetch(file, groups, columns) {
    const share = (g) => columns.reduce((n, c) => n + (g.bytes[c] ?? 0), 0) / (g.to - g.from);
    const sorted = groups.filter((g) => share(g) >= 0.8).sort((a, b) => a.from - b.from);
    for (let i = 0; i < sorted.length; ) {
      let j = i;
      while (j + 1 < sorted.length && sorted[j + 1].from <= sorted[j].to) j++;
      if (sorted[i].from < sorted[j].to) file.prefetch(sorted[i].from, sorted[j].to);
      i = j + 1;
    }
  }

  // Row groups read in parallel, rows in file order, 64-bit integers as numbers; columns the
  // file lacks (an older run's) are left out
  async function readGroups(name, groups, columns) {
    const { file, metadata, names } = await open(name);
    columns = columns.filter((c) => names.has(c));
    prefetch(file, groups, columns);
    const parts = await Promise.all(groups.map((g) => parquetReadObjects({
      file, metadata, columns, compressors, rowStart: g.start, rowEnd: g.end,
    })));
    return parts.flat().map(numbers);
  }

  return {
    // The whole features table (one small file)
    async features() {
      const { file, metadata } = await open("features.parquet");
      return (await parquetReadObjects({ file, metadata, compressors })).map(numbers);
    },

    // The class hierarchy, whole (one small file), or an empty list for a run without it
    async classes() {
      try {
        const { file, metadata } = await open("classes.parquet");
        return (await parquetReadObjects({ file, metadata, compressors })).map(numbers);
      } catch {
        return [];
      }
    },

    // A class's direct members among the coded items: [{ id: "Q…", subclass }], or [] for a
    // run without the members table
    async members(cls) {
      try {
        const { groups } = await open("members.parquet");
        const rows = await readGroups("members.parquet",
          overlapping(groups, "class", cls, cls), ["class", "id", "subclass"]);
        return rows.filter((r) => r.class === cls).map((r) => ({ id: `Q${r.id}`, subclass: r.subclass }));
      } catch {
        return [];
      }
    },

    // Fetch a file's footer ahead of its first lookup
    warm(name) { open(name).catch(() => {}); },

    // Items whose lowercased label starts with `q`: exact matches first, then coded items,
    // then by Wikipedias.
    // With several words and few such labels, also those whose label starts with the first
    // words and whose description has the rest ("transformer machine learning"). Reads at
    // most `maxGroups` row groups of names per label prefix.
    // With `aliases`, an item's aliases match too (rows with `alias` set); an item found
    // more than once is listed once.
    async search(q, { limit = 40, maxGroups = 3, aliases = false } = {}) {
      const { groups } = await open("names.parquet");
      const columns = ["key", "label", "alias", "id", "description", "wikipedias", "coded"];
      const byPrefix = async (prefix) => (await readGroups("names.parquet",
        overlapping(groups, "key", prefix, prefix + "\uffff").slice(0, maxGroups), columns))
        .filter((r) => r.key.startsWith(prefix) && (aliases || r.alias == null));
      // Exact names first, then items with a code (they can be explored), then by Wikipedias
      const rank = (rows, exact) => rows.sort((a, b) => (b.key === exact) - (a.key === exact)
        || (b.coded !== false) - (a.coded !== false) || b.wikipedias - a.wikipedias);
      const found = rank(await byPrefix(q), q);
      const words = q.split(/\s+/).filter(Boolean);
      for (let k = words.length - 1; k >= 1 && found.length < limit; k--) {
        const prefix = words.slice(0, k).join(" ");
        const rest = words.slice(k);
        const seen = new Set(found.map((r) => r.id));
        found.push(...rank((await byPrefix(prefix)).filter((r) => !seen.has(r.id) &&
          rest.every((w) => (r.description ?? "").toLowerCase().includes(w))), prefix));
      }
      const once = new Set();
      return found.filter((r) => !once.has(r.id) && once.add(r.id)).slice(0, limit);
    },

    // An item's description: names is sorted by lowercased label, so its label finds it
    async description(id, label) {
      if (!label) return null;
      const key = label.toLowerCase();
      const { groups } = await open("names.parquet");
      const rows = await readGroups("names.parquet", overlapping(groups, "key", key, key),
        ["key", "id", "description"]);
      return rows.find((r) => r.id === id)?.description ?? null;
    },

    // Items by id: one row group each (items is sorted by id)
    async items(ids, columns = ["id", "label", "kinds", "is_class", "features", "weights", "norm"]) {
      const { groups } = await open("items.parquet");
      const want = new Set(ids);
      const hit = new Set(ids.flatMap((id) => overlapping(groups, "id", id, id)));
      const rows = await readGroups("items.parquet", [...hit], columns);
      return rows.filter((r) => want.has(r.id));
    },

    // A feature's postings, as columns: `ids` (the number of each "Q…"), `units` (its weight
    // for the feature over its norm), and its `kinds`, those of posting i being
    // kinds[kindAt[i]] to kinds[kindAt[i + 1]]. Typed arrays, at about 17 bytes a posting:
    // one object per posting took ten times that, and a few items' worth of families ran a
    // browser tab out of memory. The latest are kept (filters re-rank the same postings), up
    // to KEEP_POSTINGS in all. With `limit`, the first `limit` as objects ({ id, unit, kinds }).
    // Older layouts store `weight` and `norm`, heaviest first, or a float `unit`; the current
    // one a 16-bit `unit16`, by id.
    async postings(feature, { limit = Infinity } = {}) {
      if (limit === Infinity && postingsOf.has(feature)) {
        const kept = postingsOf.get(feature);
        postingsOf.delete(feature);  // most recently used last
        postingsOf.set(feature, kept);
        return kept;
      }
      const { groups } = await open("postings.parquet");
      let hit = overlapping(groups, "feature", feature, feature);
      if (limit < Infinity) hit = hit.slice(0, 1);  // rank order: the first row group leads
      const read = readColumns("postings.parquet", hit,
        ["feature", "id", "unit16", "unit", "weight", "norm", "kinds"])
        .then((parts) => packPostings(parts, feature));
      if (limit < Infinity) {
        const p = await read;
        return Array.from({ length: Math.min(limit, p.length) }, (_, i) =>
          ({ id: `Q${p.ids[i]}`, unit: p.units[i], kinds: p.kindsOf(i) }));
      }
      postingsOf.set(feature, read);
      read.then(() => trimPostings(), () => postingsOf.delete(feature));
      return read;
    },
  };

  // The oldest postings dropped while more than KEEP_POSTINGS are kept
  async function trimPostings() {
    let total = 0;
    const sizes = [];
    for (const [f, p] of postingsOf) {
      const n = await Promise.race([p.then((x) => x.length, () => 0), 0]);
      sizes.push([f, n]);
      total += n;
    }
    for (const [f, n] of sizes) {
      if (total <= KEEP_POSTINGS || sizes.length < 2) break;
      postingsOf.delete(f);
      total -= n;
    }
  }

  // Row groups read in parallel as columns, without an object per row: for each group,
  // { column: values }, typed arrays where hyparquet decodes to them
  async function readColumns(name, groups, columns) {
    const { file, metadata, names } = await open(name);
    columns = columns.filter((c) => names.has(c));
    prefetch(file, groups, columns);
    return Promise.all(groups.map(async (g) => {
      const chunks = Object.fromEntries(columns.map((c) => [c, []]));
      await parquetRead({
        file, metadata, columns, compressors, rowStart: g.start, rowEnd: g.end,
        onChunk: ({ columnName, columnData, rowStart, rowEnd }) => {
          const from = Math.max(g.start, rowStart), to = Math.min(g.end, rowEnd);
          if (from < to) chunks[columnName].push([from, columnData.slice(from - rowStart, to - rowStart)]);
        },
      });
      const part = {};
      for (const c of columns) {
        const list = chunks[c].sort((a, b) => a[0] - b[0]).map(([, d]) => d);
        part[c] = list.length === 1 ? list[0] : list.flatMap((d) => Array.from(d));
      }
      return part;
    }));
  }
}

// The most postings kept between lookups, over all features (about 100 MB)
const KEEP_POSTINGS = 6_000_000;

// One feature's rows from row groups of postings, packed into typed arrays
function packPostings(parts, feature) {
  let n = 0, nk = 0;
  for (const p of parts) {
    for (let i = 0; i < p.feature.length; i++) {
      if (Number(p.feature[i]) !== feature || p.id[i] == null) continue;  // a property: no Q number
      n++;
      nk += p.kinds?.[i]?.length ?? 0;
    }
  }
  const ids = new Uint32Array(n), units = new Float32Array(n);
  const kindAt = new Uint32Array(n + 1), kinds = new Uint32Array(nk);
  let r = 0, k = 0;
  for (const p of parts) {
    for (let i = 0; i < p.feature.length; i++) {
      if (Number(p.feature[i]) !== feature || p.id[i] == null) continue;  // a property: no Q number
      const id = p.id[i];
      ids[r] = typeof id === "string" ? +id.slice(1) : Number(id);
      units[r] = p.unit16 ? Number(p.unit16[i]) / 65535
        : p.unit ? Number(p.unit[i]) : Number(p.weight[i]) / Number(p.norm[i]);
      kindAt[r] = k;
      for (const c of p.kinds?.[i] ?? []) kinds[k++] = Number(c);
      r++;
    }
  }
  kindAt[n] = k;
  return {
    length: n, ids, units, kindAt, kinds,
    kindsOf: (i) => Array.from(kinds.subarray(kindAt[i], kindAt[i + 1])),
  };
}

function numbers(row) {
  for (const k in row) {
    const v = row[k];
    if (typeof v === "bigint") row[k] = Number(v);
    else if (Array.isArray(v) && typeof v[0] === "bigint") row[k] = v.map(Number);
  }
  return row;
}

// The items most like `item`, by cosine of the weighted features, over the item's features
// that `use` accepts (e.g. those on at most so many items: a feature on a million items
// costs megabytes to read and weighs little), at most `top` of them, heaviest first, among
// the candidates `keep` accepts (each with its `kinds`). Returns those features ([feature,
// weight]) and the neighbours, each with `shared`: which of them it has. `onRead(done, total)`
// is called as each feature's postings arrive.
export async function neighbours(data, item,
  { use = () => true, keep = () => true, top = 12, limit = 30, onRead = () => {} } = {}) {
  const seed = item.features.map((f, i) => [f, item.weights[i]])
    .filter(([f]) => use(f))
    .sort((a, b) => b[1] - a[1]).slice(0, top);
  let done = 0;
  onRead(0, seed.length);
  const lists = await Promise.all(seed.map(([f]) =>
    data.postings(f).then((list) => (onRead(++done, seed.length), list))));
  // Each candidate's dot product, which seed features it has (a bit each), and where its
  // kinds are (the first list it is in, and its row there), in typed arrays: hundreds of
  // thousands of candidates are common
  const self = +item.id.slice(1);
  const most = lists.reduce((n, p) => n + p.length, 0);
  const slot = new Map();
  const ids = new Uint32Array(most), dot = new Float64Array(most), mask = new Uint32Array(most);
  const from = new Uint8Array(most), row = new Uint32Array(most);
  let n = 0;
  seed.forEach(([, w], i) => {
    const p = lists[i];
    for (let r = 0; r < p.length; r++) {
      const id = p.ids[r];
      if (id === self) continue;
      let s = slot.get(id);
      if (s === undefined) {
        slot.set(id, s = n++);
        ids[s] = id;
        from[s] = i;
        row[s] = r;
      }
      dot[s] += p.units[r] * w;
      mask[s] |= 1 << i;
    }
  });
  // A light object per candidate; its id, kinds and shared features are worked out when read
  class Candidate {
    constructor(s) { this.s = s; this.similarity = dot[s] / item.norm; }
    get id() { return `Q${ids[this.s]}`; }
    get kinds() { return lists[from[this.s]].kindsOf(row[this.s]); }
    get shared() { return seed.map((_, i) => i).filter((i) => mask[this.s] & (1 << i)); }
  }
  const order = new Uint32Array(n).map((_, s) => s)
    .sort((a, b) => dot[b] - dot[a] || ids[a] - ids[b]);
  const near = [];
  for (const s of order) {
    const c = new Candidate(s);
    if (keep(c)) near.push(c);
    if (near.length >= limit) break;
  }
  return { seed, near };
}