import streamlit as st
from transformers import pipeline

# ---------------------------------------------------
# PAGE CONFIG
# ---------------------------------------------------
st.set_page_config(page_title="FitPlan AI", layout="centered")

st.title("🏋️ FitPlan AI – Personalized Fitness Plan")

# ---------------------------------------------------
# LOAD MODEL (Cached)
# ---------------------------------------------------
@st.cache_resource
def load_model():
    return pipeline(
        "text-generation",
        model="google/flan-t5-base"
    )

generator = load_model()

# ---------------------------------------------------
# USER DETAILS
# ---------------------------------------------------
st.subheader("👤 Personal Information")

name = st.text_input("Enter Your Name")

col1, col2 = st.columns(2)

with col1:
    height = st.number_input("Height (in cm)", min_value=100, max_value=250, value=170)

with col2:
    weight = st.number_input("Weight (in kg)", min_value=30, max_value=200, value=70)

# ---------------------------------------------------
# BMI CALCULATION
# ---------------------------------------------------
bmi = None
bmi_category = ""

if height > 0 and weight > 0:
    height_m = height / 100
    bmi = weight / (height_m ** 2)

    st.write(f"### 📊 Your BMI: {bmi:.2f}")

    if bmi < 18.5:
        bmi_category = "Underweight"
        st.info("Category: Underweight")
    elif 18.5 <= bmi < 24.9:
        bmi_category = "Normal weight"
        st.success("Category: Normal weight")
    elif 25 <= bmi < 29.9:
        bmi_category = "Overweight"
        st.warning("Category: Overweight")
    else:
        bmi_category = "Obese"
        st.error("Category: Obese")

# ---------------------------------------------------
# GOAL
# ---------------------------------------------------
st.subheader("🎯 Fitness Goal")

goal = st.selectbox(
    "",
    [
        "Flexible",
        "Weight Loss",
        "Build Muscle",
        "Strength Gaining",
        "Abs Building"
    ]
)

# ---------------------------------------------------
# EQUIPMENT
# ---------------------------------------------------
st.subheader("🏋️ Available Equipment")

col1, col2 = st.columns(2)

with col1:
    dumbbells = st.checkbox("Dumbbells")
    resistance_band = st.checkbox("Resistance Band")
    yoga_mat = st.checkbox("Yoga Mat")
    no_equipment = st.checkbox("No Equipment")
    inclined_bench = st.checkbox("Inclined Bench")
    treadmill = st.checkbox("Treadmill")
    cycle = st.checkbox("Cycle")

with col2:
    skipping_rope = st.checkbox("Skipping Rope")
    hand_gripper = st.checkbox("Hand Gripper")
    pullups_bar = st.checkbox("Pullups Bar")
    weight_plates = st.checkbox("Weight Plates")
    hula_hoop = st.checkbox("Hula Hoop Ring")
    bosu_ball = st.checkbox("Bosu Ball")

equipment = []

equipment_map = {
    "Dumbbells": dumbbells,
    "Resistance Band": resistance_band,
    "Yoga Mat": yoga_mat,
    "No Equipment": no_equipment,
    "Inclined Bench": inclined_bench,
    "Treadmill": treadmill,
    "Cycle": cycle,
    "Skipping Rope": skipping_rope,
    "Hand Gripper": hand_gripper,
    "Pullups Bar": pullups_bar,
    "Weight Plates": weight_plates,
    "Hula Hoop Ring": hula_hoop,
    "Bosu Ball": bosu_ball,
}

for item, selected in equipment_map.items():
    if selected:
        equipment.append(item)

# ---------------------------------------------------
# FITNESS LEVEL
# ---------------------------------------------------
st.subheader("📈 Fitness Level")

fitness_level = st.radio(
    "",
    ["Beginner", "Intermediate", "Advanced"],
    horizontal=True
)

# ---------------------------------------------------
# PROMPT BUILDER
# ---------------------------------------------------
def build_prompt(name, goal, fitness_level, equipment, bmi, bmi_category):
    equipment_text = ", ".join(equipment)

    prompt = f"""
    Generate a detailed 5-day personalized workout plan.

    Name: {name}
    Goal: {goal}
    Fitness Level: {fitness_level}
    BMI: {bmi:.2f}
    BMI Category: {bmi_category}
    Available Equipment: {equipment_text}

    Adjust workout intensity according to BMI and fitness level.

    Include:
    - Day-wise workout split
    - Exercises
    - Sets and reps
    - Rest time
    - Warm-up and cool-down
    """

    return prompt


# ---------------------------------------------------
# GENERATE PLAN BUTTON
# ---------------------------------------------------
if st.button("🚀 Generate Personalized Plan"):

    if not name:
        st.error("Please enter your name.")
    elif not equipment:
        st.error("Please select at least one equipment option.")
    elif bmi is None:
        st.error("Please enter valid height and weight.")
    else:
        with st.spinner("Generating your AI workout plan..."):

            prompt = build_prompt(
                name,
                goal,
                fitness_level,
                equipment,
                bmi,
                bmi_category
            )

            result = generator(
                prompt,
                max_length=600,
                do_sample=True,
                temperature=0.7
            )

            plan = result[0]["generated_text"]

        st.success(f"✅ Plan Generated Successfully, {name}!")
        st.markdown("## 📝 Your Personalized Workout Plan")
        st.write(plan)