Added code to main.c