#!/usr/bin/env python3

"""

Tool for asking a question about the simulation results.

"""

import logging

from typing import Type, Annotated, Literal

import hydra

import basico

import pandas as pd

from pydantic import BaseModel, Field

from langchain_core.tools.base import BaseTool

from langchain_experimental.agents import create_pandas_dataframe_agent

from langgraph.prebuilt import InjectedState

# Initialize logger

logging.basicConfig(level=logging.INFO)

logger = logging.getLogger(__name__)

class AskQuestionInput(BaseModel):

    """

    Input schema for the AskQuestion tool.

    """

    question: str = Field(description="question about the simulation and steady state results")

    experiment_name: str = Field(description="""Name assigned to the simulation

                                            or steady state analysis when the tool 

                                            simulate_model or steady_state is invoked.""")

    question_context: Literal["simulation", "steady_state"] = Field(

        description="Context of the question")

    state: Annotated[dict, InjectedState]

# Note: It's important that every field has type hints.

# BaseTool is a Pydantic class and not having type hints

# can lead to unexpected behavior.

class AskQuestionTool(BaseTool):

    """

    Tool for asking a question about the simulation or steady state results.

    """

    name: str = "ask_question"

    description: str = """A tool to ask question about the

                        simulation or steady state results."""

    args_schema: Type[BaseModel] = AskQuestionInput

    return_direct: bool = False

    def _run(self,

             question: str,

             experiment_name: str,

             question_context: Literal["simulation", "steady_state"],

             state: Annotated[dict, InjectedState]) -> str:

        """

        Run the tool.

        Args:

            question (str): The question to ask about the simulation or steady state results.

            state (dict): The state of the graph.

            experiment_name (str): The name assigned to the simulation or steady state analysis.

        Returns:

            str: The answer to the question.

        """

        logger.log(logging.INFO,

                   "Calling ask_question tool %s, %s, %s",

                   question,

                   question_context,

                   experiment_name)

        # Load hydra configuration

        with hydra.initialize(version_base=None, config_path="../configs"):

            cfg = hydra.compose(config_name='config',

                                overrides=['tools/ask_question=default'])

            cfg = cfg.tools.ask_question

        # Get the context of the question

        # and based on the context, get the data

        # and prompt content to ask the question

        if question_context == "steady_state":

            dic_context = state["dic_steady_state_data"]

            prompt_content = cfg.steady_state_prompt

        else:

            dic_context = state["dic_simulated_data"]

            prompt_content = cfg.simulation_prompt

        # Extract the

        dic_data = {}

        for data in dic_context:

            for key in data:

                if key not in dic_data:

                    dic_data[key] = []

                dic_data[key] += [data[key]]

        # Create a pandas dataframe of the data

        df_data = pd.DataFrame.from_dict(dic_data)

        # Extract the data for the experiment

        # matching the experiment name

        df = pd.DataFrame(

            df_data[df_data['name'] == experiment_name]['data'].iloc[0]

        )

        logger.log(logging.INFO, "Shape of the dataframe: %s", df.shape)

        # # Extract the model units

        # model_units = basico.model_info.get_model_units()

        # Update the prompt content with the model units

        prompt_content += "Following are the model units:\n"

        prompt_content += f"{basico.model_info.get_model_units()}\n\n"

        # Create a pandas dataframe agent

        df_agent = create_pandas_dataframe_agent(

                        state['llm_model'],

                        allow_dangerous_code=True,

                        agent_type='tool-calling',

                        df=df,

                        max_iterations=5,

                        include_df_in_prompt=True,

                        number_of_head_rows=df.shape[0],

                        verbose=True,

                        prefix=prompt_content)

        # Invoke the agent with the question

        llm_result = df_agent.invoke(question, stream_mode=None)

        # print (llm_result)

        return llm_result["output"]