#!/usr/bin/env python3

"""

Tool for parameter scan.

"""

import logging

from typing import Type, Annotated

import basico

from pydantic import BaseModel, Field

from langgraph.types import Command

from langgraph.prebuilt import InjectedState

from langchain_core.tools import BaseTool

from langchain_core.messages import ToolMessage

from langchain_core.tools.base import InjectedToolCallId

from .load_biomodel import ModelData, load_biomodel

from .load_arguments import ArgumentData, add_rec_events

# Initialize logger

logging.basicConfig(level=logging.INFO)

logger = logging.getLogger(__name__)

def run_steady_state(model_object,

                     dic_species_to_be_analyzed_before_experiment):

    """

    Run the steady state analysis.

    Args:

        model_object: The model object.

        dic_species_to_be_analyzed_before_experiment: Dictionary of species data.

    Returns:

        DataFrame: The results of the steady state analysis.

    """

    # Update the fixed model species and parameters

    # These are the initial conditions of the model

    # set by the user

    model_object.update_parameters(dic_species_to_be_analyzed_before_experiment)

    logger.log(logging.INFO, "Running steady state analysis")

    # Run the steady state analysis

    output = basico.task_steadystate.run_steadystate(model=model_object.copasi_model)

    if output == 0:

        logger.error("Steady state analysis failed")

        raise ValueError("A steady state was not found")

    logger.log(logging.INFO, "Steady state analysis successful")

    # Store the steady state results in a DataFrame

    df_steady_state = basico.model_info.get_species(model=model_object.copasi_model).reset_index()

    # print (df_steady_state)

    # Rename the column name to species_name

    df_steady_state.rename(columns={'name': 'species_name'},

                           inplace=True)

    # Rename the column concentration to steady_state_concentration

    df_steady_state.rename(columns={'concentration': 'steady_state_concentration'},

                           inplace=True)

    # Rename the column transition_time to steady_state_transition_time

    df_steady_state.rename(columns={'transition_time': 'steady_state_transition_time'},

                           inplace=True)

    # Drop some columns

    df_steady_state.drop(columns=

                         [

                            'initial_particle_number',

                            'initial_expression',

                            'expression',

                            'particle_number',

                            'type',

                            'particle_number_rate',

                            'key',

                            'sbml_id',

                            'display_name'],

                            inplace=True)

    logger.log(logging.INFO, "Steady state results with shape %s", df_steady_state.shape)

    return df_steady_state

class SteadyStateInput(BaseModel):

    """

    Input schema for the steady state tool.

    """

    sys_bio_model: ModelData = Field(description="model data",

                                     default=None)

    arg_data: ArgumentData = Field(

        description="time, species, and reocurring data"

                " that must be set before the steady state analysis"

                " as well as the experiment name", default=None)

    tool_call_id: Annotated[str, InjectedToolCallId]

    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 SteadyStateTool(BaseTool):

    """

    Tool to bring a model to steady state.

    """

    name: str = "steady_state"

    description: str = "A tool to bring a model to steady state."

    args_schema: Type[BaseModel] = SteadyStateInput

    def _run(self,

        tool_call_id: Annotated[str, InjectedToolCallId],

        state: Annotated[dict, InjectedState],

        sys_bio_model: ModelData = None,

        arg_data: ArgumentData = None

    ) -> Command:

        """

        Run the tool.

        Args:

            tool_call_id (str): The tool call ID. This is injected by the system.

            state (dict): The state of the tool.

            sys_bio_model (ModelData): The model data.

            arg_data (ArgumentData): The argument data.

        Returns:

            Command: The updated state of the tool.

        """

        logger.log(logging.INFO, "Calling the steady_state tool %s, %s",

                   sys_bio_model, arg_data)

        # print (f'Calling steady_state tool {sys_bio_model}, {arg_data}, {tool_call_id}')

        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)

        # Prepare the dictionary of species data

        # that will be passed to the simulate method

        # of the BasicoModel class

        dic_species_to_be_analyzed_before_experiment = {}

        if arg_data:

            # Prepare the dictionary of species data

            if arg_data.species_to_be_analyzed_before_experiment is not None:

                dic_species_to_be_analyzed_before_experiment = dict(

                    zip(arg_data.species_to_be_analyzed_before_experiment.species_name,

                        arg_data.species_to_be_analyzed_before_experiment.species_concentration))

            # Add reocurring events (if any) to the model

            if arg_data.reocurring_data is not None:

                add_rec_events(model_object, arg_data.reocurring_data)

        # Run the parameter scan

        df_steady_state = run_steady_state(model_object,

                                           dic_species_to_be_analyzed_before_experiment)

        print (df_steady_state)

        # Prepare the dictionary of scanned data

        # that will be passed to the state of the graph

        dic_steady_state_data = {

            'name': arg_data.experiment_name,

            'source': sys_bio_model.biomodel_id if sys_bio_model.biomodel_id else 'upload',

            'tool_call_id': tool_call_id,

            'data': df_steady_state.to_dict(orient='records')

        }

        # Prepare the dictionary of updated state for the model

        dic_updated_state_for_model = {}

        for key, value in {

            "model_id": [sys_bio_model.biomodel_id],

            "sbml_file_path": [sbml_file_path],

            "dic_steady_state_data": [dic_steady_state_data]

            }.items():

            if value:

                dic_updated_state_for_model[key] = value

        # Return the updated state

        return Command(

                update=dic_updated_state_for_model|{

                # Update the message history

                "messages": [

                ToolMessage(

                        content=f"Steady state analysis of"

                                f" {arg_data.experiment_name}"

                                " was successful.",

                        tool_call_id=tool_call_id,

                        artifact={'dic_data': df_steady_state.to_dict(orient='records')}

                        )

                    ],

                }

            )