from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser

insights_mind_map_prompt_template = """
Based on the provided excerpt, create a detailed outline in Markdown format that organizes the content into distinct sections for the experiment setup and separate branches for each analysis mentioned. 
This structure should provide a clear and comprehensive overview of the study, emphasizing the setup, methodologies, findings, and implications of each analysis independently. 
Structure the outline as follows:

- **Experiment Setup**:
    - Description of the general setup for the experiments, including any datasets and baselines used.

For each analysis mentioned in the excerpt, create individual branches as follows:

- **Analysis [Name/Number]**:
    1. **Objectives**: 
        - The main goal or purpose of this specific analysis.
    2. **Methods**: 
        - Detailed overview of the approach, techniques, or methodologies employed in this analysis.
    3. **Results and Conclusions**: 
        - Key findings, data interpretation, and the implications drawn from this analysis.
        - Combine this section if the results and conclusions are closely aligned to streamline content.
    4. **Figures and Tables**:
        - Mention and describe any figures or tables referenced, linking them to their corresponding insights.

Ensure that the outline is structured in Markdown format for clarity, facilitating its integration into documents or presentations. This structured approach aims to provide a comprehensive view of the experimental framework and the individual analyses, highlighting their unique contributions to the overall study."

excerpt: {paper}
"""

insights_mind_map_output_parser = StrOutputParser()
insights_mind_map_prompt = PromptTemplate(
    template=insights_mind_map_prompt_template,
    input_variables=["paper"],
)
insights_mind_map_chain = (
    lambda model: insights_mind_map_prompt | model | insights_mind_map_output_parser
)