Refact code to make it more strong