const OP_TYPES = ["F", "B", "VF", "VB", "GF", "GB"];

const state = {
  schedulers: [],
  durationSpecs: defaultDurationSpecs(),
  limits: {
    maxMicrobatches: null,
  },
  result: null,
  chromeTrace: null,
  perfettoTrace: null,
  selectedOpId: null,
  scale: 72,
  showLabels: true,
  compare: {
    resultA: null,
    resultB: null,
    comparison: null,
    scale: 72,
    showLabels: true,
    selected: null,
  },
};

const colors = {
  F: ["#7db7f0", "#2f78c4", "#174f96"],
  B: ["#8bd2a2", "#32a463", "#197a43"],
  VF: ["#f2c66d", "#e49c31", "#b86f18"],
  VB: ["#f08d80", "#d95749", "#a9362f"],
  GF: ["#c8a2f0", "#8f62cc", "#653b9d"],
  GB: ["#ee9fc5", "#d65c93", "#aa356e"],
};

const els = {
  tabSimulate: document.getElementById("tab-simulate"),
  tabCompare: document.getElementById("tab-compare"),
  simulateView: document.getElementById("simulate-view"),
  compareView: document.getElementById("compare-view"),
  form: document.getElementById("config-form"),
  scheduler: document.getElementById("scheduler"),
  ppSize: document.getElementById("pp-size"),
  vppSize: document.getElementById("vpp-size"),
  numMicrobatches: document.getElementById("num-microbatches"),
  veForwardLimit: document.getElementById("ve-forward-limit"),
  seed: document.getElementById("seed"),
  durationFields: document.getElementById("duration-fields"),
  summary: document.getElementById("summary-strip"),
  timeline: document.getElementById("timeline"),
  timelineWrap: document.getElementById("timeline-wrap"),
  tooltip: document.getElementById("tooltip"),
  details: document.getElementById("details-panel"),
  status: document.getElementById("status"),
  zoom: document.getElementById("zoom-slider"),
  fit: document.getElementById("fit-button"),
  showLabels: document.getElementById("show-labels"),
  downloadChrome: document.getElementById("download-chrome"),
  downloadPerfetto: document.getElementById("download-perfetto"),
  compareForm: document.getElementById("compare-form"),
  comparePpSize: document.getElementById("compare-pp-size"),
  compareVppSize: document.getElementById("compare-vpp-size"),
  compareNumMicrobatches: document.getElementById("compare-num-microbatches"),
  compareVeForwardLimit: document.getElementById("compare-ve-forward-limit"),
  compareSeed: document.getElementById("compare-seed"),
  compareIncludeEncoder: document.getElementById("compare-include-encoder"),
  compareIncludeGenerator: document.getElementById("compare-include-generator"),
  compareSchedulerA: document.getElementById("compare-scheduler-a"),
  compareSchedulerB: document.getElementById("compare-scheduler-b"),
  compareDurationFields: document.getElementById("compare-duration-fields"),
  compareSummary: document.getElementById("compare-summary-strip"),
  compareTimelineA: document.getElementById("compare-timeline-a"),
  compareTimelineB: document.getElementById("compare-timeline-b"),
  compareTimelineAWrap: document.getElementById("compare-timeline-a-wrap"),
  compareTimelineBWrap: document.getElementById("compare-timeline-b-wrap"),
  compareTitleA: document.getElementById("compare-title-a"),
  compareTitleB: document.getElementById("compare-title-b"),
  compareDetails: document.getElementById("compare-details-panel"),
  compareZoom: document.getElementById("compare-zoom-slider"),
  compareFit: document.getElementById("compare-fit-button"),
  compareShowLabels: document.getElementById("compare-show-labels"),
  errorModal: document.getElementById("error-modal"),
  errorTitle: document.getElementById("error-title"),
  errorMessage: document.getElementById("error-message"),
  errorClose: document.getElementById("error-close"),
  errorOk: document.getElementById("error-ok"),
};

let gestureStartScale = null;
let compareGestureStartScale = null;
let syncingCompareScroll = false;

init();

async function init() {
  renderDurationFields(els.durationFields, OP_TYPES);
  if (els.compareDurationFields) {
    renderDurationFields(els.compareDurationFields, OP_TYPES);
  }
  bindEvents();
  await loadSchedulers();
  await simulate();
}

function bindEvents() {
  if (els.tabSimulate && els.tabCompare) {
    els.tabSimulate.addEventListener("click", () => setActiveTab("simulate"));
    els.tabCompare.addEventListener("click", () => setActiveTab("compare"));
  }

  els.form.addEventListener("submit", async (event) => {
    event.preventDefault();
    await simulate();
  });

  els.zoom.addEventListener("input", () => {
    state.scale = Number(els.zoom.value);
    renderTimeline();
  });

  els.showLabels.addEventListener("change", () => {
    state.showLabels = els.showLabels.checked;
    renderTimeline();
  });

  els.fit.addEventListener("click", () => fitTimeline());
  els.downloadChrome.addEventListener("click", () => downloadJson("pp-simulation-chrome-trace.json", state.chromeTrace));
  els.downloadPerfetto.addEventListener("click", () => downloadJson("pp-simulation-perfetto-trace.json", state.perfettoTrace));
  els.errorClose.addEventListener("click", clearError);
  els.errorOk.addEventListener("click", clearError);
  els.errorModal.addEventListener("click", (event) => {
    if (event.target === els.errorModal) clearError();
  });
  document.addEventListener("keydown", (event) => {
    if (event.key === "Escape" && !els.errorModal.hidden) clearError();
  });

  els.scheduler.addEventListener("change", updateSchedulerControls);
  els.timelineWrap.addEventListener("wheel", handleTimelineWheel, { passive: false });
  els.timelineWrap.addEventListener("gesturestart", handleGestureStart, { passive: false });
  els.timelineWrap.addEventListener("gesturechange", handleGestureChange, { passive: false });

  if (!els.compareForm) return;

  els.compareForm.addEventListener("submit", async (event) => {
    event.preventDefault();
    await compareSchedulers();
  });
  els.compareZoom.addEventListener("input", () => {
    state.compare.scale = Number(els.compareZoom.value);
    renderCompare();
  });
  els.compareShowLabels.addEventListener("change", () => {
    state.compare.showLabels = els.compareShowLabels.checked;
    renderCompare();
  });
  els.compareFit.addEventListener("click", () => fitCompareTimelines());
  for (const input of [
    els.compareIncludeEncoder,
    els.compareIncludeGenerator,
  ].filter(Boolean)) {
    input.addEventListener("change", updateCompareControls);
  }
  els.compareTimelineAWrap.addEventListener("wheel", handleCompareWheel, { passive: false });
  els.compareTimelineBWrap.addEventListener("wheel", handleCompareWheel, { passive: false });
  els.compareTimelineAWrap.addEventListener("scroll", () => syncCompareScroll(els.compareTimelineAWrap, els.compareTimelineBWrap));
  els.compareTimelineBWrap.addEventListener("scroll", () => syncCompareScroll(els.compareTimelineBWrap, els.compareTimelineAWrap));
  els.compareTimelineAWrap.addEventListener("gesturestart", handleCompareGestureStart, { passive: false });
  els.compareTimelineBWrap.addEventListener("gesturestart", handleCompareGestureStart, { passive: false });
  els.compareTimelineAWrap.addEventListener("gesturechange", handleCompareGestureChange, { passive: false });
  els.compareTimelineBWrap.addEventListener("gesturechange", handleCompareGestureChange, { passive: false });
}

async function loadSchedulers() {
  const response = await fetch("/api/schedulers");
  const payload = await response.json();
  state.schedulers = payload.schedulers || [];
  state.durationSpecs = payload.duration_specs || defaultDurationSpecs();
  state.limits.maxMicrobatches = payload.limits ? payload.limits.max_microbatches : null;
  applyInputLimits();
  populateSchedulerSelect(els.scheduler, state.schedulers);
  els.scheduler.value = "UnifiedBigMacVPP";
  renderDurationFields(els.durationFields, visibleOpTypes());
  updateSchedulerControls();
  if (els.compareForm) updateCompareControls();
}

function applyInputLimits() {
  const maxMicrobatches = Number(state.limits.maxMicrobatches);
  for (const input of [els.numMicrobatches, els.compareNumMicrobatches].filter(Boolean)) {
    if (Number.isFinite(maxMicrobatches) && maxMicrobatches > 0) {
      input.max = String(maxMicrobatches);
      input.title = `Maximum ${maxMicrobatches} microbatches in this deployment`;
      if (Number(input.value) > maxMicrobatches) {
        input.value = String(maxMicrobatches);
      }
    } else {
      input.removeAttribute("max");
      input.removeAttribute("title");
    }
  }
}

function populateSchedulerSelect(select, schedulers) {
  select.innerHTML = "";
  for (const item of schedulers) {
    const option = document.createElement("option");
    option.value = item.name;
    option.textContent = item.description || item.name;
    select.appendChild(option);
  }
}

function setActiveTab(tab) {
  if (!els.simulateView || !els.compareView) return;
  const isCompare = tab === "compare";
  els.tabSimulate.classList.toggle("active", !isCompare);
  els.tabCompare.classList.toggle("active", isCompare);
  els.simulateView.classList.toggle("active", !isCompare);
  els.compareView.classList.toggle("active", isCompare);
  if (isCompare && !state.compare.resultA) {
    renderCompareSummary();
  }
}

function updateSchedulerControls() {
  const option = state.schedulers.find((item) => item.name === els.scheduler.value);
  const usesVpp = option ? option.uses_vpp : true;
  const usesVeForwardLimit = option ? option.uses_ve_forward_limit : true;
  els.vppSize.disabled = !usesVpp;
  els.veForwardLimit.disabled = !usesVeForwardLimit;
  renderDurationFields(els.durationFields, visibleOpTypes());
}

function updateCompareControls() {
  if (!els.compareForm) return;
  if (els.compareIncludeGenerator.checked) {
    els.compareIncludeEncoder.checked = true;
  }
  els.compareVeForwardLimit.disabled = !els.compareIncludeEncoder.checked;

  const compatible = compatibleSchedulers();
  const previousA = els.compareSchedulerA.value;
  const previousB = els.compareSchedulerB.value;
  populateSchedulerSelect(els.compareSchedulerA, compatible);
  populateSchedulerSelect(els.compareSchedulerB, compatible);
  if (compatible.some((item) => item.name === previousA)) {
    els.compareSchedulerA.value = previousA;
  }
  if (compatible.some((item) => item.name === previousB)) {
    els.compareSchedulerB.value = previousB;
  } else if (compatible.length > 1) {
    els.compareSchedulerB.value = compatible[1].name;
  }
  renderDurationFields(els.compareDurationFields, compareVisibleOpTypes());
}

function compatibleSchedulers() {
  const shape = compareShape();
  return state.schedulers.filter((item) => {
    return Boolean(item.include_encoder) === shape.include_encoder
      && Boolean(item.include_generator) === shape.include_generator;
  });
}

function compareShape() {
  return {
    include_encoder: els.compareIncludeEncoder.checked,
    include_generator: els.compareIncludeGenerator.checked,
  };
}

function renderDurationFields(container, opTypes) {
  container.innerHTML = "";
  for (const opType of opTypes) {
    const spec = state.durationSpecs[opType] || { mean: 1, variance: 0 };
    const row = document.createElement("div");
    row.className = "duration-grid";
    row.innerHTML = `
      <span class="op-name">${opType}</span>
      <input data-op="${opType}" data-field="mean" type="number" min="0" step="0.01" value="${spec.mean}">
      <input data-op="${opType}" data-field="variance" type="number" min="0" step="0.01" value="${spec.variance}">
    `;
    container.appendChild(row);
  }
}

async function simulate() {
  setStatus("Simulating...");
  clearError();
  try {
    const response = await fetch("/api/simulate", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify(readConfig()),
    });
    const payload = await readApiPayload(response);
    if (!response.ok) {
      throw new Error(payload.error || "Simulation failed");
    }

    state.result = payload.simulation;
    state.chromeTrace = payload.chrome_trace;
    state.perfettoTrace = payload.perfetto_trace;
    state.selectedOpId = null;
    renderSummary();
    fitTimeline();
    renderDetails(null);
    setStatus("Ready");
  } catch (error) {
    showError("Simulation failed", error);
  }
}

async function compareSchedulers() {
  const compatible = compatibleSchedulers();
  if (compatible.length < 2 && els.compareSchedulerA.value === els.compareSchedulerB.value) {
    showError("Compare failed", new Error("Need two compatible schedulers"));
    return;
  }

  setStatus("Comparing...");
  clearError();
  try {
    const response = await fetch("/api/compare", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify(readCompareConfig()),
    });
    const payload = await readApiPayload(response);
    if (!response.ok) {
      throw new Error(payload.error || "Compare failed");
    }
    state.compare.resultA = payload.result_a;
    state.compare.resultB = payload.result_b;
    state.compare.comparison = payload.comparison;
    state.compare.selected = null;
    renderCompareSummary();
    fitCompareTimelines();
    renderCompareDetails(null, null);
    setStatus("Ready");
  } catch (error) {
    showError("Compare failed", error);
  }
}

async function readApiPayload(response) {
  const text = await response.text();
  if (!text) return {};
  try {
    return JSON.parse(text);
  } catch (error) {
    return {
      error: response.ok
        ? "Server returned an invalid JSON response"
        : `Server returned HTTP ${response.status}: ${text.slice(0, 240)}`,
    };
  }
}

function readConfig() {
  return {
    scheduler: els.scheduler.value,
    pp_size: Number(els.ppSize.value),
    vpp_size: Number(els.vppSize.value),
    num_microbatches: Number(els.numMicrobatches.value),
    ve_forward_limit: Number(els.veForwardLimit.value),
    seed: els.seed.value === "" ? null : Number(els.seed.value),
    duration_specs: readDurationSpecs(els.durationFields),
  };
}

function readCompareConfig() {
  return {
    scheduler_a: els.compareSchedulerA.value,
    scheduler_b: els.compareSchedulerB.value,
    seed: els.compareSeed.value === "" ? null : Number(els.compareSeed.value),
    workload: {
      pp_size: Number(els.comparePpSize.value),
      vpp_size: Number(els.compareVppSize.value),
      num_microbatches: Number(els.compareNumMicrobatches.value),
      ve_forward_limit: Number(els.compareVeForwardLimit.value),
      include_encoder: els.compareIncludeEncoder.checked,
      include_generator: els.compareIncludeGenerator.checked,
    },
    duration_specs: readDurationSpecs(els.compareDurationFields),
  };
}

function readDurationSpecs(container) {
  const specs = JSON.parse(JSON.stringify(state.durationSpecs));
  for (const opType of OP_TYPES) {
    if (!specs[opType]) specs[opType] = { mean: 0, variance: 0 };
  }
  for (const input of container.querySelectorAll("input")) {
    specs[input.dataset.op][input.dataset.field] = Number(input.value);
  }
  return specs;
}

function visibleOpTypes() {
  const option = state.schedulers.find((item) => item.name === els.scheduler.value);
  return option && option.op_types ? option.op_types : OP_TYPES;
}

function compareVisibleOpTypes() {
  const opTypes = ["F", "B"];
  if (els.compareIncludeEncoder.checked) opTypes.push("VF", "VB");
  if (els.compareIncludeGenerator.checked) opTypes.push("GF", "GB");
  return opTypes;
}

function renderSummary() {
  const result = state.result;
  if (!result) {
    els.summary.innerHTML = "";
    return;
  }
  els.summary.innerHTML = summaryMetrics(result).map(renderMetric).join("");
}

function renderCompareSummary() {
  const comparison = state.compare.comparison;
  if (!comparison) {
    els.compareSummary.innerHTML = "";
    return;
  }
  const delta = comparison.makespan_delta;
  const speedup = comparison.speedup_a_over_b;
  const metrics = [
    ["A Makespan", formatNumber(comparison.makespan_a)],
    ["B Makespan", formatNumber(comparison.makespan_b)],
    ["B - A", formatSigned(delta)],
    ["B/A", speedup == null ? "-" : `${speedup.toFixed(3)}x`],
  ];
  els.compareSummary.innerHTML = metrics.map(renderMetric).join("");
}

function summaryMetrics(result) {
  const summary = result.summary;
  const utilValues = Object.values(summary.rank_utilization || {});
  const avgUtil = utilValues.length ? utilValues.reduce((a, b) => a + b, 0) / utilValues.length : 0;
  return [
    ["Makespan", formatNumber(summary.makespan)],
    ["Total Wait", formatNumber(summary.total_wait_time)],
    ["Avg Utilization", `${(avgUtil * 100).toFixed(1)}%`],
    ["Ops", String(result.metadata.op_count)],
  ];
}

function renderMetric([label, value]) {
  return `
    <div class="metric">
      <div class="label">${escapeHtml(label)}</div>
      <div class="value">${escapeHtml(value)}</div>
    </div>
  `;
}

function fitTimeline() {
  if (!state.result) return;
  const available = Math.max(260, els.timelineWrap.clientWidth - 150);
  const makespan = Math.max(1, state.result.summary.makespan);
  const scale = clamp(available / makespan, 20, 180);
  setScale(scale);
}

function fitCompareTimelines() {
  if (!state.compare.resultA || !state.compare.resultB) return;
  const available = Math.max(260, Math.min(
    els.compareTimelineAWrap.clientWidth,
    els.compareTimelineBWrap.clientWidth,
  ) - 150);
  const makespan = Math.max(
    1,
    state.compare.resultA.summary.makespan,
    state.compare.resultB.summary.makespan,
  );
  setCompareScale(clamp(available / makespan, 20, 180));
}

function handleTimelineWheel(event) {
  if (!state.result || !event.ctrlKey) return;

  event.preventDefault();
  const factor = Math.exp(-event.deltaY * 0.01);
  zoomAroundClientX(factor, event.clientX);
}

function handleCompareWheel(event) {
  if (!state.compare.resultA || !event.ctrlKey) return;

  event.preventDefault();
  const factor = Math.exp(-event.deltaY * 0.01);
  zoomCompareAroundClientX(factor, event.clientX, event.currentTarget);
}

function handleGestureStart(event) {
  if (!state.result) return;
  event.preventDefault();
  gestureStartScale = state.scale;
}

function handleGestureChange(event) {
  if (!state.result || gestureStartScale == null) return;
  event.preventDefault();
  const targetScale = gestureStartScale * event.scale;
  zoomAroundClientX(targetScale / state.scale, event.clientX);
}

function handleCompareGestureStart(event) {
  if (!state.compare.resultA) return;
  event.preventDefault();
  compareGestureStartScale = state.compare.scale;
}

function handleCompareGestureChange(event) {
  if (!state.compare.resultA || compareGestureStartScale == null) return;
  event.preventDefault();
  const targetScale = compareGestureStartScale * event.scale;
  zoomCompareAroundClientX(targetScale / state.compare.scale, event.clientX, event.currentTarget);
}

function zoomAroundClientX(factor, clientX) {
  const oldScale = state.scale;
  const newScale = clamp(oldScale * factor, Number(els.zoom.min), Number(els.zoom.max));
  if (Math.abs(newScale - oldScale) < 0.001) return;

  const rect = els.timelineWrap.getBoundingClientRect();
  const localX = clientX - rect.left + els.timelineWrap.scrollLeft;
  const timeAtPointer = Math.max(0, (localX - 92) / oldScale);

  setScale(newScale);

  const newLocalX = 92 + timeAtPointer * newScale;
  els.timelineWrap.scrollLeft = Math.max(0, newLocalX - (clientX - rect.left));
}

function zoomCompareAroundClientX(factor, clientX, wrap) {
  const oldScale = state.compare.scale;
  const newScale = clamp(oldScale * factor, Number(els.compareZoom.min), Number(els.compareZoom.max));
  if (Math.abs(newScale - oldScale) < 0.001) return;

  const rect = wrap.getBoundingClientRect();
  const localX = clientX - rect.left + wrap.scrollLeft;
  const timeAtPointer = Math.max(0, (localX - 92) / oldScale);

  setCompareScale(newScale);

  const newLocalX = 92 + timeAtPointer * newScale;
  const scrollLeft = Math.max(0, newLocalX - (clientX - rect.left));
  els.compareTimelineAWrap.scrollLeft = scrollLeft;
  els.compareTimelineBWrap.scrollLeft = scrollLeft;
}

function setScale(scale) {
  state.scale = clamp(scale, Number(els.zoom.min), Number(els.zoom.max));
  els.zoom.value = String(Math.round(state.scale));
  renderTimeline();
}

function setCompareScale(scale) {
  state.compare.scale = clamp(scale, Number(els.compareZoom.min), Number(els.compareZoom.max));
  els.compareZoom.value = String(Math.round(state.compare.scale));
  renderCompare();
}

function syncCompareScroll(source, target) {
  if (syncingCompareScroll) return;
  syncingCompareScroll = true;
  target.scrollLeft = source.scrollLeft;
  requestAnimationFrame(() => {
    syncingCompareScroll = false;
  });
}

function renderTimeline() {
  renderPipelineTimeline({
    svg: els.timeline,
    result: state.result,
    scale: state.scale,
    showLabels: state.showLabels,
    maxTime: state.result ? state.result.summary.makespan : 1,
    selectedOpId: state.selectedOpId,
    side: "single",
    onSelect: (op) => {
      state.selectedOpId = op.id;
      renderTimeline();
      renderDetails(op);
    },
  });
}

function renderCompare() {
  const resultA = state.compare.resultA;
  const resultB = state.compare.resultB;
  const maxTime = Math.max(
    1,
    resultA ? resultA.summary.makespan : 1,
    resultB ? resultB.summary.makespan : 1,
  );
  els.compareTitleA.textContent = resultA ? resultA.metadata.scheduler : "Scheduler A";
  els.compareTitleB.textContent = resultB ? resultB.metadata.scheduler : "Scheduler B";

  renderPipelineTimeline({
    svg: els.compareTimelineA,
    result: resultA,
    scale: state.compare.scale,
    showLabels: state.compare.showLabels,
    maxTime,
    selectedOpId: state.compare.selected && state.compare.selected.side === "A"
      ? state.compare.selected.id
      : null,
    side: "A",
    onSelect: (op) => {
      state.compare.selected = { side: "A", id: op.id };
      renderCompare();
      renderCompareDetails(op, "A");
    },
  });
  renderPipelineTimeline({
    svg: els.compareTimelineB,
    result: resultB,
    scale: state.compare.scale,
    showLabels: state.compare.showLabels,
    maxTime,
    selectedOpId: state.compare.selected && state.compare.selected.side === "B"
      ? state.compare.selected.id
      : null,
    side: "B",
    onSelect: (op) => {
      state.compare.selected = { side: "B", id: op.id };
      renderCompare();
      renderCompareDetails(op, "B");
    },
  });
}

function renderPipelineTimeline({
  svg,
  result,
  scale,
  showLabels,
  maxTime,
  selectedOpId,
  side,
  onSelect,
}) {
  if (!result) {
    svg.innerHTML = "";
    return;
  }

  const opsByRank = groupByRank(result.ops);
  const ranks = Object.keys(opsByRank).map(Number).sort((a, b) => a - b);
  const left = 92;
  const top = 34;
  const laneHeight = 38;
  const laneGap = 8;
  const rightPad = 28;
  const bottomPad = 24;
  const width = Math.ceil(left + Math.max(1, maxTime) * scale + rightPad);
  const height = Math.ceil(top + ranks.length * (laneHeight + laneGap) + bottomPad);

  svg.setAttribute("width", width);
  svg.setAttribute("height", height);
  svg.setAttribute("viewBox", `0 0 ${width} ${height}`);
  svg.innerHTML = "";

  renderAxis(svg, left, top, width, height, maxTime, scale);

  ranks.forEach((rank, rankIndex) => {
    const y = top + rankIndex * (laneHeight + laneGap);
    addRect(svg, "rank-band", 0, y, width, laneHeight);
    addText(svg, "rank-label", 14, y + 24, rankLabel(result, rank));
    renderRankOps(svg, opsByRank[rank], left, y, laneHeight, {
      scale,
      showLabels,
      selectedOpId,
      side,
      onSelect,
    });
  });
}

function rankLabel(result, rank) {
  const layout = result && result.metadata ? result.metadata.pipeline_layout : null;
  if (!layout) return `Rank ${rank}`;
  if (rank === layout.encoder_stage) return "Encoder";
  if (rank === layout.generator_stage) return "Generator";
  const llmStart = layout.llm_stage_offset;
  const llmEnd = llmStart + layout.llm_pp_size - 1;
  if (rank >= llmStart && rank <= llmEnd) return `LLM ${rank - llmStart}`;
  return `Rank ${rank}`;
}

function renderAxis(svg, left, top, width, height, maxTime, scale) {
  const tickCount = Math.min(12, Math.max(4, Math.ceil(maxTime / 5)));
  for (let i = 0; i <= tickCount; i++) {
    const time = (maxTime * i) / tickCount;
    const x = left + time * scale;
    addLine(svg, "grid-line", x, top - 18, x, height - 12);
    addText(svg, "axis-text", x + 3, 18, formatNumber(time));
  }
  addLine(svg, "axis-line", left, top - 8, width - 18, top - 8);
}

function renderRankOps(svg, ops, left, y, laneHeight, options) {
  let previousEnd = 0;
  for (const op of ops) {
    const wait = Math.max(0, op.start_time - previousEnd);
    if (wait > 1e-9) {
      addRect(svg, "wait-rect", left + previousEnd * options.scale, y + 5, wait * options.scale, laneHeight - 10);
    }

    const x = left + op.start_time * options.scale;
    const width = Math.max(1.5, op.duration * options.scale);
    const rect = addRect(svg, "op-rect", x, y + 5, width, laneHeight - 10);
    rect.setAttribute("fill", opColor(op));
    rect.dataset.opId = op.id;
    rect.dataset.side = options.side;
    rect.dataset.matchKey = matchKey(op);
    if (op.id === options.selectedOpId) rect.classList.add("selected");

    rect.addEventListener("mouseenter", (event) => {
      showTooltip(event, op);
      highlightMatch(matchKey(op), true);
    });
    rect.addEventListener("mousemove", (event) => moveTooltip(event));
    rect.addEventListener("mouseleave", () => {
      hideTooltip();
      highlightMatch(matchKey(op), false);
    });
    rect.addEventListener("click", () => options.onSelect(op));

    if (options.showLabels && width >= 28) {
      const label = width >= 54 ? op.label : compactLabel(op);
      const text = addText(svg, "op-label", x + width / 2, y + 25, label);
      text.setAttribute("text-anchor", "middle");
    }

    previousEnd = op.end_time;
  }
}

function renderDetails(op) {
  if (!op) {
    els.details.innerHTML = `<div class="details-empty">No op selected</div>`;
    return;
  }

  const depRows = Object.entries(op.dep_reasons || {}).map(([dep, reasons]) => {
    return renderDependencyDetail(dep, reasons, state.result);
  }).join("");

  els.details.innerHTML = renderOpDetails(op, depRows);
}

function renderCompareDetails(op, side) {
  if (!op) {
    els.compareDetails.innerHTML = `<div class="details-empty">No op selected</div>`;
    return;
  }
  const result = side === "A" ? state.compare.resultA : state.compare.resultB;
  const peer = findMatchedCompareOp(op, side);
  const depRows = Object.entries(op.dep_reasons || {}).map(([dep, reasons]) => {
    return renderDependencyDetail(dep, reasons, result);
  }).join("");
  const peerRow = peer
    ? `<div class="dep-row"><strong>Matched op</strong><div>${escapeHtml(peer.label)} · ${escapeHtml(peer.metadata)} · duration ${formatNumber(peer.duration)}</div></div>`
    : "";
  els.compareDetails.innerHTML = renderOpDetails(op, peerRow + depRows, `Scheduler ${side}`);
}

function renderOpDetails(op, depRows, prefix = null) {
  return `
    <div class="details-grid">
      ${detail("Source", prefix)}
      ${detail("Label", op.label)}
      ${detail("Type", op.op_type)}
      ${detail("Rank", op.rank)}
      ${detail("Microbatch", op.microbatch_id)}
      ${detail("Chunk", op.chunk_id)}
      ${detail("Wait", formatNumber(op.wait_time))}
      ${detail("Start", formatNumber(op.start_time))}
      ${detail("End", formatNumber(op.end_time))}
      ${detail("Duration", formatNumber(op.duration))}
      ${detail("Deps", op.deps.length)}
    </div>
    <div class="deps">${depRows || "No dependencies"}</div>
  `;
}

function findMatchedCompareOp(op, side) {
  const result = side === "A" ? state.compare.resultB : state.compare.resultA;
  if (!result) return null;
  const key = matchKey(op);
  const peer = result.ops.find((candidate) => matchKey(candidate) === key);
  if (!peer) return null;
  return {
    ...peer,
    metadata: side === "A" ? "Scheduler B" : "Scheduler A",
  };
}

function renderDependencyDetail(depId, reasons, result) {
  const depOp = opById(depId, result);
  const title = depOp
    ? `${depOp.label} on Rank ${depOp.rank}`
    : depId;
  const timing = depOp
    ? `ends at ${formatNumber(depOp.end_time)}`
    : "source op not found";
  const reasonText = reasons.map(reasonExplanation).join("; ");

  return `
    <div class="dep-row">
      <div><strong>${escapeHtml(title)}</strong> <span>${escapeHtml(timing)}</span></div>
      <div>${escapeHtml(reasonText)}</div>
      <code>${escapeHtml(depId)}</code>
    </div>
  `;
}

function opById(opId, result = state.result) {
  if (!result) return null;
  return result.ops.find((op) => op.id === opId) || null;
}

function reasonExplanation(reason) {
  const explanations = {
    rank_order: "same-rank program order: this rank must finish the previous scheduled op first",
    encoder_forward_ready: "LLM forward waits for the matching encoder forward output",
    ve_forward_ready: "LLM forward waits for the matching VE forward output",
    ve_forward_unit_ready: "LLM forward waits until the VE forward unit for this batch is ready",
    llm_forward_from_previous_rank: "pipeline forward dependency: this rank consumes activation from the previous PP rank",
    vpp_forward_from_previous_chunk: "VPP forward dependency: this chunk starts after the previous chunk reaches the last PP rank",
    backward_consumes_forward_activation: "backward must wait for its matching forward activation on this rank",
    llm_backward_from_next_rank: "pipeline backward dependency: gradient arrives from the next PP rank",
    vpp_backward_from_next_chunk: "VPP backward dependency: previous chunk waits for backward from the next chunk",
    llm_output_unit_ready: "generator forward waits for the LLM output unit",
    generator_backward_consumes_forward: "generator backward consumes the matching generator forward activation on this rank",
    generator_backward_ready: "LLM backward waits for the matching generator backward gradient",
    generator_backward_unit_ready: "LLM backward waits for generator backward gradients",
    llm_output_ready: "generator forward waits for the matching LLM output",
    llm_input_gradient_unit_ready: "VE backward waits for LLM input gradients for this unit",
    llm_input_gradient_ready: "encoder backward waits for the matching LLM input gradient",
  };
  return explanations[reason] || reason;
}

function detail(key, value) {
  if (value == null) return "";
  return `
    <div class="detail-item">
      <div class="key">${escapeHtml(key)}</div>
      <div class="val">${escapeHtml(value == null ? "-" : String(value))}</div>
    </div>
  `;
}

function showTooltip(event, op) {
  els.tooltip.innerHTML = `
    <strong>${escapeHtml(op.label)}</strong>
    <div>rank ${op.rank} · ${escapeHtml(op.op_type)} · mb ${op.microbatch_id ?? "-"}</div>
    <div>start ${formatNumber(op.start_time)} · end ${formatNumber(op.end_time)}</div>
    <div>duration ${formatNumber(op.duration)} · wait ${formatNumber(op.wait_time)}</div>
    <div>deps ${op.deps.length}</div>
  `;
  els.tooltip.style.display = "block";
  moveTooltip(event);
}

function moveTooltip(event) {
  els.tooltip.style.left = `${event.clientX + 14}px`;
  els.tooltip.style.top = `${event.clientY + 14}px`;
}

function hideTooltip() {
  els.tooltip.style.display = "none";
}

function highlightMatch(key, enabled) {
  for (const rect of document.querySelectorAll(".op-rect")) {
    if (rect.dataset.matchKey === key) {
      rect.classList.toggle("matched", enabled);
    }
  }
}

function groupByRank(ops) {
  const grouped = {};
  for (const op of ops) {
    if (!grouped[op.rank]) grouped[op.rank] = [];
    grouped[op.rank].push(op);
  }
  for (const rank of Object.keys(grouped)) {
    grouped[rank].sort((a, b) => a.index - b.index);
  }
  return grouped;
}

function opColor(op) {
  const palette = colors[op.op_type] || ["#cbd5e1", "#94a3b8", "#64748b"];
  const chunk = op.chunk_id == null ? 0 : Number(op.chunk_id);
  return palette[chunk % palette.length];
}

function compactLabel(op) {
  if (op.microbatch_id == null) return op.op_type;
  return `${op.op_type}${op.microbatch_id}`;
}

function matchKey(op) {
  const allowCrossVpp = state.compare.comparison
    && state.compare.comparison.workload
    && state.compare.comparison.workload.allow_cross_vpp;
  const chunk = allowCrossVpp && (op.op_type === "F" || op.op_type === "B")
    ? 0
    : (op.chunk_id == null ? 0 : Number(op.chunk_id));
  return `${op.rank}|${op.op_type}|${op.microbatch_id ?? "none"}|${chunk}`;
}

function addRect(svg, className, x, y, width, height) {
  const rect = document.createElementNS("http://www.w3.org/2000/svg", "rect");
  rect.setAttribute("class", className);
  rect.setAttribute("x", x);
  rect.setAttribute("y", y);
  rect.setAttribute("width", Math.max(0, width));
  rect.setAttribute("height", Math.max(0, height));
  svg.appendChild(rect);
  return rect;
}

function addLine(svg, className, x1, y1, x2, y2) {
  const line = document.createElementNS("http://www.w3.org/2000/svg", "line");
  line.setAttribute("class", className);
  line.setAttribute("x1", x1);
  line.setAttribute("y1", y1);
  line.setAttribute("x2", x2);
  line.setAttribute("y2", y2);
  svg.appendChild(line);
  return line;
}

function addText(svg, className, x, y, value) {
  const text = document.createElementNS("http://www.w3.org/2000/svg", "text");
  text.setAttribute("class", className);
  text.setAttribute("x", x);
  text.setAttribute("y", y);
  text.textContent = value;
  svg.appendChild(text);
  return text;
}

function downloadJson(filename, payload) {
  if (!payload) return;
  const blob = new Blob([JSON.stringify(payload, null, 2)], { type: "application/json" });
  const url = URL.createObjectURL(blob);
  const anchor = document.createElement("a");
  anchor.href = url;
  anchor.download = filename;
  anchor.click();
  URL.revokeObjectURL(url);
}

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

function showError(title, error) {
  const message = error && error.message ? error.message : String(error || "Unknown error");
  setStatus("Error");
  els.errorTitle.textContent = title;
  els.errorMessage.textContent = message;
  els.errorModal.hidden = false;
  els.errorOk.focus();
}

function clearError() {
  els.errorModal.hidden = true;
  els.errorMessage.textContent = "";
}

function formatNumber(value) {
  return Number(value).toFixed(2);
}

function formatSigned(value) {
  const number = Number(value);
  return `${number >= 0 ? "+" : ""}${formatNumber(number)}`;
}

function clamp(value, min, max) {
  return Math.max(min, Math.min(max, value));
}

function escapeHtml(value) {
  return String(value)
    .replaceAll("&", "&amp;")
    .replaceAll("<", "&lt;")
    .replaceAll(">", "&gt;")
    .replaceAll('"', "&quot;")
    .replaceAll("'", "&#039;");
}

function defaultDurationSpecs() {
  return {
    F: { mean: 1.0, variance: 0.0 },
    B: { mean: 2.0, variance: 0.0 },
    VF: { mean: 0.8, variance: 0.01 },
    VB: { mean: 0.9, variance: 0.01 },
    GF: { mean: 0.6, variance: 0.01 },
    GB: { mean: 0.7, variance: 0.01 },
  };
}