Updated and cleaned up code