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 vectorSqMag = v => vectorDot(v, v);
const vectorMag = v => Math.sqrt(vectorSqMag(v));
const vectorSqDist = (u, v) => vectorSqMag(u.map((x, i) => x - v[i]));
const vectorDist = (u, v) => Math.sqrt(vectorSqDist(u, 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 rad2deg = 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);  

  return {
    ...data,
    yJABHex: d3.jab(...data.yJAB).formatHex(),
    yJABNorm, yJABHat,
    yRGBHex: yRGBColor.formatHex(),
    yRGBNorm, yRGBHat,
    yHueAngleJAB: rad2deg * Math.atan2(data.yJAB[2], data.yJAB[1]),
    yHueAngleRGB: 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 [ qRGBNorm, qRGBHat ] = vectorNorm(qRGB);
  const qRGBNormSq = qRGBNorm * qRGBNorm;

  return {
    qHex: rgb.formatHex(),
    qJAB, qJABHat, qJABNorm, qJABNormSq,
    qRGB, qRGBHat, qRGBNorm, qRGBNormSq,
    qHueAngleJAB: rad2deg * Math.atan2(b, a),
    qHueAngleRGB: d3.hsl(rgb).h,
  };
};

// State
const state = {
  metric: null,
  includeX: null,
  normQY: null,
  closeCoeff: null,
  numPoke: null,
  searchTerm: null,
  targetColor: null,
  searchResults: 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),
  ],
  ({ yJAB, yRGB }) => [
    vectorSqDist(state.targetColor.qJAB, yJAB),
    vectorSqDist(state.targetColor.qRGB, yRGB),
  ],
  ({ yHueAngleJAB, yHueAngleRGB }) => [
    angleDiff(state.targetColor.qHueAngleJAB, yHueAngleJAB),
    angleDiff(state.targetColor.qHueAngleRGB, yHueAngleRGB),
  ],
  ({ 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, yJAB, yRGB, yJABHat, yJABNorm, yRGBHat, yRGBNorm, yHueAngleJAB, yHueAngleRGB,
}) => {
  // 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: rad2deg * Math.acos(cosAngleJAB),
    angleRGB: rad2deg * Math.acos(cosAngleRGB),
    meanDistJAB: vectorDist(state.targetColor.qJAB, yJAB),
    meanDistRGB: vectorDist(state.targetColor.qRGB, yRGB),
    hueAngleJAB: angleDiff(state.targetColor.qHueAngleJAB, yHueAngleJAB),
    hueAngleRGB: angleDiff(state.targetColor.qHueAngleRGB, yHueAngleRGB),
  };
};

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

const mathDefinitions = {
  "x-definition": String.raw`
    X\left(P\right) = \frac{1}{\left|P\right|}\sum_{p\in P}{\left|\left|\vec{p}\right|\right|^2}
  `,
  "y-definition": String.raw`
    \vec{Y}\left(P\right) = \frac{1}{\left|P\right|}\sum_{p\in P}{\vec{p}}
  `,
  "v-perp-definition": String.raw`
    \vec{v}_{\perp} = \text{oproj}_{\left\{\vec{J}, \vec{L}\right\}}{\vec{v}}
  `,
  "del-h-definition": String.raw`
    \Delta{H} = \angle \left(\vec{q}_{\perp}, \vec{Y}_{\perp}\left(P\right) \right)
  `,
  "cluster-definition": String.raw`
    \left\{P_1, P_2, P_3\right\} = \arg\min_{\left\{P_1, P_2, P_3\right\}} \sum_{i=1}^3 \sum_{p\inP_i} \left|\left| \vec{p} - \vec{Y}\left(P_i\right) \right|\right|^2
  `,
  "z-best-definition": String.raw`
    \vec{Z}_{\text{best}}\left(P\right) = \vec{Y}\left(\arg\min_{P_i} \left|\left| \vec{q} - \vec{Y}\left(P_i\right) \right|\right| \right)
  `,
  "z-worst-definition": String.raw`
    \vec{Z}_{\text{worst}}\left(P\right) = \vec{Y}\left(\arg\max_{P_i} \left|\left| \vec{q} - \vec{Y}\left(P_i\right) \right|\right| \right)
  `,
  "result-definition": String.raw`
    \left(
      \text{RMS}_P\left(q\right), 
      \angle \left(\vec{q}, \vec{Y}\left(P\right)\right), 
      \left|\left| \vec{q} - \vec{Y}\left(P\right) \right|\right|,
      \Delta{H}
    \right)
  `,
};

const metricText = [
  String.raw`\text{RMS}_{P}\left(q\right) ~ \arg\min_{P}\left[X\left(P\right) - 2\vec{q}\cdot \vec{Y}\left(P\right)\right]`,
  String.raw`\angle \left(\vec{q}, \vec{Y}\left(P\right)\right) ~ \arg\max_{P}\left[\cos\left(\angle \left(\vec{q}, \vec{Y}\left(P\right)\right)\right)\right]`,
  String.raw`\left|\left| \vec{q} - \vec{Y}\left(P\right) \right|\right| ~ \arg\min_{P}\left[\left|\left| \vec{q} - \vec{Y}\left(P\right) \right|\right|^2\right]`,
  String.raw`\Delta{H}`,
].map(s => TeXZilla.toMathML(s));

const renderVec = math => String.raw`\vec{${math.charAt(0)}}${math.substr(1)}`;
const renderNorm = vec => String.raw`\frac{${vec}}{\left|\left|${vec}\right|\right|}`;
const updateObjective = () => {
  let tex = metricText?.[state.metric];
  if (!tex) {
    const { includeX, normQY, closeCoeff } = state;
    if (!includeX && closeCoeff === 0) {
      tex = TeXZilla.toMathML(String.raw`\text{Empty Metric}`);
    } else {
      const qyMod = normQY ? c => renderNorm(renderVec(c)) : renderVec;
      tex = TeXZilla.toMathML(String.raw`
        \arg
        \m${includeX ? "in" : "ax"}_{P}
        \left[
          ${includeX ? String.raw`X\left(P\right)` : ""}
          ${closeCoeff === 0 ? "" : String.raw`
              ${includeX ? "-" : ""}
              ${(includeX && closeCoeff !== 1) ? closeCoeff : ""}
              ${qyMod("q")}
              \cdot
              ${qyMod(String.raw`Y\left(P\right)`)}
          `}
        \right]
      `);
    }
  }
  const objFnNode = getObjFnDisplay();
  clearNodeContents(objFnNode);
  objFnNode.appendChild(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 { labelClass = "", rgbClass = "", jabClass = "", tileClass = "" } = classes;
  let { resultsClass = "" } = classes;
  let displayMetrics = {};
  if (!state.targetColor) {
    // no color selected need to skip scores
    resultsClass = "hide";
  } else {
    displayMetrics = calcDisplayMetrics(data);
  }
  const {
    stdDevJAB = 0, stdDevRGB = 0,
    angleJAB = 0, angleRGB = 0,
    meanDistJAB = 0, meanDistRGB,
    hueAngleJAB = 0, hueAngleRGB = 0,
  } = displayMetrics;

  const titleName = name.split("-").map(part => part.charAt(0).toUpperCase() + part.substr(1)).join(" ");
  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 ${tileClass}`);
  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 ${resultsClass}">
          <span class="${jabClass}">
            (${stdDevJAB.toFixed(2)}, ${angleJAB.toFixed(2)}&deg;, ${meanDistJAB.toFixed(2)}, ${hueAngleJAB.toFixed(2)}&deg;)
          </span>
          <span class="${rgbClass}">
            (${stdDevRGB.toFixed(2)}, ${angleRGB.toFixed(2)}&deg;, ${meanDistRGB.toFixed(2)}, ${hueAngleRGB.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);
};

// Update Search Results
const renderSearch = () => {
  const resultsNode = getSearchListNode();
  const append = getPokemonAppender(resultsNode);
  clearNodeContents(resultsNode);
  state.searchResults?.forEach(pkmn => append(pkmn));
};

// Scoring
const rescore = () => {
  if (!state.targetColor) {
    return;
  }

  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", tileClass: "pokemon_tile--smaller" }));

  // 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", tileClass: "pokemon_tile--smaller" }));

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

// 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 === 4) { // 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 = () => {
  state.searchTerm = getNameInputNode()?.value?.toLowerCase() ?? "";
  if (state.searchTerm.length === 0) {
    state.searchResults = [];
  } else {
    state.searchResults = pokemonLookup
      .search(state.searchTerm, { limit: 10 })
      .map(({ item }) => item);
  }
  renderSearch();
};

const onRandomPokemon = () => {
  getNameInputNode().value = "";
  state.searchResults = Array.from({ length: 10 }, () => pokemonColorData[Math.floor(Math.random() * pokemonColorData.length)]);
  renderSearch();
};

const onPageLoad = () => {
  // render static explanations
  Object.entries(mathDefinitions).forEach(([id, tex]) => {
    document.getElementById(id).appendChild(TeXZilla.toMathML(tex));
  });

  // fake some events but don't do any scoring
  onColorChanged(true);
  onMetricChanged(true);
  onLimitChanged(true);
  // then do a rescore directly, which will do nothing unless old data was loaded
  rescore();
  // finally render search in case rescore didn't
  onSearchChanged();
};