Add restart button
Some fixes to animation