Feature/code cleanup