Cleaning up; adding missing header file