#importing dependencies
import os

import streamlit as st
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain.tools import WikipediaQueryRun
from langchain.utilities import WikipediaAPIWrapper


wikipedia = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())

#app framework
st.title("🛕Gurubot AI") 
guru = st.radio(
    "Choose your Guru",
    ["A. C. Bhaktivedanta Prabhupada", "Laotse", "Master Minch Yoda of Dagobah", "Karl Marx", "Plato", "Wikipedia"],
    captions = ["ISKCON", "Taoism", "Jedi", "Socialism", "Classic greek philosophy", "Facts, just facts"])
if guru == "A. C. Bhaktivedanta Prabhupada":
    church = "ISKCON"
elif guru == "Laotse":
    church = "Tao"
elif guru == "Master Minch Yoda of Dagobah":
    church = "Jedi"
elif guru == "Karl Marx":
    church = "Social"
elif guru == "Plato":
    church = "Greek"
elif guru == "Wikipedia":
    church = "Science"

if church == "Jedi":
    welcome = "Welcome you are, young padawan. Ask or tell me anything you may."
elif church == "Science":
    welcome = "What do you want to know?"
else:
    welcome = "Welcome, my dear disciple. Ask or tell me anything."

prompt = st.text_input(welcome)

#prompt templates and system prompt context
if church == "ISKCON":
    persona = "the great spiritual leader A. C. Bhaktivedanta Prabhupada"
    scripture = "the Bhagavad Gita"
elif church == "Tao":
    persona = "the great spiritual leader Laozi"
    scripture = "the Tao Te Ching"
elif church == "Jedi":
    persona = "Master Yoda, the character from the Star Wars movie franchise"
    scripture = "the Jedi Path of the Force"
elif church == "Social":
    persona = "the philosopher Karl Marx"
    scripture = "Karl Marx' book ""Das Kapital"""
elif church == "Greek":
    persona = "the ancient Greek philosopher Plato"
    scripture = "Plato's famous writing ""The Republic"""
elif church == "Science":
    persona = "an omniscient talking encyclopedia"
    scripture = "Wikipedia, precisely the explanations you find in {wikipedia_results}"

context = PromptTemplate(
    input_variables = ["persona", "scripture", "topic"],
    template="The user will provide you with an input text \
    in form of a question or \
    in form of a description of a situation or \
    in form of a description of the user's current personal feelings. \
    Act as {persona} \
    and answer the user's question or give advice to the user. \
    Choose your words in accordance to the teachings of {scripture}. \
    And please remember, that you are in fact {persona}, so do not talk in third person about {persona}. \
    Input text: {topic}"
)

find_keyword = PromptTemplate(
    input_variables = ['topic'],
    template="The user will provide you with an input text. \
    Please return the main keyword of this input text. Input text: {topic}"
)

yoda_grammar = PromptTemplate(
    input_variables = ["answer"],
    template="You will be given an input text. \
    Take the object and the first verb of each sentence and move these to the end of the sentence. \
    Example input text: ""I am a powerful jedi master.""\
    Example output text: ""A powerful jedi master I am.""\
    Input text: {answer} \
    Output text:"
)


#llms
llm = OpenAI(temperature=0.9)
llm_facts = OpenAI(temperature=0)
main_chain = LLMChain(llm=llm, prompt=context, verbose=True)
yoda_grammar_chain = LLMChain(llm=llm_facts, prompt=yoda_grammar)
keyword_chain = LLMChain(llm=llm_facts, prompt=find_keyword)
wiki_chain = LLMChain(llm=llm_facts, prompt=context, verbose=True)


#answer on screen if prompt is entered
if prompt:
    if church == "Science":
        keyword = keyword_chain.run(topic=prompt)
        wikipedia_results = wikipedia.run(keyword)
        response = wiki_chain.run(persona=persona, scripture=scripture, topic=prompt)
    elif church == "Jedi":
        answer = main_chain.run(persona=persona, scripture=scripture, topic=prompt)
        response = yoda_grammar_chain.run(answer=answer)
        
    else:
        response = main_chain.run(persona=persona, scripture=scripture, topic=prompt)
    st.write(response)