restructure code to make more sense.
label potential bug