Add header file
Refactor code