improve code & bug fix