import os

from dotenv import load_dotenv
from langchain import hub
from langchain.agents import AgentExecutor, AgentType, initialize_agent, load_tools
from langchain.agents.format_scratchpad import format_log_to_str
from langchain.agents.output_parsers import (
    JSONAgentOutputParser,
    ReActSingleInputOutputParser,
)
from langchain.chains.conversation.memory import ConversationBufferWindowMemory
from langchain.chat_models import ChatOpenAI
from langchain.llms import OpenAI
from langchain.tools import ArxivQueryRun, WikipediaQueryRun, tool
from langchain.tools.render import render_text_description_and_args, format_tool_to_openai_function
from langchain.utilities import ArxivAPIWrapper, WikipediaAPIWrapper
from langchain.prompts import MessagesPlaceholder
from langchain.schema import ChatMessage, SystemMessage

from llamp.mp.agents import (
    MPSummaryExpert,
    MPThermoExpert,
    MPElasticityExpert,
    MPDielectricExpert,
    MPPiezoelectricExpert,
    MPMagnetismExpert,
    MPElectronicExpert,
)

load_dotenv()

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", None)

# OPENAI_GPT_MODEL = "gpt-4-1106-preview"
OPENAI_GPT_MODEL = "gpt-3.5-turbo-1106"

import re

mp_llm = ChatOpenAI(
    temperature=0.1,
    model="gpt-3.5-turbo-1106",
    openai_api_key=OPENAI_API_KEY,
    openai_organization=None,
    max_retries=3,
    # streaming=True
)

llm = ChatOpenAI(
    temperature=0.5,
    model="gpt-4-1106-preview",
    openai_api_key=OPENAI_API_KEY,
    openai_organization=None,
    # streaming=True
)

wikipedia = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
arxiv = ArxivQueryRun(api_wrapper=ArxivAPIWrapper())

tools = [
    MPThermoExpert(llm=mp_llm).as_tool(agent_kwargs=dict(return_intermediate_steps=True)),
    MPElasticityExpert(llm=mp_llm).as_tool(agent_kwargs=dict(return_intermediate_steps=True)),
    MPSummaryExpert(llm=mp_llm).as_tool(agent_kwargs=dict(return_intermediate_steps=True)),
    MPDielectricExpert(llm=mp_llm).as_tool(agent_kwargs=dict(return_intermediate_steps=True)),
    MPMagnetismExpert(llm=mp_llm).as_tool(agent_kwargs=dict(return_intermediate_steps=True)),
    MPElectronicExpert(llm=mp_llm).as_tool(agent_kwargs=dict(return_intermediate_steps=True)),
    # arxiv,
    # wikipedia,
]

prompt = hub.pull("hwchase17/react-multi-input-json")
prompt.messages[0].prompt.template = re.sub(
    r"\s+", " ",
    """You are a data-aware agent that can consult materials-related
    data through Materials Project (MP) database, arXiv, and Wikipedia. Ask 
    user to clarify their queries if needed. Please note that you don't have 
    direct control over MP but through multiple assistant agents to help you. 
    You need to provide complete context in the input for them to do their job.
    """).replace("\n", " ") + prompt.messages[0].prompt.template

prompt = prompt.partial(
    tools=render_text_description_and_args(tools),
    tool_names=", ".join([t.name for t in tools]),
)

agent = (
    {
        "input": lambda x: x["input"],
        "agent_scratchpad": lambda x: format_log_to_str(x["intermediate_steps"]),
    }
    | prompt
    | llm.bind(stop=["Observation"])
    # | map_reduce_chain  # TODO: Add map-reduce after LLM
    | JSONAgentOutputParser()
)


conversational_memory = ConversationBufferWindowMemory(
    memory_key='chat_history',
    k=5,
    return_messages=True
)

agent_kwargs = {
    "handle_parsing_errors": True,
    "extra_prompt_messages": [
        MessagesPlaceholder(variable_name="chat_history"),
        # SystemMessage(content=re.sub(
        #     r"\s+", " ",
        #     """You are a helpful data-aware agent that can consult materials-related
        #     data through Materials Project (MP) database, arXiv, and Wikipedia. Ask 
        #     user to clarify their queries if needed. Please note that you don't have 
        #     direct control to MP but through multiple assistant agents to help you. 
        #     You need to provide complete context for them to do their job. 
        #     """).replace("\n", " ")
        # )
        ],
    # "early_stopping_method": 'generate',
    # "extra_prompt_messages": 
    # )
}

agent_executor = initialize_agent(
    agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    tools=tools,
    llm=llm,
    verbose=True,
    max_iterations=5,
    memory=conversational_memory,
    agent_kwargs=agent_kwargs,
    handle_parsing_errors=True,
)

# agent_executor = initialize_agent(
#     tools=tools,
#     llm=llm,
#     agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
#     verbose=True,
#     max_iterations=5,
# )

/home/cyrus/miniconda3/envs/llamp/lib/python3.11/site-packages/mp_api/client/mprester.py:230: UserWarning: mpcontribs-client not installed. Install the package to query MPContribs data, or construct pourbaix diagrams: 'pip install mpcontribs-client'
  warnings.warn(

agent_executor.invoke({
    "input": "What's the stiffest materials with the lowest formation energy in Si-O system?"
})


> Entering new AgentExecutor chain...
Thought: To answer this question, I need to find materials in the Si-O system with the lowest formation energy and the highest stiffness. I will use the MPThermoExpert tool to search for materials in the Si-O system and sort them by formation energy. After finding candidates, I will need to use the MPElasticityExpert tool to determine the stiffness of these materials.

Action:
```json
{
  "action": "MPThermoExpert",
  "action_input": {
    "input": "What are the materials with the lowest formation energy in the Si-O system?"
  }
}
```

> Entering new AgentExecutor chain...

/home/cyrus/miniconda3/envs/llamp/lib/python3.11/site-packages/mp_api/client/mprester.py:230: UserWarning: mpcontribs-client not installed. Install the package to query MPContribs data, or construct pourbaix diagrams: 'pip install mpcontribs-client'
  warnings.warn(
/home/cyrus/miniconda3/envs/llamp/lib/python3.11/site-packages/mp_api/client/mprester.py:230: UserWarning: mpcontribs-client not installed. Install the package to query MPContribs data, or construct pourbaix diagrams: 'pip install mpcontribs-client'
  warnings.warn(

Action:
```json
{
  "action": "search_materials_thermo__get",
  "action_input": {
    "formula": "Si-O",
    "sort_fields": "formation_energy_per_atom",
    "limit": 5,
    "fields": "material_id,formula_pretty,formation_energy_per_atom"
  }
}
```{"formula": "Si-O", "sort_fields": "formation_energy_per_atom", "limit": 5, "fields": "material_id,formula_pretty,formation_energy_per_atom"}

/home/cyrus/miniconda3/envs/llamp/lib/python3.11/site-packages/mp_api/client/core/client.py:821: UserWarning: The server does not support the request made to https://api.materialsproject.org/materials/thermo/?_fields=material_id%2Cformula_pretty%2Cformation_energy_per_atom&formula=Si-O&_limit=1000. This may be due to an outdated mp-api package, or a problem with the query.
  warnings.warn(

Error on search_materials_thermo__get: REST query returned with error status code 400 on URL https://api.materialsproject.org/materials/thermo/?_fields=material_id%2Cformula_pretty%2Cformation_energy_per_atom&formula=Si-O&_limit=1000 with message:
Problem processing one or more provided formulas.. Please provide more information or try smaller request.
Error on search_materials_thermo__get: REST query returned with error status code 400 on URL https://api.materialsproject.org/materials/thermo/?_fields=material_id%2Cformula_pretty%2Cformation_energy_per_atom&formula=Si-O&_limit=1000 with message:
Problem processing one or more provided formulas.. Please provide more information or try smaller request.Thought: The query may be too broad or ambiguous. I will refine the search by specifying the chemical system more precisely.

Action:
```json
{
  "action": "search_materials_thermo__get",
  "action_input": {
    "chemsys": "Si-O",
    "sort_fields": "formation_energy_per_atom",
    "limit": 5,
    "fields": "material_id,formula_pretty,formation_energy_per_atom"
  }
}
```
{"chemsys": "Si-O", "sort_fields": "formation_energy_per_atom", "limit": 5, "fields": "material_id,formula_pretty,formation_energy_per_atom"}

Retrieving ThermoDoc documents:   0%|          | 0/836 [00:00<?, ?it/s]

[{'formula_pretty': 'Si2O5', 'material_id': 'mp-862998', 'formation_energy_per_atom': -2.289431383571429}, {'formula_pretty': 'SiO2', 'material_id': 'mp-733790', 'formation_energy_per_atom': -3.039244523055555}, {'formula_pretty': 'SiO2', 'material_id': 'mp-6922', 'formation_energy_per_atom': -3.036741518333334}, {'formula_pretty': 'SiO2', 'material_id': 'mp-556985', 'formation_energy_per_atom': -2.992733753333333}, {'formula_pretty': 'SiO2', 'material_id': 'mp-556994', 'formation_energy_per_atom': -2.9575402525000007}]```
{
  "action": "Final Answer",
  "action_input": "The materials with the lowest formation energy in the Si-O system are: Si2O5 (mp-862998) with a formation energy per atom of -2.289 eV/atom, SiO2 (mp-733790) with a formation energy per atom of -3.039 eV/atom, SiO2 (mp-6922) with a formation energy per atom of -3.037 eV/atom, SiO2 (mp-556985) with a formation energy per atom of -2.993 eV/atom, and SiO2 (mp-556994) with a formation energy per atom of -2.958 eV/atom."
}
```

> Finished chain.

Observation: {'input': 'What are the materials with the lowest formation energy in the Si-O system?', 'output': 'The materials with the lowest formation energy in the Si-O system are: Si2O5 (mp-862998) with a formation energy per atom of -2.289 eV/atom, SiO2 (mp-733790) with a formation energy per atom of -3.039 eV/atom, SiO2 (mp-6922) with a formation energy per atom of -3.037 eV/atom, SiO2 (mp-556985) with a formation energy per atom of -2.993 eV/atom, and SiO2 (mp-556994) with a formation energy per atom of -2.958 eV/atom.', 'intermediate_steps': [(AgentAction(tool='search_materials_thermo__get', tool_input={'formula': 'Si-O', 'sort_fields': 'formation_energy_per_atom', 'limit': 5, 'fields': 'material_id,formula_pretty,formation_energy_per_atom'}, log='Action:\n```json\n{\n  "action": "search_materials_thermo__get",\n  "action_input": {\n    "formula": "Si-O",\n    "sort_fields": "formation_energy_per_atom",\n    "limit": 5,\n    "fields": "material_id,formula_pretty,formation_energy_per_atom"\n  }\n}\n```'), 'Error on search_materials_thermo__get: REST query returned with error status code 400 on URL https://api.materialsproject.org/materials/thermo/?_fields=material_id%2Cformula_pretty%2Cformation_energy_per_atom&formula=Si-O&_limit=1000 with message:\nProblem processing one or more provided formulas.. Please provide more information or try smaller request.'), (AgentAction(tool='search_materials_thermo__get', tool_input={'chemsys': 'Si-O', 'sort_fields': 'formation_energy_per_atom', 'limit': 5, 'fields': 'material_id,formula_pretty,formation_energy_per_atom'}, log='Thought: The query may be too broad or ambiguous. I will refine the search by specifying the chemical system more precisely.\n\nAction:\n```json\n{\n  "action": "search_materials_thermo__get",\n  "action_input": {\n    "chemsys": "Si-O",\n    "sort_fields": "formation_energy_per_atom",\n    "limit": 5,\n    "fields": "material_id,formula_pretty,formation_energy_per_atom"\n  }\n}\n```\n'), [{'formula_pretty': 'Si2O5', 'material_id': 'mp-862998', 'formation_energy_per_atom': -2.289431383571429}, {'formula_pretty': 'SiO2', 'material_id': 'mp-733790', 'formation_energy_per_atom': -3.039244523055555}, {'formula_pretty': 'SiO2', 'material_id': 'mp-6922', 'formation_energy_per_atom': -3.036741518333334}, {'formula_pretty': 'SiO2', 'material_id': 'mp-556985', 'formation_energy_per_atom': -2.992733753333333}, {'formula_pretty': 'SiO2', 'material_id': 'mp-556994', 'formation_energy_per_atom': -2.9575402525000007}])]}
Thought:The previous work has identified materials in the Si-O system with the lowest formation energy. The next step is to determine the stiffness of these materials. I will use the MPElasticityExpert tool to find the stiffness (typically represented by the bulk modulus or Young's modulus) for each of these materials.

Action:
```json
{
  "action": "MPElasticityExpert",
  "action_input": {
    "input": "What are the bulk and Young's moduli for Si2O5 (mp-862998), SiO2 (mp-733790), SiO2 (mp-6922), SiO2 (mp-556985), and SiO2 (mp-556994)?"
  }
}
```

> Entering new AgentExecutor chain...

/home/cyrus/miniconda3/envs/llamp/lib/python3.11/site-packages/mp_api/client/mprester.py:230: UserWarning: mpcontribs-client not installed. Install the package to query MPContribs data, or construct pourbaix diagrams: 'pip install mpcontribs-client'
  warnings.warn(
/home/cyrus/miniconda3/envs/llamp/lib/python3.11/site-packages/mp_api/client/mprester.py:230: UserWarning: mpcontribs-client not installed. Install the package to query MPContribs data, or construct pourbaix diagrams: 'pip install mpcontribs-client'
  warnings.warn(

Action:
```json
{
  "action": "search_materials_elasticity__get",
  "action_input": {
    "formula": "Si2O5,SiO2,SiO2,SiO2,SiO2",
    "fields": "material_id,formula_pretty,bulk_modulus,young_modulus"
  }
}
```{"formula": "Si2O5,SiO2,SiO2,SiO2,SiO2", "fields": "material_id,formula_pretty,bulk_modulus,young_modulus"}

Retrieving SummaryDoc documents:   0%|          | 0/1285 [00:00<?, ?it/s]