Add SaveState into certain functions.  Slightly bugged atm