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 pymatgen.core.periodic_table import _pt_data

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

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,
    model="gpt-4-1106-preview",
    openai_api_key=OPENAI_API_KEY,
    openai_organization=None,
    max_retries=5,
    # streaming=True
)

llm = ChatOpenAI(
    temperature=0.7,
    model="gpt-4",
    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=False)),
    MPElasticityExpert(llm=mp_llm).as_tool(agent_kwargs=dict(return_intermediate_steps=False)),
    MPDielectricExpert(llm=mp_llm).as_tool(agent_kwargs=dict(return_intermediate_steps=False)),
    MPMagnetismExpert(llm=mp_llm).as_tool(agent_kwargs=dict(return_intermediate_steps=False)),
    MPElectronicExpert(llm=mp_llm).as_tool(agent_kwargs=dict(return_intermediate_steps=False)),
    MPSummaryExpert(llm=mp_llm).as_tool(agent_kwargs=dict(return_intermediate_steps=False)),
    MPSynthesisExpert(llm=mp_llm).as_tool(agent_kwargs=dict(return_intermediate_steps=False)),
    ArxivAgent(llm=mp_llm).as_tool(agent_kwargs=dict(return_intermediate_steps=False)),
    # 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=20,
    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 are the bandgaps of the following materials: Ba(PdS2)2, FePO4, DyBi2IO4, NaFe(SO4)2, Eu2B5BrO9, Sr5As3H, TlGaS2, AgSbS2, Pr3S3N, LiLaO2. Give me one value of the most stable structure for each material.",
})


> Entering new AgentExecutor chain...
Thought: The human inquiry requires information about the bandgaps of several materials. The best tool to fetch this data is the MPElectronicExpert from the Materials Project. I will seek the bandgap data for each material one by one starting with Ba(PdS2)2.
Action:
```
{
  "action": "MPElectronicExpert",
  "action_input": "What is the bandgap of Ba(PdS2)2?"
}
```

> 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:
```
{
  "action": "search_materials_electronic_structure__get",
  "action_input": {
    "formula": "Ba(PdS2)2"
  }
}
```{"formula": "Ba(PdS2)2"}

Retrieving ElectronicStructureDoc documents:   0%|          | 0/1 [00:00<?, ?it/s]

[{'band_gap': 0.7791999999999999, 'formula_pretty': 'Ba(PdS2)2', 'material_id': 'mp-28967'}]Action:
```
{
  "action": "Final Answer",
  "action_input": "The band gap of Ba(PdS2)2 is approximately 0.779 eV."
}
```

> Finished chain.

Observation: {'input': 'What is the bandgap of Ba(PdS2)2?', 'output': 'The band gap of Ba(PdS2)2 is approximately 0.779 eV.'}
Thought:I found the bandgap for Ba(PdS2)2. Now, I'll proceed with the next material, which is FePO4.
Action:
```
{
  "action": "MPElectronicExpert",
  "action_input": "What is the bandgap of FePO4?"
}
```

> 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:
```
{
  "action": "search_materials_electronic_structure__get",
  "action_input": {
    "formula": "FePO4"
  }
}
```{"formula": "FePO4"}

Retrieving ElectronicStructureDoc documents:   0%|          | 0/30 [00:00<?, ?it/s]

[{'band_gap': 2.8055000000000003, 'formula_pretty': 'FePO4', 'material_id': 'mp-504149'}, {'band_gap': 2.1932, 'formula_pretty': 'FePO4', 'material_id': 'mp-753395'}, {'band_gap': 2.7245, 'formula_pretty': 'FePO4', 'material_id': 'mp-766442'}, {'band_gap': 2.1857, 'formula_pretty': 'FePO4', 'material_id': 'mp-752872'}, {'band_gap': 2.7331000000000003, 'formula_pretty': 'FePO4', 'material_id': 'mp-705201'}, {'band_gap': 2.5353000000000003, 'formula_pretty': 'FePO4', 'material_id': 'mp-2236419'}, {'band_gap': 2.4257, 'formula_pretty': 'FePO4', 'material_id': 'mp-504276'}, {'band_gap': 1.7056, 'formula_pretty': 'FePO4', 'material_id': 'mp-561445'}, {'band_gap': 2.4147, 'formula_pretty': 'FePO4', 'material_id': 'mp-2234217'}, {'band_gap': 2.7961, 'formula_pretty': 'FePO4', 'material_id': 'mp-26792'}, {'band_gap': 2.5497, 'formula_pretty': 'FePO4', 'material_id': 'mp-753636'}, {'band_gap': 1.6668000000000003, 'formula_pretty': 'FePO4', 'material_id': 'mp-776288'}, {'band_gap': 2.4109, 'formula_pretty': 'FePO4', 'material_id': 'mp-753756'}, {'band_gap': 1.2423000000000002, 'formula_pretty': 'FePO4', 'material_id': 'mp-26612'}, {'band_gap': 0.9954999999999998, 'formula_pretty': 'FePO4', 'material_id': 'mp-25303'}, {'band_gap': 0.8201, 'formula_pretty': 'FePO4', 'material_id': 'mp-1048924'}, {'band_gap': 2.6991, 'formula_pretty': 'FePO4', 'material_id': 'mp-19109'}, {'band_gap': 2.7823, 'formula_pretty': 'FePO4', 'material_id': 'mp-760343'}, {'band_gap': 2.7489, 'formula_pretty': 'FePO4', 'material_id': 'mp-2216699'}, {'band_gap': 2.6721000000000004, 'formula_pretty': 'FePO4', 'material_id': 'mp-559295'}, {'band_gap': 0.6198000000000001, 'formula_pretty': 'FePO4', 'material_id': 'mp-2234625'}, {'band_gap': 2.6866, 'formula_pretty': 'FePO4', 'material_id': 'mp-704299'}, {'band_gap': 1.6047, 'formula_pretty': 'FePO4', 'material_id': 'mp-1225011'}, {'band_gap': 2.4351, 'formula_pretty': 'FePO4', 'material_id': 'mp-704220'}, {'band_gap': 2.4307, 'formula_pretty': 'FePO4', 'material_id': 'mp-694541'}, {'band_gap': 1.3138000000000003, 'formula_pretty': 'FePO4', 'material_id': 'mp-20361'}, {'band_gap': 2.5686999999999998, 'formula_pretty': 'FePO4', 'material_id': 'mp-19752'}, {'band_gap': 2.4187000000000003, 'formula_pretty': 'FePO4', 'material_id': 'mp-766460'}, {'band_gap': 0.36719999999999997, 'formula_pretty': 'FePO4', 'material_id': 'mp-540111'}, {'band_gap': 1.1156000000000001, 'formula_pretty': 'FePO4', 'material_id': 'mp-772512'}]Action:
```
{
  "action": "Final Answer",
  "action_input": "The bandgap of FePO4 varies depending on the specific material entry. The reported bandgaps range from as low as 0.3672 eV to as high as 2.8055 eV across different entries in the Materials Project database."
}
```

> Finished chain.

Observation: {'input': 'What is the bandgap of FePO4?', 'output': 'The bandgap of FePO4 varies depending on the specific material entry. The reported bandgaps range from as low as 0.3672 eV to as high as 2.8055 eV across different entries in the Materials Project database.'}
Thought:I found the bandgap for FePO4. Now, I'll proceed with the next material, which is DyBi2IO4.
Action:
```
{
  "action": "MPElectronicExpert",
  "action_input": "What is the bandgap of DyBi2IO4?"
}
```

> 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(

print(agent_executor.memory.chat_memory.messages[1].content)

Here are the bandgaps for the most stable structures of the requested materials: 
- Ba(PdS2)2: 0.779 eV 
- FePO4: 0.3672 eV 
- DyBi2IO4: 1.367 eV 
- NaFe(SO4)2: 2.094 eV 
- Eu2B5BrO9: 0.8801 eV 
- Sr5As3H: 0.676 eV 
- TlGaS2: 1.7885 eV 
- AgSbS2: 0.0 eV (Metallic) 
- Pr3S3N: 1.15 eV 
- LiLaO2: 3.6258 eV

print("To synthesize YMnO3, a solid-state reaction method can be employed. Here is a summarized procedure based on the synthesis details provided in the literature with DOI: 10.1021/ic502027k:\n\n1. Dissolve Y(NO3)3·6H2O, Mn(NO3)2·4H2O, and a source of [OH-] in water to create a homogeneous solution.\n2. Add citric acid to the solution to act as a mixing media.\n3. Heat the mixture at 250°C for 2 hours to achieve pre-calcination.\n4. Finally, heat the pre-calcined mixture at a temperature range of 700°C to 930°C for 10 to 18 hours in an air atmosphere to form YMnO3.")

To synthesize YMnO3, a solid-state reaction method can be employed. Here is a summarized procedure based on the synthesis details provided in the literature with DOI: 10.1021/ic502027k:

1. Dissolve Y(NO3)3·6H2O, Mn(NO3)2·4H2O, and a source of [OH-] in water to create a homogeneous solution.
2. Add citric acid to the solution to act as a mixing media.
3. Heat the mixture at 250°C for 2 hours to achieve pre-calcination.
4. Finally, heat the pre-calcined mixture at a temperature range of 700°C to 930°C for 10 to 18 hours in an air atmosphere to form YMnO3.