Cleaned up code in main function and moved that to class.