Minor cleanup
Add controller