#!/usr/bin/env python3

"""
Tool for plotting a custom y-axis of a simulation plot.
"""

import logging
from typing import Type, Annotated, List, Tuple, Union, Literal
from pydantic import BaseModel, Field
import hydra
import pandas as pd
from langchain_core.tools import BaseTool
from langchain_core.prompts import ChatPromptTemplate
from langgraph.prebuilt import InjectedState
from .load_biomodel import ModelData, load_biomodel
from .utils import get_model_units

# Initialize logger
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def extract_relevant_species(question, species_names, state):
    """
    Extract the relevant species from the user question.

    Args:
        question (str): The user question.
        species_names (list): The species names available in the simulation results.
        state (dict): The state of the graph.

    Returns:
        CustomHeader: The relevant species
    """
    # In the following code, we extract the species
    # from the user question. We use Literal to restrict
    # the species names to the ones available in the
    # simulation results.
    class CustomHeader(BaseModel):
        """
        A list of species based on user question.

        This is a Pydantic model that restricts the species
        names to the ones available in the simulation results.
        
        If no species is relevant, set the attribute
        `relevant_species` to None.
        """
        relevant_species: Union[None, List[Literal[*species_names]]] = Field(
                description="This is a list of species based on the user question."
                "It is restricted to the species available in the simulation results."
                "If no species is relevant, set this attribute to None."
                "If the user asks for very specific species (for example, using the"
                "keyword `only` in the question), set this attribute to correspond "
                "to the species available in the simulation results, otherwise set it to None."
                )
    # Load hydra configuration
    with hydra.initialize(version_base=None, config_path="../configs"):
        cfg = hydra.compose(config_name='config',
                            overrides=['tools/custom_plotter=default'])
        cfg = cfg.tools.custom_plotter
    # Get the system prompt
    system_prompt = cfg.system_prompt_custom_header
    # Create an instance of the LLM model
    logging.log(logging.INFO, "LLM model: %s", state['llm_model'])
    llm = state['llm_model']
    llm_with_structured_output = llm.with_structured_output(CustomHeader)
    prompt = ChatPromptTemplate.from_messages([("system", system_prompt),
                                               ("human", "{input}")])
    few_shot_structured_llm = prompt | llm_with_structured_output
    return few_shot_structured_llm.invoke(question)

class CustomPlotterInput(BaseModel):
    """
    Input schema for the custom plotter tool.
    """
    question: str = Field(description="Description of the plot")
    sys_bio_model: ModelData = Field(description="model data",
                                     default=None)
    simulation_name: str = Field(description="Name assigned to the simulation")
    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.
# 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 CustomPlotterTool(BaseTool):
    """
    Tool for custom plotting the y-axis of a plot.
    """
    name: str = "custom_plotter"
    description: str = '''A visualization tool designed to extract and display a subset
                        of the larger simulation plot generated by the simulate_model tool.
                        It allows users to specify particular species for the y-axis, 
                        providing a more targeted view of key species without the clutter 
                        of the full plot.'''
    args_schema: Type[BaseModel] = CustomPlotterInput
    response_format: str = "content_and_artifact"

    def _run(self,
             question: str,
             sys_bio_model: ModelData,
             simulation_name: str,
             state: Annotated[dict, InjectedState]
             ) -> Tuple[str, Union[None, List[str]]]:
        """
        Run the tool.

        Args:
            question (str): The question about the custom plot.
            sys_bio_model (ModelData): The model data.
            simulation_name (str): The name assigned to the simulation.
            state (dict): The state of the graph.

        Returns:
            str: The answer to the question
        """
        logger.log(logging.INFO, "Calling custom_plotter tool %s, %s", question, sys_bio_model)
        # Load the model
        sbml_file_path = state['sbml_file_path'][-1] if len(state['sbml_file_path']) > 0 else None
        model_object = load_biomodel(sys_bio_model, sbml_file_path=sbml_file_path)
        dic_simulated_data = {}
        for data in state["dic_simulated_data"]:
            for key in data:
                if key not in dic_simulated_data:
                    dic_simulated_data[key] = []
                dic_simulated_data[key] += [data[key]]
        # Create a pandas dataframe from the dictionary
        df = pd.DataFrame.from_dict(dic_simulated_data)
        # Get the simulated data for the current tool call
        df = pd.DataFrame(
                df[df['name'] == simulation_name]['data'].iloc[0]
                )
        # df = pd.DataFrame.from_dict(state['dic_simulated_data'])
        species_names = df.columns.tolist()
        # Exclude the time column
        species_names.remove('Time')
        logging.log(logging.INFO, "Species names: %s", species_names)
        # Extract the relevant species from the user question
        results = extract_relevant_species(question, species_names, state)
        print (results)
        if results.relevant_species is None:
            raise ValueError("No species found in the simulation results \
                             that matches the user prompt.")
        extracted_species = []
        # Extract the species from the results
        # that are available in the simulation results
        for species in results.relevant_species:
            if species in species_names:
                extracted_species.append(species)
        logging.info("Extracted species: %s", extracted_species)
        # Include the time column
        extracted_species.insert(0, 'Time')
        return f"Custom plot {simulation_name}",{
                            'dic_data': df[extracted_species].to_dict(orient='records')
                            }| get_model_units(model_object)