#!/usr/bin/env python2

# This script generates an erwiz file, which in turn generates a graphical
# ER diagram using the `erwiz` program. The columns and types of each table
# are written by querying nfldb's information schema. Relationships between
# entities are written out manually.
#
# Note that this script supports a `--condense` option that writes a smaller
# erwiz file. Namely, it excludes many of the play and player statistical
# columns. This provides a more compact ER diagram that is a useful
# introduction to the nfldb schema.

from __future__ import absolute_import, division, print_function
import sys

import nfldb

template_erwiz = '''
{{title: "nfldb Entity-Relationship diagram{condensed}"; title-size: 20}}

# Entities
{entities}

# Relationships
{relationships}
'''

template_entity = '''
[{name}] {{color:{color}}}
{columns}
'''

template_column = '  {pk}{column}{fk}{type}'

tables = ['player', 'team', 'game', 'drive', 'play', 'play_player', 'meta']

condense_whitelist = {
    'player': ['player_id', 'full_name', 'team', 'position', 'status'],
    'team': ['team_id', 'city', 'name'],
    'game': ['gsis_id', 'start_time', 'week', 'season_year', 'season_type',
             'finished', 'home_team', 'home_score', 'away_team', 'away_score'],
    'drive': ['gsis_id', 'drive_id', 'start_field', 'start_time', 'end_field',
              'end_time', 'pos_team', 'pos_time'],
    'play': ['gsis_id', 'drive_id', 'play_id', 'time', 'pos_team', 'yardline',
             'down', 'yards_to_go'],
    'play_player': ['gsis_id', 'drive_id', 'play_id', 'player_id', 'team'],
    'meta': ['version', 'season_type', 'season_year', 'week'],
}

colors = {
    'player': 'green', 'team': 'green', 'meta': 'red',
    'game': 'blue', 'drive': 'blue', 'play': 'blue', 'play_player': 'blue',
}


def mk_type(row):
    if row['domain_name']:
        type = row['domain_name']
    elif row['data_type'] == 'USER-DEFINED' and row['udt_name']:
        type = row['udt_name']
    else:
        type = row['data_type']

    if type == 'character varying':
        type = 'varchar'

    if row['is_nullable'] == 'YES':
        type += ', null'
    else:
        type += ', not null'
    return ' [%s]' % type


def is_pk(table, col):
    if col == 'profile_id':
        return False
    return col.endswith('_id')


def is_fk(table, col):
    if table == 'game' and col == 'gsis_id':
        return False
    if col == 'profile_id':
        return False
    return col.endswith('_id') and not col.startswith(table)


def in_whitelist(row):
    return row['column_name'] in condense_whitelist[row['table_name']]


def mk_column(row):
    col = row['column_name']
    table = row['table_name']
    pk = '*' if is_pk(table, col) else ''
    fk = '*' if is_fk(table, col) else ''
    return template_column.format(pk=pk, fk=fk, column=col, type=mk_type(row))


condense = len(sys.argv) == 2 and sys.argv[1] == '--condense'
db = nfldb.connect()
entities = []
with nfldb.Tx(db) as cursor:
    for table in tables:
        cursor.execute('''
            SELECT
                table_name, column_name, udt_name, data_type, domain_name,
                is_nullable
            FROM information_schema.columns
            WHERE table_name = %s
        ''', (table,))
        columns = []
        for row in cursor.fetchall():
            if not condense or in_whitelist(row):
                columns.append(mk_column(row))

        entity = template_entity.format(name=table, columns='\n'.join(columns),
                                        color=colors[table])
        entities.append(entity)

# No magic for relationships. Just gotta write'em out manually.
relationships = '''
[player] *--1 [team] <belongs-to->
[game] *--1 [team] <has-one->
[game] *--1 [team] <has-one->
[drive] *--1 [team] <has-one->
[play] *--1 [team] <has-one->
[play_player] *--1 [team] <has-one->

[game] 1--* [drive] <has-many->
[game] 1--* [play] <has-many->
[game] 1--* [play_player] <has-many->

[drive] 1--* [play] <has-many->
[drive] 1--* [play_player] <has-many->

[play] 1--* [play_player] <has-many->

[player] 1--* [play_player] <has-many->
'''

print(template_erwiz.format(entities='\n'.join(entities),
                            relationships=relationships,
                            condensed=' (condensed)' if condense else ''))
