Merged #4 and fixed some bugs