pokemon-color-search/nearest.js

// Selectors + DOM Manipulation
const getColorInputNode = () => document.getElementById("color-input");
const getMetricDropdownNode = () => document.getElementById("metric");
const getIncludeXToggleNode = () => document.getElementById("include-x");
const getNormQYToggleNode = () => document.getElementById("norm-q-y");
const getCloseCoeffSliderNode = () => document.getElementById("close-coeff");
const getCloseCoeffDisplayNode = () => document.getElementById("close-coeff-display");
const getLimitSliderNode = () => document.getElementById("num-poke");
const getLimitDisplayNode = () => document.getElementById("num-poke-display");
const getNameInputNode = () => document.getElementById("pokemon-name");
const getScoreListJABNode = () => document.getElementById("best-list-jab");
const getScoreListRGBNode = () => document.getElementById("best-list-rgb");
const getSearchListNode = () => document.getElementById("search-list");
const getHideableControlNodes = () => document.querySelectorAll(".hideable_control");
const getQJABDisplay = () => document.getElementById("q-vec-jab");
const getQRGBDisplay = () => document.getElementById("q-vec-rgb");
const getObjFnDisplay = () => document.getElementById("obj-fn");

const clearNodeContents = node => { node.innerHTML = ""; };
const hideCustomControls = () => getHideableControlNodes()
  .forEach(n => n.setAttribute("class", "hideable_control hideable_control--hidden"));
const showCustomControls = () => getHideableControlNodes()
  .forEach(n => n.setAttribute("class", "hideable_control"));

// Vector Math
const vectorDot = (u, v) => u.map((x, i) => x * v[i]).reduce((x, y) => x + y);
const vectorMag = v => Math.sqrt(vectorDot(v, v));
const vectorNorm = v => { const n = vectorMag(v); return [ n, v.map(c => c / n) ]; };

// Angle Math
const angleDiff = (a, b) => { const raw = Math.abs(a - b); return raw < 180 ? raw : (360 - raw); };
const acosDeg = v => Math.acos(v) * 180 / Math.PI;

// Pre-Compute Y Data
const pokemonColorData = database.map(data => {
  const yRGBColor = d3.rgb(...data.yRGB);
  const [ yJABNorm, yJABHat ] = vectorNorm(data.yJAB);
  const [ yRGBNorm, yRGBHat ] = vectorNorm(data.yRGB);
  const [ _, yChromaHat ] = vectorNorm(data.yJAB.slice(1));

  return {
    ...data,
    yJABHex: d3.jab(...data.yJAB).formatHex(),
    yJABNorm, yJABHat,
    yRGBHex: yRGBColor.formatHex(),
    yRGBNorm, yRGBHat,
    yChromaHat,
    yHueAngle: d3.hsl(yRGBColor).h,
  }
});
const pokemonLookup = new Fuse(pokemonColorData, { keys: [ "name" ] });

// Color Calculations
const getContrastingTextColor = rgb => vectorDot(rgb, [0.3, 0.6, 0.1]) >= 128 ? "#222" : "#ddd";

const readColorInput = () => {
  const colorInput = "#" + (getColorInputNode()?.value?.replace("#", "") ?? "FFFFFF");

  if (colorInput.length !== 7) {
    return;
  }

  const rgb = d3.color(colorInput);
  const { J, a, b } = d3.jab(rgb);

  const qJAB = [ J, a, b ];
  const qRGB = [ rgb.r, rgb.g, rgb.b ];

  const [ qJABNorm, qJABHat ] = vectorNorm(qJAB);
  const qJABNormSq = qJABNorm * qJABNorm; 
  const [ _, qChromaHat ] = vectorNorm(qJAB.slice(1));

  const [ qRGBNorm, qRGBHat ] = vectorNorm(qRGB);
  const qRGBNormSq = qRGBNorm * qRGBNorm;
  const qHueAngle = d3.hsl(rgb).h;

  return {
    qHex: rgb.formatHex(),
    qJAB, qJABHat, qJABNorm, qJABNormSq, qChromaHat,
    qRGB, qRGBHat, qRGBNorm, qRGBNormSq, qHueAngle,
  };
};

// State
const state = {
  metric: null,
  includeX: null,
  normQY: null,
  closeCoeff: null,
  numPoke: null,
  searchTerm: null,
  targetColor: null,
};

// Metrics
const scoringMetrics = [
  ({ xJAB, xRGB, yJAB, yRGB }) => [ 
    xJAB - 2 * vectorDot(yJAB, state.targetColor.qJAB),
    xRGB - 2 * vectorDot(yRGB, state.targetColor.qRGB),
  ],
  ({ yJABHat, yRGBHat }) => [ 
    -vectorDot(yJABHat, state.targetColor.qJABHat),
    -vectorDot(yRGBHat, state.targetColor.qRGBHat),
  ],
  ({ yChromaHat, yHueAngle }) => [
    acosDeg(vectorDot(state.targetColor.qChromaHat, yChromaHat)),
    angleDiff(state.targetColor.qHueAngle, yHueAngle),
  ],
  ({ xJAB, xRGB, yJAB, yRGB, yJABHat, yRGBHat }) => [ 
    (state.includeX ? xJAB : 0) - state.closeCoeff * vectorDot(
      state.normQY ? yJABHat : yJAB, 
      state.normQY ? state.targetColor.qJABHat : state.targetColor.qJAB
    ),
    (state.includeX ? xRGB : 0) - state.closeCoeff * vectorDot(
      state.normQY ? yRGBHat : yRGB, 
      state.normQY ? state.targetColor.qRGBHat : state.targetColor.qRGB
    ),
  ]
];

const calcDisplayMetrics = ({
  xJAB, xRGB, yJABHat, yJABNorm, yRGBHat, yRGBNorm, yChromaHat, yHueAngle,
}) => {
  // TODO - case on state.metric to avoid recalculation of subterms?

  const cosAngleJAB = vectorDot(state.targetColor.qJABHat, yJABHat);
  const yTermJAB = cosAngleJAB * yJABNorm * state.targetColor.qJABNorm;

  const cosAngleRGB = vectorDot(state.targetColor.qRGBHat, yRGBHat);
  const yTermRGB = cosAngleRGB * yRGBNorm * state.targetColor.qRGBNorm;

  return {
    stdDevRGB: Math.sqrt(xRGB - 2 * yTermRGB + state.targetColor.qRGBNormSq),
    stdDevJAB: Math.sqrt(xJAB - 2 * yTermJAB + state.targetColor.qJABNormSq),
    angleJAB: acosDeg(cosAngleJAB), 
    angleRGB: acosDeg(cosAngleRGB), 
    chromaAngle: acosDeg(vectorDot(state.targetColor.qChromaHat, yChromaHat)),
    hueAngle: angleDiff(state.targetColor.qHueAngle, yHueAngle),
  };
};

// Math Rendering
const renderQVec = (q, node, sub) => {
  node.innerHTML = TeXZilla.toMathMLString(`\\vec{q}_{\\text{${sub}}} = \\left(${q.join(", ")}\\right)`);
};

const renderVec = math => `\\vec{${math.charAt(0)}}${math.substr(1)}`;
const renderNorm = vec => `\\frac{${vec}}{\\left|\\left|${vec}\\right|\\right|}`;
const metricText = [
  "\\text{RMS}_{P} ~ \\arg\\min_{P}\\left[X\\left(P\\right) - 2\\vec{q}\\cdot \\vec{Y}\\left(P\\right)\\right]", 
  `\\angle \\left(\\vec{q}, \\vec{Y}\\left(P\\right)\\right) ~ \\arg\\min_{P}\\left[-${renderNorm(renderVec("q"))}\\cdot ${renderNorm(renderVec("Y\\left(P\\right)"))}\\right]`, 
  "\\angle \\left(\\vec{q}_{\\perp}, \\vec{Y}\\left(P\\right)_{\\perp} \\right)", 
];
const updateObjective = () => {
  let tex = metricText?.[state.metric];
  if (!tex) {
    const { includeX, normQY, closeCoeff } = state;
    const xTerm = includeX ? "X\\left(P\\right)" : "";
    const qyMod = normQY ? c => renderNorm(renderVec(c)) : renderVec;
    tex = `\\arg\\min_{P}\\left[${xTerm}-${closeCoeff === 1 ? "" : closeCoeff}${qyMod("q")}\\cdot ${qyMod("Y\\left(P\\right)")}\\right]`;
  }
  const objFnNode = getObjFnDisplay();
  clearNodeContents(objFnNode);
  objFnNode.appendChild(TeXZilla.toMathML(tex));
};

// Pokemon Rendering
const stripForm = ["flabebe", "floette", "florges", "vivillon", "basculin", "furfrou", "magearna"];

const getSprite = pokemon => {
  pokemon = pokemon
    .replace("-alola", "-alolan")
    .replace("-galar", "-galarian")
    .replace("darmanitan-galarian", "darmanitan-galarian-standard");
  if (stripForm.find(s => pokemon.includes(s))) {
    pokemon = pokemon.replace(/-.*$/, "");
  }
  return `https://img.pokemondb.net/sprites/sword-shield/icon/${pokemon}.png`;
};

const renderPokemon = (data, classes = {}) => {
  const { name, yJAB, yJABHex, yRGB, yRGBHex } = data;
  const { stdDevJAB, stdDevRGB, angleJAB, angleRGB, chromaAngle, hueAngle } = calcDisplayMetrics(data)
  const { labelClass = "", rgbClass = "", jabClass = "", scoreClass = "" } = classes;

  const titleName = name.charAt(0).toUpperCase() + name.substr(1);
  const textHex = getContrastingTextColor(yRGB);
  const rgbVec = yRGB.map(c => c.toFixed()).join(", ");
  const jabVec = yJAB.map(c => c.toFixed(1)).join(", ");

  const pkmn = document.createElement("div");
  pkmn.setAttribute("class", "pokemon_tile");
  pkmn.innerHTML = `
    <div class="pokemon_tile-image-wrapper">
      <img src="${getSprite(name)}" />
    </div>
    <div class="pokemon_tile-info_panel">
      <span class="pokemon_tile-pokemon_name">${titleName}</span>
      <div class="pokemon_tile-results">
        <div class="pokemon_tile-labels ${labelClass}">
          <span class="${jabClass}">Jab: </span>
          <span class="${rgbClass}">RGB: </span>
        </div>
        <div class="pokemon_tile-score_column ${scoreClass}">
          <span class="pokemon_tile-no_flex ${jabClass}">
            (${stdDevJAB.toFixed(2)}, ${angleJAB.toFixed(2)}&deg;, ${chromaAngle.toFixed(2)}&deg;)
          </span>
          <span class="pokemon_tile-no_flex ${rgbClass}">
            (${stdDevRGB.toFixed(2)}, ${angleRGB.toFixed(2)}&deg;, ${hueAngle.toFixed(2)}&deg;)
          </span>
        </div>
        <div class="pokemon_tile-hex_column">
          <div class="pokemon_tile-hex_color ${jabClass}" style="background-color: ${yJABHex}; color: ${textHex}">
            <span>${yJABHex}</span><span class="pokemon_tile-vector">(${jabVec})</span>
          </div>
          <div class="pokemon_tile-hex_color ${rgbClass}" style="background-color: ${yRGBHex}; color: ${textHex}">
            <span>${yRGBHex}</span><span class="pokemon_tile-vector">(${rgbVec})</span>
          </div>
        </div>
      </div>
    </div>
  `;
  return pkmn;
};

const getPokemonAppender = targetList => (pokemonData, classes) => {
  const li = document.createElement("li");
  li.appendChild(renderPokemon(pokemonData, classes));
  targetList.appendChild(li);
};

// Search
const search = () => {
  const resultsNode = getSearchListNode();
  const append = getPokemonAppender(resultsNode);
  let found;
  if (state.searchTerm.trim().toLowerCase() === "!random") {
    found = Array.from({ length: 10 }, () => pokemonColorData[Math.floor(Math.random() * pokemonColorData.length)]);
  } else {
    found = pokemonLookup.search(state.searchTerm, { limit: 10 }).map(({ item }) => item);
  }
  clearNodeContents(resultsNode);
  found.forEach(item => append(item));
};

// Scoring
const rescore = () => {
  const metricFn = scoringMetrics[state.metric ?? 0];
  // TODO might like to save this somewhere instead of recomputing when limit changes
  const scores = pokemonColorData.map(data => ({ ...data, scores: metricFn(data) }));

  const jabList = getScoreListJABNode();
  const appendJAB = getPokemonAppender(jabList);
  const rgbList = getScoreListRGBNode();
  const appendRGB = getPokemonAppender(rgbList);

  // extract best CIECAM02 results
  const bestJAB = scores
    .sort((a, b) => a.scores[0] - b.scores[0])
    .slice(0, state.numPoke);
  clearNodeContents(jabList);
  bestJAB.forEach(data => appendJAB(data, { labelClass: "hide", rgbClass: "hide" }));

  // extract best RGB results
  const bestRGB = scores
    .sort((a, b) => a.scores[1] - b.scores[1])
    .slice(0, state.numPoke);
  clearNodeContents(rgbList);
  bestRGB.forEach(data => appendRGB(data, { labelClass: "hide", jabClass: "hide" }));

  // update the rendered search results as well
  search();
};

// Listeners
const onColorChanged = skipScore => { 
  const readColor = readColorInput();
  if (readColor) {
    state.targetColor = readColor;
    
    renderQVec(state.targetColor.qJAB.map(c => c.toFixed(2)), getQJABDisplay(), "Jab");
    renderQVec(state.targetColor.qRGB.map(c => c.toFixed()), getQRGBDisplay(), "RGB");

    const textColor = getContrastingTextColor(state.targetColor.qRGB);
    document.querySelector("body").setAttribute("style", `background: ${state.targetColor.qHex}; color: ${textColor}`);
    state.targetColor
    if (!skipScore) {
      rescore();
    }
  }
};

const onRandomColor = () => {
  const color = [Math.random(), Math.random(), Math.random()].map(c => c * 255);
  getColorInputNode().value = d3.rgb(...color).formatHex();
  onColorChanged(); // triggers rescore
};

const onCustomControlsChanged = skipScore => {
  state.includeX = getIncludeXToggleNode()?.checked ?? false;
  state.normQY = getNormQYToggleNode()?.checked ?? false;
  state.closeCoeff = parseFloat(getCloseCoeffSliderNode()?.value ?? 2);

  getCloseCoeffDisplayNode().innerHTML = state.closeCoeff;
  updateObjective();

  if (!skipScore) {
    rescore();
  }
}

const onMetricChanged = skipScore => {
  const metric = getMetricDropdownNode()?.selectedIndex ?? 0;
  if (metric === state.metric) {
    return;
  }
  state.metric = metric;
  if (state.metric === 3) { // Custom
    showCustomControls();
    onCustomControlsChanged(skipScore); // triggers rescore
  } else {
    hideCustomControls();
    updateObjective();
    if (!skipScore) {
      rescore();
    }
  }
};

const onLimitChanged = skipScore => {
  state.numPoke = parseInt(getLimitSliderNode()?.value ?? 10);

  getLimitDisplayNode().textContent = state.numPoke;

  if (!skipScore) {
    // TODO don't need to rescore just need to expand
    rescore();
  }
};

const onSearchChanged = skipSearch => {
  state.searchTerm = getNameInputNode()?.value?.toLowerCase() ?? "";
  if (state.searchTerm.length === 0) {
    clearNodeContents(getSearchListNode());
  } else if (!skipSearch) {
    search();
  }
};

const onPageLoad = () => { 
  // fake some events but don't do any searching or scoring
  onSearchChanged(true);
  onColorChanged(true);
  onMetricChanged(true);
  onLimitChanged(true);
  // then do a rescore directly, which will itself trigger a search
  rescore();
};