Refactor code

Just to satisfy myself