Cleaned up and simplified code for final submission