Refactored code to use python functions instread of classes