typestats.py

"""Analyze statistics of movepool typing"""

from __future__ import division

import pokedex.db
import pokedex.db.tables as t
from sqlalchemy import sql

import argparse
from itertools import groupby
from operator import itemgetter

session = pokedex.db.connect()

# what we want to know is the diversity of the types of moves learned by
# different types of pokemon by different methods
#
# version group => pokemon type => move method => move type => count
# no, we want to collapse identical moves learned by different methods
#
# WHAT WE HAVE: pokemon => version group => move method => move
# WHAT WE WANT: pokemon type => move type
#          AND: pokemon type => move method => move type
# 
# but in the UI we probably also want to list the pokemon and show a histogram
# of the moves learned
#
# Zhorken and Butterfree's claim is that Normal pokemon learn more varied moves
# than non-Normal pokemon. so let's start with that. we should restrict
# ourselves to pure-typed pokemon only, since part of the claim is that
# Normal/Flying pokemon are different from Normal-type pokemon
#
# so we just want a histogram of move type, broken down by pokemon type

VERSION_GROUP_ID = 14 # B2/W2
BARSIZE = 80

parser = argparse.ArgumentParser()
parser.add_argument('--types', default='', choices=['', 'pure', 'dual'])
    # '': select all pokemon
    # pure: select only singly-typed pokemon
    # dual: select only dual-typed pokemon
parser.add_argument('--flying', default='', choices=['', 'normal', 'flying', 'none'])
    # '': nothing
    # normal: treat normal/flying as pure flying
    # flying: treat x/flying as pure flying
    # none: treat x/flying as pure x
parser.add_argument('--damage', default='', choices=['', 'damage', 'none'])
    # '': nothing
    # damage: count only damaging moves (physical or special)
    # none: count only non-damaging moves
parser.add_argument('--method', default='')
parser.add_argument('--version-group-id', type=int, default=VERSION_GROUP_ID)

pokemon = t.Pokemon.__table__
pokemon_moves = t.PokemonMove.__table__
pokemon_types = t.PokemonType.__table__
moves = t.Move.__table__
types = t.Type.__table__

# SELECT pokemon_type, move_type, count(moves.id)
# FROM types as pokemon_type, types as move_type, pokemon_moves pm
# JOIN pokemon p ON pm.pokemon_id = p.id
# JOIN moves m ON pm.move_id = m.id
# WHERE p.type_id = pokemon_types.id AND m.type_id = move_types.id
# GROUP BY pokemon_type, move_type

def implies(c1, c2):
    return sql.or_(sql.not_(sql.func.coalesce(c1, False)), c2)

def get_method_id(method):
    q = session.query(t.PokemonMoveMethod.id).filter_by(identifier=method)
    return q.one()
def get_damage_class_id(identifier):
    q = session.query(t.MoveDamageClass.id).filter_by(identifier=identifier)
    return q.one()
def get_type_id(identifier):
    q = session.query(t.Type.id).filter_by(identifier=identifier)
    return q.one()

def get_pokemon_query(type_config, flying_config):
    if type_config:
        pt = get_pokemon_type_query(flying_config).alias()
        count = sql.select([pt.c.type_id],
                           pt.c.pokemon_id==pokemon.c.id)
        count = count.correlate(pokemon).alias().count()
        if type_config == 'pure':
            s = sql.select([pokemon], count.as_scalar() == 1)
        elif type_config == 'dual':
            s = sql.select([pokemon], count.as_scalar() == 2)
        else:
            raise ValueError(type_config)
    else:
        s = pokemon.select()
    return s

def get_pokemon_type_query(flying_config):
    s = pokemon_types.select()
    if flying_config:
        # We're going to ignore certain rows if the pokemon has certain types
        # First we need to join to one or two new copies of pokemon_types so we
        # can query on both types at once
        f = pokemon_types
        normal = get_type_id(u'normal')
        flying = get_type_id(u'flying')
        real_type_1 = pokemon_types.alias('real_type_1')
        real_type_2 = pokemon_types.alias('real_type_2')
        def addtype(f, real_type, slot):
            return f.outerjoin(real_type,
               sql.and_(pokemon_types.c.pokemon_id == real_type.c.pokemon_id,
                        real_type.c.slot == slot)
            )
        if flying_config == 'normal':
            # normal/flying -> flying.
            f = addtype(f, real_type_1, 1)
            f = addtype(f, real_type_2, 2)
            s = s.where(implies(sql.and_(real_type_1.c.type_id == normal,
                                         real_type_2.c.type_id == flying),
                                pokemon_types.c.type_id == flying))
        elif flying_config == 'flying':
            # x/flying -> flying
            f = addtype(f, real_type_2, 2)
            s = s.where(implies(real_type_2.c.type_id == flying,
                                pokemon_types.c.type_id == flying))
        elif flying_config == 'none':
            # x/flying -> x
            f = addtype(f, real_type_2, 2)
            s = s.where(implies(real_type_2.c.type_id == flying,
                                pokemon_types.c.type_id != flying))
        else:
            raise ValueError(flying_config)
        s = s.select_from(f).correlate(None)
    return s

def get_type_query(config):
    ptype = types.alias(u'ptype')
    mtype = types.alias(u'mtype')
    pokemon = get_pokemon_query(config.types, config.flying).correlate(None).alias(u'pokemon')
    pokemon_types = get_pokemon_type_query(config.flying).alias()
    s = sql.select([ptype.c.identifier, mtype.c.identifier,
                    sql.func.count(sql.text('*')).label(u'count')],
            sql.and_(
                pokemon_moves.c.version_group_id == config.version_group_id,
                pokemon_moves.c.move_id == moves.c.id,
                pokemon_moves.c.pokemon_id == pokemon.c.id,
                ptype.c.id == pokemon_types.c.type_id,
                mtype.c.id == moves.c.type_id,
            ),
    )
    if config.method:
        method_id = get_method_id(unicode(config.method))
        s = s.where(pokemon_moves.c.pokemon_move_method_id == method_id)
    if config.damage:
        damage_class_id = get_damage_class_id(u'non-damaging')
        if config.damage == 'none':
            s = s.where(moves.c.damage_class_id == damage_class_id)
        elif config.damage == 'damage':
            s = s.where(moves.c.damage_class_id != damage_class_id)
        else:
            raise ValueError(config.damage)
    s = s.select_from(ptype.join(pokemon_types).join(pokemon))
    s = s.select_from(mtype.join(moves))
    s = s.group_by(ptype.c.id, mtype.c.id, ptype.c.identifier, mtype.c.identifier)
    s = s.order_by(ptype.c.id, mtype.c.id)
    return s

def get_pokemon_count_query(config):
    # FIXME i don't take method or version into account
    ptype = types.alias(u'ptype')
    pokemon = get_pokemon_query(config.types, config.flying).correlate(None).alias(u'pure_pokemon')
    pokemon_types = get_pokemon_type_query(config.flying).alias()
    s = sql.select([ptype.c.identifier, sql.func.count(pokemon.c.id)],
            from_obj=[ptype.join(pokemon_types).join(pokemon)])
    s = s.group_by(ptype.c.id, ptype.c.identifier)
    s = s.order_by(ptype.c.id)
    return s

def keysort(d, key=None):
    """sort a dict by its keys"""
    items = d.items()
    if key:
        items.sort(key=lambda (k, v): key(k))
    else:
        items.sort()
    return items
def valuesort(d, key=None):
    """sort a dict by its values"""
    items = d.items()
    if key:
        items.sort(key=lambda (k, v): key(v))
    else:
        items.sort(key=itemgetter(1))
    return items
def typesort(d):
    typeorder = [u'normal', u'fighting', u'flying', u'poison', u'ground', u'rock', u'bug', u'ghost', u'steel', u'fire', u'water', u'grass', u'electric', u'psychic', u'ice', u'dragon', u'dark', u'???']
    return keysort(d, typeorder.index)

def main():
    args = parser.parse_args()
    conn = session.connection()
    results = conn.execute(get_type_query(args)).fetchall()
    pcounts = {type: count for type, count in
                conn.execute(get_pokemon_count_query(args)).fetchall()}

    d = {}
    for ptype, group in groupby(results, itemgetter(0)):
        group = list(group)
        d[ptype] = {mtype: count for _, mtype, count in group}

    for ptype, moves in reversed(valuesort(d, lambda m: sum(m.values()))):
        print u"%s (%d)" % (ptype, pcounts[ptype])

        total = sum(moves.values())

        print u"%9s %4d" % (u'total', total)
        for mtype, count in reversed(valuesort(moves)):
            p = count / total
            if mtype == ptype:
                bar = u"%" * int(round(BARSIZE * p))
            else:
                bar = u"#" * int(round(BARSIZE * p))
            print u"%9s %4d %4s %s" % (mtype, count, (u"%01.3f" % p)[1:], bar)
        print

#session.bind.echo = 1
main()