Cleaned up some of the code around overriding the main class, to make it simpler and more configurable