Cleaned up and organized code better