Update category menu behaviour on smaller screens and clean up code.