SImplify logic further