Code cleanup/improvements
Bug fix & optimization