Add specs and break up rebalance method