Cleaned up and reorganized code to make it more readable