Refactor & Fix Minor Bugs