ocsparse.py

import xml.etree.cElementTree as ET
import operator
from pprint import pprint, pformat

# def etree_to_dict(t):
#     d = {t.tag : list(map(etree_to_dict, iter(t)))}
#     d.update(('@' + k, v) for k, v in t.attrib.items())
#     d['text'] = t.text
#     return d

def etree_to_dict(t):
    if len(t) > 0:
        return {t.tag: [etree_to_dict(x) for x in t]}
    else:
        return {t.tag: t.text}


def can_have_successors(action):
    for step in action.find("Steps"):
        if step.findtext("PrimitiveAction") == "DIE":
            return False
    return True


def can_be_successor(parent, child, default_value=True, optimistic_chain_detection=False):
    """Determine if child can be performed after parent.

    Always returns False on definite exclusion. Returns default_value if
    possibility is not definitively excluded. If optimistic_chain_detection is
    True, returns True if child has any positive conditions on U_ACTION, and
    can be satisfied."""
    if not can_have_successors(parent):
        return False

    # TODO: handle conditions on things other than U_ACTION
    possible_chain = False
    for condition in child.find("Conditions"):
        resource_name = condition.findtext("Resource")
        if resource_name == "U_ACTION":
            comparison = condition.findtext("Comparison")
            if comparison != "NOT_EQUAL":
                possible_chain = True
            value = int(condition.findtext("Value"))
            parent_num = int(parent.findtext("ActionNumber"))
            comparator = {
                "EQUAL": operator.eq,
                "GEQUAL": operator.ge,
                "LEQUAL": operator.le,
                "NOT_EQUAL": operator.ne
            }[comparison]
            if not comparator(parent_num, value):
                return False
    if optimistic_chain_detection and possible_chain:
        return True

    return default_value


def succeeds_in_action_chain(parent, child):
    return can_have_successors(child) and can_be_successor(parent, child, False, True)


def prettify_action_graph(graph, action_map):
    prettified = {}
    for action, successors in graph.items():
        action_name = action_map[action].findtext("Identifier")
        new_successors = [action_map[s].findtext("Identifier") for s in successors]
        prettified[action_name] = new_successors
    return prettified


def prettify_relatives_graph(graph, action_map):
    prettified = {}
    for action, neighbors in graph.items():
        action_name = "%s (%d)" % (action_map[action].findtext("Identifier"), action)
        prettified[action_name] = prettify_relatives_graph(neighbors, action_map)
    return prettified


def get_action_parents(child, actions):
    parents = []
    for action in actions:
        if succeeds_in_action_chain(action, child):
            parents.append(int(action.findtext("ActionNumber")))
    return parents


def get_action_children(parent, actions):
    children = []
    for action in actions:
        if succeeds_in_action_chain(parent, action):
            children.append(int(action.findtext("ActionNumber")))
    return children


def get_action_ancestors(child, actions, action_map):
    parents = get_action_parents(child, actions)
    ancestors = {int(action_map[parent].findtext("ActionNumber")): get_action_ancestors(action_map[parent], actions, action_map) for parent in parents}
    return ancestors


def get_action_descendants(child, actions, action_map):
    children = get_action_children(child, actions)
    descendants = {int(action_map[child].findtext("ActionNumber")): get_action_descendants(action_map[child], actions, action_map) for child in children}
    return descendants


def get_action_tree_duration(tree, action_map):
    dur = 0
    for action_num, children in tree.items():
        dur = max(dur, int(action_map[action_num].findtext("TimeTaken")) + get_action_tree_duration(children, action_map))
    return dur

def dumpinfo(objectclass):
    print(objectclass.findtext("Name"))

    action_graph = {}
    action_map = {}
    for action in objectclass.find("Actions"):
        number = int(action.findtext("ActionNumber"))
        action_graph[number] = set()
        action_map[number] = action

    attack_chain_members = set()
    for action in objectclass.find("Actions"):
        # pprint(etree_to_dict(action))
        for step in action.find("Steps"):
            if step.findtext("PrimitiveAction") == "ATTACK_POSITION":
                attack_chain_members.add(int(action.findtext("ActionNumber")))
                break

    for member in attack_chain_members:
        # This may go into an infinite loop if the OCS XML has loops.
        # Members can apparently have multiple children. TODO: investigate.
        # Currently assuming that means they run concurrently.

        ancestors = get_action_ancestors(action_map[member], objectclass.find("Actions"), action_map)
        print("Attack chain with %s (%d):" % (action_map[member].findtext("Identifier"), member))
        self_dur = int(action_map[member].findtext("TimeTaken"))
        print("Self duration: %d" % self_dur)
        print("Ancestors: ", prettify_relatives_graph(ancestors, action_map))
        ancestor_dur = get_action_tree_duration(ancestors, action_map)
        print("Duration: %d" % ancestor_dur)
        descendants = get_action_descendants(action_map[member], objectclass.find("Actions"), action_map)
        print("Descendants: ", prettify_relatives_graph(descendants, action_map))
        descendant_dur = get_action_tree_duration(descendants, action_map)
        print("Duration: %d" % descendant_dur)
        print("Total duration: %d\n" % (ancestor_dur + descendant_dur + self_dur))

    # pprint(prettify_action_graph(action_graph, action_map))

    print("----")

if __name__ == "__main__":
    tree = ET.parse('Achron.ocs.xml')
    objectclasses = {oc.findtext("Name"): oc for oc in tree.find("ObjectClasses")}
    # dumpinfo(objectclasses["Marine"])
    for oc in sorted(objectclasses.values(), key=lambda x: int(x.findtext("ClassNumber"))):
        dumpinfo(oc)