Updating code, small changes, fixes #3