Adding logging, cleaning up