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 {Worker, isMainThread} = require('worker_threads');

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

Dex.includeModData();

const SEED = [1, 2, 3, 4];
const N = +process.argv[2] || 100_000;
const NUM_WORKERS = 4;
const FORMAT_ARG = process.argv[3];
const SKIP = ['factory', 'unrated', 'cup', 'staff', 'cap', 'monotype', 'blitz', 'generated'];
const ORDER = ['hp', 'atk', 'def', 'spa', 'spd', 'spe'];
const FIELDS = ['role', 'movepool', 'teraTypes'];
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 = sets => {
  for (const species in sets) {
    for (const set of sets[species].sets) {
      for (const field of Object.keys(set)) {
        if (!FIELDS.includes(field)) {
          throw new Error(`Unexpected field in set: '${field}'`);
        }
      }
    }
  }
  return sets;
};

// gen9randomdoublebattle currently uses the same sets as singles
const getSets = format => format.id.startsWith('gen9')
  ? verify(require('./vendor/pokemon-showdown/data/random-sets.json'))
  : undefined;

const getForme = (dex, pool, set) => {
  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 = (species, doubles, sets) => {
  const moves = new Map();
  if (sets) {
    for (const set of (sets[species.id]?.sets || [])) {
      set.movepool.forEach(move =>  moves.set(Dex.toID(move), 0));
    }
  } if (doubles && species.randomDoubleBattleMoves) {
    species.randomDoubleBattleMoves.forEach(m => moves.set(m, 0));
  } else {
    if (species.essentialMove) moves.set(species.essentialMove, 0);
    if (species.exclusiveMoves) species.exclusiveMoves.forEach(m => moves.set(m, 0));
    if (species.comboMoves) species.comboMoves.forEach(m => moves.set(m, 0));
    if (species.randomBattleMoves) species.randomBattleMoves.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) {
    roles[set.role] = {
      n: 0,
      abilities: new Map(),
      items: new Map(),
      moves: new Map(),
      teraTypes: new Map(),
      ivs: {},
      evs: {},
    },
    set.movepool.forEach(move => roles[set.role].moves.set(Dex.toID(move), 0));
    set.teraTypes.forEach(type => roles[set.role].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 () => {
  if(!isMainThread) {
    return;
  }
  fs.rmSync('data', {recursive: true, force: true});
  fs.mkdirSync('data/js', {recursive: true});
  fs.mkdirSync('data/stats');

  const index = {options: {}, stats: {}};
  const formatsList = FORMAT_ARG ? [Dex.formats.get(FORMAT_ARG)] : Dex.formats.all();
  for (const format of formatsList) {
    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 doubles = format.gameType === 'doubles';
    const sets = getSets(format);

    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(species, 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 workers = Array(NUM_WORKERS);
    for(let i = 0; i < NUM_WORKERS; i++) {
      workers[i] = new Promise((resolve, reject) => {
        const numteams = Math.floor(N / NUM_WORKERS) + (i === NUM_WORKERS ? N % NUM_WORKERS : 0);
        const w = new Worker(__dirname + '/update-worker', 
          {workerData: pool, argv: [numteams, format.id]});
        w.on('message', resolve);
      });
    }
    const formesList = await Promise.all(workers);
    const formes = formesList[0];
    const attributesToIterate = ['abilities', 'items', 'moves', 'teraTypes'];
    for(let forme of formesList.slice(1)) {
      for(let pkmn in forme) {
        if(!formes[pkmn]) {
          formes[pkmn] = forme[pkmn];
          continue;
        }
        const currPkmn = formes[pkmn];
        const updatePkmn = forme[pkmn];
        currPkmn.n += updatePkmn.n;
        for(const attribute of attributesToIterate) {
          if(!currPkmn[attribute]) {
            continue;
          }
          for(const [k, count] of updatePkmn[attribute][Symbol.iterator]()) {
            currPkmn[attribute].set(k, count + (currPkmn[attribute].get(k) || 0));
          }
        }

        if(currPkmn.roles) {
          for(const role in currPkmn.roles) {
            const roleSet = currPkmn.roles[role];
            roleSet.n += updatePkmn.roles[role].n;
            for(const attribute of attributesToIterate) {
              for(const [k, count] of updatePkmn.roles[role][attribute][Symbol.iterator]()) {
                roleSet[attribute].set(k, count + (roleSet[attribute].get(k) || 0));
              }
            }
          }
        }
      }
    }

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

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

module.exports = {getSets, getForme, getMoves, getRoles, increment};