// Selectors + DOM Manipulation
const getColorInputNode = () => document.getElementById("color-input");
const getMetricDropdownNode = () => document.getElementById("metric");
const getMeanArgumentDropdownNode = () => document.getElementById("image-summary");
const getClusterMeanWarning = () => document.getElementById("cluster-mean-warning");
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;

// Conversions
const jab2hex = jab => d3.jab(...jab).formatHex();
const rgb2hex = rgb => d3.rgb(...rgb).formatHex();
const jab2hue = ([, a, b]) => rad2deg * Math.atan2(b, a);
const rgb2hue = rgb => d3.hsl(d3.rgb(...rgb)).h || 0;

// Arg Compare
const argComp = comp => ra => ra.map((x, i) => [x, i]).reduce((a, b) => comp(a[0], b[0]) > 0 ? b : a)[1];
const argMin = argComp((a, b) => a - b);
const argMax = argComp((a, b) => b - a);

// Pre-Compute Data
const computeVectorData = (vector, toHex, toHue) => {
  const [ magnitude, unit ] = vectorNorm(vector);
  return {
    vector,
    magnitude,
    magSq: magnitude * magnitude,
    unit,
    hex: toHex(vector),
    hue: toHue(vector),
  };
};

const computeStats = (varFromZero, trueMeanVec, kMeanStruct, toHex, toHue) => ({
  varFromZero,
  trueMean: computeVectorData(trueMeanVec, toHex, toHue),
  kMeans: kMeanStruct.slice(0, 3).map(z => computeVectorData(z, toHex, toHue)),
  kWeights: kMeanStruct[3],
  largestCluster: argMax(kMeanStruct[3]),
  smallestCluster: argMin(kMeanStruct[3]),
});

const pokemonColorData = database.map(({
  name, xJAB, xRGB, yJAB, yRGB, zJAB, zRGB,
}) => ({
  name,
  jabStats: computeStats(xJAB, yJAB, zJAB, jab2hex, jab2hue),
  rgbStats: computeStats(xRGB, yRGB, zRGB, rgb2hex, rgb2hue),
}));

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);

  return {
    jabData: computeVectorData([ J, a, b ], jab2hex, jab2hue),
    rgbData: computeVectorData([ rgb.r, rgb.g, rgb.b ], rgb2hex, rgb2hue),
  };
};

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

// Metrics
const summarySelectors = [
  // true mean
  stats => stats.trueMean,
  // largest cluster
  stats => stats.kMeans[stats.largestCluster],
  // smallest cluster
  stats => stats.kMeans[stats.smallestCluster],
  // best fit cluster
  (stats, q) => stats.kMeans[argMin(stats.kMeans.map((z, i) => vectorSqDist(z.vector, q.vector) / stats.kWeights[i]))],
  // worst fit cluster
  (stats, q) => stats.kMeans[argMax(stats.kMeans.map((z, i) => vectorSqDist(z.vector, q.vector) / stats.kWeights[i]))],
];

const selectedSummary = (stats, q) => summarySelectors[state.meanArgument](stats, q);

const metrics = [
  // RMS
  (stats, q) => stats.varFromZero - 2 * vectorDot(selectedSummary(stats, q).vector, q.vector),
  // mean angle
  (stats, q) => -vectorDot(selectedSummary(stats, q).unit, q.unit),
  // mean dist
  (stats, q) => vectorSqDist(selectedSummary(stats, q).vector, q.vector),
  // hue angle
  (stats, q) => angleDiff(selectedSummary(stats, q).hue, q.hue),
  // custom
  (stats, q) => (state.includeX ? stats.varFromZero : 0) - state.closeCoeff * vectorDot(
    selectedSummary(stats, q)[state.normQY ? "unit" : "vector"], 
    state.normQY ? q.unit : q.vector,
  ),
];

const scorePokemon = pkmn => ({
  jab: metrics[state.metric](pkmn.jabStats, state.targetColor.jabData),
  rgb: metrics[state.metric](pkmn.rgbStats, state.targetColor.rgbData),
});

const calcDisplayMetrics = ({ jabStats, rgbStats }) => {
  // TODO - case on metric and meanArgument to avoid recalculation
  // TODO - is there ever any value to computing these around the selected summary instead?
  // obviously that has no mathematical value, and screws up the sqrts, but maybe?

  const cosAngleJAB = vectorDot(state.targetColor.jabData.unit, jabStats.trueMean.unit);
  const yTermJAB = cosAngleJAB * jabStats.trueMean.magnitude * state.targetColor.jabData.magnitude;

  const cosAngleRGB = vectorDot(state.targetColor.rgbData.unit, rgbStats.trueMean.unit);
  const yTermRGB = cosAngleRGB * rgbStats.trueMean.magnitude * state.targetColor.rgbData.magnitude;

  return {
    stdDevJAB: Math.sqrt(jabStats.varFromZero - 2 * yTermJAB + state.targetColor.jabData.magSq),
    stdDevRGB: Math.sqrt(rgbStats.varFromZero - 2 * yTermRGB + state.targetColor.rgbData.magSq),
    angleJAB: rad2deg * Math.acos(cosAngleJAB),
    angleRGB: rad2deg * Math.acos(cosAngleRGB),
    meanDistJAB: vectorDist(state.targetColor.jabData.vector, jabStats.trueMean.vector),
    meanDistRGB: vectorDist(state.targetColor.rgbData.vector, rgbStats.trueMean.vector),
    hueAngleJAB: angleDiff(state.targetColor.jabData.hue, jabStats.trueMean.hue),
    hueAngleRGB: angleDiff(state.targetColor.rgbData.hue, rgbStats.trueMean.hue),
  };
};

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

const mathArgBest = (mxn, arg) => `\\underset{${arg}}{\\arg\\${mxn}}`;

const mathDefinitions = {
  "main-definition": String.raw`
    \begin{aligned}
      \vec{\mu}\left(P\right) &= \frac{1}{\left|P\right|}\sum_{p\in P}{\vec{p}} \\
      I\left(P\right) &= \frac{1}{\left|P\right|}\sum_{p\in P}{\left|\left|\vec{p}\right|\right|^2} \\
      \vec{x}_{\perp} &= \text{oproj}_{\left\{\vec{J}, \vec{L}\right\}}{\vec{x}} \\
      \Delta{\theta}\left(P\right) &= \angle \left(\vec{q}_{\perp}, \vec{\mu}\left(P\right)_{\perp} \right)
    \end{aligned}
  `,
  "k-definition": String.raw`
    \begin{aligned}
      K_{\text{big}}\left(P\right) &= ${mathArgBest("max", "P_i")} \frac{\left|P_i\right|}{\left|P\right|} \\
      K_{\text{small}}\left(P\right) &= ${mathArgBest("min", "P_i")} \frac{\left|P_i\right|}{\left|P\right|} \\
      K_{\text{best}}\left(P\right) &= ${mathArgBest("min", "P_i")} \frac{\left|P\right|}{\left|P_i\right|} \left|\left| \vec{q} - \vec{\mu}\left(P_i\right) \right|\right| \\
      K_{\text{worst}}\left(P\right) &= ${mathArgBest("max", "P_i")} \frac{\left|P\right|}{\left|P_i\right|} \left|\left| \vec{q} - \vec{\mu}\left(P_i\right) \right|\right|
    \end{aligned}
  `,
  "cluster-definition": String.raw`
    \begin{aligned}
      \left\{P_1, P_2, P_3\right\} = ${mathArgBest("max", String.raw`\left\{P_1, P_2, P_3\right\}`)} \sum_{i=1}^3 \sum_{p\inP_i} \left|\left| \vec{p} - \vec{\mu}\left(P_i\right) \right|\right|^2
    \end{aligned}
  `,
  "rms-definition": String.raw`
    \text{RMS}_{P}\left(q\right) = \sqrt{E\left[\left|\left|\vec{q} - \vec{p}\right|\right|^2\right]} = \sqrt{\frac{1}{|P|}\sum_{p \in P}{\left|\left|\vec{p} - \vec{q}\right|\right|^2}}
  `,
  "result-definition": String.raw`
    \left(
      \text{RMS}_P\left(q\right),
      \angle \left(\vec{q}, \vec{\mu}\left(P\right)\right),
      \left|\left| \vec{q} - \vec{\mu}\left(P\right) \right|\right|,
      \Delta{\theta}\left(P\right)
    \right)
  `,
};

const metricText = [
  muArg => String.raw`${mathArgBest("min", "P")}\left[I\left(P\right) - 2\vec{q}\cdot \vec{\mu}\left(${muArg}\right)\right]`,
  muArg => String.raw`${mathArgBest("max", "P")}\left[\cos\left(\angle \left(\vec{q}, \vec{\mu}\left(${muArg}\right)\right)\right)\right]`,
  muArg => String.raw`${mathArgBest("min", "P")}\left[\left|\left| \vec{q} - \vec{\mu}\left(${muArg}\right) \right|\right|^2\right]`,
  muArg => String.raw`${mathArgBest("min", "P")} \left[\angle \left(\vec{q}_{\perp}, \vec{\mu}\left(${muArg}\right)_{\perp} \right)\right]`,
].map(s => muArg => TeXZilla.toMathML(s(muArg)));

const muArgs = [
  "P", 
  String.raw`K_{\text{big}}\left(P\right)`, 
  String.raw`K_{\text{small}}\left(P\right)`, 
  String.raw`K_{\text{best}}\left(P\right)`, 
  String.raw`K_{\text{worst}}\left(P\right)`,
];

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 = () => {
  const muArg = muArgs[state.meanArgument];
  let tex = metricText?.[state.metric]?.(muArg);
  if (!tex) {
    const { includeX, normQY, closeCoeff } = state;
    if (!includeX && closeCoeff === 0) {
      tex = TeXZilla.toMathML(String.raw`\text{Empty Metric}`);
    } else {
      const qyMod = normQY ? renderNorm : c => c;
      tex = TeXZilla.toMathML(String.raw`
        ${mathArgBest(includeX ? "min" : "max", "P")}
        \left[
          ${includeX ? String.raw`I\left(P\right)` : ""}
          ${closeCoeff === 0 ? "" : String.raw`
              ${includeX ? "-" : ""}
              ${(includeX && closeCoeff !== 1) ? closeCoeff : ""}
              ${qyMod("\\vec{q}")}
              \cdot
              ${qyMod(String.raw`\vec{\mu}\left(${muArg}\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, jabStats, rgbStats, scores } = 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 = 0,
    hueAngleJAB = 0, hueAngleRGB = 0,
  } = displayMetrics;

  const titleName = name.split("-").map(part => part.charAt(0).toUpperCase() + part.substr(1)).join(" ");
  const textHex = getContrastingTextColor(rgbStats.trueMean.vector);
  const rgbVec = rgbStats.trueMean.vector.map(c => c.toFixed()).join(", ");
  const jabVec = jabStats.trueMean.vector.map(c => c.toFixed(1)).join(", ");

  // TODO Z dists, Z colors

  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} ${scores?.jab?.toFixed(2) ?? ""} ${scores?.rgb?.toFixed(2) ?? ""}
      </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: ${jabStats.trueMean.hex}; color: ${textHex}">
            <span>${jabStats.trueMean.hex}</span><span class="pokemon_tile-vector">(${jabVec})</span>
          </div>
          <div class="pokemon_tile-hex_color ${rgbClass}" style="background-color: ${rgbStats.trueMean.hex}; color: ${textHex}">
            <span>${rgbStats.trueMean.hex}</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;
  }

  // TODO might like to save this somewhere instead of recomputing when limit changes
  const scores = pokemonColorData.map(data => ({ ...data, scores: scorePokemon(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.jab - b.scores.jab)
    .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.rgb - b.scores.rgb)
    .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.jabData.vector.map(c => c.toFixed(2)), getQJABDisplay(), "Jab");
    renderQVec(state.targetColor.rgbData.vector.map(c => c.toFixed()), getQRGBDisplay(), "RGB");

    const textColor = getContrastingTextColor(state.targetColor.rgbData.vector);
    document.querySelector("body").setAttribute("style", `background: ${state.targetColor.rgbData.hex}; 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 checkClusterMeanWarning = () => {
  const warning = getClusterMeanWarning();
  const unhidden = warning.getAttribute("class").replaceAll("hide", "");
  if (state.meanArgument !== 0 && state.metric === 0) {
    warning.setAttribute("class", unhidden);
  } else {
    warning.setAttribute("class", unhidden + " hide");
  }
}

const onMeanArgumentChanged = skipScore => {
  const meanArgument = getMeanArgumentDropdownNode()?.selectedIndex ?? 0;
  if (meanArgument === state.meanArgument) {
    return;
  }
  state.meanArgument = meanArgument;
  checkClusterMeanWarning();
  updateObjective();
  if (!skipScore) {
    rescore();
  }
}

const onMetricChanged = skipScore => {
  const metric = getMetricDropdownNode()?.selectedIndex ?? 0;
  if (metric === state.metric) {
    return;
  }
  state.metric = metric;
  checkClusterMeanWarning();
  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);
  onMeanArgumentChanged(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();
};