'''

Test cases for Talk2Biomodels.

'''

import pandas as pd

from langchain_core.messages import HumanMessage, ToolMessage

from langchain_openai import ChatOpenAI

from ..agents.t2b_agent import get_app

LLM_MODEL = ChatOpenAI(model='gpt-4o-mini', temperature=0)

def test_integration():

    '''

    Test the integration of the tools.

    '''

    unique_id = 1234567

    app = get_app(unique_id, llm_model=LLM_MODEL)

    config = {"configurable": {"thread_id": unique_id}}

    # ##########################################

    # ## Test simulate_model tool

    # ##########################################

    prompt = '''Simulate the model BIOMD0000000537 for 100 hours and time intervals

    100 with an initial concentration of `DoseQ2W` set to 300 and `Dose`

    set to 0. Reset the concentration of `Ab{serum}` to 100 every 25 hours.'''

    # Test the tool get_modelinfo

    response = app.invoke(

                        {"messages": [HumanMessage(content=prompt)]},

                        config=config

                    )

    assistant_msg = response["messages"][-1].content

    print (assistant_msg)

    # Check if the assistant message is a string

    assert isinstance(assistant_msg, str)

    ##########################################

    # Test ask_question tool when simulation

    # results are available

    ##########################################

    # Update state

    app.update_state(config, {"llm_model": LLM_MODEL})

    prompt = """What is the concentration of CRP in serum after 100 hours?

    Round off the value to 2 decimal places."""

    # Test the tool get_modelinfo

    response = app.invoke(

                        {"messages": [HumanMessage(content=prompt)]},

                        config=config

                    )

    assistant_msg = response["messages"][-1].content

    # print (assistant_msg)

    # Check if the assistant message is a string

    assert '211' in assistant_msg

    ##########################################

    # Test the custom_plotter tool when the

    # simulation results are available but

    # the species is not available

    ##########################################

    prompt = """Call the custom_plotter tool to make a plot

        showing only species 'Infected cases'. Let me

        know if these species were not found. Do not

        invoke any other tool."""

    # Update state

    app.update_state(config, {"llm_model": LLM_MODEL}

                    )

    # Test the tool get_modelinfo

    response = app.invoke(

                        {"messages": [HumanMessage(content=prompt)]},

                        config=config

                    )

    assistant_msg = response["messages"][-1].content

    current_state = app.get_state(config)

    # Get the messages from the current state

    # and reverse the order

    reversed_messages = current_state.values["messages"][::-1]

    # Loop through the reversed messages until a

    # ToolMessage is found.

    predicted_artifact = []

    for msg in reversed_messages:

        if isinstance(msg, ToolMessage):

            # Work on the message if it is a ToolMessage

            # These may contain additional visuals that

            # need to be displayed to the user.

            if msg.name == "custom_plotter":

                predicted_artifact = msg.artifact

                break

    # Check if the the predicted artifact is `None`

    assert predicted_artifact is None

    ##########################################

    # Test custom_plotter tool when the

    # simulation results are available

    ##########################################

    prompt = "Plot only CRP related species."

    # Update state

    app.update_state(config, {"llm_model": LLM_MODEL}

                    )

    # Test the tool get_modelinfo

    response = app.invoke(

                        {"messages": [HumanMessage(content=prompt)]},

                        config=config

                    )

    assistant_msg = response["messages"][-1].content

    current_state = app.get_state(config)

    # Get the messages from the current state

    # and reverse the order

    reversed_messages = current_state.values["messages"][::-1]

    # Loop through the reversed messages

    # until a ToolMessage is found.

    expected_header = ['Time', 'CRP{serum}', 'CRPExtracellular']

    expected_header += ['CRP Suppression (%)', 'CRP (% of baseline)']

    expected_header += ['CRP{liver}']

    predicted_artifact = []

    for msg in reversed_messages:

        if isinstance(msg, ToolMessage):

            # Work on the message if it is a ToolMessage

            # These may contain additional visuals that

            # need to be displayed to the user.

            if msg.name == "custom_plotter":

                predicted_artifact = msg.artifact['dic_data']

                break

    # Convert the artifact into a pandas dataframe

    # for easy comparison

    df = pd.DataFrame(predicted_artifact)

    # Extract the headers from the dataframe

    predicted_header = df.columns.tolist()

    # Check if the header is in the expected_header

    # assert expected_header in predicted_artifact

    assert set(expected_header).issubset(set(predicted_header))