Small code cleanup
Re-add screen buffer