Added state machine back into main.c