Bug fixes and some code re-factoring.