Cleaned up and modified some code