Fix remove button
Bunch of cleanup