Refactoring and small bug fix regarding element selection