Feature/clean up