Cleaned up main class