import streamlit as st
import cohere
import os


# 'growth', 'remove_negations', 'self_affirmation'
EXAMPLES = {
    "I lost my Job": ("I lost my Job.", 'growth'),
    "I hate when you show off your expensive motorbike": ("I hate when you show off your expensive motorbike.", 'remove_negations'),
    "I have a lot of work": ("I have a lot of work.", 'self_affirmation'),
    "Are you saying you are better than me": ("Are you saying you are better than me?", 'detect_rhetoric'),
    "I am worried about upcomig exams": ("I am worried about upcomig exams.", 'self_affirmation'),
}
categories = ['growth', 'remove_negations', 'self_affirmation', 'detect_rhetoric']

co = cohere.Client(os.getenv('COHERE_API_KEY')) 

# Initialization
def fill_example_caption(**kwargs):
    print(f"Called with {kwargs}")
    print(f"selectbox {prompt}")

def generate_hashtags(input):
    if len(input) == 0:
        return None
    
    if category == "growth":
        input_prompt = "This program will take a sentence as input and reframe it with positive and growth mindset.\n--\nInput: I have a lot of work to do today.\nOutput: I have a lot of work to do today. It's better for me to make a list and break it down into smaller chunks and finish them off one by one.\n--\nInput: I am barely able to lift 10 pound weights.\nOutput: I should exercise more consistently and gradually increase my weight limit. I should also try eating healthy foods.\n--\nInput: {}\nOutput:".format(input.strip())
    if category == "remove_negations":
        input_prompt = "This program takes an input and reframes it by removing negative words and rephrasing with with positivity while preserving the meaning of the original sentence. The word 'no' doesn't appear in the output\n--\nInput: I wouldn’t say I don’t want to go.\nOutput: I would like to go.\n--\nInput: That’s not a bad idea.\nOutput: That is a good idea.\n--\nInput: {}\nOutput:".format(input.strip())
    if category == "self_affirmation":
        input_prompt = "This program takes an input and reframes it with positive affirmations.\n-- \nInput: I am worthless\nOutput: I have value and my actions impact the world around me.\n--\nInput: I am a failure\nOutput: I never fail, only my attempts do. I will keep trying and improving\n--\nInput:  I don’t like myself\nOutput: I accept myself and improve what I can.\n--\nInput: {}\nOutput:".format(input.strip())
    if category == "detect_rhetoric":
        input_prompt = "This program takes an input and tells if it is rhetoric or not.\n--\nInput: Are you kidding me?\nOutput: rhetoric\n--\nInput: How should I know?\nOutput: rhetoric\n--\nInput:  How much does this cost?\nOutput: not rhetoric\n--\nInput: Do you know the answer to this problem?\nOutput: not rhetoric\n--\nInput: Is this for real?\nOutput: rhetoric\n--\nInput: Can you please finish this task by monday?\nOutput: not rhetoric\n--\nInput:{}\nOutput:".format(input.strip())
    response = co.generate( 
    model='xlarge', 
    prompt=input_prompt, 
    max_tokens=100, 
    temperature=0.9, 
    k=0, 
    p=1, 
    frequency_penalty=0, 
    presence_penalty=0, 
    stop_sequences=["--"], 
    return_likelihoods='NONE') 
    
    output = response.generations[0].text
    return output

st.title('ShubhMitra शुभ-मित्रा')
st.write('''AI assistance for positive reframing of situation''')

prompts = ["Choose"] + list(EXAMPLES.keys())
prompt = st.selectbox(
    'Examples (select from this list)',
    prompts,
    index=0,
)
if prompt == "Choose":
    input_val = ""
    input_cat = categories[0]
else:
    input_val = EXAMPLES[prompt][0]
    input_cat = EXAMPLES[prompt][1]
input = st.text_area('Enter your post title caption or select from examples', input_val, height=100)

category = st.radio(
    "Please select a category",
    categories,
    categories.index(input_cat),
    )

if st.button('Reframe'):
    output = generate_hashtags(input)
    st.write(output)

'''
---
Growth Mindset: Viewing a challenging event as an opportunity to grow and improve

Remove Negations: Positive way of saying things. It's not a bad idea -> It's a good idea.

Self Affirmation: Reprhase the sentence in a affirmation 

Detect Rhetoric: Detect if the questions asked are sincere or just rhetorical. eg: Are you kidding me?

This is a Cohere API powered implementation, taking inspiration from a custom model around positive reframing available [here](https://huggingface.co/spaces/Ella2323/Positive-Reframing)
'''