Started to restructure the code