Fixed some more bugs

did not quite get how the dicts worked, but found out eventually. Forgot to import itertools. Also added my test