// Runtime-agnostic prediction pipeline: `ort` is onnxruntime-web (browser) or
// onnxruntime-node (tests); `getSession(task)` returns an InferenceSession.
import {
  prepareFeatures, buildMembers, prepareClassification, prepareRegression,
  trimmedMean, softmax, makeRng, isMissing, parseNumber, isNumericColumn,
} from "./kumo.js";

export const LIMITS = { context: 5000, query: 20000, classes: 10, chunk: 256 };

export async function predict({ ort, getSession, ctx, qry, target, task, numEstimators, seed = 0, onProgress = () => {} }) {
  if (!ctx.columns.includes(target)) throw new Error(`Target "${target}" not found`);
  const features = ctx.columns.filter((c) => c !== target);
  if (!features.length) throw new Error("Need at least one feature column besides the target");

  const labeled = ctx.rows.filter((r) => !isMissing(r[target]));
  const unlabeled = ctx.rows.filter((r) => isMissing(r[target]));
  const queryRows = (qry ? qry.rows : []).concat(unlabeled);
  if (!queryRows.length) throw new Error("Nothing to predict: add a query CSV or leave the target blank on some rows");
  if (labeled.length < 2) throw new Error("Need at least 2 labeled context rows");
  if (labeled.length > LIMITS.context) throw new Error(`Browser demo limit: ${LIMITS.context} labeled context rows`);
  if (queryRows.length > LIMITS.query) throw new Error(`Browser demo limit: ${LIMITS.query} query rows`);

  const yVals = labeled.map((r) => r[target]);
  if (task === "auto") {
    const numeric = isNumericColumn(yVals);
    task = numeric && new Set(yVals.map(Number)).size > 10 ? "regression" : "classification";
  }
  if (task === "regression" && yVals.some((v) => !Number.isFinite(parseNumber(v)))) throw new Error("Regression needs a numeric target");

  const E = numEstimators;
  const rng = makeRng(seed + 1);
  // RNG order mirrors sdm: target first, then feature flips, then column shuffle.
  const tgt = task === "classification" ? prepareClassification(yVals, rng, E) : prepareRegression(yVals, rng, E);
  if (task === "classification") {
    if (tgt.K < 2) throw new Error("Classification needs at least 2 distinct target values");
    if (tgt.K > LIMITS.classes) throw new Error(`Browser demo supports up to ${LIMITS.classes} classes (got ${tgt.K})`);
  }
  onProgress("Preprocessing…", 0);
  const prep = prepareFeatures(features, labeled, queryRows);
  if (!prep.columns.length) throw new Error("All feature columns are constant in the context rows");
  const members = buildMembers(prep, E, rng);
  const session = await getSession(task);

  const n = prep.n, nq = queryRows.length, chunk = LIMITS.chunk;
  const nChunks = Math.ceil(nq / chunk);
  const outW = task === "classification" ? tgt.K : 999;
  const perMember = members.map(() => new Float64Array(nq * outW));
  let step = 0;
  for (let m = 0; m < E; m++) {
    const { x, F, mask } = members[m];
    const yT = task === "classification"
      ? new ort.Tensor("int64", tgt.yFor(m), [n])
      : new ort.Tensor("float32", tgt.yFor(m), [n]);
    const maskT = new ort.Tensor("bool", mask, [F]);
    for (let c = 0; c < nChunks; c++) {
      const q0 = c * chunk, q1 = Math.min(nq, q0 + chunk), R = n + (q1 - q0);
      const xb = new Float32Array(R * F);
      xb.set(x.subarray(0, n * F), 0);
      xb.set(x.subarray((n + q0) * F, (n + q1) * F), n * F);
      const out = await session.run({ x: new ort.Tensor("float32", xb, [R, F]), y: yT, categorical_mask: maskT });
      const o = out.logits.data, W = out.logits.dims[1];
      for (let i = 0; i < q1 - q0; i++) {
        const row = o.subarray(i * W, (i + 1) * W);
        const v = task === "classification" ? tgt.unshift(m, row) : tgt.invert(m, row);
        perMember[m].set(v.subarray(0, outW), (q0 + i) * outW);
      }
      step++;
      onProgress(`Estimator ${m + 1}/${E} · chunk ${c + 1}/${nChunks}`, step / (E * nChunks));
    }
  }

  // Post-process
  const columns = [], rows = [];
  const truthPresent = queryRows.some((r) => target in r && !isMissing(r[target]));
  let metrics = "";
  if (task === "classification") {
    const names = tgt.classes;
    columns.push(`predicted_${target}`, ...(truthPresent ? [`true_${target}`] : []), "confidence", ...features, ...names.map((c) => `p(${c})`));
    let correct = 0, scored = 0;
    for (let i = 0; i < nq; i++) {
      const avg = Array.from({ length: tgt.K }, (_, k) => members.reduce((s, _m, m) => s + perMember[m][i * tgt.K + k], 0) / E);
      const p = softmax(avg);
      let best = 0; for (let k = 1; k < p.length; k++) if (p[k] > p[best]) best = k;
      const row = { [`predicted_${target}`]: names[best], confidence: +p[best].toFixed(4) };
      if (truthPresent) {
        const t = queryRows[i][target] ?? "";
        row[`true_${target}`] = t;
        if (!isMissing(t)) { scored++; if (String(t) === names[best] || (Number(t) === Number(names[best]) && t !== "")) correct++; }
      }
      for (const f of features) row[f] = queryRows[i][f] ?? "";
      names.forEach((c, k) => (row[`p(${c})`] = +p[k].toFixed(4)));
      rows.push(row);
    }
    if (scored) metrics = `accuracy ${(correct / scored).toFixed(3)} on ${scored} labeled query rows`;
  } else {
    columns.push(`predicted_${target}`, ...(truthPresent ? [`true_${target}`] : []), "p10", "p90", ...features);
    const pick = (i, qi) => trimmedMean(members.map((_m, m) => perMember[m][i * 999 + qi]));
    let se = 0, ae = 0, cov = 0, scored = 0; const ts = [];
    for (let i = 0; i < nq; i++) {
      const med = pick(i, 499), p10 = pick(i, 99), p90 = pick(i, 899);
      const row = { [`predicted_${target}`]: +med.toPrecision(6), p10: +p10.toPrecision(6), p90: +p90.toPrecision(6) };
      if (truthPresent) {
        const t = queryRows[i][target] ?? "";
        row[`true_${target}`] = t;
        const tv = parseNumber(t);
        if (Number.isFinite(tv)) { scored++; se += (med - tv) ** 2; ae += Math.abs(med - tv); cov += tv >= p10 && tv <= p90; ts.push([tv, med]); }
      }
      for (const f of features) row[f] = queryRows[i][f] ?? "";
      rows.push(row);
    }
    if (scored) {
      const mt = ts.reduce((a, [t]) => a + t, 0) / scored, sst = ts.reduce((a, [t]) => a + (t - mt) ** 2, 0);
      metrics = `RMSE ${Math.sqrt(se / scored).toFixed(3)} · MAE ${(ae / scored).toFixed(3)} · R² ${(1 - se / sst).toFixed(3)} · 80% interval coverage ${(cov / scored).toFixed(2)} on ${scored} rows`;
    }
  }
  return { task, columns, rows, metrics, nContext: n, nQuery: nq, nFeatures: prep.columns.length };
}