updated code to correctly handle changing the UI when the screen size changes