code reorg and fix stupid bug