Closed #9 & Some cleanup