Fix/update headers a little bit