# Commonsense memory for the agent
# Written by Jiageng Mao & Junjie Ye

class CommonSenseMemory:
    def __init__(self) -> None:
        self.common_sense = {
            "Traffic Rules": TRAFFIC_RULES,
        }

    def retrieve(self, knowledge_types: list = None):
        commonsense_prompt = "\n"
        if knowledge_types is not None:
            for knowledge_type in knowledge_types:
                commonsense_prompt += ("*"*5 + knowledge_type + ":" + "*"*5 + "\n")
                for rule in self.common_sense[knowledge_type]:
                    commonsense_prompt += ("- " + rule + "\n")
        else: # fetch all knowledge
            for knowledge_type in self.common_sense.keys():
                commonsense_prompt += ("*"*5 + knowledge_type + ":" + "*"*5 + "\n")
                for rule in self.common_sense[knowledge_type]:
                    commonsense_prompt += ("- " + rule + "\n")
        return commonsense_prompt

TRAFFIC_RULES = [
    "Avoid collision with other objects.",
    "Always drive on drivable regions.",
    "Avoid driving on occupied regions.",
    "Pay attention to your ego-states and historical trajectory when planning.",
    "Maintain a safe distance from the objects in front of you.",
]