changed code a bit and added code for main menu