UPDATE - minor refactor and changes