Refactor the whole code and rearrange classes