Add missing `#include` statements
Update python wrapper