Clean up code

Cleaned up the code before adding main flow