tracker

#!/usr/bin/env python3

from datetime import datetime, timedelta
import math
import os
import sys
import sqlite3
import subprocess

# Data stored in sqlite so it can be shared easily
db_file = os.path.expanduser("~") + "/.tracker_data.db"
tracker_db = sqlite3.connect(db_file)
db = tracker_db.cursor()
db.execute('''CREATE TABLE IF NOT EXISTS trackers
                (name text, day date, value real)''')
tracker_db.commit()


def error(message = None):
    if message is not None:
        print(message)
    exit(1)

def help(command = None):
    """Display usage information, exit program.
    """
    if command is None:
        option = "   {:12} {}"
        help_text = [
            "Usage: tracker <command> [<args>]",
            "",
            "Available commands:",
            option.format("help",   "display this dialog"),
            option.format("update", "save data to tracker"),
            option.format("list",   "list available trackers"),
            option.format("show",   "display raw tracker data"),
            option.format("rename", "rename tracker"),
            option.format("delete", "remove tracker"),
            option.format("stats",  "show statistics"),
            option.format("plot",   "show graph"),
            "",
            "Use 'tracker help <command>' for a command's detailed usage."
        ]
        print("\n".join(help_text))
    else:
        # commands = ["update", "list", "show", "rename", "delete", "stats", "plot"]
        usage = "   {}"
        desc  = "      {}"
        if command == "update":
            help_text = [
                "Update: command which adds (numerical) data to a tracker.",
                "",
                "Usage:",
                usage.format("tracker update <tracker> <data>"),
                usage.format("tracker update <tracker>"),
                "",
                "Description:",
                usage.format("tracker update <tracker> <data>"),
                desc.format("This form is shorthand for saving <data> to " +
                    "<tracker> for today's date."),
                "",
                usage.format("tracker update <tracker>"),
                desc.format("This form is used to set the value for an " +
                    "arbitrary date for <tracker>."),
                desc.format("The date must be in the format YYYY-MM-DD."),
                "",
                "Options:",
                usage.format("<tracker>"),
                desc.format("The name of the tracker to update, converted to lowercase."),
                desc.format("If <tracker> does not exist, you will be prompted to create it."),
                "",
                usage.format("<data>"),
                desc.format("The value to save to the tracker to update, must be numerical.")
            ]
        elif command == "list":
            help_text = [
                "List: displays a list of trackers which have been created",
                "",
                "Usage:",
                usage.format("tracker list")
            ]
        elif command == "show":
            help_text = [
                "Show: displays raw dates and values for a tracker",
                "",
                "Usage:",
                usage.format("tracker show <tracker>"),
                "",
                "Description:",
                usage.format("tracker show <tracker>"),
                desc.format("Displays all data for <tracker> in the form '<date> | <value>'."),
                desc.format("Note: <date> is formatted as 'YYYY-MM-DD'."),
                "",
                "Options:",
                usage.format("<tracker>"),
                desc.format("The name of the tracker to show, converted to lowercase."),
            ]
        elif command == "rename":
            help_text = [
                "Rename: change name of a tracker",
                "",
                "Usage:",
                usage.format("tracker rename <tracker> <new_tracker>"),
                "",
                "Description:",
                usage.format("tracker rename <tracker> <new_tracker>"),
                desc.format("All <tracker> entries will not be <new_tracker> entries."),
                "",
                "Options:",
                usage.format("<tracker>"),
                desc.format("The name of the existing tracker to change, converted to lowercase."),
                "",
                usage.format("<new_tracker>"),
                desc.format("The name of the new tracker (must not already exist), converted to lowercase.")
            ]
        elif command == "delete":
            help_text = [
                "Delete: permanently removes all data entries for a given tracker",
                "",
                "Usage:",
                usage.format("tracker delete <tracker>"),
                "",
                "Description:",
                usage.format("tracker delete <tracker>"),
                desc.format("All sqlite entries associated with <tracker> are deleted."),
                "",
                "Options:",
                usage.format("<tracker>"),
                desc.format("The name of the tracker to delete, converted to lowercase.")
            ]
        elif command == "stats":
            help_text = [
                "Stats: show statistics for tracker(s)",
                "",
                "Usage:",
                usage.format("tracker stats <tracker> <tracker>"),
                usage.format("tracker stats <tracker>"),
                usage.format("tracker stats"),
                "",
                "Description:",
                usage.format("tracker stats <tracker> <tracker>"),
                desc.format("Show correlation coefficient between two trackers."),
                "",
                usage.format("tracker stats <tracker>"),
                desc.format("Display information for each weekday and entire time period."),
                desc.format("Stats included: total, mean, min, max."),
                "",
                usage.format("tracker stats"),
                desc.format("Displays information about all trackers."),
                desc.format("Stats included: total entries, entries per tracker."),
                "",
                "Options:",
                usage.format("<tracker>"),
                desc.format("The name of the tracker to show stats for, converted to lowercase.")
            ]
        elif command == "plot":
            help_text = [
                "Plot: show graph for tracker",
                "",
                "Usage:",
                usage.format("tracker plot <tracker>"),
                "",
                "Description:",
                usage.format("tracker stats <tracker>"),
                desc.format("Displays graph for <tracker> from first entry to last entry."),
                "",
                "Options:",
                usage.format("<tracker>"),
                desc.format("The name of the tracker to graph, converted to lowercase.")
            ]
        else:
            error("Invalid command: '{}'".format(command))
        print("\n".join(help_text))
    sys.exit(1)


def str_to_date(date_string, fmt="%Y-%m-%d"):
    """Converts string in given format to date

    Input:
        date_string: string representing a date
        fmt: format string for date (optional)
    Output:
        date: datetime.date() object corresponding to date_string
    """
    return datetime.strptime(date_string, fmt).date()


def tracker_list():
    """Returns list of tracker names"""
    trackers = db.execute("SELECT DISTINCT name FROM trackers")
    names = [tup[0] for tup in trackers.fetchall()]
    return names


def tracker_exists(tracker, trackers = None):
    """Return whether given tracker exists

    Input:
        tracker: name of tracker to check
        trackers: optional list of trackers to determine existance with
    """
    if trackers is None:
        trackers = tracker_list()
    return tracker in trackers


def fetch_all(tracker = None):
    """Returns all data entries for the given tracker

    Input:
        tracker: name of tracker to return data for
    Output:
        List of entries for tracker. If no tracker, returns all entries
    """
    if tracker:
        query = "SELECT * FROM TRACKERS where name=:tracker ORDER BY day ASC"
        cursor = db.execute(query, {"tracker": tracker})
    else:
        query = "SELECT * FROM TRACKERS ORDER BY day ASC"
        cursor = db.execute(query)
    return cursor.fetchall()


def count_entries(tracker = None):
    """Returns number of data entries for the given tracker

    Input:
        tracker: optional name of tracker to return data for
    Output:
        Number of entries for tracker. If no tracker, returns total entry count
    """
    if tracker:
        query = "SELECT Count(*) FROM TRACKERS where name=:tracker"
        cursor = db.execute(query, {"tracker": tracker})
    else:
        query = "SELECT Count(*) FROM TRACKERS"
        cursor = db.execute(query)
    return cursor.fetchone()[0]


def show_tracker(tracker_info):
    """Show data for list of trackers

    Input:
        tracker_info: list in the form [tracker, ...]
    """
    if not tracker_info:
        error("Must provide tracker to show")
    entry = "{} | {}"
    for t in tracker_info:
        if not tracker_exists(t):
            error("Cannot show data for unknown tracker '" + t + "'")
        results = fetch_all(t)
        for r in results:
            print(entry.format(r[1], r[2]))


def rename_tracker(tracker_info):
    """Rename the given tracker

    Input:
        tracker_info: list in the form [tracker, new_tracker]
    """
    if len(tracker_info) == 0:
        error("Must supply tracker to rename")
    elif len(tracker_info) != 2:
        error("Must suppply tracker and new tracker name")
    tracker = tracker_info[0]
    new_tracker = tracker_info[1]
    if not tracker_exists(tracker):
        error("Tracker '" + tracker + "' not found")
    if tracker_exists(new_tracker):
        error("Tracker '" + new_tracker + "' already exists")
    db.execute('''UPDATE trackers SET name = :new_name
                    WHERE name = :tracker''',
                    {"new_name": new_tracker, "tracker": tracker})
    tracker_db.commit()
    print("'" + tracker + "' renamed to '" + new_tracker + "'")


def update_tracker(tracker_info):
    """Update tracker with new data

    Input:
        tracker_info: list in the form [tracker, data]
    """
    # Ensure tracker and data are provided
    if len(tracker_info) == 0:
        error("Must supply tracker and data to update")
    elif len(tracker_info) > 2:
        error("Too many arguments given for update command")

    # Ensure tracker exists
    tracker = tracker_info[0]
    if not tracker_exists(tracker):
        print("Tracker '" + tracker + "' not found")
        confirm = input("Confirm '" + tracker + "' to create: ").strip()
        if confirm != tracker:
            error("Update failed")

    if len(tracker_info) == 1:
        date = None
        while not date:
            date = input("Provide date YYYY-MM-DD (leave blank for today): ").strip()
            if not date:
                date = datetime.today().date()
                break
            try:
                date = datetime.strptime(date, "%Y-%m-%d").date()
            except ValueError:
                print("Error: Invalid date provided")
                date = None
        data = input("Value: ").strip()
    else:
        data = tracker_info[1]
        date = datetime.today().date()

    # Ensure data is a number
    try:
        data = float(data)
    except ValueError:
        error("'" + data + "' is not numerical, update failed")
    # Update existing row for today's value if possible
    db.execute('''UPDATE trackers SET value = value + :data
                    WHERE name = :tracker AND day = :date''',
                    {"data": data, "tracker": tracker, "date": date})
    # Otherwise, insert a new row
    if not db.rowcount:
        db.execute("INSERT INTO trackers VALUES (?, ?, ?)",
                (tracker, date, data))
    tracker_db.commit()
    print("'" + tracker + "' updated")


def display_list(args = None):
    """Lists available trackers"""
    trackers = tracker_list()
    if len(trackers) == 0:
        print("No trackers found")
    else:
        print("Current Trackers:")
        tracker_option = "   {}"
        for t in trackers:
            print(tracker_option.format(t))


def delete_tracker(tracker):
    """Delete existing tracker

    Input:
        tracker_info: existing tracker
    """
    # Ensure tracker is provided
    if len(tracker) != 1:
        error("Must supply tracker to delete")
    tracker = tracker[0]
    if not tracker_exists(tracker):
        error("Cannot delete unknown tracker '" + tracker + "'")

    # User must confirm tracker to delete all entries
    confirm = input("Type '" + tracker + "' to confirm: ").strip()
    if confirm == tracker:
        db.execute("DELETE FROM trackers WHERE name=:tracker", {"tracker": tracker})
        tracker_db.commit()
        print("Tracker '" + tracker + "' deleted")
    else:
        print("Deletion cancelled")


def longest_streak(tracker_data):
    """Find longest streak for a tracker

    Input:
        tracker_data: list of entries for a tracker
    Output:
        longest number of consecutive entries
    """
    max_streak = 0
    streak = 0
    streak_start = None
    streak_end = None
    current_start = None
    last_day = None
    for entry in tracker_data:
        entry_day = str_to_date(entry[1])
        next_day = (entry_day - timedelta(days=1))
        if last_day == next_day:
            streak += 1
        else:
            streak = 1
            current_start = entry_day
        if streak > max_streak:
            max_streak = streak
            streak_start = current_start
            streak_end = entry_day
        last_day = entry_day
    return max_streak, streak_start, streak_end


def longest_break(tracker_data):
    """Find longest streak for a tracker

    Input:
        tracker_data: list of entries for a tracker
    Output:
        longest number of consecutive entries
    """
    max_gap = 0
    gap_start = None
    gap_end = None
    previous = None
    for entry in tracker_data:
        entry_day = str_to_date(entry[1])
        if previous:
            gap = (entry_day - previous).days
            if gap > max_gap:
                max_gap = gap
                gap_start = previous
                gap_end = entry_day
        previous = entry_day
    return max_gap, gap_start, gap_end


def correlation(x, y):
    """Find correlation coefficient for two sets of data

    Input:
        x: list of tuples in the form [(date, value), ...]
        y: list of tuples in the form [(date, value), ...]
    Output:
        r: correlation coefficient
        error: data has no variance
    """
    # Use data in natural key-value form
    xs = {}
    for (_, date, value) in x:
        xs[date] = value
    ys = {}
    for (_, date, value) in y:
        ys[date] = value

    # Fill 0s for missing dates
    for d in set(ys.keys()) - set(xs.keys()):
        xs[d] = 0
    for d in set(xs.keys()) - set(ys.keys()):
        ys[d] = 0

    x_avg = sum(xs.values()) / len(xs.values())
    y_avg = sum(ys.values()) / len(ys.values())

    # Pearson correlation coefficient for given sample
    covariance = 0
    x_variance = 0
    y_variance = 0
    for d in xs.keys():
        x_diff = xs[d] - x_avg
        y_diff = ys[d] - y_avg
        covariance += x_diff * y_diff
        x_variance += math.pow(x_diff, 2)
        y_variance += math.pow(y_diff, 2)
    if x_variance == 0:
        return -1
    elif y_variance == 0:
        return -2
    return covariance / (math.sqrt(x_variance) * math.sqrt(y_variance))


def justify_column(column):
    """Justifies floats based on longest string in column

    Input:
        column: list of float values
    Output:
        column_strings: list of justified, formatted floats as strings
        justify: amount to justify by to match formatted data
    """
    formatted_strings = ['{:,.3f}'.format(v) for v in column]
    justify = len(max(formatted_strings, key=len))
    column_strings = [s.rjust(justify) for s in formatted_strings]
    return (column_strings, justify)


def display_stats(trackers):
    """Show stats for trackers

    Input:
        trackers: list of tracker names, empty implies global stats
    """
    if len(trackers) == 0:
        # number of trackers
        trackers = tracker_list()
        print("{} tracker(s)".format(len(trackers)))

        # total number of entries
        entries = count_entries()
        print("{} total entries".format(entries))

        # number of entries per tracker
        for t in trackers:
            print()
            print(t)
            t_entries = count_entries(t)
            print("   {} entries".format(t_entries))

    elif len(trackers) == 2:
        t1 = trackers[0]
        t2 = trackers[1]
        if not tracker_exists(t1):
            error("Cannot show stats for unknown tracker '" + t1 + "'")
        if not tracker_exists(t2):
            error("Cannot show stats for unknown tracker '" + t2 + "'")
        r = correlation(fetch_all(t1), fetch_all(t2))
        if r == -1:
            error("Error: '{}' data has no variance".format(t1))
        elif r == -2:
            error("Error: '{}' data has no variance".format(t2))
        print("Correlation coefficient: {}".format(r))
    else:
        days = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
        for t in trackers:
            if not tracker_exists(t):
                error("Cannot show stats for unknown tracker '" + t + "'")
            results = fetch_all(t)
            num_entries = len(results)
            counts = [0] * len(days)
            tally = [0] * len(days)
            mins = [None] * len(days)
            maxs = [None] * len(days)
            min_value = None
            max_value = None
            for r in results:
                count = r[2]
                day = str_to_date(r[1])
                day_index = day.weekday()
                counts[day_index] += count
                tally[day_index] += 1
                if not min_value or min_value[0] > count:
                    min_value = (count, day)
                if not max_value or max_value[0] < count:
                    max_value = (count, day)
                if not mins[day_index] or count < mins[day_index]:
                    mins[day_index] = count
                if not maxs[day_index] or count > maxs[day_index]:
                    maxs[day_index] = count


            print("{}: {} entries\n".format(t, num_entries))

            # TODO: decide how to format output for better readability
            # Give general tracker stats
            total = sum(counts)
            avg = total/sum(tally)
            gen = [total, avg, min_value[0], max_value[0]]
            gen = ['{:,.3f}'.format(v) for v in gen]
            gen_len = len(max(gen, key=len))

            # format string for single data point
            point = "   {:3} -- {}"
            print(point.format("sum", gen[0].rjust(gen_len)))
            print(point.format("avg", gen[1].rjust(gen_len)))
            print(point.format("min", gen[2].rjust(gen_len)) +
                    " (" + str(min_value[1]) + ")")
            print(point.format("max", gen[3].rjust(gen_len)) +
                    " (" + str(max_value[1]) + ")")
            print()

            # Give stats about consistency
            streak = longest_streak(results)
            gap = longest_break(results)

            # format string for streak data
            alt_point = "   {:14} -- {:2} day(s)  ({} to {})"
            print(alt_point.format("longest streak", streak[0], streak[1], streak[2]))
            print(alt_point.format("longest break", gap[0], gap[1], gap[2]))
            print()

            # prepare min/max data for printing
            mins = [v if v else 0 for v in mins]
            maxs = [v if v else 0 for v in maxs]


            # Format data into justified columns
            counts_str, count_len = justify_column(counts)
            avgs = [None] * len(counts)
            for i in range(-1,6):
                avgs[i] = 0 if tally[i] is 0 else counts[i]/tally[i]
            avgs_str, avgs_len = justify_column(avgs)
            mins_str, mins_len = justify_column(mins)
            maxs_str, maxs_len = justify_column(maxs)

            row = "   {:8}{:<22}{:<22}{:<22}{:<22}"
            title_row = '\033[1m' + row + '\033[0m'
            print(title_row.format(
                "",
                "Total".rjust(count_len),
                "Avg".rjust(avgs_len),
                "Min".rjust(mins_len),
                "Max".rjust(maxs_len)))

            for i in range(-1,6):
                print(row.format(days[i] + " --",
                    counts_str[i],
                    avgs_str[i],
                    mins_str[i],
                    maxs_str[i]))


def display_plot(tracker):
    """Show plot for a tracker

    Input:
        tracker: tracker name to plot
    """
    if len(tracker) != 1:
        error("Only a single tracker can be plotted")
    else:
        tracker = tracker[0]
        if not tracker_exists(tracker):
            error("Cannot show stats for unknown tracker '" + tracker + "'")

        results = fetch_all(tracker)

        # extract data for plotting
        dates = [r[1] for r in results]
        values = [v[2] for v in results]

        minY = min(values) * 0.75
        maxY = max(values) * 1.1

        minX = str_to_date(min(dates)) + timedelta(days=-1)
        maxX = str_to_date(max(dates)) + timedelta(days=1)

        data = [
            # give graph a size and title
            "set term dumb 79 25\n",
            "set title '{}'\n".format(tracker),

            # set date format for gnuplot
            "set xdata time\n",
            'set timefmt "%Y-%m-%d\n',

            # lazy scaling, won't work nicely for all data ranges
            'set yrange[{}:{}]\n'.format(minY, maxY),
            'set xrange["{}":"{}"]\n'.format(minX, maxX),

            "plot '-' using 1:2 title '{}' with impulses\n".format(tracker)
        ]
        # plot data
        data.extend(["%s %f" % (d, v) for d, v in zip(dates, values)])
        data.append("exit")

        gnuplot = subprocess.Popen(["gnuplot"], stdin=subprocess.PIPE)
        gnuplot.communicate("\n".join(data).encode("utf-8"))


def handle_commands(args_in):
    """Ensures user input contains valid commands

    Input:
        args:
    """
    args = [arg.lower() for arg in args_in]

    # available commands
    # TODO consider consolidating with help()
    commands = ["update", "list", "show", "rename", "delete", "stats", "plot"]

    if len(args) == 0:
        help()

    command = args[0]
    if command in ["help", "usage"]:
        if len(args) > 2:
            error("Too many args provided for 'help'")
        elif len(args) == 2:
            help(args[1])
        else:
            help()
    if command not in commands:
        error("Invalid command: '{}'".format(args[0]))

    # Map to command's handler function
    # Remaining args are passed regardless, dealt with in handler
    handlers = {
        "update" : update_tracker,
        "list"   : display_list,
        "show"   : show_tracker,
        "rename" : rename_tracker,
        "delete" : delete_tracker,
        "stats"  : display_stats,
        "plot"   : display_plot
    }
    handlers[command](args[1:])
    

if __name__ == '__main__':
    handle_commands(sys.argv[1:])