Add restart function to finish screen

Also added some comments