Refactor the code, and added JavaDocs