Refactoring, tidyup and improve layering