Cleaning up code, trying to get everything to work.