Feature/code enhance