const HF_REPO_BASE = "https://huggingface.co/datasets";
const LARGE_FILE_BYTES = 15 * 1024 * 1024;
// Pairplot panels grow quadratically; past 12 dims the plot is unreadable
// and canvas rendering gets slow. Cells never shrink below MIN_CELL_PX —
// the canvas grows instead and the wrapper scrolls.
const MAX_DIMS = 12;
const MIN_CELL_PX = 110;

// Canvas colors per theme; the page chrome is themed via CSS variables and
// the [data-theme] attribute, but canvas drawing needs explicit values.
const THEMES = {
  light: {
    canvasBg: "#ffffff",
    text: "#667085",
    grid: "#d9deea",
    hist: "rgba(49, 91, 232, 0.55)",
    point: "rgba(49, 91, 232, 0.25)",
    truth: "#e5484d",
  },
  dark: {
    canvasBg: "#1e222d",
    text: "#9aa3b5",
    grid: "#3a4257",
    hist: "rgba(126, 153, 255, 0.6)",
    point: "rgba(126, 153, 255, 0.35)",
    truth: "#ff6b70",
  },
};

const state = {
  manifest: null,
  task: null,
  samples: [],
  trueTheta: null,
  loadedObservation: null,
  themeName: "light",
  canvasMessage: "Click Load to fetch posterior samples.",
};

const params = new URLSearchParams(window.location.search);
if (params.get("embedded") === "1") {
  document.body.classList.add("embedded");
}

function theme() {
  return THEMES[state.themeName] || THEMES.light;
}

function applyTheme(name) {
  state.themeName = name === "dark" ? "dark" : "light";
  document.documentElement.dataset.theme = state.themeName;
  if (state.samples.length) {
    drawPlot();
  } else {
    clearCanvas(state.canvasMessage);
  }
}

// Initial theme: explicit ?theme= wins (set by the embedding docs page),
// otherwise follow the OS preference. Only the attribute is set here —
// nothing is drawn yet, and applyTheme would touch elements not yet created.
const themeParam = params.get("theme");
const prefersDark = window.matchMedia("(prefers-color-scheme: dark)");
state.themeName = (themeParam || (prefersDark.matches ? "dark" : "light")) === "dark" ? "dark" : "light";
document.documentElement.dataset.theme = state.themeName;
if (!themeParam) {
  prefersDark.addEventListener("change", (event) => applyTheme(event.matches ? "dark" : "light"));
}

// The embedding page notifies us when the docs palette toggles.
window.addEventListener("message", (event) => {
  const data = event.data;
  if (data && data.type === "sbibm-theme" && (data.theme === "dark" || data.theme === "light")) {
    applyTheme(data.theme);
  }
});

const els = {
  taskSelect: document.getElementById("taskSelect"),
  observationSelect: document.getElementById("observationSelect"),
  pointLimitSelect: document.getElementById("pointLimitSelect"),
  loadButton: document.getElementById("loadButton"),
  taskInfo: document.getElementById("taskInfo"),
  dimensionGrid: document.getElementById("dimensionGrid"),
  status: document.getElementById("status"),
  canvas: document.getElementById("plotCanvas"),
};

function formatBytes(bytes) {
  if (!Number.isFinite(bytes)) return "unknown size";
  if (bytes < 1024 * 1024) return `${Math.round(bytes / 1024)} KB`;
  return `${(bytes / (1024 * 1024)).toFixed(1)} MB`;
}

function taskByName(name) {
  return state.manifest.tasks.find((task) => task.task === name) || state.manifest.tasks[0];
}

function hfResolveUrl(path) {
  return `${HF_REPO_BASE}/${state.manifest.repo}/resolve/main/${path}`;
}

function hfTreeUrl(path) {
  return `${HF_REPO_BASE}/${state.manifest.repo}/tree/main/${path}`;
}

function apiTreeUrl(path) {
  return `https://huggingface.co/api/datasets/${state.manifest.repo}/tree/main/${path}`;
}

function setStatus(message) {
  els.status.textContent = message;
}

function selectedDims() {
  return Array.from(els.dimensionGrid.querySelectorAll("input:checked"))
    .map((input) => Number(input.value))
    .filter(Number.isInteger)
    .slice(0, MAX_DIMS);
}

function renderTaskOptions() {
  els.taskSelect.innerHTML = state.manifest.tasks
    .map((task) => `<option value="${task.task}">${task.display_name}</option>`)
    .join("");
}

function renderObservationOptions(task) {
  const count = task.num_observations || 100;
  let html = "";
  for (let i = 0; i < count; i += 1) {
    html += `<option value="${i}">${i}</option>`;
  }
  els.observationSelect.innerHTML = html;
}

function renderTaskInfo(task) {
  els.taskInfo.innerHTML = [
    `<span class="pill">${task.dim_parameters} θ dims</span>`,
    `<span class="pill">${task.dim_data} x dims</span>`,
    `<span class="pill">${task.sample_count.toLocaleString()} samples</span>`,
    `<span class="pill">${formatBytes(task.sample_file_size_bytes)} / observation</span>`,
    `<span>Reference potential: <code>${task.reference_potential_type}</code></span>`,
    `<a href="${hfTreeUrl(`${task.task}/reference_samples`)}" target="_blank" rel="noreferrer">HF files</a>`,
  ].join(" ");
}

function renderDimensionOptions(task) {
  const requested = (params.get("dims") || "")
    .split(",")
    .map((value) => Number(value.trim()))
    .filter((value) => Number.isInteger(value) && value >= 0 && value < task.dim_parameters)
    .slice(0, MAX_DIMS);
  const defaultCount = Math.min(task.dim_parameters, task.dim_parameters <= 5 ? task.dim_parameters : 4);
  const defaults = requested.length >= 2 ? requested : Array.from({ length: defaultCount }, (_, i) => i);

  let html = "";
  for (let dim = 0; dim < task.dim_parameters; dim += 1) {
    const checked = defaults.includes(dim) ? " checked" : "";
    html += `<label class="dim-check"><input type="checkbox" value="${dim}"${checked}> θ${dim}</label>`;
  }
  els.dimensionGrid.innerHTML = html;

  els.dimensionGrid.querySelectorAll("input").forEach((input) => {
    input.addEventListener("change", () => {
      const checkedCount = els.dimensionGrid.querySelectorAll("input:checked").length;
      if (input.checked && checkedCount > MAX_DIMS) {
        input.checked = false;
        setStatus(`At most ${MAX_DIMS} dimensions can be plotted at once.`);
        return;
      }
      drawPlot();
    });
  });
}

function applyTask(task) {
  state.task = task;
  state.samples = [];
  state.trueTheta = null;
  state.loadedObservation = null;
  renderObservationOptions(task);
  renderTaskInfo(task);
  renderDimensionOptions(task);
  clearCanvas("Click Load to fetch posterior samples.");
  setStatus("Pick an observation, then click Load.");
}

async function importHyparquet() {
  return import("https://cdn.jsdelivr.net/npm/hyparquet@1.14.0/+esm");
}

function vectorFromRow(row, name) {
  const value = row[name];
  if (value == null) return null;
  if (Array.isArray(value)) return value.map(Number);
  if (ArrayBuffer.isView(value)) return Array.from(value, Number);
  if (typeof value === "object") return Object.values(value).map(Number);
  return null;
}

function selectParquetFiles(entries) {
  const files = entries
    .filter((entry) => entry.type === "file" && entry.path.endsWith(".parquet"))
    .sort((a, b) => a.path.localeCompare(b.path));
  const single = files.find((entry) => /data-00000-of-00001\.parquet$/.test(entry.path));
  if (single) return [single];
  return files;
}

async function listSampleFiles(task, observationId) {
  const path = `${task.task}/reference_samples/observation_${observationId}`;
  const response = await fetch(apiTreeUrl(path));
  if (!response.ok) throw new Error(`Could not list sample files: HTTP ${response.status}`);
  const entries = await response.json();
  const files = selectParquetFiles(entries);
  if (!files.length) throw new Error("No parquet sample files found for this observation.");
  return files;
}

async function readParquetRows(url) {
  const { asyncBufferFromUrl, parquetReadObjects } = await importHyparquet();
  const file = await asyncBufferFromUrl({ url });
  return parquetReadObjects({ file });
}

async function loadSamples(task, observationId) {
  const files = await listSampleFiles(task, observationId);
  const totalBytes = files.reduce((sum, file) => sum + (file.size || 0), 0);
  if (totalBytes > LARGE_FILE_BYTES) {
    const ok = window.confirm(
      `${task.display_name} observation ${observationId} is ${formatBytes(totalBytes)}. Continue loading it in the browser?`,
    );
    if (!ok) return null;
  }

  const samples = [];
  for (const file of files) {
    const rows = await readParquetRows(hfResolveUrl(file.path));
    for (const row of rows) {
      const theta = vectorFromRow(row, "theta");
      if (theta) samples.push(theta);
    }
  }
  return samples;
}

async function loadTrueTheta(task, observationId) {
  try {
    const rows = await readParquetRows(hfResolveUrl(`${task.task}/reference_observations/data-00000-of-00001.parquet`));
    const row = rows.find((item) => Number(item.observation_id) === observationId) || rows[observationId];
    return row ? vectorFromRow(row, "true_theta") : null;
  } catch (_) {
    return null;
  }
}

function downsampleRows(rows, limit) {
  if (rows.length <= limit) return rows;
  const step = rows.length / limit;
  const out = [];
  for (let i = 0; i < limit; i += 1) {
    out.push(rows[Math.floor(i * step)]);
  }
  return out;
}

async function loadCurrentObservation() {
  const task = state.task;
  const observationId = Number(els.observationSelect.value);
  els.loadButton.disabled = true;
  els.status.classList.add("loading");
  setStatus("Loading parquet files from Hugging Face...");
  try {
    const samples = await loadSamples(task, observationId);
    if (!samples) {
      setStatus("Load cancelled.");
      return;
    }
    state.samples = samples;
    state.trueTheta = await loadTrueTheta(task, observationId);
    state.loadedObservation = observationId;
    setStatus(`Loaded ${samples.length.toLocaleString()} samples for observation ${observationId}.`);
    drawPlot();
  } catch (error) {
    state.samples = [];
    state.trueTheta = null;
    setStatus(error && error.message ? error.message : String(error));
    clearCanvas("Could not load samples.");
  } finally {
    els.loadButton.disabled = false;
    els.status.classList.remove("loading");
  }
}

function clearCanvas(message) {
  state.canvasMessage = message;
  const canvas = els.canvas;
  // Undo any explicit sizing a previous large pairplot applied.
  canvas.style.width = "100%";
  canvas.style.height = "auto";
  const ctx = canvas.getContext("2d");
  const ratio = window.devicePixelRatio || 1;
  const width = Math.max(320, Math.floor(canvas.clientWidth || 900));
  const height = Math.min(900, Math.max(320, width));
  canvas.width = Math.floor(width * ratio);
  canvas.height = Math.floor(height * ratio);
  ctx.setTransform(ratio, 0, 0, ratio, 0, 0);
  ctx.clearRect(0, 0, width, height);
  ctx.fillStyle = theme().canvasBg;
  ctx.fillRect(0, 0, width, height);
  ctx.fillStyle = theme().text;
  ctx.textAlign = "center";
  ctx.font = "14px system-ui, sans-serif";
  ctx.fillText(message, width / 2, height / 2);
}

function extent(values) {
  let lo = Infinity;
  let hi = -Infinity;
  for (const value of values) {
    if (!Number.isFinite(value)) continue;
    lo = Math.min(lo, value);
    hi = Math.max(hi, value);
  }
  if (!Number.isFinite(lo) || !Number.isFinite(hi)) return [-1, 1];
  if (lo === hi) return [lo - 1, hi + 1];
  const pad = (hi - lo) * 0.05;
  return [lo - pad, hi + pad];
}

function scale(value, domain, range) {
  return range[0] + ((value - domain[0]) / (domain[1] - domain[0])) * (range[1] - range[0]);
}

function drawHistogram(ctx, values, x0, y0, size, domain) {
  const bins = 28;
  const counts = Array.from({ length: bins }, () => 0);
  for (const value of values) {
    const idx = Math.max(0, Math.min(bins - 1, Math.floor(((value - domain[0]) / (domain[1] - domain[0])) * bins)));
    counts[idx] += 1;
  }
  const maxCount = Math.max(...counts, 1);
  ctx.fillStyle = theme().hist;
  counts.forEach((count, idx) => {
    const h = (count / maxCount) * (size - 22);
    const x = x0 + 10 + (idx / bins) * (size - 20);
    const w = (size - 20) / bins - 1;
    ctx.fillRect(x, y0 + size - 10 - h, w, h);
  });
}

function drawPlot() {
  if (!state.samples.length) {
    clearCanvas("Click Load to fetch posterior samples.");
    return;
  }
  const dims = selectedDims();
  if (dims.length < 2) {
    clearCanvas("Select at least two dimensions.");
    return;
  }

  const limit = Number(els.pointLimitSelect.value);
  const rows = downsampleRows(state.samples, limit);
  const canvas = els.canvas;
  const ctx = canvas.getContext("2d");
  const ratio = window.devicePixelRatio || 1;
  const n = dims.length;
  const margin = 34;
  const gap = 7;

  // Fit the grid to the container width; once cells would drop below
  // MIN_CELL_PX the canvas grows instead and .plot-scroll scrolls.
  const available = Math.max(340, Math.floor(canvas.parentElement.clientWidth || 900));
  const fitCell = (Math.min(available, 980) - margin * 2 - gap * (n - 1)) / n;
  const cell = Math.max(fitCell, MIN_CELL_PX);
  const cssSize = Math.ceil(margin * 2 + cell * n + gap * (n - 1));
  canvas.style.width = cssSize + "px";
  canvas.style.height = cssSize + "px";
  canvas.width = Math.floor(cssSize * ratio);
  canvas.height = Math.floor(cssSize * ratio);
  ctx.setTransform(ratio, 0, 0, ratio, 0, 0);
  ctx.clearRect(0, 0, cssSize, cssSize);
  ctx.fillStyle = theme().canvasBg;
  ctx.fillRect(0, 0, cssSize, cssSize);
  const domains = new Map();
  for (const dim of dims) {
    domains.set(dim, extent(state.samples.map((row) => row[dim])));
  }

  ctx.font = "11px system-ui, sans-serif";
  ctx.textAlign = "center";
  ctx.textBaseline = "middle";

  dims.forEach((yDim, rowIdx) => {
    dims.forEach((xDim, colIdx) => {
      const x0 = margin + colIdx * (cell + gap);
      const y0 = margin + rowIdx * (cell + gap);
      ctx.strokeStyle = theme().grid;
      ctx.lineWidth = 1;
      ctx.strokeRect(x0, y0, cell, cell);

      if (rowIdx === colIdx) {
        drawHistogram(ctx, rows.map((row) => row[xDim]), x0, y0, cell, domains.get(xDim));
      } else {
        ctx.fillStyle = theme().point;
        for (const row of rows) {
          const x = scale(row[xDim], domains.get(xDim), [x0 + 8, x0 + cell - 8]);
          const y = scale(row[yDim], domains.get(yDim), [y0 + cell - 8, y0 + 8]);
          ctx.fillRect(x, y, 1.8, 1.8);
        }

        if (state.trueTheta) {
          const tx = scale(state.trueTheta[xDim], domains.get(xDim), [x0 + 8, x0 + cell - 8]);
          const ty = scale(state.trueTheta[yDim], domains.get(yDim), [y0 + cell - 8, y0 + 8]);
          ctx.strokeStyle = theme().truth;
          ctx.lineWidth = 1.8;
          ctx.beginPath();
          ctx.moveTo(tx - 5, ty);
          ctx.lineTo(tx + 5, ty);
          ctx.moveTo(tx, ty - 5);
          ctx.lineTo(tx, ty + 5);
          ctx.stroke();
        }
      }

      if (rowIdx === n - 1) {
        ctx.fillStyle = theme().text;
        ctx.fillText(`θ${xDim}`, x0 + cell / 2, y0 + cell + 15);
      }
      if (colIdx === 0) {
        ctx.save();
        ctx.translate(x0 - 17, y0 + cell / 2);
        ctx.rotate(-Math.PI / 2);
        ctx.fillStyle = theme().text;
        ctx.fillText(`θ${yDim}`, 0, 0);
        ctx.restore();
      }
    });
  });
}

async function init() {
  const response = await fetch("manifest.json", { cache: "no-cache" });
  state.manifest = await response.json();
  renderTaskOptions();
  const initialTask = taskByName(params.get("task"));
  els.taskSelect.value = initialTask.task;
  applyTask(initialTask);

  const requestedObservation = Number(params.get("observation"));
  if (Number.isInteger(requestedObservation) && requestedObservation >= 0 && requestedObservation < initialTask.num_observations) {
    els.observationSelect.value = String(requestedObservation);
  }

  els.taskSelect.addEventListener("change", () => applyTask(taskByName(els.taskSelect.value)));
  els.loadButton.addEventListener("click", loadCurrentObservation);
  els.pointLimitSelect.addEventListener("change", drawPlot);
  window.addEventListener("resize", drawPlot);

  if (params.get("autoload") === "1") {
    loadCurrentObservation();
  }
}

init().catch((error) => {
  setStatus(error && error.message ? error.message : String(error));
});