[IMP] Refactoring + notes for improvement