import streamlit as st
import re
from ast import literal_eval
import openai
import plotly.graph_objects as go

api_key = st.secrets["openai_key"]
openai.api_key = api_key
st.header('GPT-based Recipe Generation and Nutrition Analysis')
st.markdown("Generate recipe from given ingredients and food name. Analyze the healthiness of the generated recipe. Tips: to make the recipe healthier, indicate healthiness in the food name (e.g., healthy chicken soup, very healthy fried rice, low-carb noodles)")
name = st.text_input('Food name (Optional)')
ingr = st.text_input('Ingredients')
prompt = open('prompt.txt', 'r').read()

if st.button('Generate'):

    prompt = prompt.replace('[FOOD_NAME]', name).replace('[INGREDIENTS]', ingr)
    max_token = 3500 - len(prompt)//4
    with st.spinner('Please wait for completion'):
        response = openai.Completion.create(model="text-davinci-003", prompt=prompt, temperature=0, max_tokens=max_token)
        response = response['choices'][0]['text']
#         response = '''Ingredients: ['1 lb ground beef', '1 onion, diced', '2 cloves garlic, minced', '1 red bell pepper, diced', '1 green bell pepper, diced', '1 jalapeño, diced', '1/2 cup tomato sauce', '1/4 cup Worcestershire sauce', '1/4 cup soy sauce', '1/4 cup brown sugar', '1/4 cup red wine vinegar', '1/4 cup olive oil', 'Salt and pepper, to taste']
# Instructions: ['In a large skillet over medium-high heat, cook the ground beef until browned. Drain off any excess fat.', 'Add the onion, garlic, bell peppers, and jalapeño to the skillet and cook until softened, about 5 minutes.', 'Stir in the tomato sauce, Worcestershire sauce, soy sauce, brown sugar, red wine vinegar, and olive oil. Season with salt and pepper, to taste.', 'Reduce heat to low and simmer for 10 minutes, stirring occasionally.', 'Serve over cooked rice or noodles.']
# Image of final dish: ['The final dish is a savory and flavorful beef stir-fry. The beef is cooked with onions, garlic, bell peppers, and jalapeño, and then simmered in a mixture of tomato sauce, Worcestershire sauce, soy sauce, brown sugar, red wine vinegar, and olive oil. The stir-fry is served over cooked rice or noodles.']
# Estimated calories in number: [500, 600]
# Healthiness score out of 100: [70]
# Estimated percentage of food categories: {'Vegetables': 40, 'Whole grains': 0, 'Healthy protein': 30,  'Fruits': 0,  'Healthy oils': 10}
# '''
        print(response)
        full_ingredients = literal_eval(re.findall('Ingredients: (\[.+\])', response)[0])
        full_ingredients = '\n• ' + '\n• '.join(full_ingredients)

        instructions = literal_eval(re.findall('Instructions: (\[.+\])', response)[0])
        instructions = ['\n{}. {}'.format(str(i + 1), x) for i, x in enumerate(instructions)]
        instructions = ''.join(instructions)

        if name != '':
            name += '. '
        description = name + re.findall('Image of final dish: \[(.+)\]', response)[0]

        est_cal = literal_eval(re.findall('Estimated calories in number: (\[.+\])', response)[0])
        est_cal = list(map(int, est_cal))
        compo = literal_eval(re.findall('Estimated percentage of food categories: ({.+})', response)[0])
        if len(est_cal) == 1:
            est_cal = est_cal[0]
        elif len(est_cal) == 2:
            est_cal = sum(est_cal) // 2
        else:
            print("Wrong calories output")
            est_cal = est_cal[0]
        healthiness = literal_eval(re.findall('Healthiness score out of 100: (\[.+\])', response)[0])
        try:
            healthiness = healthiness[0]
        except:
            print("Wrong healthiness estimation")
        image_url = openai.Image.create(prompt=description, n=1, size="1024x1024")['data'][0]['url']
        col1, col2 = st.columns(2)

        with col1:
            st.subheader('Ingredients')
            st.text(full_ingredients)

        with col2:
            st.image(image_url, caption="Illustration of the final dish")

        st.subheader('Instructions')
        st.markdown(instructions)


        st.subheader("Plate Analysis")
        compo = {i: compo[i] for i in compo if compo[i] != 0}
        labels = list(compo.keys())
        values = list(map(int, list(compo.values())))
        if sum(values) != 100:
            labels.append('Others')
            values.append(100 - sum(values))

        col3, col4 = st.columns(2)

        with col3:
            col3.markdown(
                """
                Estimated calories
                <h1>🔥 {}</h1>
                """.format(str(est_cal)), unsafe_allow_html=True
            )
        with col4:
            col4.markdown(
                """
                Healthiness score
                <h1>🌿 {}/100</h1>
                """.format(str(healthiness)), unsafe_allow_html=True
            )
        st.subheader("")
        st.subheader("")
        col1, col2 = st.columns(2)
        with col1:
            col1.markdown(
                """
                Plate composition
                """, unsafe_allow_html=True
            )
            fig = go.Figure(data=[go.Pie(labels=labels, values=values, hole=.3)])
            st.plotly_chart(fig, use_container_width=True)

        with col2:
            col2.markdown(
                """
                Recommended composition
                """, unsafe_allow_html=True
            )
            st.header("")
            st.image('plate_comp.png')