added some functionality, rewrite some code as 'pythonic'