const operations = [
  {
    id: "add",
    label: "Add two matrices (element-wise)",
    description: "A + B. Shapes of A and B should be equal or broadcast-compatible.",
    inputs: 2,
    code: "output = A + B",
    match: ["add", "sum", "plus", "element-wise add"],
    run: (A, B) => elementWise(A, B, (a, b) => a + b, "add"),
  },
  {
    id: "subtract",
    label: "Subtract two matrices (element-wise)",
    description: "A - B. Shapes of A and B should be equal or broadcast-compatible.",
    inputs: 2,
    code: "output = A - B",
    match: ["subtract", "minus", "difference"],
    run: (A, B) => elementWise(A, B, (a, b) => a - b, "subtract"),
  },
  {
    id: "multiply",
    label: "Multiply two matrices (element-wise)",
    description: "A * B. Shapes of A and B should be equal or broadcast-compatible.",
    inputs: 2,
    code: "output = A * B",
    match: ["multiply", "times", "element-wise multiply"],
    run: (A, B) => elementWise(A, B, (a, b) => a * b, "multiply"),
  },
  {
    id: "matmul",
    label: "Matrix multiplication (np.matmul)",
    description: "np.matmul(A, B). Works for 2D and stacked arrays (up to 3D here).",
    inputs: 2,
    code: "output = np.matmul(A, B)",
    match: ["matmul", "matrix multiplication", "dot product", "@"],
    run: (A, B) => matmul(A, B),
  },
  {
    id: "transpose",
    label: "Transpose matrix (np.transpose)",
    description: "Swap axes. For 2D, rows and columns are flipped.",
    inputs: 1,
    code: "output = np.transpose(A)",
    match: ["transpose", "swap axes", "flip rows columns"],
    run: (A) => transpose(A),
  },
  {
    id: "reshape",
    label: "Reshape matrix (np.reshape)",
    description: "Reshape A into the target shape with same total number of elements.",
    inputs: 1,
    code: "output = np.reshape(A, new_shape)",
    match: ["reshape", "change shape"],
    requiresTarget: true,
    run: (A, _unused, targetShape) => reshape(A, targetShape),
  },
  {
    id: "concat",
    label: "Concatenate two arrays (np.concatenate)",
    description: "Join A and B along axis 0. Non-concat dimensions must match.",
    inputs: 2,
    code: "output = np.concatenate([A, B], axis=0)",
    match: ["concatenate", "concat", "join arrays"],
    run: (A, B) => concatAxis0(A, B),
  },
  {
    id: "stack",
    label: "Stack two arrays (np.stack)",
    description: "Stack A and B along a new axis 0. A and B shapes must be identical.",
    inputs: 2,
    code: "output = np.stack([A, B], axis=0)",
    match: ["stack", "new axis"],
    run: (A, B) => stackAxis0(A, B),
  },
  {
    id: "sum",
    label: "Sum matrix values (np.sum)",
    description: "Sum values on an optional axis with optional keepdims.",
    inputs: 1,
    code: "output = np.sum(A)",
    match: ["sum all", "total", "np.sum", "sum"],
    supportsReductionOptions: true,
    run: (A, _unused, _unused2, options) =>
      reduceArray(A, {
        mode: "sum",
        axis: options.axis,
        keepdims: options.keepdims,
      }),
  },
  {
    id: "mean",
    label: "Mean matrix values (np.mean)",
    description: "Compute mean on an optional axis with optional keepdims.",
    inputs: 1,
    code: "output = np.mean(A)",
    match: ["mean", "average", "np.mean"],
    supportsReductionOptions: true,
    run: (A, _unused, _unused2, options) =>
      reduceArray(A, {
        mode: "mean",
        axis: options.axis,
        keepdims: options.keepdims,
      }),
  },
  {
    id: "ones",
    label: "Create ones matrix (np.ones)",
    description: "Generate an array filled with ones for the given shape.",
    inputs: 1,
    code: "output = np.ones(shape)",
    match: ["ones", "all ones", "np.ones"],
    generatorOnly: true,
    run: (A) => clone(A),
  },
  {
    id: "zeros",
    label: "Create zeros matrix (np.zeros)",
    description: "Generate an array filled with zeros for the given shape.",
    inputs: 1,
    code: "output = np.zeros(shape)",
    match: ["zeros", "all zeros", "np.zeros"],
    generatorOnly: true,
    run: (A) => clone(A),
  },
];

const operationSelect = document.getElementById("operationSelect");
const nlInput = document.getElementById("nlInput");
const matchBtn = document.getElementById("matchBtn");
const runBtn = document.getElementById("runBtn");
const operationInfo = document.getElementById("operationInfo");
const shapeInputs = document.getElementById("shapeInputs");
const operationOptions = document.getElementById("operationOptions");
const codeOutput = document.getElementById("codeOutput");
const inputViz = document.getElementById("inputViz");
const outputViz = document.getElementById("outputViz");
const detailViz = document.getElementById("detailViz");
let lastRunMeta = null;

function init() {
  operations.forEach((op) => {
    const option = document.createElement("option");
    option.value = op.id;
    option.textContent = op.label;
    operationSelect.appendChild(option);
  });

  operationSelect.value = "add";
  renderShapeInputs();
  bindEvents();
}

function bindEvents() {
  operationSelect.addEventListener("change", renderShapeInputs);

  matchBtn.addEventListener("click", () => {
    const query = nlInput.value.trim().toLowerCase();
    if (!query) return;

    let best = operations[0];
    let bestScore = 0;

    for (const op of operations) {
      let score = 0;
      for (const key of op.match) {
        if (query.includes(key)) score += key.length;
      }
      if (score > bestScore) {
        best = op;
        bestScore = score;
      }
    }

    operationSelect.value = best.id;
    renderShapeInputs();
  });

  runBtn.addEventListener("click", () => {
    try {
      lastRunMeta = null;
      const op = currentOperation();
      const parsed = parseAllShapes(op);
      const inputArrays = buildInputArrays(op, parsed);
      const options = parsed.options || {};
      const output = op.run(
        inputArrays[0],
        inputArrays[1],
        parsed.targetShape || null,
        options
      );

      renderCode(op, parsed);
      renderInputs(inputArrays, op);
      renderOutput(output);
      renderComputationDetails(op, inputArrays, output, parsed, lastRunMeta);
    } catch (err) {
      codeOutput.textContent = "Error: " + err.message;
      inputViz.innerHTML = "";
      outputViz.innerHTML = `<p class=\"error\">${escapeHtml(err.message)}</p>`;
      detailViz.innerHTML = "";
    }
  });
}

function currentOperation() {
  return operations.find((op) => op.id === operationSelect.value);
}

function renderShapeInputs() {
  const op = currentOperation();

  operationInfo.innerHTML = `<strong>${op.label}</strong><br/>${escapeHtml(op.description)}`;

  const cards = [];
  for (let i = 0; i < op.inputs; i += 1) {
    cards.push(shapeCard(`inputShape${i + 1}`, `Input ${i + 1} shape`, "2,3"));
  }

  if (op.requiresTarget) {
    cards.push(shapeCard("targetShape", "Target shape", "3,2"));
  }

  shapeInputs.innerHTML = cards.join("");
  operationOptions.innerHTML = op.supportsReductionOptions
    ? reductionOptionsCard()
    : "";

  const s1 = document.getElementById("inputShape1");
  const s2 = document.getElementById("inputShape2");

  if (op.id === "matmul") {
    s1.value = "2,3";
    if (s2) s2.value = "3,2";
  } else if (op.id === "concat" || op.id === "stack") {
    s1.value = "2,2";
    if (s2) s2.value = "2,2";
  } else if (op.id === "transpose") {
    s1.value = "2,3";
  } else if (op.id === "reshape") {
    s1.value = "2,3";
    document.getElementById("targetShape").value = "3,2";
  } else if (op.id === "ones" || op.id === "zeros") {
    s1.value = "3,3";
  } else if (op.id === "sum" || op.id === "mean") {
    s1.value = "2,3";
  }

  if (op.supportsReductionOptions) {
    const axisInput = document.getElementById("reduceAxis");
    const keepdimsInput = document.getElementById("reduceKeepdims");

    const setAxisHint = () => {
      try {
        const dims = parseShape(document.getElementById("inputShape1").value);
        axisInput.placeholder = `axis (optional): 0 to ${dims.length - 1}`;
      } catch (_err) {
        axisInput.placeholder = "axis (optional): 0";
      }
    };

    document.getElementById("inputShape1").addEventListener("input", setAxisHint);
    setAxisHint();
    keepdimsInput.checked = false;
  }

  detailViz.innerHTML = "";
}

function shapeCard(inputId, label, placeholder) {
  return `
    <div class="shape-card">
      <h3>${label}</h3>
      <input id="${inputId}" type="text" placeholder="${placeholder}" />
      <small>Use comma-separated dimensions, max 3D, e.g. 2,3 or 2,2,3</small>
    </div>
  `;
}

function reductionOptionsCard() {
  return `
    <div class="shape-card options-card">
      <h3>Reduction options</h3>
      <div class="inline-fields">
        <label class="mini-label" for="reduceAxis">axis</label>
        <input id="reduceAxis" type="text" placeholder="axis (optional): 0" />
      </div>
      <div class="inline-fields">
        <label class="mini-label" for="reduceKeepdims">keepdims</label>
        <input id="reduceKeepdims" type="checkbox" />
      </div>
      <small>Leave axis empty to reduce all dimensions. keepdims keeps reduced axes as size 1.</small>
    </div>
  `;
}

function parseAllShapes(op) {
  const shapes = [];

  for (let i = 0; i < op.inputs; i += 1) {
    const input = document.getElementById(`inputShape${i + 1}`);
    shapes.push(parseShape(input.value));
  }

  const parsed = { shapes };

  if (op.requiresTarget) {
    parsed.targetShape = parseShape(document.getElementById("targetShape").value);
  }

  if (op.supportsReductionOptions) {
    const axisRaw = document.getElementById("reduceAxis").value.trim();
    const keepdims = document.getElementById("reduceKeepdims").checked;
    let axis = null;

    if (axisRaw !== "") {
      axis = Number(axisRaw);
      if (!Number.isInteger(axis)) {
        throw new Error("axis must be an integer or left empty.");
      }
      if (axis < 0 || axis >= shapes[0].length) {
        throw new Error(`axis out of range for input rank ${shapes[0].length}.`);
      }
    }

    parsed.options = { axis, keepdims };
  }

  return parsed;
}

function parseShape(raw) {
  if (!raw || !raw.trim()) {
    throw new Error("Shape cannot be empty.");
  }

  let normalized = raw.trim();
  if (normalized.startsWith("(") && normalized.endsWith(")")) {
    normalized = normalized.slice(1, -1);
  }
  normalized = normalized.trim();

  const dims = normalized
    .split(",")
    .map((x) => x.trim())
    .filter((x) => x.length > 0)
    .map((x) => Number(x));

  if (dims.length < 1 || dims.length > 3) {
    throw new Error("Each shape must have 1 to 3 dimensions.");
  }

  dims.forEach((d) => {
    if (!Number.isFinite(d) || !Number.isInteger(d) || d < 1 || d > 6) {
      throw new Error("Dimensions must be integers between 1 and 6.");
    }
  });

  return dims;
}

function buildInputArrays(op, parsed) {
  if (op.id === "ones") {
    return [fillArray(parsed.shapes[0], 1)];
  }
  if (op.id === "zeros") {
    return [fillArray(parsed.shapes[0], 0)];
  }

  const [shapeA, shapeB] = parsed.shapes;
  const A = randomArray(shapeA);
  if (op.inputs === 1) return [A];

  let B;
  if (["add", "subtract", "multiply"].includes(op.id)) {
    B = randomArray(shapeB || shapeA);
  } else if (op.id === "matmul") {
    B = randomArray(shapeB);
  } else if (op.id === "concat" || op.id === "stack") {
    B = randomArray(shapeB);
  } else {
    B = randomArray(shapeB || shapeA);
  }

  return [A, B];
}

function randomArray(shape) {
  return createByShape(shape, () => Math.floor(Math.random() * 9) + 1);
}

function fillArray(shape, val) {
  return createByShape(shape, () => val);
}

function createByShape(shape, valueFn, level = 0) {
  const len = shape[level];
  const arr = new Array(len);

  for (let i = 0; i < len; i += 1) {
    arr[i] =
      level === shape.length - 1
        ? valueFn()
        : createByShape(shape, valueFn, level + 1);
  }

  return arr;
}

function shapeOf(arr) {
  if (!Array.isArray(arr)) return [];
  return [arr.length, ...shapeOf(arr[0])];
}

function formatShape(shape) {
  return `(${shape.join(", ")})`;
}

function broadcastShapes(shapeA, shapeB) {
  const maxRank = Math.max(shapeA.length, shapeB.length);
  const out = new Array(maxRank);

  for (let i = 0; i < maxRank; i += 1) {
    const a = shapeA[shapeA.length - 1 - i] ?? 1;
    const b = shapeB[shapeB.length - 1 - i] ?? 1;

    if (a !== b && a !== 1 && b !== 1) {
      throw new Error(
        `Broadcast mismatch at dimension ${maxRank - i - 1}: ${a} vs ${b}.`
      );
    }
    out[maxRank - 1 - i] = Math.max(a, b);
  }

  return out;
}

function getAtIndices(arr, indices) {
  let cur = arr;
  for (let i = 0; i < indices.length; i += 1) {
    cur = cur[indices[i]];
  }
  return cur;
}

function createByShapeIndexed(shape, valueFn, idx = []) {
  if (shape.length === 0) return valueFn(idx);
  const dim = shape[idx.length];
  const out = new Array(dim);
  for (let i = 0; i < dim; i += 1) {
    const nextIdx = idx.concat(i);
    if (nextIdx.length === shape.length) {
      out[i] = valueFn(nextIdx);
    } else {
      out[i] = createByShapeIndexed(shape, valueFn, nextIdx);
    }
  }
  return out;
}

function projectBroadcastIndices(outputIndices, sourceShape) {
  const offset = outputIndices.length - sourceShape.length;
  const mapped = [];
  for (let i = 0; i < sourceShape.length; i += 1) {
    const srcDim = sourceShape[i];
    const outIndex = outputIndices[offset + i];
    mapped.push(srcDim === 1 ? 0 : outIndex);
  }
  return mapped;
}

function elementWise(A, B, fn, opName) {
  const shapeA = shapeOf(A);
  const shapeB = shapeOf(B);
  const outputShape = broadcastShapes(shapeA, shapeB);

  const output = createByShapeIndexed(outputShape, (outIdx) => {
    const idxA = projectBroadcastIndices(outIdx, shapeA);
    const idxB = projectBroadcastIndices(outIdx, shapeB);
    return fn(getAtIndices(A, idxA), getAtIndices(B, idxB));
  });

  const usedBroadcasting =
    shapeA.length !== shapeB.length ||
    shapeA.some((dim, i) => dim !== shapeB[i]) ||
    shapeA.join(",") !== outputShape.join(",") ||
    shapeB.join(",") !== outputShape.join(",");

  lastRunMeta = {
    kind: "elementwise",
    opName,
    shapeA,
    shapeB,
    outputShape,
    usedBroadcasting,
  };

  return output;
}

function elementWiseDeep(A, B, fn) {
  if (!Array.isArray(A) && !Array.isArray(B)) return fn(A, B);
  return A.map((v, i) => elementWiseDeep(v, B[i], fn));
}

function transpose(A) {
  const shape = shapeOf(A);
  if (shape.length === 1) return clone(A);

  if (shape.length === 2) {
    const [rows, cols] = shape;
    const out = [];
    for (let c = 0; c < cols; c += 1) {
      const row = [];
      for (let r = 0; r < rows; r += 1) {
        row.push(A[r][c]);
      }
      out.push(row);
    }
    return out;
  }

  if (shape.length === 3) {
    const [d0, d1, d2] = shape;
    const out = [];
    for (let i = 0; i < d2; i += 1) {
      const level2 = [];
      for (let j = 0; j < d1; j += 1) {
        const row = [];
        for (let k = 0; k < d0; k += 1) {
          row.push(A[k][j][i]);
        }
        level2.push(row);
      }
      out.push(level2);
    }
    return out;
  }

  throw new Error("Transpose supports up to 3D in this interface.");
}

function reshape(A, targetShape) {
  const flat = flatten(A);
  const totalA = flat.length;
  const totalTarget = targetShape.reduce((x, y) => x * y, 1);

  if (totalA !== totalTarget) {
    throw new Error(
      `reshape needs same number of elements. Got ${totalA} and ${totalTarget}.`
    );
  }

  return unflatten(flat, targetShape);
}

function flatten(arr) {
  if (!Array.isArray(arr)) return [arr];
  return arr.flatMap((x) => flatten(x));
}

function unflatten(flat, shape) {
  let idx = 0;
  function build(level = 0) {
    const len = shape[level];
    const out = [];
    for (let i = 0; i < len; i += 1) {
      if (level === shape.length - 1) {
        out.push(flat[idx]);
        idx += 1;
      } else {
        out.push(build(level + 1));
      }
    }
    return out;
  }
  return build();
}

function matmul(A, B) {
  const sA = shapeOf(A);
  const sB = shapeOf(B);
  const leftVec = sA.length === 1;
  const rightVec = sB.length === 1;

  const leftBatch = sA.length === 3 ? sA[0] : 1;
  const rightBatch = sB.length === 3 ? sB[0] : 1;
  const outBatch = Math.max(leftBatch, rightBatch);

  if (leftBatch !== rightBatch && leftBatch !== 1 && rightBatch !== 1) {
    throw new Error(
      `matmul batch broadcast mismatch: ${leftBatch} vs ${rightBatch}.`
    );
  }

  const leftRows = leftVec ? 1 : sA[sA.length - 2];
  const leftInner = sA[sA.length - 1];
  const rightInner = rightVec ? sB[0] : sB[sB.length - 2];
  const rightCols = rightVec ? 1 : sB[sB.length - 1];

  if (leftInner !== rightInner) {
    throw new Error(
      `matmul shape mismatch on core dims: ${formatShape(sA)} @ ${formatShape(
        sB
      )} (inner ${leftInner} vs ${rightInner}).`
    );
  }

  const getLeft = (batch, row, k) => {
    if (leftVec) return A[k];
    if (sA.length === 2) return A[row][k];
    const batchIdx = leftBatch === 1 ? 0 : batch;
    return A[batchIdx][row][k];
  };

  const getRight = (batch, k, col) => {
    if (rightVec) return B[k];
    if (sB.length === 2) return B[k][col];
    const batchIdx = rightBatch === 1 ? 0 : batch;
    return B[batchIdx][k][col];
  };

  const matrixForBatch = (batch) => {
    const out = [];
    for (let r = 0; r < leftRows; r += 1) {
      const row = [];
      for (let c = 0; c < rightCols; c += 1) {
        let sum = 0;
        for (let k = 0; k < leftInner; k += 1) {
          sum += getLeft(batch, r, k) * getRight(batch, k, c);
        }
        row.push(sum);
      }
      out.push(row);
    }
    return out;
  };

  const hasBatchAxis = sA.length > 2 || sB.length > 2;
  const matrices = hasBatchAxis
    ? Array.from({ length: outBatch }, (_, batch) => matrixForBatch(batch))
    : [matrixForBatch(0)];

  let output;
  if (leftVec && rightVec) {
    output = hasBatchAxis ? matrices.map((m) => m[0][0]) : matrices[0][0][0];
  } else if (leftVec) {
    output = hasBatchAxis ? matrices.map((m) => m[0].slice()) : matrices[0][0].slice();
  } else if (rightVec) {
    output = hasBatchAxis
      ? matrices.map((m) => m.map((row) => row[0]))
      : matrices[0].map((row) => row[0]);
  } else {
    output = hasBatchAxis ? matrices : matrices[0];
  }

  const outputShape = shapeOf(output);
  lastRunMeta = {
    kind: "matmul",
    shapeA: sA,
    shapeB: sB,
    outputShape,
    leftBatch,
    rightBatch,
    outBatch,
    usedBroadcasting:
      leftBatch !== rightBatch || sA.length !== sB.length || leftVec || rightVec,
    leftVectorPromoted: leftVec,
    rightVectorPromoted: rightVec,
  };

  return output;
}

function matmul2D(A, B) {
  const rowsA = A.length;
  const colsA = A[0].length;
  const rowsB = B.length;
  const colsB = B[0].length;

  if (colsA !== rowsB) {
    throw new Error(
      `matmul shape mismatch: (${rowsA},${colsA}) x (${rowsB},${colsB})`
    );
  }

  const out = [];
  for (let r = 0; r < rowsA; r += 1) {
    const row = [];
    for (let c = 0; c < colsB; c += 1) {
      let sum = 0;
      for (let k = 0; k < colsA; k += 1) {
        sum += A[r][k] * B[k][c];
      }
      row.push(sum);
    }
    out.push(row);
  }
  return out;
}

function concatAxis0(A, B) {
  const sA = shapeOf(A);
  const sB = shapeOf(B);
  if (sA.length !== sB.length) {
    throw new Error("concatenate requires same rank.");
  }
  for (let i = 1; i < sA.length; i += 1) {
    if (sA[i] !== sB[i]) {
      throw new Error(
        "concatenate axis=0 requires other dimensions to be identical."
      );
    }
  }
  return [...clone(A), ...clone(B)];
}

function stackAxis0(A, B) {
  const sA = JSON.stringify(shapeOf(A));
  const sB = JSON.stringify(shapeOf(B));
  if (sA !== sB) {
    throw new Error("stack requires A and B to have the same shape.");
  }
  return [clone(A), clone(B)];
}

function reduceArray(A, { mode, axis, keepdims }) {
  const rank = shapeOf(A).length;

  if (axis === null) {
    const flat = flatten(A);
    let scalar;
    if (mode === "sum") {
      scalar = flat.reduce((acc, x) => acc + x, 0);
    } else {
      scalar = Number(
        (flat.reduce((acc, x) => acc + x, 0) / flat.length).toFixed(4)
      );
    }

    if (!keepdims) return scalar;

    let wrapped = scalar;
    for (let i = 0; i < rank; i += 1) wrapped = [wrapped];
    return wrapped;
  }

  const reduced = reduceAlongAxis(A, axis, mode);
  if (!keepdims) return reduced;
  return insertAxisDimension(reduced, axis);
}

function reduceAlongAxis(arr, axis, mode) {
  if (axis === 0) {
    if (arr.length === 0) throw new Error("Cannot reduce empty array.");
    let accum = clone(arr[0]);
    for (let i = 1; i < arr.length; i += 1) {
      accum = elementWiseDeep(accum, arr[i], (a, b) => a + b);
    }

    if (mode === "sum") return accum;
    return elementWiseDeep(accum, accum, (a) => Number((a / arr.length).toFixed(4)));
  }

  return arr.map((sub) => reduceAlongAxis(sub, axis - 1, mode));
}

function insertAxisDimension(value, axis) {
  if (axis === 0) return [value];
  if (!Array.isArray(value)) return [value];
  return value.map((v) => insertAxisDimension(v, axis - 1));
}

function clone(value) {
  return JSON.parse(JSON.stringify(value));
}

function renderCode(op, parsed) {
  const lines = ["import numpy as np", ""];

  if (op.id === "ones" || op.id === "zeros") {
    lines.push(`shape = (${parsed.shapes[0].join(", ")})`);
  } else {
    lines.push(`A = np.random.randint(1, 10, size=(${parsed.shapes[0].join(", ")}))`);
    if (op.inputs === 2) {
      lines.push(
        `B = np.random.randint(1, 10, size=(${parsed.shapes[1].join(", ")}))`
      );
    }
    if (op.requiresTarget) {
      lines.push(`new_shape = (${parsed.targetShape.join(", ")})`);
    }
  }

  if (op.supportsReductionOptions) {
    const axisPart = parsed.options.axis === null ? "None" : parsed.options.axis;
    lines.push(
      `output = np.${op.id}(A, axis=${axisPart}, keepdims=${parsed.options.keepdims})`
    );
  } else {
    lines.push(op.code);
  }

  lines.push("print(output)");

  codeOutput.textContent = lines.join("\n");
}

function renderInputs(arrays, op) {
  inputViz.innerHTML = "";

  if (op.id === "ones" || op.id === "zeros") {
    inputViz.appendChild(
      buildMatrixCard("Generated Array", arrays[0], "from requested shape")
    );
    return;
  }

  arrays.forEach((arr, i) => {
    inputViz.appendChild(buildMatrixCard(`Input ${i + 1}`, arr));
  });
}

function renderOutput(output) {
  outputViz.innerHTML = "";
  outputViz.appendChild(buildMatrixCard("Output", output));
}

function renderComputationDetails(op, inputArrays, output, parsed, meta) {
  detailViz.innerHTML = "";

  if (meta && (meta.kind === "elementwise" || meta.kind === "matmul")) {
    detailViz.appendChild(buildBroadcastUsageCard(meta));
  }

  if (op.id === "matmul") {
    detailViz.appendChild(
      buildMatmulDetail(inputArrays[0], inputArrays[1], output, meta)
    );
  } else if (op.id === "reshape") {
    detailViz.appendChild(buildReshapeDetail(inputArrays[0], parsed.targetShape));
  } else if (op.supportsReductionOptions) {
    detailViz.appendChild(buildReductionDetail(op, inputArrays[0], parsed.options, output));
  } else {
    const note = document.createElement("p");
    note.className = "slice-label";
    note.textContent = "No extra computation walkthrough for this operation yet.";
    detailViz.appendChild(note);
  }
}

function buildBroadcastUsageCard(meta) {
  const card = document.createElement("article");
  card.className = "matrix-card";

  const title = document.createElement("p");
  title.className = "matrix-title";
  title.textContent = meta.usedBroadcasting
    ? "Broadcasting detected"
    : "No broadcasting needed";
  card.appendChild(title);

  const line1 = document.createElement("p");
  line1.className = "formula-line";
  line1.textContent = `Input shapes: A${formatShape(meta.shapeA)}, B${formatShape(
    meta.shapeB
  )}`;
  card.appendChild(line1);

  const line2 = document.createElement("p");
  line2.className = "formula-line";
  line2.textContent = `Output shape: ${formatShape(meta.outputShape)}`;
  card.appendChild(line2);

  const explain = document.createElement("p");
  explain.className = "slice-label";

  if (!meta.usedBroadcasting) {
    explain.textContent = "Inputs already align directly; operation runs without dimension expansion.";
  } else if (meta.kind === "elementwise") {
    explain.textContent =
      "Element-wise broadcasting aligns dimensions from the right. Any dimension with size 1 is repeated to match the other input.";
  } else {
    explain.textContent =
      "For np.matmul, only batch dimensions are broadcast; core matrix dimensions still follow (..., m, k) @ (..., k, n).";
  }
  card.appendChild(explain);

  return card;
}

function buildMatmulDetail(A, B, output, meta) {
  const wrapper = document.createElement("article");
  wrapper.className = "matrix-card";

  const title = document.createElement("p");
  title.className = "matrix-title";
  title.textContent = "np.matmul interactive computation breakdown";
  wrapper.appendChild(title);

  const sA = shapeOf(A);
  const sB = shapeOf(B);

  if (sA.length === 2 && sB.length === 2) {
    wrapper.appendChild(buildMatmulInteractive2D(A, B, output));
    return wrapper;
  }

  if (sA.length === 3 && sB.length === 3 && sA[0] === sB[0]) {
    const help = document.createElement("p");
    help.className = "slice-label";
    help.textContent =
      "Select a batch and output cell to highlight A row, B column, and formula.";
    wrapper.appendChild(help);

    const controls = document.createElement("div");
    controls.className = "inline-fields";
    controls.innerHTML = `
      <label class="mini-label" for="batchSelect">batch</label>
      <select id="batchSelect"></select>
    `;
    wrapper.appendChild(controls);

    const batchSelect = controls.querySelector("#batchSelect");
    for (let i = 0; i < sA[0]; i += 1) {
      const opt = document.createElement("option");
      opt.value = String(i);
      opt.textContent = `batch ${i}`;
      batchSelect.appendChild(opt);
    }

    const host = document.createElement("div");
    wrapper.appendChild(host);

    const renderBatch = () => {
      const idx = Number(batchSelect.value);
      host.innerHTML = "";
      const label = document.createElement("p");
      label.className = "slice-label";
      label.textContent = `output[${idx}] = A[${idx}] @ B[${idx}]`;
      host.appendChild(label);
      host.appendChild(buildMatmulInteractive2D(A[idx], B[idx], output[idx]));
    };

    batchSelect.addEventListener("change", renderBatch);
    renderBatch();
    return wrapper;
  }

  const mixedNote = document.createElement("p");
  mixedNote.className = "slice-label";
  mixedNote.textContent =
    "Mixed-rank/broadcasted matmul: each output batch uses A_batch @ B_batch after NumPy batch broadcasting.";
  wrapper.appendChild(mixedNote);

  if (meta && meta.outBatch > 1) {
    const mapLine = document.createElement("p");
    mapLine.className = "formula-line";
    mapLine.textContent = `Batch mapping: output batch i uses A[${meta.leftBatch === 1 ? "0" : "i"}] and B[${
      meta.rightBatch === 1 ? "0" : "i"
    }].`;
    wrapper.appendChild(mapLine);
  }

  wrapper.appendChild(buildMatrixCard("Input A", A));
  wrapper.appendChild(buildMatrixCard("Input B", B));
  wrapper.appendChild(buildMatrixCard("Output", output));

  return wrapper;
}

function buildMatmulInteractive2D(A, B, out) {
  const box = document.createElement("div");
  const helper = document.createElement("p");
  helper.className = "slice-label";
  helper.textContent =
    "Click a value in the output matrix. The corresponding row/column will be highlighted.";
  box.appendChild(helper);

  const state = { row: 0, col: 0 };
  const pickerHost = document.createElement("div");
  const matrixHost = document.createElement("div");
  matrixHost.className = "viz-grid";
  const formulaLine = document.createElement("p");
  formulaLine.className = "formula-line";

  function renderPicker() {
    pickerHost.innerHTML = "";
    const table = document.createElement("table");
    table.className = "matrix-table";

    out.forEach((rowValues, r) => {
      const tr = document.createElement("tr");
      rowValues.forEach((value, c) => {
        const td = document.createElement("td");
        const btn = document.createElement("button");
        btn.type = "button";
        btn.className = "matrix-pick-btn";
        if (r === state.row && c === state.col) btn.classList.add("active");
        btn.textContent = String(value);
        btn.addEventListener("click", () => {
          state.row = r;
          state.col = c;
          renderAll();
        });
        td.appendChild(btn);
        tr.appendChild(td);
      });
      table.appendChild(tr);
    });

    pickerHost.appendChild(table);
  }

  function build2DCard(title, matrix, subtitle, highlight) {
    const card = document.createElement("article");
    card.className = "matrix-card";

    const shape = shapeOf(matrix);
    const titleEl = document.createElement("p");
    titleEl.className = "matrix-title";
    titleEl.innerHTML = `${escapeHtml(title)} <span class="shape-badge">shape: (${shape.join(
      ", "
    )})</span>`;
    card.appendChild(titleEl);

    const sub = document.createElement("p");
    sub.className = "slice-label";
    sub.textContent = subtitle;
    card.appendChild(sub);

    card.appendChild(render2DTableWithHighlights(matrix, highlight));
    return card;
  }

  function renderAll() {
    renderPicker();
    matrixHost.innerHTML = "";
    matrixHost.appendChild(
      build2DCard(
        "Input A",
        A,
        `highlighted row: ${state.row}`,
        { highlightRow: state.row }
      )
    );
    matrixHost.appendChild(
      build2DCard(
        "Input B",
        B,
        `highlighted column: ${state.col}`,
        { highlightCol: state.col }
      )
    );
    matrixHost.appendChild(
      build2DCard(
        "Output",
        out,
        `selected cell: [${state.row}, ${state.col}]`,
        { highlightCell: { row: state.row, col: state.col } }
      )
    );

    const terms = [];
    const values = [];
    for (let k = 0; k < A[0].length; k += 1) {
      terms.push(`${A[state.row][k]}*${B[k][state.col]}`);
      values.push(A[state.row][k] * B[k][state.col]);
    }
    formulaLine.textContent = `output[${state.row}, ${state.col}] = ${terms.join(
      " + "
    )} = ${values.join(" + ")} = ${out[state.row][state.col]}`;
  }

  box.appendChild(pickerHost);
  box.appendChild(formulaLine);
  box.appendChild(matrixHost);
  renderAll();

  return box;
}

function buildReshapeDetail(A, targetShape) {
  const wrapper = document.createElement("article");
  wrapper.className = "matrix-card";

  const title = document.createElement("p");
  title.className = "matrix-title";
  title.textContent = "np.reshape step-by-step (dynamic)";
  wrapper.appendChild(title);

  const flat = flatten(A);
  const total = flat.length;

  const helper = document.createElement("p");
  helper.className = "slice-label";
  helper.textContent =
    "Use Play to animate reshaping; slider also works for manual inspection.";
  wrapper.appendChild(helper);

  const controls = document.createElement("div");
  controls.className = "slider-row";
  controls.innerHTML = `
    <label for="reshapeStep">step</label>
    <input class="reshape-step" id="reshapeStep" type="range" min="0" max="${total}" value="0" />
    <span id="reshapeStepLabel">0 / ${total}</span>
    <button type="button" id="reshapePlayBtn">Play</button>
    <button type="button" id="reshapeResetBtn">Reset</button>
  `;
  wrapper.appendChild(controls);

  const flatLine = document.createElement("p");
  flatLine.className = "formula-line";
  flatLine.textContent = `flat(A) = [${flat.join(", ")}]`;
  wrapper.appendChild(flatLine);

  const previewHost = document.createElement("div");
  wrapper.appendChild(previewHost);

  const slider = controls.querySelector("#reshapeStep");
  const stepLabel = controls.querySelector("#reshapeStepLabel");
  const playBtn = controls.querySelector("#reshapePlayBtn");
  const resetBtn = controls.querySelector("#reshapeResetBtn");
  let timer = null;
  let playing = false;

  const renderStep = () => {
    const step = Number(slider.value);
    stepLabel.textContent = `${step} / ${total}`;
    const partialFlat = new Array(total).fill("·");
    for (let i = 0; i < step; i += 1) partialFlat[i] = flat[i];
    const partial = unflatten(partialFlat, targetShape);

    previewHost.innerHTML = "";
    previewHost.appendChild(
      buildMatrixCard(
        "Current reshaped output",
        partial,
        `first ${step} element(s) assigned in row-major order`
      )
    );
  };

  const stopAnimation = () => {
    if (timer) {
      clearInterval(timer);
      timer = null;
    }
    playing = false;
    playBtn.textContent = "Play";
  };

  const startAnimation = () => {
    stopAnimation();
    playing = true;
    playBtn.textContent = "Pause";
    timer = setInterval(() => {
      if (!document.body.contains(wrapper)) {
        stopAnimation();
        return;
      }
      const current = Number(slider.value);
      if (current >= total) {
        stopAnimation();
        return;
      }
      slider.value = String(current + 1);
      renderStep();
    }, 500);
  };

  playBtn.addEventListener("click", () => {
    if (playing) {
      stopAnimation();
    } else {
      startAnimation();
    }
  });

  resetBtn.addEventListener("click", () => {
    stopAnimation();
    slider.value = "0";
    renderStep();
  });

  slider.addEventListener("input", () => {
    if (playing) stopAnimation();
    renderStep();
  });

  renderStep();
  startAnimation();

  return wrapper;
}

function buildReductionDetail(op, A, options, output) {
  const wrapper = document.createElement("article");
  wrapper.className = "matrix-card";

  const title = document.createElement("p");
  title.className = "matrix-title";
  title.textContent = `np.${op.id} options summary`;
  wrapper.appendChild(title);

  const summary = document.createElement("p");
  summary.className = "formula-line";
  summary.textContent = `axis=${options.axis === null ? "None" : options.axis}, keepdims=${options.keepdims}`;
  wrapper.appendChild(summary);

  const explain = document.createElement("p");
  explain.className = "slice-label";
  explain.textContent =
    options.axis === null
      ? "Reducing over all dimensions."
      : `Reducing along axis ${options.axis}; values on that axis are aggregated.`;
  wrapper.appendChild(explain);

  wrapper.appendChild(buildMatrixCard("Input A", A));
  wrapper.appendChild(buildMatrixCard("Reduced output", output));

  return wrapper;
}

function buildMatrixCard(title, arr, subtitle = "") {
  const card = document.createElement("article");
  card.className = "matrix-card";

  const shape = shapeOf(arr);
  const titleEl = document.createElement("p");
  titleEl.className = "matrix-title";
  titleEl.innerHTML = `${escapeHtml(title)} <span class="shape-badge">shape: (${shape.join(
    ", "
  ) || "scalar"})</span>`;

  card.appendChild(titleEl);

  if (subtitle) {
    const sub = document.createElement("p");
    sub.className = "slice-label";
    sub.textContent = subtitle;
    card.appendChild(sub);
  }

  if (!Array.isArray(arr)) {
    const scalar = document.createElement("p");
    scalar.textContent = String(arr);
    scalar.style.fontFamily = "Courier New, monospace";
    scalar.style.fontWeight = "700";
    card.appendChild(scalar);
    return card;
  }

  const rank = shape.length;

  if (rank === 1) {
    card.appendChild(render2DTable([arr]));
  } else if (rank === 2) {
    card.appendChild(render2DTable(arr));
  } else if (rank === 3) {
    arr.forEach((slice, idx) => {
      const lbl = document.createElement("p");
      lbl.className = "slice-label";
      lbl.textContent = `slice ${idx} (axis 0)`;
      card.appendChild(lbl);
      card.appendChild(render2DTable(slice));
    });
  }

  return card;
}

function render2DTable(matrix2d) {
  const table = document.createElement("table");
  table.className = "matrix-table";

  matrix2d.forEach((row) => {
    const tr = document.createElement("tr");
    row.forEach((value) => {
      const td = document.createElement("td");
      td.textContent = String(value);
      tr.appendChild(td);
    });
    table.appendChild(tr);
  });

  return table;
}

function render2DTableWithHighlights(matrix2d, highlight = {}) {
  const table = document.createElement("table");
  table.className = "matrix-table";

  matrix2d.forEach((row, r) => {
    const tr = document.createElement("tr");
    row.forEach((value, c) => {
      const td = document.createElement("td");
      td.textContent = String(value);

      if (highlight.highlightRow === r) td.classList.add("hl-row");
      if (highlight.highlightCol === c) td.classList.add("hl-col");
      if (
        highlight.highlightCell &&
        highlight.highlightCell.row === r &&
        highlight.highlightCell.col === c
      ) {
        td.classList.add("hl-cell");
      }

      tr.appendChild(td);
    });
    table.appendChild(tr);
  });

  return table;
}

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

init();