veekun_pokedex/pokedex/util/movesets.py

#! /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')

    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

        `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()
        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
            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)
        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 egg groups
        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,
            )))
        query = query.order_by(eg1.egg_group_id != None)
        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:
                self.egg_groups[evolution_chain] = (g1, g2) if g2 else (g1, )
            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

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