Refactor code and prepare of assigment 2