#!/usr/bin/env python3
# vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4

"""
Simulate button presses for assignment.
"""
import sys
import random
import os
import matplotlib.pyplot as pyplot
import math

class Person:
    def __init__(self):
        # Rat's remaining lifespan in weeks
        self.rat = 0
        # Number of weeks until out of fertilizer
        self.fertilizer = 0

    def update(self, weeks):
        # Reduce the remaining rat/fertilizer
        self.rat = max(self.rat - weeks, 0)
        self.fertilizer = max(self.fertilizer - weeks, 0)

    def choose_button(self):
        if self.rat == 0 and self.fertilizer == 0:
            return "c"
        elif self.rat == 0:
            return "b"
        elif self.fertilizer == 0:
            return "a"
        else:
            return None

    def add_weeks(self, item, weeks):
        # Add weeks to the item received from machine
        if item == "rat":
            self.rat += weeks
        elif item == "fertilizer":
            self.fertilizer += weeks

if __name__ == "__main__":
    # Assume a rat lifespan of 130 weeks
    rat = 30
    fertilizer = int(sys.argv[1])
    machine = int(sys.argv[2])
    visits = int(sys.argv[3]) // machine
    if os.path.exists(sys.argv[4]):
        print("Output file {} exists. Exiting.".format(sys.argv[4]))
        sys.exit(1)
    else:
        fig_file = sys.argv[4]

    people = [Person() for i in range(100)]
    buttons = []
    items = {'rat': rat, 'fertilizer': fertilizer}

    for i in range(visits):
        presses = {"a": 0, "b": 0, "c": 0}
        for p in people:
            p.update(machine)
            press = p.choose_button()
            if press == "c":
                item = random.choice(list(items.keys()))
            elif press == "b":
                item = 'rat'
            elif press == "a":
                item = 'fertilizer'
            else:
                continue
            weeks = items[item]
            p.add_weeks(item, weeks)
            presses[press] += 1
        buttons.append(presses)

    # Make the lists to plot, leading zeroes to get cumulative total
    a, b, c, = [0], [0], [0]
    for button in buttons:
        a.append(a[-1] + button["a"])
        b.append(b[-1] + button["b"])
        c.append(c[-1] + button["c"])
    # drop the leading zero
    a.pop(0)
    b.pop(0)
    c.pop(0)

    pyplot.plot(a)
    pyplot.plot(b)
    pyplot.plot(c)
    pyplot.legend(["fertilizer", "rat", "random"])
    pyplot.xlabel("Visits to machine")
    pyplot.ylabel("Cumulative button presses")
    pyplot.title("Visits every {} weeks, fertilizer lasts {} weeks".format(
        machine, fertilizer))

    pyplot.savefig(fig_file)