Code refactoring. Removed all the warnings from IDEA where possible