veekun_pokedex/scripts/split-movesets.py

# coding: utf-8
import itertools
from operator import itemgetter
from collections import OrderedDict

import pokedex.db
import sqlalchemy
from sqlalchemy import (Column, Integer, ForeignKey, Table, Unicode)
from sqlalchemy import ForeignKeyConstraint, PrimaryKeyConstraint, UniqueConstraint

meta = sqlalchemy.MetaData()
Base = sqlalchemy.ext.declarative.declarative_base(metadata=meta)

class OldPokemonMove(Base):
    u"""Record of a move a Pokémon can learn."""
    __tablename__ = 'pokemon_moves'
    pokemon_id = Column(Integer, ForeignKey('pokemon.id'), nullable=False, index=True,
        doc=u"ID of the Pokémon")
    version_group_id = Column(Integer, ForeignKey('version_groups.id'), nullable=False, index=True,
        doc=u"ID of the version group this applies to")
    move_id = Column(Integer, ForeignKey('moves.id'), nullable=False, index=True,
        doc=u"ID of the move")
    pokemon_move_method_id = Column(Integer, ForeignKey('pokemon_move_methods.id'), nullable=False, index=True,
        doc=u"ID of the method this move is learned by")
    level = Column(Integer, nullable=True, index=True, autoincrement=False,
        doc=u"Level the move is learned at, if applicable")
    order = Column(Integer, nullable=True,
        doc=u"The order which moves learned at the same level are learned in")

    __table_args__ = (
        PrimaryKeyConstraint('pokemon_id', 'version_group_id', 'move_id', 'pokemon_move_method_id', 'level'),
    )

# The pokemon move table has much redundancy across versions, but also many small variations.
# The general strategy here is to break movesets into chunks, such that a moveset is the union
# of a number of an arbitrary number of chunks, and those chunks can be shared across versions.

pokemon_moveset_table = Table('pokemon_movesets', Base.metadata,
    Column('id', Integer, primary_key=True, autoincrement=True),
    #Column('hint', Unicode(79), unique=True, index=True),
    Column('pokemon_id', Integer, ForeignKey('pokemon.id')),
    Column('method_id', Integer, ForeignKey('pokemon_move_methods.id'), nullable=False, index=True),
)

pokemon_moveset_moves_table = Table('pokemon_moveset_moves', Base.metadata,
    Column('moveset_id', Integer),
    Column('move_id', Integer, ForeignKey('moves.id'), nullable=False, index=True),
    Column('level', Integer, nullable=True, index=True, autoincrement=False),
    Column('order', Integer, nullable=True, index=False),
    ForeignKeyConstraint(['moveset_id'], ['pokemon_movesets.id']),
    ForeignKeyConstraint(['move_id'], ['moves.id']),
    PrimaryKeyConstraint('moveset_id', 'move_id', 'level'),
)

pokemon_moveset_version_groups_table = Table('pokemon_moveset_version_groups', Base.metadata,
    Column('moveset_id', Integer),
    Column('version_group_id', Integer, ForeignKey('version_groups.id'), nullable=False, index=True),
    ForeignKeyConstraint(['moveset_id'], ['pokemon_movesets.id']),
    PrimaryKeyConstraint('moveset_id', 'version_group_id'),
)

def upgrade():
    engine = sqlalchemy.create_engine('postgresql:///pokedex')
    meta.reflect(bind=engine)

    pokemon_moveset_version_groups_table.drop(engine, checkfirst=True)
    pokemon_moveset_moves_table.drop(engine, checkfirst=True)
    pokemon_moveset_table.drop(engine, checkfirst=True)

    pokemon_moveset_table.create(engine, checkfirst=True)
    pokemon_moveset_moves_table.create(engine, checkfirst=True)
    pokemon_moveset_version_groups_table.create(engine, checkfirst=True)

    old_table = OldPokemonMove.__table__
    pokemon_table = meta.tables['pokemon']
    method_table = meta.tables['pokemon_move_methods']
    version_group_table = meta.tables['version_groups']
    generation_table = meta.tables['generations']
    print "all systems go"
    conn = engine.connect()
    for pokemon_id, in conn.execute(sqlalchemy.select([pokemon_table.c.id])):
        moves = conn.execute(
            sqlalchemy.select([
                old_table.c.pokemon_move_method_id,
                version_group_table.c.generation_id,
                old_table.c.version_group_id,
                old_table.c.move_id,
                old_table.c.level,
                old_table.c.order,
            ])
            .where(old_table.c.version_group_id == version_group_table.c.id)
            .where(old_table.c.pokemon_id == pokemon_id)
            .order_by("pokemon_move_method_id", "generation_id", version_group_table.c.order)
        ).fetchall()
        print pokemon_id

        for method_id, group in itertools.groupby(moves, itemgetter(0)):
            if method_id == 4:
                # TMs stayed more or less same across Gens 3&4 and 5&6
                def grouper(x):
                    return [0, 1, 2, 3, 3, 5, 5][x[1]]
            #elif method_id == 2:
                # Egg moves stick around forever
                #grouper = lambda _: None
            else:
                grouper = itemgetter(1)
            for _, group in itertools.groupby(group, grouper):
                movesets = OrderedDict()
                for version_group_id, group in itertools.groupby(group, itemgetter(2)):
                    movesets[version_group_id] = set(map(itemgetter(3, 4, 5), group))
                common = reduce(set.intersection, movesets.values())
                if common:
                    create_moveset(conn, pokemon_id, method_id, common, movesets.keys())
                deltas = []
                for version_group_id in movesets:
                    moveset = movesets[version_group_id] - common
                    delta_ids = []
                    for moveset_id, delta in reversed(deltas):
                        if delta <= moveset:
                            delta_ids.append(moveset_id)
                            moveset -= delta
                    if delta_ids:
                        add_to_movesets(conn, pokemon_id, method_id, delta_ids, version_group_id)
                    if moveset:
                        moveset_id = create_moveset(conn, pokemon_id, method_id, moveset, [version_group_id])
                        deltas.append((moveset_id, moveset))


def create_moveset(conn, pokemon_id, method_id, moves, versions):
    result = conn.execute(pokemon_moveset_table.insert(), pokemon_id=pokemon_id, method_id=method_id)
    moveset_id, = result.inserted_primary_key
    conn.execute(pokemon_moveset_moves_table.insert()
        .values(moveset_id=moveset_id),
        [dict(move_id=move_id, level=level, order=order) for move_id, level, order in moves])
    conn.execute(pokemon_moveset_version_groups_table.insert()
        .values(moveset_id=moveset_id),
        [dict(version_group_id=version_group_id) for version_group_id in versions])
    return moveset_id

def add_to_movesets(conn, pokemon_id, method_id, moveset_ids, version_group_id):
    conn.execute(pokemon_moveset_version_groups_table.insert()
        .values(version_group_id=version_group_id),
        [dict(moveset_id=moveset_id) for moveset_id in moveset_ids])

if __name__ == '__main__':
    upgrade()