drinks-scanner-analysis

#!/usr/bin/env python3
# Loads drinks purchase statistics with according user database and plots it in a meaningful way.

import argparse
import datetime
import logging
import os
import psycopg2
import requests
import sys

import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

from pathlib import Path
from bs4 import BeautifulSoup # TODO Fix version in requirements.txt
from sshtunnel import SSHTunnelForwarder, HandlerSSHTunnelForwarderError

LDAP_SCRAPE_HOST = '192.168.3.231' # power-pi.fd ?
DRINKS_DATA_HOST = '192.168.3.231' # drinks.fd

SSH_PORT = 22
SSH_USERNAME = '' # Please see:
SSH_PASSWORD = '' # https://www.raspberrypi.org/documentation/linux/usage/users.md
SSH_REMOTE_BIND = ('127.0.0.1', 5432)
SSH_LOCAL_BIND = SSH_REMOTE_BIND

DATA_DIRECTORY = 'data'
CSV_SEPARATOR = ';'
LDAP_DATA_FILE = '{}/users.csv'.format(DATA_DIRECTORY)
DRINKS_DATA_FILE = '{}/drinks.csv'.format(DATA_DIRECTORY)
SCAN_DATA_FILE = '{}/scans.csv'.format(DATA_DIRECTORY)
RECHARGE_DATA_FILE = '{}/recharges.csv'.format(DATA_DIRECTORY)

LOG_DIRECTORY = 'logs'
DATE_STR = datetime.date.today().isoformat()
FILE_STR = os.path.basename(__file__ or '').replace('.py', '')
LOG_FILE = os.path.join(LOG_DIRECTORY, "{}-{}.log".format(FILE_STR, DATE_STR))

LOG_LEVELS = {
    'debug': logging.DEBUG,
    'info': logging.INFO,
    'warning': logging.WARNING,
    'error': logging.ERROR,
    'critical': logging.CRITICAL,
}


def get_users(reload=False):
    # Load local data if it exists and reloading is disabled
    ldap_data_file = Path(LDAP_DATA_FILE)
    if not reload and ldap_data_file.is_file():
        return pd.read_csv(LDAP_DATA_FILE, sep=CSV_SEPARATOR, dtype={
            'USER_NAME': str,
            'USER_ID': str
        })

    # Request scraping content
    r = requests.get('http://{}/'.format(LDAP_SCRAPE_HOST))
    if r.status_code != 200:
        raise ConnectionError('Could not connect to LDAP scrape host ({})!'.format(LDAP_SCRAPE_HOST))
    if r.text == '':
        raise ConnectionError('Could not collect data from LDAP scrape host ({})!'.format(LDAP_SCRAPE_HOST))

    # Scrape
    soup = BeautifulSoup(r.text, 'html.parser')
    users = ['USER_NAME']
    uids = ['USER_ID']
    for option in soup.find_all('option'):
        user = option.get_text().strip()
        uid = option.get('value').strip()
        # Skip erroneous/empty entries
        if not user or not uid or user == '' or uid == '':
            continue
        users.append(user)
        uids.append(uid)

    # Create data directory if necessary
    if not os.path.exists(DATA_DIRECTORY):
        os.makedirs(DATA_DIRECTORY)

    csv_data = pd.DataFrame({users[0]: users[1:], uids[0]: uids[1:]}).drop_duplicates()

    # Save data to disk
    csv_data.to_csv(LDAP_DATA_FILE, sep=CSV_SEPARATOR, index=False)
    logger.info("Saved data to: {}/{}".format(os.getcwd(), LDAP_DATA_FILE))
    return get_users()


def postgres_query(query):
    try:
        conn = psycopg2.connect("dbname='drinks' user='postgres' host='localhost'")
    except ConnectionError as e:
        raise ConnectionError("Unable to connect to the database:\n{}".format(e))

    cur = conn.cursor()
    cur.execute(query)
    rows = cur.fetchall()
    return rows


def ssh_postgres_query(query):
    try:
        with SSHTunnelForwarder(
                (DRINKS_DATA_HOST, SSH_PORT),
                ssh_username=SSH_USERNAME,
                ssh_password=SSH_PASSWORD,
                remote_bind_address=SSH_REMOTE_BIND,
                local_bind_address=SSH_LOCAL_BIND,
        ):
            return postgres_query(query)
    except ValueError as e:
        logger.error("Error with SSH credentials: {}".format(e))
        raise e
    except HandlerSSHTunnelForwarderError as e:
        logger.error("Could not rebind local SSH {}:{}. Is it already bound by another process?".format(SSH_LOCAL_BIND[0], SSH_LOCAL_BIND[1]))
        raise e


def postgres_data(map_dict, sql_query, data_file_path, reload=False):
    # Load local data if it exists and reloading is disabled
    data_file = Path(data_file_path)
    if not reload and data_file.is_file():
        map_dtypes = {}
        for k, v in map_dict.items():
            map_dtypes[k] = v[1]
        return pd.read_csv(data_file_path, sep=CSV_SEPARATOR, dtype=map_dtypes)

    # Query data otherwise
    try:
        data_raw = ssh_postgres_query(sql_query)
    except ValueError:
        sys.exit(10)
    except HandlerSSHTunnelForwarderError:
        sys.exit(11)
    data_raw = pd.DataFrame(data_raw)
    data = {}
    for k, v in map_dict.items():
        data[k] = data_raw[v[0]]
    data = pd.DataFrame(data)

    # Save data to disk
    data.to_csv(data_file_path, sep=CSV_SEPARATOR, index=False)
    logger.info("Saved data to: {}/{}".format(os.getcwd(), data_file_path))
    return postgres_data(map_dict, sql_query, data_file_path, reload=False)


def get_drinks(reload=False):
    map_dict = {
        "EAN": (1, str),
        "NAME": (2, str),
        "SIZE": (3, float),
        "TYPE": (5, str),
    }
    sql_query = "SELECT * FROM drink;"
    data_file_path = DRINKS_DATA_FILE
    return postgres_data(map_dict, sql_query, data_file_path, reload=reload)


def get_scans(reload=False):
    map_dict = {
        "EAN": (1, str),
        "USER_ID": (2, str),
        "TIMESTAMP": (3, str),
    }
    sql_query = "SELECT * FROM scanevent;"
    data_file_path = SCAN_DATA_FILE
    data = postgres_data(map_dict, sql_query, data_file_path, reload=reload)
    data['TIMESTAMP'] = pd.to_datetime(data['TIMESTAMP'])
    return data


def get_recharges(reload=False):
    map_dict = {
        "USER_ID": (1, str),
        "HELPER_USER_ID": (2, str),
        "AMOUNT": (3, float),
        "TIMESTAMP": (4, str),
    }
    sql_query = "SELECT * FROM rechargeevent;"
    data_file_path = RECHARGE_DATA_FILE
    data = postgres_data(map_dict, sql_query, data_file_path, reload=reload)
    data['TIMESTAMP'] = pd.to_datetime(data['TIMESTAMP'])
    return data


def plot_drink_consumption_by_name(scans_full):
    """Plot drink consumption by name (total)"""
    drinks_count = scans_full.groupby('NAME').size().sort_values(ascending=False)

    # Limit by minimum count
    #drinks_count = drinks_count.where(drinks_count > 50)

    plt.xticks(rotation=90)
    ax = sns.barplot(
        x=drinks_count.index.where(drinks_count > 50),
        y=drinks_count,
        palette=sns.color_palette("Paired"))
    ax.set(
        xlabel='Getränkename',
        ylabel='Gescannte Einheiten',
        title='Gescannte Getränke nach Name',
    )
    plt.show()


def plot_drink_consumption_by_category(scans_full):
    """Plot drink consumption by category (total)"""
    drinks_typed = scans_full.groupby(['TYPE', 'NAME']).size().sort_values(ascending=False).to_frame().reset_index()
    ax = sns.swarmplot(x='TYPE', y=0, data=drinks_typed)
    ax.set(
        xlabel='Typ',
        ylabel='Gescannte Einheiten',
        title='Gescannte Getränke nach Typ',
    )

    drinks_categories = drinks_typed['TYPE'].unique()

    # TODO Find out how to get categorical x position in swarm plot
    #print(drinks_categories)
    #print(drinks_typed)
    #print(drinks_typed['TYPE'])

    #for line in range(0, drinks_typed.shape[0]):
    #    ax.text(
    #        #drinks_typed.index[line] + 0.2,
    #        drinks_typed['TYPE'] * 10,
    #        drinks_typed.get(0)[line],
    #        drinks_typed['NAME'][line],
    #        horizontalalignment='left',
    #    )

    ## Annotate highest ranking
    #drinks_highest = drinks_typed.groupby('TYPE').head(1)
    ## Annotate lowest ranking
    #drinks_lowest = drinks_typed.groupby('TYPE').tail(1)

    ##for a in range(0, len(drinks_highest)):
    ##    sp.text(drinks_highest['TYPE'][a],
    ##            drinks_highest[0][a],
    #            drinks_highest['NAME'][a])
    plt.show()


def plot_drink_total_consumption_by_person(scans_full):
    scans_by_person = scans_full.sort_values(['USER_NAME', 'TYPE']).get(['USER_NAME', 'TYPE'])
    scans_by_person = scans_by_person.groupby('USER_NAME').size().sort_values(ascending=False)

    plt.xticks(rotation=90)
    ax = sns.barplot(
        x=scans_by_person.index.where(scans_by_person > 20),
        y=scans_by_person,
        palette=sns.color_palette("Paired")
    )
    ax.set(
        xlabel='Member',
        ylabel='Gescannte Getränke',
        title='Gescannte Getränke nach Member',
    )
    plt.show()


def plot_drink_categorized_consumption_by_person(scans_full):
    scans_by_person = scans_full.sort_values(['USER_NAME', 'TYPE']).get(['USER_NAME', 'TYPE'])
    scans_by_person = pd.crosstab(scans_by_person['USER_NAME'], scans_by_person['TYPE'])
    scans_by_person = scans_by_person.stack().reset_index().rename(columns={0: 'value'})

    fig, ax = plt.subplots(figsize=(18, 6))
    plt.xticks(rotation=90)
    sns.barplot(
        x=scans_by_person.USER_NAME,
        y=scans_by_person.value,
        hue=scans_by_person.TYPE,
        ax=ax,
    )
    ax.set(
        xlabel='Member',
        ylabel='Gescannte Getränke',
        title='Gescannte Getränke nach Member und Typ',
    )
    plt.show()


def plot_drink_categorized_consumption_by_person_and_time_of_day(scans_full):
    scans_by_person = scans_full.sort_values(['USER_NAME', 'TYPE']).get(['USER_NAME', 'TYPE', 'TIMESTAMP'])
    scans_by_person['TIMESTAMP'] = scans_by_person['TIMESTAMP'].apply(lambda x: x.hour)
    scans_by_person = scans_by_person.groupby(['USER_NAME', 'TYPE', 'TIMESTAMP']).size()


def plot_recharge_vouching_by_person(recharges_full):
    recharges_names = recharges_full.get(['HELPER_USER_NAME', 'USER_NAME'])
    recharges_names = pd.crosstab(recharges_names['USER_NAME'], recharges_names['HELPER_USER_NAME'])

    # Create custom palette
    max_value = recharges_names.max().max()
    cutoff = 7
    cpal = sns.hls_palette(max_value - 1 + cutoff, h=.5)
    del cpal[-cutoff:]
    cpal.insert(0, (.95, .95, .95))

    fig, ax = plt.subplots(figsize=(8, 8))
    ax = sns.heatmap(
        recharges_names,
        cmap=cpal,
        square=True,
        linewidths=1,
        linecolor='white',
    )
    ax.set(
        xlabel='Bürge',
        ylabel='Aufladender',
        title='Gebürgte Aufladungen unter Membern',
    )

    plt.show()


def plot_drink_consumption_by_name_over_time(scans_full):
    df = scans_full
    df.set_index(pd.DatetimeIndex(df['TIMESTAMP']), inplace=True)
    #ts = df.resample('M').sum()
    ts = df.groupby(['USER_NAME', 'TYPE']).resample('M')['TIMESTAMP'].count()


    df = ts.to_frame()
    #df2 = df.reset_index()
    #print(ts)
    print(df.head())

    print(df.reset_index(level='TYPE').reset_index(level='USER_NAME'))

    #print(df["USER_NAME"])

    return


    for k, v in ts.groupby(['USER_NAME']).sum().items():
        print(k, v)
    return

    #print(scans_full.groupby('USER_NAME').count())

    #for d in dir(ts):
    #    print(d)
    print(ts.head())
    #print(ts['TIMESTAMP'])
    print(ts.values)
    print(ts.index)


    #sns.kdeplot(ts.values)
    #sns.lineplot(hue=ts['USER_NAME'], data=ts.values)
    sns.lineplot(hue=ts.index, data=ts.values)
    plt.show()


    return


    print(ts['USER_NAME'])
    print(ts['TYPE'])

    sns.lineplot(
        x=ts['TIMESTAMP'],
        y=ts.values,
    )
    return

    print(ts)
    for k, v in ts.items():
        print(k, v)
    return


    ############################################################

    print("XXXX")
    print(type(df['TIMESTAMP'][0]))
    #df['TIMESTAMP'].replace(day=1)

    print(pd.DatetimeIndex(df['TIMESTAMP']))
    print(df.head())

    #df = df.assign(TIMESTAMP = [mdates.date2num(x) for x in df['TIMESTAMP']])
    df = df.assign(TIMESTAMP = [datetime.datetime(x) for x in df['TIMESTAMP']])
    print(df.head())


    return
    #df = df.assign(YEAR_MONTH = ['{}-{:0=2d}'.format(x.year, x.month) for x in df['TIMESTAMP']])
    # TODO
    # Fix datetime objects to date objects with same day but different month
    # After that, convert to matplotlib's own date format
    # https://stackoverflow.com/a/31262531
    print(type(df['TIMESTAMP'][0]))
    df = df.assign(YEAR_MONTH = [mdates.date2num(x) for x in df['TIMESTAMP']])
    #df = df.groupby(['USER_NAME', 'TYPE', 'YEAR_MONTH']).size()
    df = df.groupby(['USER_NAME', 'TYPE', 'TIMESTAMP']).size()
    print(df.head())
    #df = df.get(['USER_NAME', 'TYPE', 'YEAR_MONTH'])
    print(df.values)
    print(type(df))
    #print(df['TYPE'])
    return
    print(type(df['TIMESTAMP']))
    fig, ax = plt.subplots()
    #sns.kdeplot(df.values, shade=True)
    #sns.lineplot(x=df['YEAR_MONTH'], y=df.values, time='YEAR_MONTH', value='values', ax=ax)
    sns.lineplot(x=df['TIMESTAMP'], y=df.values, value='values', ax=ax)


    # assign locator and formatter for the xaxis ticks.
    ax.xaxis.set_major_locator(mdates.AutoDateLocator())
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y.%m.%d'))

    # put the labels at 45deg since they tend to be too long
    fig.autofmt_xdate()
    plt.show()

    return
    #df = df.assign(YEAR = [x.year for x in df['TIMESTAMP']])
    #df = df.assign(MONTH = [x.month for x in df['TIMESTAMP']])
    #df = df.get(['USER_NAME', 'TYPE', 'YEAR', 'MONTH'])
    #df = df.groupby(['USER_NAME', 'YEAR', 'MONTH', 'TYPE']).size()
    #plt.show()


if __name__ == '__main__':
    # Create empty directories if they don't exist
    for directory in (LOG_DIRECTORY, DATA_DIRECTORY):
        if not os.path.exists(directory):
            os.makedirs(directory)

    # Argument parsing
    parser = argparse.ArgumentParser()
    parser.add_argument('--disable-stdout', action='store_true',
                        help='Disable output to stdout')
    parser.add_argument('--disable-logging', action='store_true',
                        help='Disable logging')
    parser.add_argument('--reload-data', action='store_true',
                        help='Enable reloading locally cached data')
    parser.add_argument('--log-level', choices=LOG_LEVELS,
                        help='Set log level accordingly')
    args = parser.parse_args()

    # Create logger
    logger = logging.getLogger()
    formatter = logging.Formatter('%(asctime)s %(levelname)-8s %(message)s')
    if not args.log_level:
        logger.setLevel(logging.INFO)
    else:
        logger.setLevel(LOG_LEVELS[args.log_level])

    if not args.disable_stdout:
        handler_stream = logging.StreamHandler()
        handler_stream.setFormatter(formatter)
        logger.addHandler(handler_stream)

    if not args.disable_logging:
        handler_file = logging.FileHandler(LOG_FILE)
        handler_file.setFormatter(formatter)
        logger.addHandler(handler_file)

    # Load and preprocess data
    users = get_users(reload=args.reload_data)
    drinks = get_drinks(reload=args.reload_data)

    scans = get_scans(reload=args.reload_data)
    scans_full = drinks.merge(users.merge(scans, 'inner', 'USER_ID'), 'inner', 'EAN')

    recharges = get_recharges(reload=args.reload_data)
    recharges_full = users.merge(users.merge(recharges, 'inner', 'USER_ID'), 'inner', left_on='USER_ID', right_on='HELPER_USER_ID')
    recharges_full = recharges_full.drop(columns=('HELPER_USER_ID'))
    recharges_full = recharges_full.rename(columns={
        'USER_NAME_x' : 'HELPER_USER_NAME',
        'USER_ID_x' : 'HELPER_USER_ID',
        'USER_NAME_y' : 'USER_NAME',
        'USER_ID_y' : 'USER_ID',
    })

    # Print options for pandas
    pd.set_option('display.max_columns', 1000)
    #pd.set_option('display.max_rows', 1000)
    pd.set_option('display.expand_frame_repr', False)

    # Actual plotting
    plot_drink_consumption_by_name(scans_full)
    plot_drink_consumption_by_category(scans_full)
    plot_drink_total_consumption_by_person(scans_full)
    plot_drink_categorized_consumption_by_person(scans_full)
    plot_drink_categorized_consumption_by_person_and_time_of_day(scans_full)
    plot_recharge_vouching_by_person(recharges_full)
    plot_drink_consumption_by_name_over_time(scans_full)


    #plot_recharge_vouching_correlation_by_person(recharges_full)