Cleaned up code; fixed button removal.