import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

st.set_page_config(page_title="FitPlan AI", layout="centered")

# LOAD MODEL 
@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
    return tokenizer, model

tokenizer, model = load_model()

# TITLE
st.title(" FitPlan AI – User Fitness Profile")

# PERSONAL INFORMATION
st.subheader(" Personal Information")

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

gender = st.radio(
    "Gender",
    ["Male", "Female"],
    horizontal=True
)

# ---------------------------------------------------
# HEIGHT & WEIGHT
# ---------------------------------------------------
col1, col2 = st.columns(2)

with col1:
    height = st.number_input(
        "Height (in cm)",
        min_value=0.0,
        max_value=250.0,
        value=0.0,
        step=0.1
    )

with col2:
    weight = st.number_input(
        "Weight (in kg)",
        min_value=0.0,
        max_value=200.0,
        value=0.0,
        step=0.1
    )

#BMI Function
def bmi_category(bmi):
    if bmi < 18.5:
        return "Underweight"
    elif bmi < 25:
        return "Normal weight"
    elif bmi < 30:
        return "Overweight"
    else:
        return "Obese"

# BMI CALCULATION
bmi = None

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

    st.metric("📊 Your BMI", f"{bmi:.2f}")
    st.info(f"BMI Category: {bmi_category(bmi)}")

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

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

# EQUIPMENT

st.subheader(" Available Equipment")

equipment_map = {}

col1, col2, col3 = st.columns(3)

with col1:
    equipment_map["No Equipment"] = st.checkbox("No Equipment")
    equipment_map["Pull-up Bar"] = st.checkbox("Pull-up Bar")
    equipment_map["Dip Bars"] = st.checkbox("Dip Bars")
    equipment_map["Push-up Handles"] = st.checkbox("Push-up Handles")
    equipment_map["Dumbbells"] = st.checkbox("Dumbbells")
    equipment_map["Adjustable Dumbbells"] = st.checkbox("Adjustable Dumbbells")

with col2:
    equipment_map["Barbell"] = st.checkbox("Barbell")
    equipment_map["Weight Plates"] = st.checkbox("Weight Plates")
    equipment_map["Kettlebells"] = st.checkbox("Kettlebells")
    equipment_map["Medicine Ball"] = st.checkbox("Medicine Ball")
    equipment_map["Yoga Mat"] = st.checkbox("Yoga Mat")
    equipment_map["Resistance Band"] = st.checkbox("Resistance Band")

with col3:
    equipment_map["Bosu Ball"] = st.checkbox("Bosu Ball")
    equipment_map["Stability Ball"] = st.checkbox("Stability Ball")
    equipment_map["Foam Roller"] = st.checkbox("Foam Roller")
    equipment_map["Treadmill"] = st.checkbox("Treadmill")
    equipment_map["Exercise Bike"] = st.checkbox("Exercise Bike")
    equipment_map["Skipping Rope"] = st.checkbox("Skipping Rope")

equipment = [item for item, selected in equipment_map.items() if selected]


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

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

# SUBMIT BUTTON
if st.button(" Submit Profile"):

    if not name:
        st.error("Please enter your name.")

    elif height <= 0 or weight <= 0:
        st.error("Please enter valid height and weight.")

    elif not equipment:
        st.error("Please select at least one equipment option.")

    else:
        st.success("✅ Profile Submitted Successfully!")

        bmi_status = bmi_category(bmi)
        equipment_list = ", ".join(equipment)

        prompt = f"""
You are a certified professional fitness trainer.

Generate a structured 5-day workout plan.

You MUST follow the format EXACTLY.
Do NOT add any extra text.
Do NOT add explanations.
Do NOT change spacing.
Do NOT remove indentation.

Return ONLY the workout plan.

Strict Format Example:

Day 1: Upper Body
1. Push-Ups
   Sets: 3
   Reps: 10-12
   Rest: 60

Now generate for:

Gender: {gender}
BMI: {bmi:.2f} ({bmi_status})
Goal: {goal}
Fitness Level: {fitness_level}
Equipment: {equipment_list}
"""

        with st.spinner("Generating your AI workout plan..."):

            inputs = tokenizer(prompt, return_tensors="pt", truncation=True)

            outputs = model.generate(
                inputs,
                max_new_tokens=600,
                temperature=0.3,
                do_sample=False,
                top_p=0.9,
                repetition_penalty=1.2
            )

            result = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

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