#importing dependencies
import os
from keys import token

import streamlit as st
from langchain import HuggingFaceHub
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())

os.environ["HUGGINGFACEHUB_API_TOKEN"] = token
repo_id = "tiiuae/falcon-40b"

#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)
llm = HuggingFaceHub(
    repo_id=repo_id, model_kwargs={"temperature": 0.9, "max_length": 500}
)
llm_facts = HuggingFaceHub(
    repo_id=repo_id, model_kwargs={"temperature": 0.0, "max_length": 500}
)
main_chain = LLMChain(prompt=context, llm=llm)
yoda_grammar_chain = LLMChain(prompt=yoda_grammar, llm=llm_facts)
keyword_chain = LLMChain(prompt=find_keyword, llm=llm_facts)
wiki_chain = LLMChain(prompt=context, llm=llm_facts)


#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)