* Remove goofy closure thing to do the argmax. Just use a list comprehension instead