// In-browser analysis engine, ported from src/artifex/predict.py, retrieve.py and uncertainty.py.
// Everything runs on the visitor's device; nothing about an upload is sent anywhere.

import * as store from "./store.js";
import { toModelArray } from "./preprocess.js";
import { randomId } from "./util.js";

let ort = null;
let sessionPromise = null;

async function loadOrt() {
  if (!ort) {
    ort = await import("../vendor/ort/ort.wasm.min.mjs");
    ort.env.wasm.wasmPaths = new URL("../vendor/ort/", import.meta.url).href;
    ort.env.wasm.numThreads = self.crossOriginIsolated ? Math.max(1, Math.min(4, navigator.hardwareConcurrency || 1)) : 1;
  }
  return ort;
}

async function fetchModelBytes(url, expected, onProgress) {
  const cacheName = "artifex-model";
  let cache = null;
  try {
    cache = await caches.open(cacheName);
    const hit = await cache.match(url);
    if (hit) {
      const buf = new Uint8Array(await hit.arrayBuffer());
      if (buf.length === expected) {
        onProgress?.(1, true);
        return buf;
      }
      await cache.delete(url);
    }
  } catch {
    cache = null; // Cache Storage unavailable (e.g. private mode): fall back to the HTTP cache
  }
  const r = await fetch(url);
  if (!r.ok) throw new Error(`model download failed (HTTP ${r.status})`);
  const buf = new Uint8Array(expected);
  const reader = r.body.getReader();
  let got = 0;
  for (;;) {
    const { done, value } = await reader.read();
    if (done) break;
    if (got + value.length > expected) throw new Error("model file is larger than expected");
    buf.set(value, got);
    got += value.length;
    onProgress?.(got / expected, false);
  }
  if (got !== expected) throw new Error(`model download incomplete (${got} of ${expected} bytes)`);
  if (cache) {
    try {
      // Old versions are removed so that only the current model occupies storage.
      for (const req of await cache.keys()) if (req.url !== new URL(url, location.href).href) await cache.delete(req);
      await cache.put(url, new Response(buf.slice(), { headers: { "content-type": "application/octet-stream" } }));
    } catch {
      /* storage quota: keep working without persistent caching */
    }
  }
  return buf;
}

export function loadSession(onProgress) {
  if (!sessionPromise) {
    sessionPromise = (async () => {
      const [m, o] = await Promise.all([store.model(), loadOrt()]);
      const bytes = await fetchModelBytes(m.onnx_file, m.onnx_bytes, onProgress);
      return o.InferenceSession.create(bytes, { executionProviders: ["wasm"], graphOptimizationLevel: "all" });
    })().catch((e) => {
      sessionPromise = null;
      throw e;
    });
  }
  return sessionPromise;
}

export async function embed(img) {
  const [m, session] = await Promise.all([store.model(), loadSession()]);
  const spec = m.encoder.transform;
  const x = toModelArray(img, spec);
  const out = await session.run({ pixel_values: new ort.Tensor("float32", x, [1, 3, spec.image_size, spec.image_size]) });
  const q = out.embedding.data;
  for (const v of q) if (!Number.isFinite(v)) throw new Error("encoder produced non-finite values");
  return Float32Array.from(q);
}

// ---------------------------------------------------------------- classifier heads

function softmax(z, t) {
  const s = z.map((v) => v / t);
  const mx = Math.max(...s);
  const e = s.map((v) => Math.exp(v - mx));
  const sum = e.reduce((a, b) => a + b, 0);
  return e.map((v) => v / sum);
}

export function headProba(head, q) {
  const d = q.length;
  const x = new Float64Array(d);
  for (let i = 0; i < d; i++) x[i] = (q[i] - head.mean[i]) / head.scale[i];
  const logits = head.coef.map((w, c) => {
    let s = head.intercept[c];
    for (let i = 0; i < d; i++) s += w[i] * x[i];
    return s;
  });
  return softmax(logits, head.temperature);
}

function normalizedEntropy(p) {
  if (p.length < 2) return 0;
  let h = 0;
  for (const v of p) {
    const c = Math.min(1, Math.max(1e-12, v));
    h -= c * Math.log(c);
  }
  return h / Math.log(p.length);
}

const NO_ABSTAIN_NOTE =
  "Validation accuracy already exceeded the target without abstaining, so no confidence threshold applies to this head; " +
  "accuracy is dominated by the largest classes, and small classes can still have low recall.";
const TARGET_MISSED_NOTE = "The validation accuracy target was not reachable for this head; the threshold enforces minimum coverage only.";

function headPrediction(name, head, q) {
  const p = headProba(head, q);
  const order = p.map((v, i) => i).sort((a, b) => p[b] - p[a] || a - b);
  const ranked = order.map((i) => ({
    label: head.classes[i],
    display_name: head.display[i],
    score: Math.min(1, Math.max(0, p[i])),
    train_count: head.train_counts[head.classes[i]] || 0,
  }));
  const thr = head.abstention && head.abstention.threshold !== undefined ? head.abstention.threshold : null;
  const top = p[order[0]];
  const second = order.length > 1 ? p[order[1]] : 0;
  let note = null;
  if (head.abstention && head.abstention.target_met === false) note = TARGET_MISSED_NOTE;
  else if (thr !== null && thr <= 0) note = NO_ABSTAIN_NOTE;
  return {
    pred: {
      head: name, ranked, calibrated: true, calibration_method: "temperature scaling (validation split)",
      abstained: thr !== null && top < thr, abstention_threshold: thr, top_score: top, margin: top - second, note,
    },
    p,
  };
}

// ---------------------------------------------------------------- retrieval

export function similarities(E, q) {
  const { data, n, dim } = E;
  const out = new Float32Array(n);
  for (let r = 0; r < n; r++) {
    let s = 0;
    const o = r * dim;
    for (let i = 0; i < dim; i++) s += data[o + i] * q[i];
    out[r] = s;
  }
  return out;
}

function rowDot(E, a, b) {
  const { data, dim } = E;
  let s = 0;
  for (let i = 0; i < dim; i++) s += data[a * dim + i] * data[b * dim + i];
  return s;
}

// ReferenceIndex.search: top-k, optional eligibility mask, near-identical results collapsed.
export function search(E, sims, k, { mask = null, collapse = null, exclude = null } = {}) {
  const cand = [];
  for (let i = 0; i < sims.length; i++) if ((!mask || mask[i]) && (exclude === null || i !== exclude)) cand.push(i);
  cand.sort((a, b) => sims[b] - sims[a] || a - b);
  const pre = cand.slice(0, Math.max(k * 8, 50));
  const hits = [];
  for (const i of pre) {
    if (collapse !== null && hits.some((h) => rowDot(E, i, h.index) >= collapse)) continue;
    hits.push({ index: i, similarity: sims[i] });
    if (hits.length === k) break;
  }
  return hits;
}

function votes(E, sims, recs, column, k) {
  const hits = search(E, sims, k);
  const acc = new Map();
  let total = 0;
  for (const h of hits) {
    const w = Math.max(h.similarity, 0);
    const lab = recs[h.index][column];
    const key = lab === null || lab === undefined ? "unlabelled" : String(lab);
    acc.set(key, (acc.get(key) || 0) + w);
    total += w;
  }
  if (total <= 0) return {};
  const out = {};
  for (const [key, v] of [...acc.entries()].sort((a, b) => b[1] - a[1])) out[key] = v / total;
  return out;
}

export function similarRecord(rec, rank, sim, vocab) {
  return {
    artwork_id: rec.id, rank, similarity: Math.round(sim * 1e4) / 1e4, title: rec.title, artist_name: rec.artist,
    date_display: rec.date, date_start: rec.ds, date_end: rec.de, region_normalized: rec.region,
    region_display: rec.region ? vocab.region[rec.region] || rec.region : "Not recorded", culture_original: rec.culture_o,
    period_label: rec.period, institution: rec.inst, source_url: rec.url, thumbnail_path: rec.thumb, image_license: rec.license,
  };
}

export async function similar(q, sims, { region = null, period = null } = {}) {
  const [m, E, R, vocab] = await Promise.all([store.model(), store.embeddings(), store.records(), store.vocab()]);
  let mask = null;
  if (region || period) mask = R.list.map((r) => (!region || r.region === region) && (!period || r.period === period));
  return search(E, sims, m.retrieval.top_k, { mask, collapse: m.retrieval.collapse_cosine }).map((h, i) =>
    similarRecord(R.list[h.index], i + 1, h.similarity, vocab),
  );
}

// ---------------------------------------------------------------- analysis

const f2 = (x) => x.toFixed(2);
const f0 = (x) => x.toFixed(0);

export async function analyze(img, onStage) {
  const t0 = performance.now();
  const [m, E, R, vocab] = await Promise.all([store.model(), store.embeddings(), store.records(), store.vocab()]);
  onStage?.("encode");
  const q = await embed(img);
  onStage?.("score");
  const region = headPrediction("region", m.heads.region, q);
  const period = headPrediction("period", m.heads.period, q);
  let subregion = null;
  const subHead = m.heads[`subregion:${region.pred.ranked[0].label}`];
  if (subHead && !region.pred.abstained) subregion = headPrediction(`subregion:${region.pred.ranked[0].label}`, subHead, q).pred;

  const sims = similarities(E, q);
  const sim = await similar(q, sims);
  const k = m.retrieval.vote_k;
  const regionVotes = votes(E, sims, R.list, "region", k);
  const periodVotes = votes(E, sims, R.list, "period", k);
  let nearestRef = -Infinity;
  let nearestTrain = -Infinity;
  for (let i = 0; i < sims.length; i++) {
    nearestRef = Math.max(nearestRef, sims[i]);
    if (R.list[i].split === "train") nearestTrain = Math.max(nearestTrain, sims[i]);
  }
  if (nearestRef === -Infinity) nearestRef = 0;
  if (nearestTrain === -Infinity) nearestTrain = 0;
  const ood = m.ood ? nearestTrain < m.ood.threshold : false;
  const unc = {
    nearest_similarity: nearestRef, ood_threshold: m.ood ? m.ood.threshold : null, out_of_distribution: ood,
    region_entropy: normalizedEntropy(region.p), period_entropy: normalizedEntropy(period.p),
  };
  const status = ood ? "out_of_scope" : region.pred.abstained && period.pred.abstained ? "insufficient_evidence" : "success";
  const result = {
    analysis_id: randomId(), status, model_version: m.model_version, dataset_version: m.dataset_version,
    encoder: `${m.encoder.checkpoint}@${(m.encoder.commit || "").slice(0, 8)}`,
    region: region.pred, subregion, period: period.pred, similar_artworks: sim,
    retrieval_region_votes: regionVotes, retrieval_period_votes: periodVotes, uncertainty: unc,
    evidence: evidence(m, vocab, region.pred, period.pred, subregion, regionVotes, periodVotes, unc),
    image_width: img.width, image_height: img.height, elapsed_ms: Math.round((performance.now() - t0) * 10) / 10,
    runtime: "browser (onnxruntime-web, WebAssembly)",
  };
  return { result, q, sims };
}

function displayName(vocab, kind, label) {
  if (kind === "region") return vocab.region[label] || label;
  if (kind === "period") return vocab.period[label] || label;
  return vocab.subregion[label] || label;
}

function evidence(m, vocab, region, period, subregion, regionVotes, periodVotes, unc) {
  const ev = [];
  const k = m.retrieval.vote_k;
  if (unc.out_of_distribution)
    ev.push({ kind: "caution", text:
      `The closest training image has similarity ${f2(unc.nearest_similarity)}, below the threshold ` +
      `(${f2(unc.ood_threshold)}) exceeded by 99% of validation images. The upload may lie outside the ` +
      "kinds of paintings the model was trained on; treat all scores below with particular caution." });
  if (unc.nearest_similarity >= m.near_identical)
    ev.push({ kind: "caution", text:
      `A reference image is nearly identical to the upload (similarity ${f2(unc.nearest_similarity)}). The upload may be ` +
      "a photograph of a work already in the reference collection; in that case the scores are not independent evidence." });

  for (const [head, votesMap, noun] of [[region, regionVotes, "region"], [period, periodVotes, "period"]]) {
    const top = head.ranked[0];
    ev.push({ kind: "info", text:
      `The ${noun} classifier ranks ${top.display_name} first with a calibrated score of ${f2(top.score)} ` +
      `(runner-up ${head.ranked[1].display_name}, ${f2(head.ranked[1].score)}).` });
    if (head.abstained)
      ev.push({ kind: "caution", text:
        `This ${noun} score is below the abstention threshold of ${f2(head.abstention_threshold)} chosen on validation data, ` +
        `so Artifex reports insufficient evidence for the ${noun}.` });
    const entries = Object.entries(votesMap);
    if (entries.length) {
      const [vtop, vshare] = entries[0];
      const nMatch = Math.round((votesMap[top.label] || 0) * 100);
      const vname = vtop === "unlabelled" ? "unlabelled works" : displayName(vocab, noun, vtop);
      if (vtop === top.label)
        ev.push({ kind: "agreement", text:
          `Retrieval agrees: ${top.display_name} accounts for ${nMatch}% of the similarity-weighted votes ` +
          `of the ${k} nearest reference works.` });
      else
        ev.push({ kind: "disagreement", text:
          `Retrieval disagrees: the ${k} nearest reference works are mostly documented as ${vname} ` +
          `(${Math.round(vshare * 100)}% of weighted votes), while ${top.display_name} accounts for ${nMatch}%.` });
    }
    if (top.train_count < m.limited_representation)
      ev.push({ kind: "caution", text: `${top.display_name} has limited representation in the training data (${top.train_count} works).` });
    const rec = (m.test_recall[noun] || {})[top.label];
    if (rec)
      ev.push({ kind: "info", text:
        `On the held-out test split, the ${noun} classifier recovered ${f0(rec.recall * 100)}% of works labelled ` +
        `${top.display_name} (n = ${rec.support}); precision when predicting it was ${f0(rec.precision * 100)}%.` });
  }
  if (subregion) {
    const s = subregion.ranked[0];
    ev.push({ kind: "info", text:
      `Within ${region.ranked[0].display_name}, the tradition classifier ranks ${s.display_name} first (${f2(s.score)}).` +
      (subregion.abstained ? " This score is below its abstention threshold." : "") });
  }
  const supported = m.heads.region.display.join(", ");
  ev.push({ kind: "caution", text:
    `Supported regions: ${supported}. Supported periods: ${m.heads.period.display.join(", ")}. Works from other traditions or periods ` +
    "cannot be recognised and will be assigned to the nearest supported category." });
  return ev;
}