import streamlit as st
import pandas as pd
import transformers
from transformers import pipeline, AutoTokenizer
import torch


# Load Hugging Face model (replace with your desired access token)
torch.manual_seed(0)
model = "tiiuae/falcon-7b-instruct"

tokenizer = AutoTokenizer.from_pretrained(model)
pipeline = pipeline(
    "text-generation", #task
    model=model,
    tokenizer=tokenizer,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map="auto",
    max_length=500,
    do_sample=True,
    top_k=10,
    num_return_sequences=1,
    eos_token_id=tokenizer.eos_token_id
)
# Knowledge base data 
knowledge_base = {
    "health_conditions": [
        "diabetes",
        "heart disease",
        "high blood pressure",
        "kidney disease",
        "liver disease",
    ],
    "dietary_restrictions": ["vegetarian", "vegan", "gluten-free", "dairy-free"],
    "food_preferences": ["spicy", "low-carb", "high-protein", "Mediterranean"],
    "fruits": ["apple", "banana", "orange", "grapefruit", "strawberry"],
    "vegetables": ["broccoli", "spinach", "kale", "carrot", "tomato"],
    "whole_grains": ["brown rice", "quinoa", "oats", "whole-wheat bread", "barley"],
    "lean_proteins": ["chicken breast", "fish", "beans", "lentils", "tofu"],
    "healthy_fats": ["avocado", "nuts", "seeds", "olive oil"],
}

def get_prompt(name, age, gender, weight, height, body_type, health_conditions, dietary_restrictions, food_preferences):
    """
    Generates a prompt for model based on user input.

    Args:
        name (str): User's name.
        age (int): User's age.
        gender (str): User's gender.
        weight (float): User's weight (kg).
        height (float): User's height (meters).
        body_type (str): User's body type.
        health_conditions (list): User's health conditions (if any).
        dietary_restrictions (list): User's dietary restrictions (if any).
        food_preferences (list): User's food preferences (if any).

    Returns:
        str: The generated prompt.
    """

    prompt = f"""Based on the information provided about {name} (age: {age}, gender: {gender}, weight: {weight} kg, height: {height} m, body type: {body_type}), who has the following health conditions: {', '.join(health_conditions) if health_conditions else 'none'} and dietary restrictions: {', '.join(dietary_restrictions) if dietary_restrictions else 'none'}, what would be a personalized diet plan that considers their food preferences for {', '.join(food_preferences) if food_preferences else 'healthy eating'}?

**Knowledge:**

* Health conditions: {', '.join(knowledge_base['health_conditions'])}
* Dietary restrictions: {', '.join(knowledge_base['dietary_restrictions'])}
* Food preferences: {', '.join(knowledge_base['food_preferences'])}
* Fruits: {', '.join(knowledge_base['fruits'])}
* Vegetables: {', '.join(knowledge_base['vegetables'])}
* Whole grains: {', '.join(knowledge_base['whole_grains'])}
* Lean proteins: {', '.join(knowledge_base['lean_proteins'])}
* Healthy fats: {', '.join(knowledge_base['healthy_fats'])}
"""

    return prompt

def predict_diet(prompt):
    input_ids = tokenizer(prompt, return_tensors="pt")
    output = pipeline(prompt)
    predicted_text = output[0]['generated_text']
    return predicted_text



st.title("Personalized Diet Recommendation")

health_conditions = "No health conditions"
dietary_restrictions = "No dietary conditions"
food_preferences = "No food preferences conditions"
# Input fields
name = st.text_input("Name")
age = st.number_input("Age", min_value=0)
gender = st.selectbox("Gender", ["Male", "Female", "Non-binary"])
weight = st.number_input("Weight (kg)", min_value=0.0)
height = st.number_input("Height (meters)", min_value=0.0)
body_type = st.selectbox("Body Type", ["Ectomorph", "Mesomorph", "Endomorph"])
health_conditions = st.text_input("health_conditions")
dietary_restrictions = st.text_input("dietary_restrictions")
food_preferences = st.text_input("food_preferences")


if st.button("Generate"):
    op_prompt = get_prompt(name, age, gender, weight, height, body_type, health_conditions, dietary_restrictions, food_preferences)
    predected_text= predict_diet(op_prompt)
    st.write(predected_text)