Cleaning up some code to fix outside of screen bugs