import streamlit as st

from transformers import pipeline



# -------------------------

# PAGE CONFIG

# -------------------------

st.set_page_config(page_title="FitPlan-AI", page_icon="💪")



# -------------------------

# BMI FUNCTIONS

# -------------------------

def calculate_bmi(weight, height_cm):

    height_m = height_cm / 100

    return round(weight / (height_m ** 2), 2)



def get_category(bmi):

    if bmi < 18.5:

        return "Underweight"

    elif bmi < 24.9:

        return "Normal"

    elif bmi < 29.9:

        return "Overweight"

    else:

        return "Obese"



# -------------------------

# LOAD AI MODEL

# -------------------------

@st.cache_resource

def load_model():

    return pipeline(

        "text2text-generation",  # Correct for FLAN-T5

        model="google/flan-t5-base"

    )



generator = load_model()



# -------------------------

# UI STARTS

# -------------------------

st.title("💪 FitPlan-AI: Personalized Fitness Profile")



# Form to collect all inputs

with st.form("fitness_form"):



    st.subheader("Personal Information")

    name = st.text_input("Full Name*", placeholder="Enter your name")



    col1, col2 = st.columns(2)

    with col1:

        height = st.number_input("Height (cm)*", min_value=1.0, step=0.1)

    with col2:

        weight = st.number_input("Weight (kg)*", min_value=1.0, step=0.1)



    st.subheader("Fitness Details")

    goal = st.selectbox(

        "Fitness Goal",

        ["Build Muscle", "Weight Loss", "Strength Gain", "Abs Building", "Flexible"]

    )

    level = st.radio("Fitness Level", ["Beginner", "Intermediate", "Advanced"])

    equipment = st.multiselect(

        "Available Equipment",

        ["Dumbbells", "Resistance Band", "Yoga Mat", "No Equipment", "Kettlebell", "Pull-up Bar"]

    )



    # -------------------------

    # SUBMIT BUTTON AT BOTTOM

    # -------------------------

    submit = st.form_submit_button("Submit Profile")
    # -------------------------
    # HANDLE SUBMISSION
    # -------------------------
    if submit:
        # Validation (Indented from here downwards)
        if not name.strip():
            st.error("Please enter your name.")
        elif height <= 0 or weight <= 0:
            st.error("Height and Weight must be positive values.")
        elif not equipment:
            st.error("Please select at least one equipment option.")
        else:
            # Calculate BMI
            user_bmi = calculate_bmi(weight, height)
            bmi_status = get_category(user_bmi)

            st.success(f"✅ Profile Created Successfully for {name}!")
            st.metric("Your BMI", user_bmi)
            st.info(f"Health Category: **{bmi_status}**")

            equipment_list = ", ".join(equipment)

            # Prepare AI prompt
            prompt = f"""
            Create a professional 5-day structured workout plan.
            User Details:
            Name: {name}
            BMI: {user_bmi} ({bmi_status})
            Goal: {goal}
            Fitness Level: {level}
            Equipment: {equipment_list}
            Include:
            - Warm-up
            - Exercises with sets and reps
            - Rest time
            - Intensity adjustment
            - Day-wise structure
            """

            with st.spinner("Generating your AI workout plan..."):
                result = generator(prompt, max_new_tokens=400)[0]["generated_text"]

            st.subheader("🏋️ Your Personalized Workout Plan")
            st.write(result)