Cleaned up ugly spaghetti code and made it work!