/* undertow -- the browser half.
 *
 * Everything here runs on the visitor's machine. The four TF-IDF models are the real
 * fitted models from `make grid`, shipped as JSON, and this file is a line-for-line port
 * of `undertow/tfidf.py`: same tokeniser, same sublinear TF, same smooth IDF, same L2
 * norm, same linear margin. A test in the suite runs this file under Node against the
 * Python implementation on the same strings and requires agreement to 1e-9, so the four
 * verdicts you see are the four verdicts the paper's tables were computed from.
 *
 * It is a port rather than an API call on purpose: the point of the page is that four
 * models trained on four corpora disagree about one sentence, and that is only convincing
 * if you can type your own sentence and get an answer with no round trip to anywhere.
 */

'use strict';

/* sklearn's default analyzer is `(?u)\b\w\w+\b` over the lowercased string. JavaScript's
 * `\w` is ASCII-only, so the Unicode property escapes are spelled out to match Python. */
const TOKEN_RE = /(?<![\p{L}\p{N}_])[\p{L}\p{N}_]{2,}(?![\p{L}\p{N}_])/gu;

function tokenizeWords(text) {
  return String(text).toLowerCase().match(TOKEN_RE) || [];
}

function ngrams(tokens, lo, hi) {
  const out = lo <= 1 ? tokens.slice() : [];
  for (let n = Math.max(2, lo); n <= hi; n += 1) {
    for (let i = 0; i + n <= tokens.length; i += 1) {
      out.push(tokens.slice(i, i + n).join(' '));
    }
  }
  return out;
}

class TfidfModel {
  constructor(payload) {
    this.terms = payload.terms;
    this.idf = payload.idf;
    this.coef = payload.coef;
    this.intercept = payload.intercept;
    this.ngramRange = payload.ngram_range || [1, 2];
    this.vocabulary = new Map();
    for (let i = 0; i < this.terms.length; i += 1) this.vocabulary.set(this.terms[i], i);
  }

  /* Sparse by construction: a review touches a few hundred of thirty thousand columns,
   * so the dense 30k vector of the Python version is never materialised here. */
  transformOne(text) {
    const counts = new Map();
    const grams = ngrams(tokenizeWords(text), this.ngramRange[0], this.ngramRange[1]);
    for (const gram of grams) {
      const idx = this.vocabulary.get(gram);
      if (idx !== undefined) counts.set(idx, (counts.get(idx) || 0) + 1);
    }
    const entries = [];
    let sumsq = 0;
    for (const [idx, count] of counts) {
      const value = (1 + Math.log(count)) * this.idf[idx];
      sumsq += value * value;
      entries.push([idx, value]);
    }
    const norm = Math.sqrt(sumsq);
    if (norm > 0) for (const e of entries) e[1] /= norm;
    return entries;
  }

  score(text) {
    const entries = this.transformOne(text);
    let margin = this.intercept;
    const contributions = [];
    for (const [idx, value] of entries) {
      const contribution = value * this.coef[idx];
      margin += contribution;
      contributions.push({ term: this.terms[idx], weight: contribution });
    }
    contributions.sort((a, b) => Math.abs(b.weight) - Math.abs(a.weight));
    return {
      margin,
      probability: 1 / (1 + Math.exp(-margin)),
      label: margin > 0 ? 'positive' : 'negative',
      matched: entries.length,
      terms: contributions.slice(0, 12),
    };
  }
}

/* ------------------------------------------------------------------ Node parity hook */
if (typeof module !== 'undefined' && module.exports) {
  module.exports = { TfidfModel, tokenizeWords, ngrams };
}

/* --------------------------------------------------------------------------- the page */
if (typeof document !== 'undefined') {
  const state = { data: null, models: {}, method: 'ft' };
  const $ = (sel) => document.querySelector(sel);
  const el = (tag, cls, text) => {
    const node = document.createElement(tag);
    if (cls) node.className = cls;
    if (text !== undefined) node.textContent = text;
    return node;
  };
  const pct = (x) => `${(x * 100).toFixed(1)}%`;
  const pp = (x) => `${x >= 0 ? '+' : ''}${(x * 100).toFixed(1)} pp`;

  async function boot() {
    const [data, models] = await Promise.all([
      fetch('data.json').then((r) => r.json()),
      fetch('models.json').then((r) => r.json()),
    ]);
    state.data = data;
    for (const [key, payload] of Object.entries(models)) {
      state.models[key] = new TfidfModel(payload);
    }

    /* ?text= and ?method= make a particular view of the page linkable, which is also how
     * the README screenshots are captured reproducibly rather than by hand. */
    const params = new URLSearchParams(window.location.search);
    if (params.has('text')) $('#text').value = params.get('text');
    if (params.has('example')) {
      const wanted = params.get('example').toLowerCase();
      const match = data.examples.find((e) => e.label.toLowerCase() === wanted);
      if (match) $('#text').value = match.text;
    }
    if (params.has('method') && data.methods.some((m) => m.key === params.get('method'))) {
      state.method = params.get('method');
    }

    renderExamples();
    renderScorer();
    renderMethodTabs();
    renderMatrix();
    renderSummary();
    renderDecomposition();
    renderShiftAxes();
    $('#loading').remove();
    $('#app').hidden = false;
  }

  /* --------------------------------------------------------------- the live scorer */

  function renderExamples() {
    const box = $('#examples');
    for (const example of state.data.examples) {
      const chip = el('button', 'chip', example.label);
      chip.type = 'button';
      chip.addEventListener('click', () => {
        $('#text').value = example.text;
        renderScorer();
      });
      box.append(chip);
    }
    const clear = el('button', 'chip chip-ghost', 'Clear');
    clear.type = 'button';
    clear.addEventListener('click', () => {
      $('#text').value = '';
      renderScorer();
    });
    box.append(clear);
  }

  function renderScorer() {
    const text = $('#text').value.trim();
    const box = $('#verdicts');
    box.textContent = '';
    const labels = {};
    for (const domain of state.data.domains) {
      const model = state.models[domain.key];
      if (!model) continue;
      const result = text ? model.score(text) : null;
      labels[domain.key] = result ? result.label : null;
      box.append(verdictCard(domain, result));
    }
    const seen = new Set(Object.values(labels).filter(Boolean));
    const note = $('#disagreement');
    if (!text) {
      note.className = 'note note-idle';
      note.textContent =
        'Type a review, or pick one of the examples. Each card is a logistic model over '
        + 'TF-IDF n-grams, trained on 6,000 reviews from one corpus and nothing else.';
    } else if (seen.size > 1) {
      const split = Object.entries(labels)
        .filter(([, v]) => v)
        .map(([k, v]) => `${state.data.domains.find((d) => d.key === k).label} says ${v}`)
        .join(', ');
      note.className = 'note note-split';
      note.textContent = `The four models disagree. ${split}. Same sentence, same architecture — the only difference is which corpus each one read.`;
    } else {
      note.className = 'note note-agree';
      note.textContent =
        `All four agree this is ${[...seen][0]}. Agreement is the common case; the transfer `
        + 'matrix below is where the disagreement shows up at scale.';
    }
  }

  function verdictCard(domain, result) {
    const card = el('article', 'verdict');
    const head = el('header', 'verdict-head');
    head.append(el('span', 'verdict-domain', domain.label));
    head.append(el('span', 'verdict-meta', `${domain.topic} · ${domain.length}`));
    card.append(head);

    if (!result) {
      card.append(el('p', 'verdict-empty', 'waiting for text'));
      return card;
    }

    const positive = result.label === 'positive';
    card.classList.add(positive ? 'is-positive' : 'is-negative');
    const verdict = el('p', 'verdict-label', positive ? 'positive' : 'negative');
    verdict.append(el('span', 'verdict-prob', ` ${pct(positive ? result.probability : 1 - result.probability)} confident`));
    card.append(verdict);

    const meter = el('div', 'meter');
    const fill = el('div', 'meter-fill');
    fill.style.width = `${Math.min(100, Math.max(0, result.probability * 100))}%`;
    meter.append(fill);
    const mid = el('div', 'meter-mid');
    meter.append(mid);
    card.append(meter);

    const terms = el('ul', 'terms');
    for (const t of result.terms.slice(0, 6)) {
      const item = el('li', t.weight > 0 ? 'term term-pos' : 'term term-neg');
      item.append(el('span', 'term-text', t.term));
      item.append(el('span', 'term-weight', t.weight.toFixed(3)));
      terms.append(item);
    }
    card.append(terms);
    card.append(el('p', 'verdict-foot', `${result.matched} of its ${state.models[domain.key].terms.length.toLocaleString()} n-grams matched`));
    return card;
  }

  /* ------------------------------------------------------------- the transfer matrix */

  function renderMethodTabs() {
    const box = $('#method-tabs');
    for (const method of state.data.methods) {
      if (!state.data.grid.some((r) => r.method === method.key)) continue;
      const tab = el('button', 'tab', method.label);
      tab.type = 'button';
      tab.dataset.method = method.key;
      tab.addEventListener('click', () => {
        state.method = method.key;
        renderMethodTabs();
        renderMatrix();
      });
      if (method.key === state.method) tab.classList.add('is-active');
      box.append(tab);
    }
    if (box.children.length > state.data.methods.length) return;
  }

  function renderMatrix() {
    const box = $('#matrix');
    box.textContent = '';
    const domains = state.data.domains;
    const cells = state.data.grid.filter((r) => r.method === state.method);
    if (!cells.length) return;
    const values = cells.map((c) => c.accuracy);
    const lo = Math.min(...values);
    const hi = Math.max(...values);

    box.append(el('div', 'mcell mcell-corner', ''));
    for (const d of domains) box.append(el('div', 'mcell mcell-head', d.label));

    for (const train of domains) {
      box.append(el('div', 'mcell mcell-head mcell-row', train.label));
      for (const ev of domains) {
        const cell = cells.find((c) => c.train === train.key && c.eval === ev.key);
        const node = el('div', 'mcell mcell-value');
        if (!cell) {
          node.textContent = '–';
          box.append(node);
          continue;
        }
        const t = hi > lo ? (cell.accuracy - lo) / (hi - lo) : 0.5;
        node.style.setProperty('--t', t.toFixed(3));
        if (t > 0.62) node.classList.add('is-dark');
        if (train.key === ev.key) node.classList.add('is-diagonal');
        node.append(el('span', 'mcell-acc', pct(cell.accuracy)));
        node.title =
          `${state.method}: trained on ${train.label}, evaluated on ${ev.label}\n`
          + `accuracy ${pct(cell.accuracy)} · AUROC ${pct(cell.auroc)}\n`
          + `best reachable by moving the threshold ${pct(cell.oracle_accuracy)}\n`
          + `predicts positive ${pct(cell.positive_rate)} of the time (truth is 50%)`;
        box.append(node);
      }
    }
    const blurb = state.data.methods.find((m) => m.key === state.method);
    $('#matrix-blurb').textContent = blurb ? blurb.blurb : '';
  }

  /* -------------------------------------------------------------------- the summary */

  function renderSummary() {
    const body = $('#summary-body');
    body.textContent = '';
    for (const row of state.data.report.methods) {
      const method = state.data.methods.find((m) => m.key === row.method);
      const tr = el('tr');
      tr.append(el('th', null, method ? method.label : row.method));
      tr.append(el('td', null, pct(row.id_accuracy)));
      tr.append(el('td', null, pct(row.ood_accuracy)));
      const gap = el('td', 'num-gap', pp(-row.ood_gap));
      tr.append(gap);
      tr.append(el('td', null, pct(row.ood_auroc)));
      tr.append(el('td', null, pct(row.ood_oracle_accuracy)));
      tr.append(el('td', null, row.ood_ece.toFixed(3)));
      body.append(tr);
    }
  }

  /* ------------------------------------------------- what the gap is actually made of */

  function renderDecomposition() {
    const box = $('#decomposition');
    box.textContent = '';
    const rows = state.data.report.methods;
    const worst = Math.max(...rows.map((r) => Math.max(0, r.ranking_drop) + Math.max(0, r.placement_drop)));
    for (const row of rows) {
      const method = state.data.methods.find((m) => m.key === row.method);
      const item = el('div', 'decomp-row');
      item.append(el('span', 'decomp-label', method ? method.label : row.method));
      const bar = el('div', 'decomp-bar');
      const ranking = Math.max(0, row.ranking_drop);
      const placement = Math.max(0, row.placement_drop);
      const scale = worst > 0 ? 100 / worst : 0;
      const a = el('div', 'seg seg-ranking');
      a.style.width = `${ranking * scale}%`;
      a.title = `ranking: ${pp(ranking)} of accuracy the representation genuinely lost`;
      const b = el('div', 'seg seg-placement');
      b.style.width = `${placement * scale}%`;
      b.title = `placement: ${pp(placement)} recoverable by moving the threshold alone`;
      bar.append(a, b);
      item.append(bar);
      item.append(el('span', 'decomp-total', pp(-(ranking + placement))));
      box.append(item);
    }
  }

  /* ------------------------------------------------------- which shifts actually hurt */

  function renderShiftAxes() {
    const rows = state.data.report.shift_axes || [];
    if (!rows.length) return;
    const body = $('#shift-body');
    const shifts = [...new Set(rows.map((r) => r.shift))].sort();
    const head = $('#shift-head');
    head.textContent = '';
    head.append(el('th', null, 'method'));
    for (const s of shifts) head.append(el('th', null, s.replace('topic-', 'topic ').replace('/length-', ' · length ')));
    body.textContent = '';
    for (const method of state.data.methods) {
      if (!rows.some((r) => r.method === method.key)) continue;
      const tr = el('tr');
      tr.append(el('th', null, method.label));
      for (const s of shifts) {
        const row = rows.find((r) => r.method === method.key && r.shift === s);
        tr.append(el('td', null, row ? pct(row.accuracy) : '–'));
      }
      body.append(tr);
    }
  }

  document.addEventListener('DOMContentLoaded', () => {
    $('#text').addEventListener('input', renderScorer);
    boot().catch((err) => {
      const box = $('#loading');
      if (box) {
        box.textContent = `Could not load the payload: ${err}. Run \`make web\` to build it.`;
      }
    });
  });
}