More minor fixes and optimized logic