Added restart

Added restart option after the game is finished. The code is a bit
strange though - further updates required to make it better