fix some part of code based on arthur's comments