bug fix

rearrange code