veekun_pokedex/pokedex/util/movesets.py
2011-04-26 03:14:32 +03:00

461 lines
19 KiB
Python
Executable file
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#! /usr/bin/env python
# Encoding: UTF-8
import sys
import argparse
from collections import defaultdict
from sqlalchemy.orm import aliased
from sqlalchemy.orm.exc import NoResultFound
from sqlalchemy.sql.expression import not_, and_, or_
from pokedex.db import connect, tables, util
from pokedex.util import querytimer
from pokedex.util.astar import a_star
class IllegalMoveCombination(ValueError): pass
class TooManyMoves(IllegalMoveCombination): pass
class NoMoves(IllegalMoveCombination): pass
class MovesNotLearnable(IllegalMoveCombination): pass
class NoParent(IllegalMoveCombination): pass
class TargetExcluded(IllegalMoveCombination): pass
class MovesetSearch(object):
def __init__(self, session, pokemon, version, moves, level=100, costs=None,
exclude_versions=(), exclude_pokemon=(), debug=False):
self.generator = self.error = None
if not moves:
self.error = NoMoves('No moves specified.')
return
elif len(moves) > 4:
self.error = NoMoves('Too many moves specified.')
return
self.debug = debug
self.session = session
self.sketch = util.get(session, tables.Move, identifier=u'sketch').id
self.no_eggs_group = util.get(session, tables.EggGroup,
identifier=u'no-eggs').id
self.ditto_group = util.get(session, tables.EggGroup,
identifier=u'ditto').id
if costs is None:
self.costs = default_costs
else:
self.costs = costs
self.excluded_families = frozenset(p.evolution_chain_id
for p in exclude_pokemon)
if pokemon:
self.goal_evolution_chain = pokemon.evolution_chain_id
if self.goal_evolution_chain in self.excluded_families:
self.error = TargetExcluded('The target pokemon was excluded.')
return
else:
self.goal_evolution_chain = None
if debug:
print 'Specified moves:', [move.id for move in moves]
self.goal_moves = frozenset(move.id for move in moves)
self.goal_version_group = version.version_group_id
self.load_pokemon()
# Fill self.generation_id_by_version_group
self.load_version_groups(version.version_group_id,
[v.version_group_id for v in exclude_versions])
self.pokemon_moves = defaultdict( # key: pokemon
lambda: defaultdict( # key: move
lambda: defaultdict( # key: version_group
lambda: defaultdict( # key: method
list)))) # list of (level, cost)
self.movepools = defaultdict(dict) # evo chain -> move -> best cost
self.learnpools = defaultdict(set) # evo chain -> move, w/o egg moves
easy_moves, non_egg_moves = self.load_pokemon_moves(
self.goal_evolution_chain, 'family')
hard_moves = self.goal_moves - easy_moves
egg_moves = self.goal_moves - non_egg_moves
if hard_moves:
# Have to breed!
self.load_pokemon_moves(self.goal_evolution_chain, 'others')
self.construct_breed_graph()
def load_version_groups(self, version, excluded):
"""Load generation_id_by_version_group
"""
query = self.session.query(tables.VersionGroup.id,
tables.VersionGroup.generation_id)
query = query.join(tables.Version.version_group)
if excluded:
query = query.filter(not_(tables.VersionGroup.id.in_(excluded)))
self.generation_id_by_version_group = dict(query)
def expand(v2):
for v1 in self.generation_id_by_version_group:
if self.trade_cost(v1, v2):
yield 0, None, v1
def is_goal(v):
return True
goal = self.goal_version_group
filtered_map = {goal: self.generation_id_by_version_group[goal]}
for result in a_star(self.goal_version_group, expand, is_goal):
for cost, transition, version in result:
filtered_map[version] = (
self.generation_id_by_version_group[version])
self.generation_id_by_version_group = filtered_map
if self.debug:
print 'Excluded version groups:', excluded
print 'Trade cost table:'
print '%03s' % '',
for g1 in sorted(self.generation_id_by_version_group):
print '%03s' % g1,
print
for g1 in sorted(self.generation_id_by_version_group):
print '%03s' % g1,
for g2 in sorted(self.generation_id_by_version_group):
print '%03s' % (self.trade_cost(g1, g2) or '---'),
print
def load_pokemon_moves(self, evolution_chain, selection):
"""Load pokemon_moves, movepools, learnpools, smeargle_families
`selection`:
'family' for loading only pokemon in evolution_chain
'others' for loading only pokemon NOT in evolution_chain
Returns: (easy_moves, non_egg_moves)
If `selection` == 'family':
easy_moves is a set of moves that are easier to obtain than by
breeding
non_egg_moves is a set of moves that don't require breeding
Otherwise, these are empty sets.
"""
if self.debug:
print 'Loading moves, c%s %s' % (evolution_chain, selection)
query = self.session.query(
tables.PokemonMove.pokemon_id,
tables.PokemonMove.move_id,
tables.PokemonMove.version_group_id,
tables.PokemonMoveMethod.identifier,
tables.PokemonMove.level,
tables.Pokemon.evolution_chain_id,
)
query = query.join(tables.PokemonMove.pokemon)
query = query.filter(tables.PokemonMoveMethod.id ==
tables.PokemonMove.pokemon_move_method_id)
query = query.filter(tables.PokemonMove.version_group_id.in_(
set(self.generation_id_by_version_group)))
query = query.filter(or_(
tables.PokemonMove.level > 100, # XXX: Chaff?
tables.PokemonMove.move_id.in_(self.goal_moves),
tables.PokemonMove.move_id == self.sketch,
tables.PokemonMove.move_id.in_(
self.evolution_moves.values()),
))
if self.excluded_families:
query = query.filter(not_(tables.Pokemon.evolution_chain_id.in_(
self.excluded_families)))
if evolution_chain:
if selection == 'family':
query = query.filter(tables.Pokemon.evolution_chain_id == (
evolution_chain))
elif selection == 'others':
query = query.filter(tables.Pokemon.evolution_chain_id != (
evolution_chain))
query = query.order_by(tables.PokemonMove.level)
easy_moves = set()
non_egg_moves = set()
self.smeargle_families = set()
for pokemon, move, vg, method, level, chain in query:
if move in self.goal_moves or move == self.sketch:
cost = self.learn_cost(method, vg)
self.movepools[chain][move] = min(
self.movepools[chain].get(move, cost), cost)
if method != 'egg':
self.learnpools[chain].add(move)
non_egg_moves.add(move)
if cost < self.costs['breed']:
easy_moves.add(move)
else:
cost = -1
if move == self.sketch:
self.smeargle_families.add(self.evolution_chains[pokemon])
self.pokemon_moves[pokemon][move][vg][method].append((level, cost))
if self.debug and selection == 'family':
print 'Easy moves:', sorted(easy_moves)
print 'Non-egg moves:', sorted(non_egg_moves)
if self.debug:
print 'Smeargle families:', sorted(self.smeargle_families)
return easy_moves, non_egg_moves
def learn_cost(self, method, version_group):
"""Return cost of learning a move by method (identifier) in ver. group
"""
if method == 'level-up':
return self.costs['level-up']
gen = self.generation_id_by_version_group[version_group]
if method == 'machine' and gen < 5:
return self.costs['machine-once']
elif method == 'tutor' and gen == 3:
return self.costs['tutor-once']
elif method == 'egg':
return self.costs['breed']
else:
return self.costs[method]
def trade_cost(self, version_group_from, version_group_to, *thing_generations):
"""Return cost of trading between versions, None if impossibble
`thing_generations` should be the generation IDs of the pokemon and
moves being traded.
"""
# XXX: this ignores HM transfer restrictions
gen_from = self.generation_id_by_version_group[version_group_from]
gen_to = self.generation_id_by_version_group[version_group_to]
if gen_from == gen_to:
return self.costs['trade']
elif any(gen > gen_to for gen in thing_generations):
return None
elif gen_from in (1, 2):
if gen_to in (1, 2):
return self.costs['trade']
else:
return None
elif gen_to in (1, 2):
return None
elif gen_from > gen_to:
return None
elif gen_from < gen_to - 1:
return None
else:
return self.costs['trade'] + self.costs['transfer']
def load_pokemon(self):
"""Load pokemon breed groups and evolutions
self.egg_groups: maps evolution chains to their sorted egg groups
(wil contain empty tuple for no-eggs or ditto)
self.evolution_chains: maps pokemon to their evolution chains
self.pokemon_by_evolution_chain: maps evolution chains to their pokemon
self.unbreedable: set of unbreedable pokemon
self.evolution_parents: maps pokemon to their pre-evolved form
self.evolutions: maps pokemon to lists of (trigger, move, level, child)
self.evolution_moves: maps evolution_chains to their evolution moves
"""
eg1 = tables.PokemonEggGroup
eg2 = aliased(tables.PokemonEggGroup)
query = self.session.query(
tables.Pokemon.id,
tables.Pokemon.evolution_chain_id,
tables.Pokemon.evolves_from_pokemon_id,
eg1.egg_group_id,
eg2.egg_group_id,
)
query = query.join((eg1, eg1.pokemon_id == tables.Pokemon.id))
query = query.outerjoin((eg2, and_(
eg2.pokemon_id == tables.Pokemon.id,
eg1.egg_group_id < eg2.egg_group_id,
)))
bad_groups = (self.no_eggs_group, self.ditto_group)
unbreedable = set()
self.evolution_parents = dict()
self.egg_groups = defaultdict(tuple)
self.evolution_chains = dict()
self.pokemon_by_evolution_chain = defaultdict(set)
for pokemon, evolution_chain, parent, g1, g2 in query:
if g1 in bad_groups:
unbreedable.add(pokemon)
else:
new_groups = (g1, g2) if g2 else (g1, )
if len(self.egg_groups.get(evolution_chain, ())) <= len(new_groups):
self.egg_groups[evolution_chain] = new_groups
self.evolution_chains[pokemon] = evolution_chain
self.pokemon_by_evolution_chain[evolution_chain].add(pokemon)
if parent:
self.evolution_parents[pokemon] = parent
self.unbreedable = frozenset(unbreedable)
self.evolutions = defaultdict(set)
self.evolution_moves = dict()
query = self.session.query(
tables.PokemonEvolution.evolved_pokemon_id,
tables.EvolutionTrigger.identifier,
tables.PokemonEvolution.known_move_id,
tables.PokemonEvolution.minimum_level,
)
query = query.join(tables.PokemonEvolution.trigger)
for child, trigger, move, level in query:
self.evolutions[self.evolution_parents[child]].add(
(trigger, move, level, child))
if move:
self.evolution_moves[self.evolution_chains[child]] = move
if self.debug:
print 'Loaded %s pokemon: %s evo; %s families: %s breedable' % (
len(self.evolution_chains),
len(self.pokemon_by_evolution_chain),
len(self.egg_groups),
len(self.evolutions),
)
print 'Evolution moves: %s' % self.evolution_moves
def construct_breed_graph(self):
# eg1_movepools[egg_group_id] = set of moves passable by pkmn in that group
eg1_movepools = defaultdict(set)
# eg2_movepools[b_g_id1, b_g_id2] = ditto for pkmn in *both* groups
eg2_movepools = defaultdict(set)
# non_egg_pools = as eg1_movepools but for *learnable* moves
learn_pools = defaultdict(set)
goal_egg_groups = self.egg_groups[self.goal_evolution_chain]
all_groups = set()
for family, groups in self.egg_groups.iteritems():
if not groups:
continue
if family == self.goal_evolution_chain:
continue
elif family in self.smeargle_families:
pool = self.goal_moves
else:
pool = self.movepools[family]
pool = set(pool) & self.goal_moves
learnpool = self.learnpools[family] & pool
for group in groups:
eg1_movepools[group].update(pool)
learn_pools[group].update(learnpool)
all_groups.add(group)
if len(groups) >= 2:
eg2_movepools[groups].update(pool)
if self.debug:
for group in sorted(all_groups):
print "%2s can pass: %s" % (group, sorted(eg1_movepools[group]))
if learn_pools[group] != eg1_movepools[group]:
print " but learn: %s" % sorted(learn_pools[group])
for g2 in sorted(all_groups):
for g1 in sorted(all_groups):
if eg2_movepools[g1, g2]:
print " %2s/%2s pass: %s" % (g1, g2, sorted(eg2_movepools[g1, g2]))
print 'Goal groups:', goal_egg_groups
default_costs = {
# Costs for learning a move in verious ways
'level-up': 20, # The normal way
'machine': 40, # Machines are slightly inconvenient.
'machine-once': 2000, # before gen. 5, TMs only work once. Avoid.
'tutor': 60, # Tutors are slightly more inconvenient than TMs can't carry them around
'tutor-once': 2100, # gen III: tutors only work once (well except Emerald frontier ones)
'sketch': 10, # Quite cheap. (Doesn't include learning Sketch itself)
# Gimmick moves we need to use this method to learn the move anyway,
# so make a big-ish dent in the score if missing
'stadium-surfing-pikachu': 100,
'light-ball-egg': 100, # …
# Ugh... I don't know?
'colosseum-purification': 100,
'xd-shadow': 100,
'xd-purification': 100,
'form-change': 100,
# Other actions.
# Breeding should cost more than 3 times than a lv-up/machine/tutor move.
'evolution': 100, # We have to do this anyway, usually.
'evolution-delayed': 50, # *in addition* to evolution. Who wants to mash B on every level.
'breed': 400, # Breeding's a pain.
'trade': 200, # Trading's a pain, but not as much as breeding.
'transfer': 200, # *in addition* to trade. For one-way cross-generation transfers
'delete': 300, # Deleting a move. (Not needed unless deleting an evolution move.)
'relearn': 150, # Also a pain, though not as big as breeding.
'per-level': 1, # Prefer less grinding. This is for all lv-ups but the final “grow”
}
def main(argv):
parser = argparse.ArgumentParser(description=
'Find out if the specified moveset is valid, and provide a suggestion '
'on how to obtain it.')
parser.add_argument('pokemon', metavar='POKEMON', type=unicode,
help='Pokemon to check the moveset for')
parser.add_argument('move', metavar='MOVE', type=unicode, nargs='*',
help='Moves in the moveset')
parser.add_argument('-l', '--level', metavar='LV', type=int, default=100,
help='Level of the pokemon')
parser.add_argument('-v', '--version', metavar='VER', type=unicode,
default='black',
help='Version to search in.')
parser.add_argument('-V', '--exclude-version', metavar='VER', type=unicode,
action='append', default=[],
help='Versions to exclude (along with their '
'counterparts, if any, e.g. `black` will also exclude White).')
parser.add_argument('-P', '--exclude-pokemon', metavar='PKM', type=unicode,
action='append', default=[],
help='Pokemon to exclude (along with their families, e.g. `pichu` '
'will also exclude Pikachu and Raichu).')
parser.add_argument('-d', '--debug', action='append_const', const=1,
default=[],
help='Output timing and debugging information (can be specified more '
'than once).')
args = parser.parse_args(argv)
args.debug = len(args.debug)
if args.debug:
print 'Connecting'
session = connect(engine_args={'echo': args.debug > 1})
if args.debug:
print 'Parsing arguments'
def _get_list(table, idents, name):
result = []
for ident in idents:
try:
result.append(util.get(session, table, identifier=ident))
except NoResultFound:
print>>sys.stderr, ('%s %s not found. Please use '
'the identifier.' % (name, ident))
return 2
return result
pokemon = _get_list(tables.Pokemon, [args.pokemon], 'Pokemon')[0]
moves = _get_list(tables.Move, args.move, 'Move')
version = _get_list(tables.Version, [args.version], 'Version')[0]
excl_versions = _get_list(tables.Version, args.exclude_version, 'Version')
excl_pokemon = _get_list(tables.Pokemon, args.exclude_pokemon, 'Pokemon')
if args.debug:
print 'Starting search'
search = MovesetSearch(session, pokemon, version, moves, args.level,
exclude_versions=excl_versions, exclude_pokemon=excl_pokemon,
debug=args.debug)
if search.error:
print 'Error:', search.error
return 1
if args.debug:
print 'Done'
if __name__ == '__main__':
sys.exit(main(sys.argv[1:]))