Correct small bugs and nearly finish refactoring