Fixed #5 and #6

Done :)