Bug fixes and cleaning up ready for live testing