changed syntax to support e1,e2 in AX6 and add convert_to_p behavior

common/lib/capa/capa/chem/chemtools.py

@@ -1,19 +1,13 @@
-from collections import OrderedDict
 import json
 import unittest
-
+import itertools
 
 def vsepr_parse_user_answer(user_input):
-    d = OrderedDict(json.loads(user_input))
-    d['atoms'] = OrderedDict(sorted(d['atoms'].items()))
-    return d
+    return json.loads(user_input)
 
 
 def vsepr_build_correct_answer(geometry, atoms):
-    correct_answer = OrderedDict()
-    correct_answer['geometry'] = geometry
-    correct_answer['atoms'] = OrderedDict(sorted(atoms.items()))
-    return correct_answer
+    return {'geometry': geometry, 'atoms': atoms}
 
 
 def vsepr_grade(user_input, correct_answer, convert_to_perepherial=False):
@@ -22,6 +16,7 @@ def vsepr_grade(user_input, correct_answer, convert_to_perepherial=False):
             c0, a, e
             c0, p
     """
+    # import ipdb; ipdb.set_trace()
     # print user_input, type(user_input)
     # print correct_answer, type(correct_answer)
     if user_input['geometry'] != correct_answer['geometry']:
@@ -33,19 +28,45 @@ def vsepr_grade(user_input, correct_answer, convert_to_perepherial=False):
     if convert_to_perepherial:
         # convert user_input from (a,e,e1,e2) to (p)
         # correct_answer must be set in (p) using this flag
+        c0 = user_input['atoms'].pop('c0')
         user_input['atoms'] = {'p' + str(i): v for i, v in enumerate(user_input['atoms'].values())}
+        user_input['atoms']['c0'] = c0
 
-    for ea_position in ['p', 'a', 'e', 'e1', 'e2']:
-        # collecting atoms:
-        a_user = [v for k, v in user_input['atoms'].items() if k.startswith(ea_position)]
-        a_correct = [v for k, v in correct_answer['atoms'].items() if k.startswith(ea_position)]
-        # print a_user, a_correct
-        if len(a_user) != len(a_correct):
-            return False
-        if sorted(a_user) != sorted(a_correct):
-            return False
+    # special case for AX6
+    if 'e10' in correct_answer['atoms']:  # need check e1x, e2x symmetry for AX6..
+        a_user = {}
+        a_correct = {}
+        for ea_position in ['a', 'e1', 'e2']:  # collecting positions:
+            a_user[ea_position] = [v for k, v in user_input['atoms'].items() if k.startswith(ea_position)]
+            a_correct[ea_position] = [v for k, v in correct_answer['atoms'].items() if k.startswith(ea_position)]
 
-    return True
+        correct = [sorted(a_correct['a'])] + [sorted(a_correct['e1'])] + [sorted(a_correct['e2'])]
+        for permutation in itertools.permutations(['a', 'e1', 'e2']):
+            if correct == [sorted(a_user[permutation[0]])] + [sorted(a_user[permutation[1]])] + [sorted(a_user[permutation[2]])]:
+                return True
+        return False
+
+    else:  # no need to checl e1x,e2x symmetry - convert them to ex
+        if 'e10' in user_input['atoms']:  # e1x exists, it is AX6.. case
+            e_index = 0
+            for k, v in user_input['atoms'].items():
+                if len(k) == 3:  # e1x
+                    del user_input['atoms'][k]
+                    user_input['atoms']['e' + str(e_index)] = v
+                    e_index += 1
+
+        # common case
+        for ea_position in ['p', 'a', 'e']:
+            # collecting atoms:
+            a_user = [v for k, v in user_input['atoms'].items() if k.startswith(ea_position)]
+            a_correct = [v for k, v in correct_answer['atoms'].items() if k.startswith(ea_position)]
+            # print a_user, a_correct
+            if len(a_user) != len(a_correct):
+                return False
+            if sorted(a_user) != sorted(a_correct):
+                return False
+
+        return True
 
 
 class Test_Grade(unittest.TestCase):