Add comments
Improve module structure