"""A survey has a memory plan that specifies what the agent should remember when answering a question."""

from collections import UserDict, defaultdict
from typing import Optional

# from edsl.surveys.Memory import Memory
# from edsl.prompts.Prompt import Prompt
# from edsl.surveys.DAG import DAG


class MemoryPlan(UserDict):
    """A survey has a memory plan that specifies what the agent should remember when answering a question.

    The basic structure of a memory plan is a dictionary of focal questions to memories.

    {focal_question1: [prior_question1, prior_question2, ...], focal_question: [prior_question3]}
    """

    def __init__(self, survey: Optional["Survey"] = None, data: Optional[dict] = None):
        """Initialize a memory plan.

        The actual 'data' attributes of the memory plan are a dictionary of focal questions to memories.
        """
        if survey is not None:
            self.survey = survey
            self.survey_question_names = [q.question_name for q in survey.questions]
            self.question_texts = [q.question_text for q in survey.questions]
        super().__init__(data or {})

    @property
    def name_to_text(self) -> dict:
        """Return a dictionary mapping question names to question texts."""
        return dict(zip(self.survey_question_names, self.question_texts))

    def add_question(self, question: "QuestionBase") -> None:
        """Add a question to the survey.

        :param question: A question to add to the survey

        """
        self.survey_question_names.append(question.question_name)
        self.question_texts.append(question.question_text)

    def _check_valid_question_name(self, question_name: str) -> None:
        """Ensure a passed question name is valid.

        :param question_name: The name of the question to check.

        """
        if question_name not in self.survey_question_names:
            raise ValueError(
                f"{question_name} is not in the survey. Current names are {self.survey_question_names}"
            )

    def get_memory_prompt_fragment(
        self, focal_question: str, answers: dict
    ) -> "Prompt":
        """Generate the prompt fragment descripting that past question and answer.

        :param focal_question: The current question being answered.
        :param answers: A dictionary of question names to answers.

        """
        from edsl.prompts.Prompt import Prompt

        self._check_valid_question_name(focal_question)

        if focal_question not in self:
            return Prompt("")

        q_and_a_pairs = [
            (self.name_to_text[question_name], answers.get(question_name, None))
            for question_name in self[focal_question]
        ]

        base_prompt_text = """
        Before the question you are now answering, you already answered the following question(s):
        """

        def gen_line(question_text, answer):
            """Return a line of memory."""
            return f"\tQuestion: {question_text}\n\tAnswer: {answer}\n"

        lines = [gen_line(*pair) for pair in q_and_a_pairs]
        if lines:
            return Prompt(
                base_prompt_text + "\n Prior questions and answers:".join(lines)
            )
        else:
            return Prompt("")

    def _check_order(self, focal_question: str, prior_question: str) -> None:
        """Ensure the prior question comes before the focal question."""
        focal_index = self.survey_question_names.index(focal_question)
        prior_index = self.survey_question_names.index(prior_question)
        if focal_index <= prior_index:
            raise ValueError(f"{prior_question} must come before {focal_question}.")

    def add_single_memory(self, focal_question: str, prior_question: str) -> None:
        """Add a single memory to the memory plan.

        :param focal_question: The current question being answered.
        :param prior_question: The question that was answered before the focal question that should be remembered.

        >>> mp = MemoryPlan.example()
        >>> mp.add_single_memory("q0", "q1")
        Traceback (most recent call last):
        ...
        ValueError: q1 must come before q0.

        >>> mp = MemoryPlan.example()
        >>> mp.add_single_memory("q0", "crap")
        Traceback (most recent call last):
        ...
        ValueError: crap is not in the survey. Current names are ['q0', 'q1', 'q2']

        >>> mp = MemoryPlan.example()
        >>> mp.add_single_memory("crap", "q0")
        Traceback (most recent call last):
        ...
        ValueError: crap is not in the survey. Current names are ['q0', 'q1', 'q2']
        """
        self._check_valid_question_name(focal_question)
        self._check_valid_question_name(prior_question)
        self._check_order(focal_question, prior_question)
        from edsl.surveys.Memory import Memory

        if focal_question not in self:
            memory = Memory()
            memory.add_prior_question(prior_question)
            self[focal_question] = memory
        else:
            self[focal_question].add_prior_question(prior_question)

    def add_memory_collection(
        self, focal_question: str, prior_questions: list[str]
    ) -> None:
        """Add a collection of prior questions to the memory plan.

        :param focal_question: The current question being answered.
        :param prior_questions: A list of questions that were answered before the focal question that should be remembered.
        """
        for question in prior_questions:
            self.add_single_memory(focal_question, question)

    def to_dict(self, add_edsl_version=True) -> dict:
        """Serialize the memory plan to a dictionary.

        >>> mp = MemoryPlan.example()
        >>> mp.to_dict()
        {'survey_question_names': ['q0', 'q1', 'q2'], 'survey_question_texts': ['Do you like school?', 'Why not?', 'Why?'], 'data': {'q1': {'prior_questions': ['q0']}}}
        """
        newdata = {}
        for question_name, memory in self.items():
            newdata[question_name] = memory.to_dict()

        return {
            "survey_question_names": self.survey_question_names,
            "survey_question_texts": self.question_texts,
            "data": newdata,
        }

    @classmethod
    def from_dict(cls, data) -> "MemoryPlan":
        """Deserialize a memory plan from a dictionary."""
        from edsl.surveys.Memory import Memory

        newdata = {}
        for question_name, memory in data["data"].items():
            newdata[question_name] = Memory.from_dict(memory)

        memory_plan = cls(survey=None, data=newdata)
        memory_plan.survey_question_names = data["survey_question_names"]
        memory_plan.question_texts = data["survey_question_texts"]
        return memory_plan

    def _indexify(self, d: dict):
        """Convert a dictionary of question names to a dictionary of question indices.

        :param d: A dictionary of question names to indices.
        """
        new_d = {}
        for k, v in d.items():
            key = self.survey_question_names.index(k)
            new_v = set({self.survey_question_names.index(q) for q in v})
            new_d[key] = new_v
        return new_d

    @property
    def dag(self) -> "DAG":
        """Return a directed acyclic graph of the memory plan.

        >>> mp = MemoryPlan.example()
        >>> mp.dag
        {1: {0}}
        """
        from edsl.surveys.DAG import DAG

        d = defaultdict(set)
        for focal_question, memory in self.items():
            for prior_question in memory:
                d[focal_question].add(prior_question)
        return DAG(self._indexify(d))

    @classmethod
    def example(cls):
        """Return an example memory plan."""
        from edsl import Survey

        mp = cls(survey=Survey.example())
        mp.add_single_memory("q1", "q0")
        return mp

    def remove_question(self, question_name: str) -> None:
        """Remove a question from the memory plan.

        :param question_name: The name of the question to remove.
        """
        self._check_valid_question_name(question_name)

        # Remove the question from survey_question_names and question_texts
        index = self.survey_question_names.index(question_name)
        self.survey_question_names.pop(index)
        self.question_texts.pop(index)

        # Remove the question from the memory plan if it's a focal question
        self.pop(question_name, None)

        # Remove the question from all memories where it appears as a prior question
        for focal_question, memory in self.items():
            memory.remove_prior_question(question_name)

        # Update the DAG
        self.dag.remove_node(index)

    def remove_prior_question(self, question_name: str) -> None:
        """Remove a prior question from the memory."""
        self.prior_questions = [q for q in self.prior_questions if q != question_name]


if __name__ == "__main__":
    import doctest

    doctest.testmod(optionflags=doctest.ELLIPSIS)
