update

#!/usr/bin/env node
'use strict';

const fs = require('fs');
const util = require('util');
const stringify = require('json-stringify-pretty-compact');
const zlib = require('zlib');
const gzip = util.promisify(zlib.gzip);

const { Dex, Teams } = require('./vendor/pokemon-showdown');
Dex.includeModData();

const SEED = [1, 2, 3, 4];
const N = +process.argv[2] || 100_000;
const SKIP = [
  'factory', 'unrated', 'cup', 'staff', 'cap', 'monotype', 'chimera',
  'blitz', 'generated', 'joltemons', 'pokebilities', 'firstblood',
];
const ORDER = ['hp', 'atk', 'def', 'spa', 'spd', 'spe'];
const FIELDS = {
  gen9: ['role', 'movepool', 'teraTypes', 'abilities'],
  roles: ['role', 'movepool', 'preferredTypes', 'abilities'],
  gen1: ['level', 'moves', 'essentialMoves', 'exclusiveMoves', 'comboMoves'],
  other: ['level', 'moves', 'doublesLevel', 'doublesMoves', 'noDynamaxMoves'],
};
const ROLES = [
  'gen7randombattle', 'gen6randombattle', 'gen5randombattle',
  'gen4randombattle', 'gen3randombattle', 'gen2randombattle',
];
const CMP = {
  stats: (a, b) => ORDER.indexOf(a[0]) - ORDER.indexOf(b[0]),
  str: (a, b) => a[0].localeCompare(b[0]),
  entries: (a, b) => b[1] - a[1] || a[0].localeCompare(b[0]),
};

const verify = (data, type) => {
  for (const species in data) {
    if (type === 'gen9' || type === 'roles') {
      for (const set of data[species].sets) {
        for (const field of Object.keys(set)) {
          if (!FIELDS[type].includes(field)) {
            throw new Error(`Unexpected field in set: '${field}'`);
          }
        }
      }
    } else {
      for (const field of Object.keys(data[species])) {
        if (!FIELDS[type].includes(field)) {
          throw new Error(`Unexpected field in data: '${field}'`);
        }
      }
    }
  }
  return data;
};

const getData = (format, doubles, baby) => format.id.startsWith('gen9')
  ? verify(require(`./vendor/pokemon-showdown/data/random-battles/${format.mod}${baby
      ? 'baby' : ''}/${doubles ? 'doubles-' : ''}sets.json`), 'gen9')
  : ROLES.includes(format.id) ?
    verify(require(`./vendor/pokemon-showdown/data/random-battles/${format.mod}/sets.json`),
      'roles')
    : format.id === 'gen7randomdoublesbattle' ?
      verify(require(`./vendor/pokemon-showdown/data/mods/${format.mod}/random-doubles-data.json`),
        'other')
      : verify(require(`./vendor/pokemon-showdown/data/random-battles/${format.mod}/data.json`),
        format.mod === 'gen1' ? 'gen1' : 'other');

const getForme = (dex, pool, set) => {
  if (dex.gen === 7 && set.species === 'Rayquaza' &&
    set.moves.includes('dragonascent') && !dex.items.get(set.item).zMove) {
    return dex.species.get('Rayquaza-Mega');
  }
  if (set.gigantamax) {
    return dex.species.get(set.species.startsWith('Pikachu')
      ? 'Pikachu-Gmax'
      : `${set.species}-Gmax`);
  }

  const item = dex.items.get(set.item);
  if (item.megaStone) return dex.species.get(item.megaStone);

  const formes = Object.keys(pool[set.species] || pool[set.name]);
  if (formes.length === 1) return dex.species.get(formes[0]);

  const matches = [];
  for (const f in pool[set.species]) {
    const forme = dex.species.get(f);
    if (forme.requiredItem) {
      if (set.item !== forme.requiredItem) continue;
      return forme;
    }
    if (forme.requiredAbility && set.ability !== forme.requiredAbility) continue;
    if (set.moves.every(m => pool[set.species][f].has(m))) {
      matches.push(forme);
    }
  }

  // Certain generators perform "Moveset modification" on specific Pokémon or moveset combinations.
  // This can be problematic if the species in question has multiple formes, if so we just give up
  // and go with the set species (mostly relevant for Aggron/Aggron-Mega in Gen 6/7).
  return dex.species.get(matches.length ? matches[0] : set.species);
};

const getMoves = (data, doubles, sets) => {
  const moves = new Map();
  if (!data) return moves;
  if (sets) {
    for (const set of data.sets) {
      set.movepool.forEach(move => moves.set(Dex.toID(move), 0));
    }
  } else if (doubles && data.doublesMoves) {
    data.doublesMoves.forEach(m => moves.set(m, 0));
  } else {
    if (data.essentialMoves) data.essentialMoves.forEach(m => moves.set(m, 0));
    if (data.exclusiveMoves) data.exclusiveMoves.forEach(m => moves.set(m, 0));
    if (data.comboMoves) data.comboMoves.forEach(m => moves.set(m, 0));
    if (data.randomBattleMoves) data.randomBattleMoves.forEach(m => moves.set(m, 0));
    if (data.moves) data.moves.forEach(m => moves.set(m, 0));
  }
  return moves;
};

const getRoles = (species, sets) => {
  if (!sets || !sets[species.id]?.sets) return undefined;
  const roles = {};
  for (const set of sets[species.id].sets) {
    const r = roles[set.role] = {
      n: 0,
      abilities: new Map(),
      items: new Map(),
      moves: new Map(),
      ivs: {},
      evs: {},
    };
    set.movepool.forEach(move => r.moves.set(Dex.toID(move), 0));
    if (set.teraTypes) {
      r.teraTypes = new Map();
      set.teraTypes.forEach(type => r.teraTypes.set(type, 0));
    }
  }
  return roles;
};

const round = v => Math.round(v * 1e4) / 1e4;

const serialize = async (data, file, factor = 1) => {
  const json = stringify(data, { maxLength: factor * 1000 });
  await fs.promises.writeFile(`data/${file}`, json);
  const compressed = await gzip(json);
  return [json.length, compressed.length];
};

const increment = (gen, s, set) => {
  s.n++;
  if (set.ability) s.abilities.set(set.ability, (s.abilities.get(set.ability) || 0) + 1);
  if (set.item) s.items.set(set.item, (s.items.get(set.item) || 0) + 1);

  const IVS = gen < 3 ? 30 : 31;
  const EVS = gen < 3 ? 255 : 85;
  for (const stat in set.ivs) {
    if (set.ivs[stat] !== IVS) {
      if (!(stat in s.ivs)) {
        s.ivs[stat] = set.ivs[stat];
      } else if (s.ivs[stat] !== set.ivs[stat]) {
        s.ivs[stat] = -1;
      }
    } else if (stat in s.ivs) {
      s.ivs[stat] = -1;
    }
    if (set.evs[stat] !== EVS) {
      if (!(stat in s.evs)) {
        s.evs[stat] = set.evs[stat];
      } else if (s.evs[stat] !== set.evs[stat]) {
        s.evs[stat] = -1;
      }
    } else if (stat in s.evs) {
      s.evs[stat] = -1;
    }
  }
};

const displayAbilities = (gen, letsgo, data, o, s) => {
  if (gen >= 3 && !letsgo) {
    o.abilities = Array.from(data.abilities.keys()).sort();
    s.abilities = Object.fromEntries(Array.from(data.abilities.entries())
      .map(e => [e[0], round(e[1] / data.n)]).sort(CMP.entries));
  }
};

const displayItems = (data, o, s) => {
  if (data.items.size) {
    o.items = Array.from(data.items.keys()).sort();
    s.items = Object.fromEntries(Array.from(data.items.entries())
      .map(e => [e[0], round(e[1] / data.n)]).sort(CMP.entries));
  }
};

const displayStats = (letsgo, data, o, s) => {
  if (!letsgo) {
    for (const [stat, val] of Object.entries(data.evs).sort((CMP.stats))) {
      if (val !== -1) {
        o.evs = o.evs || {};
        o.evs[stat] = val;
      }
    }
    s.evs = o.evs;
  }
  for (const [stat, val] of Object.entries(data.ivs).sort((CMP.stats))) {
    if (val !== -1) {
      o.ivs = o.ivs || {};
      o.ivs[stat] = val;
    }
    s.ivs = o.ivs;
  }
};

(async () => {
  fs.rmSync('data', { recursive: true, force: true });
  fs.mkdirSync('data/js', { recursive: true });
  fs.mkdirSync('data/stats');

  const index = { options: {}, stats: {} };
  for (const format of Dex.formats.all()) {
    if (!format.team || !['singles', 'doubles'].includes(format.gameType)) continue;
    if (SKIP.some(s => format.id.includes(s)) || !format.mod.startsWith('gen')) continue;
    const start = process.hrtime.bigint();

    // Of course, PS is braindead and Format#gen always reports 0
    const gen = Number(format.id.charAt(3));
    const dex = Dex.forFormat(format);
    const letsgo = format.id.includes('letsgo');
    const baby = format.id.includes('baby');
    const doubles = format.gameType === 'doubles';
    const generator = Teams.getGenerator(format, SEED);
    const data = getData(format, doubles, baby);
    const sets = format.id.startsWith('gen9') || ROLES.includes(format.id);

    const pool = {};
    for (const species of dex.species.all()) {
      if (species.gen > gen) continue;
      if (letsgo && (species.num < 1 || species.num > 151 && ![808, 809].includes(species.num))) {
        continue;
      }
      const moves = getMoves(data[species.id], doubles, sets);
      if (!moves.size) continue;
      const base = typeof species.battleOnly === 'string' ? species.battleOnly : species.name;
      pool[base] = (pool[base] || {});
      pool[base][species.name] = moves;
    }

    const formes = {};
    for (let i = 0; i < N; i++) {
      for (const set of generator.getTeam()) {
        const forme = getForme(dex, pool, set);
        const roles = getRoles(forme, data);
        const s = formes[forme.name] || (formes[forme.name] = {
          n: 0,
          level: set.level,
          roles,
          abilities: new Map(),
          items: new Map(),
          moves: roles ? undefined : getMoves(forme, doubles),
          roles,
          ivs: {},
          evs: {},
        });

        const moves = set.role ? s.roles[set.role].moves : s.moves;
        for (const m of set.moves) moves.set(m, (moves.get(m) || 0) + 1);

        increment(gen, s, set);
        if (set.role) {
          increment(gen, s.roles[set.role], set);
          if (set.teraType) {
            s.roles[set.role].teraTypes.set(set.teraType,
              (s.roles[set.role].teraTypes.get(set.teraType) || 0) + 1);
          }
        }
      }
    }

    const options = {};
    const stats = {};
    for (const [name, data] of Object.entries(formes).sort(CMP.str)) {
      const o = { level: data.level };
      const s = { level: data.level };

      displayAbilities(gen, letsgo, data, o, s);
      displayItems(data, o, s);

      if (sets) {
        o.roles = {};
        s.roles = {};
        for (const [role, d] of Object.entries(data.roles).sort((a, b) => b[1].n - a[1].n)) {
          o.roles[role] = {};
          s.roles[role] = { weight: round(d.n / data.n) };

          displayAbilities(gen, letsgo, d, o.roles[role], s.roles[role]);
          displayItems(d, o.roles[role], s.roles[role]);

          if (d.teraTypes) {
            o.roles[role].teraTypes = Array.from(d.teraTypes.keys()).sort();
            s.roles[role].teraTypes = Object.fromEntries(Array.from(d.teraTypes.entries())
              .map(e => [e[0], round(e[1] / d.n)]).sort(CMP.entries));
          }

          o.roles[role].moves = Array.from(d.moves.keys()).map(m => dex.moves.get(m).name).sort();
          s.roles[role].moves = Object.fromEntries(Array.from(d.moves.entries())
            .map(e => [dex.moves.get(e[0]).name, round(e[1] / d.n)]).sort(CMP.entries));

          displayStats(letsgo, d, o.roles[role], s.roles[role]);
        }
      } else {
        o.moves = Array.from(data.moves.keys()).map(m => dex.moves.get(m).name).sort();
        s.moves = Object.fromEntries(Array.from(data.moves.entries())
          .map(e => [dex.moves.get(e[0]).name, round(e[1] / data.n)]).sort(CMP.entries));
      }

      displayStats(letsgo, data, o, s);

      options[name] = o
      stats[name] = s;
    }

    index.options[`${format.id}.json`] = await serialize(options, `${format.id}.json`);
    index.stats[`${format.id}.json`] = await serialize(stats, `stats/${format.id}.json`, 10);
    const code = util.inspect(options, {
      depth: Infinity,
      maxArrayLength: Infinity,
      compact: true,
      breakLength: 1000,
    }).slice(1, -1);
    fs.writeFileSync(`data/js/${format.id}.js`, `var ${format.id.toUpperCase()} = {\n ${code}\n};`);
    console.log(`${format.name}: ${(Number(process.hrtime.bigint() - start) / 1e9).toFixed(3)}s`);
  }
  fs.writeFileSync('data/index.json', stringify(index.options));
  fs.writeFileSync('data/stats/index.json', stringify(index.stats));
})().catch(err => {
  console.error(err);
  process.exit(1);
});