// ---- Math and Utilities ----
// Vector Math
const vectorDot = (u, v) => u.map((x, i) => x * v[i]).reduce((x, y) => x + y);
const unitVector = v => {
  const mag = Math.hypot(...v);
  return v.map(c => c / mag);
};

// Angle Math
const rad2deg = 180 / Math.PI;
const twoPi = 2 * Math.PI;

// Color Conversion
const hex2rgb = hex => {
  hex = hex.replace("#", "");
  const red = hex.substr(0, 2);
  const grn = hex.substr(2, 2);
  const blu = hex.substr(4, 2);
  return [red, grn, blu].map(c => parseInt(c, 16));
};

// calculated from analyze.py
RGB_TO_LMS = [
  [0.4121965, 0.53627432, 0.05143268],
  [0.2119195, 0.68071831, 0.10738379],
  [0.08834911, 0.28185414, 0.63018663],
];
LMS_TO_OKLAB = [
  [0.2104542553, 0.793617785, -0.0040720468],
  [1.9779984951, -2.428592205, 0.4505937099],
  [0.0259040371, 0.7827717662, -0.808675766],
];
const hex2oklab = hex => {
  const lrgb = hex2rgb(hex).map(c => {
    const v = c / 255;
    return v <= 0.04045 ? v / 12.92 : Math.pow((v + 0.055) / 1.055, 2.4);
  });
  const lms = RGB_TO_LMS.map(row => Math.cbrt(vectorDot(row, lrgb)));
  return LMS_TO_OKLAB.map(row => vectorDot(row, lms));
};

// Misc
const clamp = (min, value, max) => Math.min(Math.max(value, min), max);
const productLift =
  (...factors) =>
  (...args) =>
    factors
      .filter(fn => !!fn)
      .map(fn => fn(...args))
      .reduce((x, y) => x * y, 1);

// Pre-computation
const getColorData = hex => {
  const lab = hex2oklab(hex);
  return {
    hex,
    vector: lab,
    unit: unitVector(lab),
    chroma: Math.hypot(lab[1], lab[2]),
    hue: (Math.atan2(lab[2], lab[1]) + twoPi) % twoPi,
  };
};

const pokemonData = database.map(({ total, clusters, ...rest }) => ({
  total: {
    ...total,
    hex: total.medianHex,
    stddevTotal: Math.hypot(...total.stddev),
    stddevL: total.stddev[0],
    stddevA: total.stddev[1],
    stddevB: total.stddev[2],
    chromaMean: total.chroma[0],
    chromaDev: total.chroma[1],
    unitCentroid: unitVector(total.centroid),
    proportionDisplay: null,
    proportion: 1,
    inverseProportion: 1,
  },
  clusters: clusters
    .sort((a, b) => b.size - a.size)
    .map(c => ({
      ...c,
      hex: c.medianHex,
      stddevTotal: Math.hypot(...c.stddev),
      stddevL: c.stddev[0],
      stddevA: c.stddev[1],
      stddevB: c.stddev[2],
      chromaMean: c.chroma[0],
      chromaDev: c.chroma[1],
      unitCentroid: unitVector(c.centroid),
      proportion: c.size / total.size,
      inverseProportion: total.size / c.size,
    }))
    .map(c => ({
      ...c,
      proportionDisplay: (100 * c.proportion).toFixed(2),
    })),
  ...rest,
}));
const pokemonLookup = new Fuse(pokemonData, { keys: ["name"] });
const pokemonDisplayLookup = Object.fromEntries(
  database.map(({ name, species }) => {
    const formattedName = name
      .split("-")
      .map(part => part.charAt(0).toUpperCase() + part.substr(1))
      .join(" ")
      .replace(/ M$/, " ♂")
      .replace(/ F$/, " ♀");
    let spriteName = name.toLowerCase();
    if (
      ["alcremie", "minior", "wobbuffet", "hippopotas", "hippowdon", "unfezant"].find(s =>
        spriteName.includes(s)
      )
    ) {
      spriteName = species.toLowerCase();
    } else if (spriteName === "basculin-white-striped") {
      // pdb has the wrong filenames lol
      spriteName = "1x/basculin-red-striped";
    } else if (!spriteName.includes("nidoran-")) {
      spriteName = spriteName
        .replace("-gmax", "-gigantamax")
        .replace("-alola", "-alolan")
        .replace("-galar", "-galarian")
        .replace("-hisui", "-hisuian")
        .replace("-paldea", "-paldean")
        .replace("-paldean-combat", "-paldean") // tauros
        .replace("-paldean-blaze", "-paldean-fire") // tauros
        .replace("-paldean-aqua", "-paldean-water") // tauros
        .replace("-pokeball", "-poke-ball") // vivillon
        .replaceAll("é", "e") // flabébé
        .replace(/darmanitan-galarian$/, "darmanitan-galarian-standard")
        .replace("urshifu-gigantamax", "urshifu-gigantamax-single-strike")
        .replace("urshifu-rapid-strike-gigantamax", "urshifu-gigantamax-rapid-strike")
        .replace("calyrex-shadow", "calyrex-shadow-rider")
        .replace("calyrex-ice", "calyrex-ice-rider")
        .replace("maushold-four", "maushold-family4")
        .replace("chienpao", "chien-pao")
        .replace("tinglu", "ting-lu")
        .replace("wochien", "wo-chien")
        .replace("chiyu", "chi-yu")
        .replace(/-m$/, "")
        .replace(/-f$/, "-female")
        .replaceAll(/['%]+/g, "")
        .replaceAll(/[:.\s]+/g, "-")
        .replace(/-+$/, "");
    }

    const linkName = species.replace(/[':.]/, "").replace(" ", "-");

    const handlerChain = [
      `https://img.pokemondb.net/sprites/sword-shield/normal/${spriteName}.png`,
      `https://img.pokemondb.net/sprites/scarlet-violet/icon/${spriteName}.png`,
      `https://img.pokemondb.net/sprites/sword-shield/icon/${spriteName}.png`,
    ].map((url, i) => ({ target }) => {
      target.removeEventListener("error", handlerChain[i]);
      if (handlerChain[i + 1]) {
        target.addEventListener("error", handlerChain[i + 1]);
      }
      pokemonDisplayLookup[name].sprite = target.src = url;
    });

    return [
      name,
      {
        formattedName,
        link: `https://pokemondb.net/pokedex/${linkName}`,
        sprite: `https://img.pokemondb.net/sprites/scarlet-violet/normal/${spriteName}.png`,
        handleSpriteError: handlerChain[0],
      },
    ];
  })
);

const calcScores = (data, target) => {
  const { centroid, unitCentroid, tilt, stddev, stddevTotal, chromaMean, chromaDev, hue } = data;
  const deltaVec = centroid.map((c, i) => c - target.vector[i]);
  const delta = Math.hypot(...deltaVec);
  const psiVec = stddev.map((s, i) => Math.hypot(s, deltaVec[i]));
  const psi = Math.hypot(...psiVec);
  const theta = Math.acos(Math.min(1, vectorDot(tilt, target.unit))) * rad2deg;
  const deltaC = chromaMean - target.chroma;
  const deltaH = Math.abs(hue - target.hue) * rad2deg;
  const likelihood = Math.exp(-Math.pow(delta / (2 * stddevTotal), 2)) / stddevTotal;
  const arcDiff = psi * theta;

  return {
    ...data,
    probArcDiff: arcDiff / likelihood, 
    arcDiff,
    likelihood,

    psi,
    psiVec,
    psiL: psiVec[0],
    psiA: psiVec[1],
    psiB: psiVec[2],
    psiC: Math.hypot(chromaDev, deltaC),

    theta,

    delta,
    deltaVec,
    deltaL: Math.abs(deltaVec[0]),
    deltaA: Math.abs(deltaVec[1]),
    deltaB: Math.abs(deltaVec[2]),
    deltaC: Math.abs(deltaC),
    deltaH: Math.min(deltaH, 360 - deltaH),

    angDiff: Math.acos(vectorDot(unitCentroid, target.unit)) * rad2deg,
    hue: hue * rad2deg,
  };
};

const sortOrders = {
  max: (a, b) => b - a,
  min: (a, b) => a - b,
};

// who needs a framework?
const makeTemplate = (id, definition = () => ({})) => {
  const content = document.getElementById(id).content;
  return (...args) => {
    const fragment = content.cloneNode(true);
    const binds = Object.fromEntries(
      Array.from(fragment.querySelectorAll("[bind-to]")).map(element => {
        const name = element.getAttribute("bind-to");
        element.removeAttribute("bind-to");
        return [name, element];
      })
    );
    const children = Array.from(content.children);
    const update = (...args) => {
      Object.entries(definition(...args))
        .map(([name, settings]) => [binds[name], settings])
        .filter(([bind]) => !!bind)
        .forEach(([bind, settings]) =>
          Object.entries(settings).forEach(([setting, value]) => {
            if (setting.startsWith("@")) {
              bind.addEventListener(setting.slice(1), value);
            } else if (setting.startsWith("--")) {
              bind.style.setProperty(setting, value);
            } else if (setting === "dataset") {
              Object.entries(value).forEach(([key, data]) => (bind.dataset[key] = data));
            } else if (setting === "append") {
              if (Array.isArray(value)) {
                bind.append(...value);
              } else {
                bind.append(value);
              }
            } else {
              bind[setting] = value;
            }
          })
        );
    };
    update(...args);
    return {
      mount: root =>
        (typeof root === "string" ? document.querySelector(root) : root).append(fragment),
      unmount: () => children.forEach(c => c.remove()),
      update,
      children,
      firstElementChild: fragment.firstElementChild,
      fragment,
    };
  };
};

// ---- Selectors ----

const rootStyle = document.querySelector(":root").style;
const colorSearchResultsTarget = document.getElementById("color-results");
const nameSearchResultsTarget = document.getElementById("name-results");
const prevColorsSidebar = document.getElementById("prevColors");
const clusterMetricSection = document.getElementById("cls-metric-mount");
const clusterFunctionSection = document.getElementById("cls-fn");

const colorCalculateForm = document.forms.colorCalculateForm;
const colorSortForm = document.forms.colorSortForm;
const targetColorElements = document.forms.targetColorForm.elements;
const colorDisplayElements = document.forms.colorDisplayForm.elements;
const filterElements = document.forms.filterControl.elements;
const nameSearchFormElements = document.forms.nameSearchForm.elements;

// ---- Add Metric Selects ----

const createMetricSelect = makeTemplate("metric-select-template");
const { firstElementChild: sortMetricForm, mount: mountSortMetricForm } =
  createMetricSelect();
sortMetricForm.elements.metricKind.value = "compare";
mountSortMetricForm("#sort-metric-mount");

const { firstElementChild: clusterMetricForm, mount: mountClusterMetricForm } =
  createMetricSelect();
clusterMetricForm.elements.metricKind.value = "stat";
mountClusterMetricForm("#cls-metric-mount");

const updateMetricSelects = form => {
  const kind = form.elements.metricKind.value;
  form.elements.compare.disabled = kind !== "compare";
  form.elements.stat.disabled = kind !== "stat";
  form.elements.metric.value = form.elements[kind].value;
};

// bit of a hack, but lets us control this all from the template
const metricSymbols = Object.fromEntries(
  Array.from(document.querySelectorAll("option")).map(el => [
    el.value,
    el.textContent.split("(")[1].split(")")[0],
  ])
);

const updateMetricDisplays = () => {
  updateMetricSelects(sortMetricForm);
  updateMetricSelects(clusterMetricForm);

  colorCalculateForm.elements.sortMetricSymbolP.value =
    colorCalculateForm.elements.sortMetricSymbolB.value =
      metricSymbols[
        sortMetricForm.elements[sortMetricForm.elements.metricKind.value].value
      ];

  colorCalculateForm.elements.clusterMetricSymbol.value =
    metricSymbols[
      clusterMetricForm.elements[clusterMetricForm.elements.metricKind.value].value
    ];
};

// ---- Styling ----

const getColorStyles = hex => {
  const rgb = hex2rgb(hex);
  const lum = vectorDot(rgb, [0.3, 0.6, 0.1]);
  const highlight = lum >= 128 ? "var(--color-dark)" : "var(--color-light)";
  return {
    "--highlight": highlight,
    "--background": hex,
    "--shadow-component": lum >= 128 ? "0" : "255",
  };
};

const setColorStyles = (style, hex) =>
  Object.entries(getColorStyles(hex)).forEach(([prop, value]) =>
    style.setProperty(prop, value)
  );

// ---- Pokemon Display ----

// pulled out bc the render uses them
const metricScores = {};
const bestClusterIndices = {};
const objectiveValues = {};

const createPokemonTooltip = makeTemplate("pkmn-data-template", data =>
  Object.fromEntries(
    Object.entries(data).map(([metric, value]) => [
      metric,
      {
        innerText: Array.isArray(value)
          ? value.map(v => v.toFixed(2)).join(", ")
          : value?.toFixed?.(3)?.replace(".000", ""),
      },
    ])
  )
);

const createPokemonTile = makeTemplate(
  "pkmn-tile-template",
  ({ name, color }, enableTotalFlags, enableClusterFlags) => {
    const { formattedName, link, sprite, handleSpriteError } = pokemonDisplayLookup[name];
    const image = {
      alt: formattedName,
      src: sprite,
      "@error": handleSpriteError,
    };

    const score = {
      innerText: objectiveValues[name].toFixed(2),
    };

    const { total, clusters } = metricScores[name];
    const buttonBinds = [
      [clusters[0], "cls1Btn", "cls1Data", "infoHover1"],
      [clusters[1], "cls2Btn", "cls2Data", "infoHover2"],
      [clusters[2], "cls3Btn", "cls3Data", "infoHover3"],
      [clusters[3], "cls4Btn", "cls4Data", "infoHover4"],
      [total, "totalBtn", "totalData", "infoHover"],
    ]
      .filter(([data]) => !!data)
      .map(([data, button, tooltip, infoHover], index) => ({
        [button]: {
          dataset: {
            included: enableClusterFlags && index === bestClusterIndices[name],
          },
          hidden: false,
          innerText:
            data.hex + (data.proportionDisplay ? ` - ${data.proportionDisplay}%` : ""),
          "@click"() {
            model.setTargetColor(data.hex);
          },
          ...getColorStyles(data.hex),
        },
        [tooltip]: {
          append: createPokemonTooltip(data).fragment,
        },
        [infoHover]: {
          "@mouseenter"({ target }) {
            target.closest(".pkmn-tile").dataset.info = `show-${infoHover}`;
          },
          "@mouseleave"({ target }) {
            target.closest(".pkmn-tile").dataset.info = "hide";
          },
        },
      }))
      .reduce((a, b) => ({ ...a, ...b }), {});
    buttonBinds.totalBtn.dataset.included = enableTotalFlags;

    return {
      name: {
        innerText: formattedName,
        title: formattedName,
      },
      image,
      color: { innerText: color },
      link: { href: link },
      score,
      tileRoot: { dataset: { info: "hide" } },
      ...buttonBinds,
    };
  }
);

const renderPokemon = (list, target) => {
  target.innerText = "";

  const {
    sortUseWholeImage,
    sortUseBestCluster,
    sortUseClusterProportion,
    sortUseInvClusterProportion,
  } = Object.fromEntries(new FormData(colorCalculateForm).entries());

  const enableClusterFlags = !!(
    sortUseBestCluster ||
    sortUseClusterProportion ||
    sortUseInvClusterProportion
  );

  target.append(
    ...list.map(
      pkmn => createPokemonTile(pkmn, !!sortUseWholeImage, enableClusterFlags).fragment
    )
  );
};

// ---- Calculation Logic ----

const model = {
  setTargetColor(newColor) {
    const hex = `#${newColor?.replace("#", "")}`;
    if (hex.length !== 7) {
      return;
    }

    setColorStyles(rootStyle, hex);

    const oldColor = this.targetColor;
    this.targetColor = hex;

    targetColorElements.colorText.value = hex;
    targetColorElements.colorText.dataset.lastValid = hex;
    targetColorElements.colorPicker.value = hex;

    if (oldColor) {
      const prevButton = document.createElement("button");
      prevButton.innerText = oldColor;
      prevButton.classList = "color-select";
      setColorStyles(prevButton.style, oldColor);
      prevButton.addEventListener("click", () => this.setTargetColor(oldColor));
      prevColorsSidebar.prepend(prevButton);
    }

    const targetData = getColorData(hex);
    targetColorElements.info.value = `
      (${(targetData.vector[0] * 100).toFixed(2)}%,
      ${targetData.chroma.toFixed(4)}, 
      ${(targetData.hue * rad2deg).toFixed(1)}°)
    `;

    pokemonData.forEach(({ name, total, clusters }) => {
      metricScores[name] = {
        total: calcScores(total, targetData),
        clusters: clusters.map(c => calcScores(c, targetData)),
      };
    });

    this.calculateObjective();
  },

  calculateObjective() {
    const {
      clusterUseClusterSize,
      clusterUseInvClusterSize,
      clusterSortOrder,
      sortUseWholeImage,
      sortUseBestCluster,
      sortUseClusterProportion,
      sortUseInvClusterProportion,
    } = Object.fromEntries(new FormData(colorCalculateForm).entries());

    const clsMetric = clusterMetricForm.elements.metric.value;
    const getClusterScore = productLift(
      cluster => cluster[clsMetric],
      clusterUseClusterSize && (cluster => cluster.size),
      clusterUseInvClusterSize && (cluster => cluster.inverseSize)
    );

    const clsSort = sortOrders[clusterSortOrder];
    const getBestClusterIndex = ({ total, clusters }) =>
      clusters
        .map((c, i) => [getClusterScore(c, total), i])
        .reduce((a, b) => (clsSort(a[0], b[0]) > 0 ? b : a))[1];

    Object.entries(metricScores).forEach(([name, scores]) => {
      bestClusterIndices[name] = getBestClusterIndex(scores);
    });

    const metric = sortMetricForm.elements.metric.value;
    const getSortScoreClusterFactors = productLift(
      sortUseBestCluster && (({ clusters }, i) => clusters[i][metric]),
      sortUseClusterProportion && (({ clusters }, i) => clusters[i].proportion),
      sortUseInvClusterProportion && (({ clusters }, i) => clusters[i].inverseProportion)
    );
    const usingCluster =
      sortUseBestCluster || sortUseClusterProportion || sortUseInvClusterProportion;
    const getSortScore = (scores, i) =>
      (sortUseWholeImage ? scores.total[metric] : 0) +
      (usingCluster ? getSortScoreClusterFactors(scores, i) : 0);

    Object.entries(metricScores).forEach(([name, scores]) => {
      objectiveValues[name] = getSortScore(scores, bestClusterIndices[name]);
    });

    this.renderNameSearchResults();
    this.rank();
  },

  rank() {
    const { sortOrder } = Object.fromEntries(new FormData(colorSortForm).entries());
    const compare = sortOrders[sortOrder];
    const sortFn = (a, b) => compare(objectiveValues[a], objectiveValues[b]);

    this.ranked = pokemonData
      .slice(0)
      .sort((a, b) => sortFn(a.name, b.name) || a.name.localeCompare(b.name));

    this.renderColorSearchResults();
  },

  setNameSearchResults(newNameResults) {
    this.nameSearchResults = newNameResults;
    this.renderNameSearchResults();
  },

  renderNameSearchResults() {
    renderPokemon(this.nameSearchResults ?? [], nameSearchResultsTarget);
  },

  resultsToDisplay: 6,

  setResultsToDisplay(count) {
    colorDisplayElements.output.value = this.resultsToDisplay = clamp(0, count, 100);
    this.renderColorSearchResults();
  },

  renderColorSearchResults() {
    const dexNums = new Set();
    const toRender = [];
    for (const pkmn of this.ranked) {
      if (toRender.length >= this.resultsToDisplay) {
        break;
      }
      if (filterElements.hideNoStart.checked && pkmn.traits.includes("nostart")) {
        continue;
      }
      if (!filterElements.allowNFE.checked && pkmn.traits.includes("nfe")) {
        continue;
      }
      if (!filterElements.allowRepeatDexNum.checked) {
        if (dexNums.has(pkmn.num)) {
          continue;
        }
        dexNums.add(pkmn.num);
      }
      toRender.push(pkmn);
    }
    renderPokemon(toRender, colorSearchResultsTarget);
  },
};

// ---- Form Controls ----

nameSearchFormElements.input.addEventListener("input", ({ target: { value } }) => {
  model.setNameSearchResults(
    pokemonLookup.search(value, { limit: 36 }).map(({ item }) => item)
  );
});

nameSearchFormElements.all.addEventListener("click", () => {
  nameSearchFormElements.input.value = "";
  model.setNameSearchResults(pokemonData.slice(0));
});

nameSearchFormElements.clear.addEventListener("click", () => {
  nameSearchFormElements.input.value = "";
  model.setNameSearchResults([]);
});

nameSearchFormElements.random.addEventListener("click", () => {
  model.setNameSearchResults(
    Array.from(
      { length: 36 },
      () => pokemonData[Math.floor(Math.random() * pokemonData.length)]
    )
  );
});

targetColorElements.colorText.addEventListener("input", ({ target }) => {
  if (target.willValidate && !target.validity.valid) {
    target.value = target.dataset.lastValid || "";
  } else {
    model.setTargetColor(target.value);
  }
});

targetColorElements.colorPicker.addEventListener("change", ({ target }) =>
  model.setTargetColor(target.value)
);

const randomizeTargetColor = () =>
  model.setTargetColor(
    [Math.random(), Math.random(), Math.random()]
      .map(component =>
        Math.floor(component * 256)
          .toString(16)
          .padStart(2, "0")
      )
      .reduce((x, y) => x + y)
  );

targetColorElements.randomColor.addEventListener("click", randomizeTargetColor);

colorDisplayElements.lessResults.addEventListener("click", () => {
  model.setResultsToDisplay(model.resultsToDisplay - 6);
});

colorDisplayElements.moreResults.addEventListener("click", () => {
  model.setResultsToDisplay(model.resultsToDisplay + 6);
});

Array.from(filterElements).forEach(el =>
  el.addEventListener("change", () => model.renderColorSearchResults())
);

Array.from(colorSortForm.elements).forEach(el =>
  el.addEventListener("change", () => model.rank())
);

Array.from(colorCalculateForm.elements).forEach(el =>
  el.addEventListener("change", () => {
    const { sortUseBestCluster, sortUseClusterProportion, sortUseInvClusterProportion } =
      Object.fromEntries(new FormData(colorCalculateForm).entries());
    clusterFunctionSection.dataset.faded = !(
      sortUseBestCluster ||
      sortUseClusterProportion ||
      sortUseInvClusterProportion
    );
    model.calculateObjective();
  })
);

sortMetricForm.addEventListener("change", () => {
  updateMetricDisplays();
  model.calculateObjective();
});

clusterMetricForm.addEventListener("change", () => {
  updateMetricDisplays();
  model.calculateObjective();
});

// ---- Initial Setup ----

updateMetricDisplays();
randomizeTargetColor();