Minor Bug Fix and Some Refactoring