import os
import streamlit as st

import os
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")


from langchain.chat_models import ChatOpenAI
from langchain.schema import (
    AIMessage,
    HumanMessage,
    SystemMessage
)

# From here down is all the StreamLit UI
st.set_page_config(page_title="Def_off_Testing", page_icon=":robot:")
st.header("Defensive_Offensive_LLM_Test")



# if "sessionMessages" not in st.session_state:
st.session_state.sessionMessages = [
SystemMessage(content="You are playing a game where you have to team up with 2 AI agents to fight a team of 3 enemy units. Your health point is x_1, AI agents' health points are x_2 and x_3. The enemy units' health are y_1, y_2, and y_3. Analyze the situation and give one float number to each agent between 0.0 to 1.0, where 0.0 represents total defense and 1.0 represents total offense as weights. Generate simulated reward functions for each agent including friendly agents and enemy agents."),

HumanMessage(content="We want to cause as much damage as possible on the enemy units and don't care about our health remaining"),
AIMessage(content="Reward: 0 * x_1 + 0 * x_2 + 0 * x_3 - 1.0 * y_1 - 1.0 * y_2 - 1.0 * y_3"),
HumanMessage(content="We want to destroy enemy unit 3 and hopefully preserve the health of the human unit"),
AIMessage(content="Reward: 0.6 * x_1 + 0 * x_2 + 0 * x_3 - 0 * y_1 - 0 * y_2 - 1.0 * y_3"),
HumanMessage(content="We want to protect the human unit and protect AI units if possible, causing damage to enemy units is secondary"),
AIMessage(content="Reward: 1.0 * x_1 + 0.6 * x_2 + 0.6 * x_3 -  0.3 * y_1 - 0.3 * y_2 - 0.3 * y_3"),
]





def load_answer(question):

    st.session_state.sessionMessages.append(HumanMessage(content=question))

    assistant_answer  = chat(st.session_state.sessionMessages )

    st.session_state.sessionMessages.append(AIMessage(content=assistant_answer.content))

    return assistant_answer.content


def get_text():
    input_text = st.text_area("Situation: ", "")
    return input_text


chat = ChatOpenAI(temperature=.25, model='gpt-3.5-turbo')
# chat = ChatOpenAI(temperature=.25, model='gpt-3.5-turbo')
st.write("Your health point is x_1, AI agents' health points are x_2 and x_3.   \n")
st.write("The enemy units' health are y_1, y_2, and y_3.")
# txt = st.text_area(
#     "Example",
#     "You have 5 units with 10 health, enemy has 1 units with 2 health"
#     )


user_input=get_text()
submit = st.button('Generate')  

if submit:
    

    response = load_answer(user_input)
    st.subheader("Answer:")

    st.write(response)

    print(response)