- add header file
- optimize code