// Lists the support results of the aws-neuron/torch-neuronx-model-support
// dataset, one tab per library, 100 models per page: for the models the latest
// sweep selected, or for the Hub's models by trending score or downloads.
import { parquetReadObjects } from "https://cdn.jsdelivr.net/npm/hyparquet@1.31.2/src/index.js";

const HUB = "https://huggingface.co";
const DATASET = "aws-neuron/torch-neuronx-model-support";
const NEFF_CACHE = "hf://buckets/aws-neuron/torch-neuronx-neff-cache";
const LIBRARIES = ["transformers", "diffusers", "sentence-transformers"];
const VIEWS = { selected: "Selected", trending: "Trending", downloads: "Downloads" };
// The task Trending and Downloads open on: all tasks would open on models the sweep never tests.
const DEFAULT_TASK = "text-generation";
// The Hub's sort, and model field, of each view but "selected".
const HUB_SORTS = { trending: "trendingScore", downloads: "downloads" };
const PAGE = 100;
// The test instance: 4 Trainium2 devices of 96 GiB of HBM each.
const DEVICES = 4;
const DEVICE_HBM = 96 * 1024 ** 3;
const COLUMNS = ["model_id", "cores", "result", "wall (s)", "Memory (GB)", "notes"];

// An element with attributes and children. Strings become text nodes, never HTML.
function h(tag, attrs = {}, ...children) {
  const el = document.createElement(tag);
  for (const [key, value] of Object.entries(attrs)) {
    if (key.startsWith("on")) el.addEventListener(key.slice(2), value);
    else if (value === true) el.setAttribute(key, "");
    else if (value !== false && value != null) el.setAttribute(key, value);
  }
  el.append(...children.flat(Infinity).filter((c) => c != null && c !== false));
  return el;
}

const encodePath = (path) => path.split("/").map(encodeURIComponent).join("/");

async function fetchOk(url) {
  const response = await fetch(url);
  if (!response.ok) throw new Error(`${url}: HTTP ${response.status}`);
  return response;
}

// Promises by key, dropped on failure so that the next render retries.
const cache = new Map();
function cached(key, load) {
  if (!cache.has(key)) cache.set(key, load().catch((e) => { cache.delete(key); throw e; }));
  return cache.get(key);
}

// The dataset's files at its latest commit: one per (neuronx-cc, library version)
// pair.
async function loadIndex() {
  const info = await (await fetchOk(`${HUB}/api/datasets/${DATASET}`)).json();
  const pairs = [];
  for (const { rfilename: path } of info.siblings) {
    const pair = path.match(/^data\/neuronxcc-([^/]+)\/(.+)-([^-/]+)\.parquet$/);
    if (pair) pairs.push({ sdk: pair[1], library: pair[2], version: pair[3], path });
  }
  return { sha: info.sha, pairs };
}

const index = await loadIndex().catch((e) => e);
const fileUrl = (path) => `${HUB}/datasets/${DATASET}/resolve/${index.sha}/${encodePath(path)}`;

// A pair's rows: its runs, the models its sweep left out, and each model's rank
// in the selection of its last sweep, best-ranked first.
const loadRows = (path) => cached(path, async () => {
  const rows = await parquetReadObjects({ file: await (await fetchOk(fileUrl(path))).arrayBuffer() });
  const ranked = rows.filter((row) => row.selection_rank != null)
    .sort((a, b) => a.selection_rank - b.selection_rank);
  return {
    runs: rows.filter((row) => row.result !== "EXCLUDED"),
    exclusions: new Map(rows.filter((row) => row.result === "EXCLUDED").map((row) => [row.model_id, row])),
    selection: new Map(ranked.map((row) => [row.model_id, row.selection_rank])),
  };
});

// A model's task on the Hub, or "unknown" if the Hub no longer has it.
const loadTask = (id) => cached(`task:${id}`, async () => {
  const query = new URLSearchParams({ "expand[]": "pipeline_tag" });
  const response = await fetch(`${HUB}/api/models/${encodePath(id)}?${query}`);
  if (response.status === 404) return "unknown";
  if (!response.ok) throw new Error(`${id}: HTTP ${response.status}`);
  return (await response.json()).pipeline_tag ?? "unknown";
}).catch(() => "unknown");

// The models of a library, of a task unless it is null, and whose id contains
// `name` unless it is empty, in the order of the Hub's model list, with their
// parameter count when the Hub knows it: at least `count` of them, or all if it
// has fewer. `next` is the Hub's link to the models after them, if any. The list
// grows one Hub page at a time, as deep as it was asked for.
const hubLists = new Map();
function hubList(library, task, name, view, count) {
  const key = `${library}:${task}:${name}:${view}`;
  const sort = HUB_SORTS[view];
  if (!hubLists.has(key)) {
    const query = new URLSearchParams({ library, sort, direction: "-1", limit: PAGE });
    if (task) query.set("pipeline_tag", task);
    if (name) query.set("search", name);
    for (const field of [sort, "safetensors", "library_name"]) query.append("expand[]", field);
    hubLists.set(key, { models: [], next: `${HUB}/api/models?${query}`, growing: Promise.resolve() });
  }
  const list = hubLists.get(key);
  list.growing = list.growing.catch(() => {}).then(async () => {
    while (list.models.length < count && list.next) {
      const response = await fetchOk(list.next);
      list.models.push(...(await response.json())
        // The Hub's library filter matches a tag that models of other libraries
        // carry too, such as diffusers' FLUX or sentence-transformers' MiniLM.
        .filter((model) => model.library_name === library)
        .map((model) => ({ id: model.id, score: model[sort] ?? 0, params: model.safetensors?.total })));
      list.next = response.headers.get("Link")?.match(/<([^>]+)>;\s*rel="next"/)?.[1];
    }
  });
  return list.growing.then(() => list);
}

// Newest first. A version that is not a release, such as a git commit pin,
// comes before the releases.
function newestFirst(a, b) {
  const release = (v) => v.match(/^\d+(?:\.\d+)+/)?.[0].split(".").map(Number);
  const [ra, rb] = [release(a), release(b)];
  if (!ra || !rb) return Boolean(ra) - Boolean(rb) || a.localeCompare(b);
  for (let i = 0; i < Math.max(ra.length, rb.length); i++) {
    const d = (rb[i] ?? 0) - (ra[i] ?? 0);
    if (d) return d;
  }
  // A pre-release, such as 5.21.0rc1, is older than its release.
  const pre = (v) => /^\d+(?:\.\d+)+(?:a|b|rc|\.?dev)/.test(v);
  return pre(a) - pre(b);
}

const badge = (result) => `${result === "PASS" ? "✅" : "❌"} ${result}`;
const memory = (run) => run.hbm_peak_gb == null ? ""
  : `${run.hbm_peak_gb.toFixed(2)} / ${Number(run.hbm_capacity_gb.toFixed(1))}`;
const day = (date) => new Date(date).toISOString().slice(0, 10);

const tabs = document.getElementById("tabs");
const panel = document.getElementById("panel");
const dialog = document.getElementById("run");
const state = new Map(LIBRARIES.map((library) =>
  [library, { view: "selected", sdk: null, version: null, tasks: {}, names: {}, page: 1 }]));
let active = LIBRARIES[0];
let rendering = 0;
// The pending list render of the name being typed.
let typing;

function renderTabs() {
  tabs.replaceChildren(...LIBRARIES.map((library) => h("button", {
    role: "tab",
    "aria-selected": String(library === active),
    onclick: () => { active = library; renderTabs(); renderPanel(); },
  }, library)));
}

async function renderPanel() {
  const render = ++rendering;
  clearTimeout(typing);
  if (index instanceof Error) {
    panel.replaceChildren(h("p", { class: "status" }, `Could not read the dataset: ${index.message}`));
    return;
  }
  const pairs = index.pairs.filter((p) => p.library === active);
  if (!pairs.length) {
    panel.replaceChildren(h("p", { class: "status" }, "No records yet."));
    return;
  }
  const s = state.get(active);
  const sdks = [...new Set(pairs.map((p) => p.sdk))].sort(newestFirst);
  const versionsOf = (sdk) => pairs.filter((p) => p.sdk === sdk).map((p) => p.version).sort(newestFirst);
  // The latest selection was swept with the newest library version of the newest compiler.
  const latest = { sdk: sdks[0], version: versionsOf(sdks[0])[0] };
  if (!sdks.includes(s.sdk)) s.sdk = latest.sdk;
  if (!versionsOf(s.sdk).includes(s.version)) s.version = versionsOf(s.sdk)[0];
  const selected = s.view === "selected";
  const shown = selected ? latest : { sdk: s.sdk, version: s.version };
  const pathOf = ({ sdk, version }) => pairs.find((p) => p.sdk === sdk && p.version === version).path;

  const select = (label, key, options, onchange) => h("label", {}, label, h("select", {
    onchange: (e) => { s[key] = e.target.value; s.page = 1; onchange?.(); renderPanel(); },
  }, options.map((o) => h("option", { value: o, selected: o === s[key] }, o))));
  const tasks = h("div", { class: "tasks", role: "group", "aria-label": "Task" }, h("span", {}, "Task"));
  const results = h("div", { class: "results" }, h("p", { class: "status" }, "Loading…"));
  // Typing renders the list again, not the panel, which would take the focus away.
  let renderResults;
  const nameField = h("label", {}, "Name", h("input", {
    type: "search",
    placeholder: "Filter by name",
    value: s.names[s.view] ?? "",
    oninput: (e) => {
      s.names[s.view] = e.target.value.trim();
      s.page = 1;
      clearTimeout(typing);
      typing = setTimeout(() => renderResults?.(), 300);
    },
  }));
  panel.replaceChildren(
    h("div", { class: "views", role: "tablist" }, Object.entries(VIEWS).map(([view, label]) => h("button", {
      role: "tab",
      "aria-selected": String(view === s.view),
      onclick: () => { s.view = view; s.page = 1; renderPanel(); },
    }, label))),
    selected
      ? h("p", { class: "pair" }, `neuronx-cc ${latest.sdk} · ${active} ${latest.version}`)
      : h("div", { class: "controls" },
        select("SDK (neuronx-cc)", "sdk", sdks, () => { s.version = null; }),
        select(active, "version", versionsOf(s.sdk)),
        nameField),
    ...(selected ? [] : [tasks]),
    results);

  let runs, rank, exclusions, taskList;
  try {
    let latestRows;
    [{ runs, exclusions }, latestRows] = await Promise.all([loadRows(pathOf(shown)), loadRows(pathOf(latest))]);
    rank = latestRows.selection;
    if (!selected) {
      // The tasks the sweep covers: those of the latest pair's and the selection's models.
      const covered = new Set([...latestRows.runs.map((run) => run.model_id), ...rank.keys()]);
      taskList = [...new Set(await Promise.all([...covered].map(loadTask)))].sort();
    }
  } catch (e) {
    if (render === rendering) results.replaceChildren(h("p", { class: "status" }, `Could not load the results: ${e.message}`));
    return;
  }
  if (render !== rendering) return;

  const runsOf = new Map();
  for (const run of runs) runsOf.set(run.model_id, [...(runsOf.get(run.model_id) ?? []), run]);
  // The task Trending or Downloads shows, or null for all of them.
  if (!selected && !(s.view in s.tasks)) s.tasks[s.view] = taskList.includes(DEFAULT_TASK) ? DEFAULT_TASK : null;
  const task = s.tasks[s.view] ?? null;

  // The current page of the selection, or of the Hub's list of the chosen task or
  // of all tasks: its models, each with its runs on the shown pair, the position
  // of the first, and whether more follow.
  async function listed() {
    const start = (s.page - 1) * PAGE;
    const end = start + PAGE;
    let models, more;
    if (selected) {
      const ids = [...rank.keys()];
      models = ids.slice(start, end).map((id) => ({ id }));
      more = ids.length > end;
    } else {
      const list = await hubList(active, task, s.names[s.view] ?? "", s.view, end);
      models = list.models.slice(start, end);
      more = list.models.length > end || Boolean(list.next);
    }
    return { start, more, models: models.map((model) => ({ ...model, runs: runsOf.get(model.id) ?? [] })) };
  }

  // Why a model was not tested: weights larger than the test instance's HBM, or
  // the reason the sweep's selection left it out.
  function untestedReason({ id, params }) {
    if (params * 2 > DEVICES * DEVICE_HBM) return `Requires more than ${DEVICES} Trainium2 devices`;
    return exclusions.get(id)?.notes ?? "";
  }

  renderResults = async () => {
    const render = ++rendering;
    let page;
    try {
      page = await listed();
    } catch (e) {
      if (render === rendering) results.replaceChildren(h("p", { class: "status" }, `Could not load the Hub's list: ${e.message}`));
      return;
    }
    if (render !== rendering) return;
    if (!page.models.length) {
      const name = s.names[s.view];
      results.replaceChildren(h("p", { class: "status" }, selected ? "No selection yet."
        : `No model of ${task ?? "any task"}${name ? ` matching "${name}"` : ""} is listed.`));
      return;
    }
    // A model passes if one of its runs passed.
    const passed = page.models.filter((m) => m.runs.some((run) => run.result === "PASS")).length;
    const untested = page.models.filter((m) => !m.runs.length).length;
    const go = async (pages) => {
      s.page += pages;
      await renderResults();
      if (results.getBoundingClientRect().top < 0) results.scrollIntoView();
    };
    const pager = () => h("div", { class: "pager" },
      h("button", { disabled: s.page === 1, onclick: () => go(-1) }, "‹ Previous"),
      h("button", { disabled: !page.more, onclick: () => go(1) }, "Next ›"));
    results.replaceChildren(
      h("div", { class: "summary" },
        h("p", {}, h("strong", {}, `${page.start + 1}-${page.start + page.models.length}`),
          ` · ✅ Pass ${passed} · ❌ Fail ${page.models.length - passed - untested} · Untested ${untested}`),
        pager()),
      h("div", { class: "scroll" }, h("table", {},
        h("thead", {}, h("tr", {}, COLUMNS.map((c) => h("th", {}, c)))),
        h("tbody", {}, page.models.map((model) => model.runs.length ? model.runs.map((run) => h("tr", {
          tabindex: 0,
          onclick: () => showRun(run),
          onkeydown: (e) => { if (e.key === "Enter") showRun(run); },
        },
        h("td", {}, h("code", {}, model.id)),
        h("td", { class: "num" }, run.neuron_cores),
        h("td", { class: "result" }, badge(run.result)),
        h("td", { class: "num" }, run.wall_s ?? ""),
        h("td", { class: "num" }, memory(run)),
        h("td", { class: "notes" }, run.notes ?? "")))
          : h("tr", { class: "untested" },
            h("td", {}, h("code", {}, model.id)), h("td"), h("td", { class: "result" }, "not tested"),
            h("td"), h("td"), h("td", { class: "notes" }, untestedReason(model))))))),
      pager());
  };
  // One task at a time, as on the Hub: clicking the chosen task again shows all of them.
  if (!selected) {
    tasks.append(...[null, ...taskList].map((t) => h("button", {
      "aria-pressed": String(task === t),
      onclick: () => { s.tasks[s.view] = task === t ? null : t; s.page = 1; renderPanel(); },
    }, t ?? "All")));
  }
  renderResults();
}

let scriptUrl;
function showRun(run) {
  URL.revokeObjectURL(scriptUrl);
  scriptUrl = URL.createObjectURL(new Blob([run.run_py], { type: "text/x-python" }));
  const pip = run.run_py.match(/^#   (pip install .*)$/m)?.[1];
  const stack = [
    ["neuronx-cc", run.neuronx_cc], ["torch-neuronx", run.torch_neuronx], ["torch", run.torch],
    ["nki", run.nki], ["Python", run.python], ["Neuron driver", run.neuron_driver],
  ].filter(([, version]) => version);
  dialog.replaceChildren(h("div", { class: "dialog" },
    h("form", { method: "dialog" }, h("button", { class: "close", "aria-label": "Close" }, "×")),
    h("h2", {}, h("a", { href: `${HUB}/${run.model_id}`, target: "_blank", rel: "noopener" }, run.model_id)),
    h("p", {}, `${badge(run.result)} · ${run.neuron_cores} NeuronCore${run.neuron_cores > 1 ? "s" : ""}`
      + ` · ${run.runner} · ${run.library} ${run.library_version}`
      + ` · tested ${day(run.tested_at)}`),
    h("h3", {}, "Run it"),
    h("p", {}, "On a Trainium2 (trn2) instance with this Neuron stack, for example in the AWS Neuron"
      + " Deep Learning Container:"),
    h("dl", {}, stack.map(([name, version]) => [h("dt", {}, name), h("dd", {}, h("code", {}, version))])),
    h("ol", {},
      pip && h("li", {}, "Install the packages the script adds:", h("pre", {}, pip)),
      h("li", {}, "Optionally, skip compilation with the graphs this compiler already built:",
        h("pre", {}, `hf buckets sync ${NEFF_CACHE}/neuronxcc-${run.neuronx_cc} /tmp/neff_cache`)),
      h("li", {}, "Download ", h("a", { href: scriptUrl, download: "run.py" }, "run.py"), " and run it",
        run.neuron_cores > 1 ? ". It launches its own torch.distributed ranks:" : ":",
        h("pre", {}, "python3 run.py"))),
    run.error && [h("h3", {}, "Error"), h("pre", { class: "error" }, run.error)],
    h("details", {}, h("summary", {}, "run.py"), h("pre", {}, run.run_py))));
  dialog.showModal();
}
dialog.addEventListener("click", (e) => { if (e.target === dialog) dialog.close(); });

renderTabs();
renderPanel();