Bug fix & clean up the code