make the code simpler and fix issue with ab