Code update in order to look better and well organized