#!/usr/bin/env python3

"""
Tool for parameter scan.
"""

import logging
from dataclasses import dataclass
from typing import Type, Union, List, Annotated, Optional
import pandas as pd
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 TimeData, SpeciesInitialData
from .utils import get_model_units

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

@dataclass
class ParameterScanData(BaseModel):
    """
    Dataclass for storing the parameter scan data.
    """
    species_names: List[str] = Field(
                    description="species to be observed after each scan."
                    " These are the species whose concentration"
                    " will be observed after the parameter scan."
                    " Do not make up this data.",
                    default=[])
    species_parameter_name: str = Field(
                    description="Species or parameter name to be scanned."
                    " This is the species or parameter whose value will be scanned"
                    " over a range of values. This does not include the species"
                    " that are to be observed after the scan."
                    "Do not make up this data.",
                    default=None)
    species_parameter_values: List[Union[int, float]] = Field(
                    description="Species or parameter values to be scanned."
                    " These are the values of the species or parameters that will be"
                    " scanned over a range of values. This does not include the "
                    "species that are to be observed after the scan."
                    "Do not make up this data.",
                    default=None)

@dataclass
class ArgumentData:
    """
    Dataclass for storing the argument data.
    """
    time_data: TimeData = Field(description="time data", default=None)
    species_to_be_analyzed_before_experiment: Optional[SpeciesInitialData] = Field(
                    description=" This is the initial condition of the model."
                    " This does not include species that reoccur or the species"
                    " whose concentration is to be determined/observed at the end"
                    " of the experiment. This also does not include the species"
                    " or the parameter that is to be scanned. Do not make up this data.",
                    default=None)
    parameter_scan_data: ParameterScanData = Field(
                    description="parameter scan data",
                    default=None)
    experiment_name: str = Field(
                    description="An AI assigned `_` separated unique name of"
                    " the parameter scan experiment based on human query."
                    " This must be unique for each experiment.")

def make_list_dic_scanned_data(dic_param_scan, arg_data, sys_bio_model, tool_call_id):
    """
    Prepare the list dictionary of scanned data
    that will be passed to the state of the graph.

    Args:
        dic_param_scan: Dictionary of parameter scan results.
        arg_data: The argument data.
        sys_bio_model: The model data.
        tool_call_id: The tool call ID.

    Returns:
        list: List of dictionary of scanned data.
    """
    list_dic_scanned_data = []
    for species_name, df_param_scan in dic_param_scan.items():
        logger.log(logging.INFO, "Parameter scan results for %s with shape %s",
                    species_name,
                    df_param_scan.shape)
        # Prepare the list dictionary of scanned data
        # that will be passed to the state of the graph
        list_dic_scanned_data.append({
            'name': arg_data.experiment_name+':'+species_name,
            'source': sys_bio_model.biomodel_id if sys_bio_model.biomodel_id else 'upload',
            'tool_call_id': tool_call_id,
            'data': df_param_scan.to_dict()
        })
    return list_dic_scanned_data

def run_parameter_scan(model_object,
                       arg_data,
                       dic_species_data,
                       duration,
                       interval) -> dict:
    """
    Run parameter scan on the model.

    Args:
        model_object: The model object.
        arg_data: The argument data.
        dic_species_data: Dictionary of species data.
        duration: Duration of the simulation.
        interval: Interval between time points in the simulation.

    Returns:
        dict: Dictionary of parameter scan results. Each key is a species name
        and each value is a DataFrame containing the results of the parameter scan.
    """
    # Extract all parameter names from the model
    df_all_parameters = basico.model_info.get_parameters(model=model_object.copasi_model)
    all_parameters = []
    if df_all_parameters is not None:
        # For example model 10 in the BioModels database
        # has no parameters
        all_parameters = df_all_parameters.index.tolist()

    # Extract all species name from the model
    df_all_species = basico.model_info.get_species(model=model_object.copasi_model)
    all_species = df_all_species['display_name'].tolist()

    # Verify if the given species or parameter names to be scanned are valid
    if arg_data.parameter_scan_data.species_parameter_name not in all_parameters + all_species:
        logger.error(
            "Invalid species or parameter name: %s",
            arg_data.parameter_scan_data.species_parameter_name)
        raise ValueError(
            "Invalid species or parameter name: "
            f"{arg_data.parameter_scan_data.species_parameter_name}.")

    # Dictionary to store the parameter scan results
    dic_param_scan_results = {}

    # Loop through the species names that are to be observed
    for species_name in arg_data.parameter_scan_data.species_names:
        # Verify if the given species name to be observed is valid
        if species_name not in all_species:
            logger.error("Invalid species name: %s", species_name)
            raise ValueError(f"Invalid species name: {species_name}.")

        # Copy the model object to avoid modifying the original model
        model_object_copy = model_object.model_copy()

        # Update the fixed model species and parameters
        # These are the initial conditions of the model
        # set by the user
        model_object_copy.update_parameters(dic_species_data)

        # Initialize empty DataFrame to store results
        # of the parameter scan
        df_param_scan = pd.DataFrame()

        # Loop through the parameter that are to be scanned
        for param_value in arg_data.parameter_scan_data.species_parameter_values:
            # Update the parameter value in the model
            model_object_copy.update_parameters(
                {arg_data.parameter_scan_data.species_parameter_name: param_value})
            # Simulate the model
            model_object_copy.simulate(duration=duration, interval=interval)
            # If the column name 'Time' is not present in the results DataFrame
            if 'Time' not in df_param_scan.columns:
                df_param_scan['Time'] = model_object_copy.simulation_results['Time']
            # Add the simulation results to the results DataFrame
            col_name = f"{arg_data.parameter_scan_data.species_parameter_name}_{param_value}"
            df_param_scan[col_name] = model_object_copy.simulation_results[species_name]

        logger.log(logging.INFO, "Parameter scan results with shape %s", df_param_scan.shape)

        # Add the results of the parameter scan to the dictionary
        dic_param_scan_results[species_name] = df_param_scan
    # return df_param_scan
    return dic_param_scan_results

class ParameterScanInput(BaseModel):
    """
    Input schema for the ParameterScan tool.
    """
    sys_bio_model: ModelData = Field(description="model data",
                                     default=None)
    arg_data: ArgumentData = Field(description=
                                   """time, species, and reocurring data
                                   as well as the parameter scan name and
                                   data""",
                                   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 ParameterScanTool(BaseTool):
    """
    Tool for parameter scan.
    """
    name: str = "parameter_scan"
    description: str = """A tool to perform scanning of a given
    parameter over a range of values and observe the effect on
    the concentration of a given species"""
    args_schema: Type[BaseModel] = ParameterScanInput

    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 parameter_scan tool %s, %s",
                   sys_bio_model, arg_data)
        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
        duration = 100.0
        interval = 10
        dic_species_data = {}
        if arg_data:
            # Prepare the dictionary of species data
            if arg_data.species_to_be_analyzed_before_experiment is not None:
                dic_species_data = 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)

            # Set the duration and interval
            if arg_data.time_data is not None:
                duration = arg_data.time_data.duration
                interval = arg_data.time_data.interval

        # Run the parameter scan
        dic_param_scan = run_parameter_scan(model_object,
                                           arg_data,
                                           dic_species_data,
                                           duration,
                                           interval)

        logger.log(logging.INFO, "Parameter scan results ready")
        # Prepare the list dictionary of scanned data
        list_dic_scanned_data = make_list_dic_scanned_data(dic_param_scan,
                                                           arg_data,
                                                           sys_bio_model,
                                                           tool_call_id)
        # 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_scanned_data": list_dic_scanned_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"Parameter scan results of {arg_data.experiment_name}",
                        tool_call_id=tool_call_id,
                        artifact=get_model_units(model_object)
                        )
                    ],
                }
            )